From 47bc09f9c3757b4c3d041202f837638d689a23c2 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sat, 8 Aug 2026 23:53:31 -0500 Subject: [PATCH 1/6] test(drift): port the rigid translation + wizard suites from the quarry Pulled from feat/seg-fast-engine (PR #122) at file level; these are the behavioral spec for the rigid drift PR. test_drift_wizard.py is pruned of its non-rigid coverage (three TestMethodStubs tests: nonrigid selectable, nonrigid field parameterisation, unknown-field fallback) -- non-rigid lands in its own later PR. test_drift_translation.py is rigid-only and comes over verbatim. --- .../tests/migrated/test_drift_translation.py | 552 +++++++++++ spyde/tests/migrated/test_drift_wizard.py | 860 ++++++++++++++++++ 2 files changed, 1412 insertions(+) create mode 100644 spyde/tests/migrated/test_drift_translation.py create mode 100644 spyde/tests/migrated/test_drift_wizard.py diff --git a/spyde/tests/migrated/test_drift_translation.py b/spyde/tests/migrated/test_drift_translation.py new file mode 100644 index 00000000..506596ce --- /dev/null +++ b/spyde/tests/migrated/test_drift_translation.py @@ -0,0 +1,552 @@ +""" +Tests for spyde.drift — rigid translation solve, warp, and DriftModel. + +The acceptance gate from DRIFT_AND_PARTICLES_PLAN.md is numerical, not +structural: recover a synthetically applied shift to better than 0.1 px, and +agree with ``skimage.registration.phase_cross_correlation`` on the same data. +That is what most of this file asserts. + +Qt-free and dask-free — these are pure-compute tests on small synthetic stacks. +""" +from __future__ import annotations + +import numpy as np +import pytest + +from spyde.drift import DriftModel, coverage_mask, frame_source, shift_frame +from spyde.drift.translation import solve_translation + +# The solver's own tolerance target. Sub-pixel ground truth on a smooth +# synthetic scene should land well inside this. +GATE_PX = 0.1 + + +# ── synthetic data ─────────────────────────────────────────────────────────── + +def _scene(h=96, w=112, seed=3, noise=0.02): + """A smooth, asymmetric, non-periodic scene. + + Asymmetric on purpose: a symmetric scene correlates equally well at several + offsets, so a sign error or an axis swap would still pass. Non-periodic on + purpose: a lattice invites the exact wrong-translation lock that ``max_shift`` + exists to prevent, which is a separate test. + + ``noise=0`` gives a band-limited scene, needed wherever a test resamples + twice — bilinear interpolation legitimately destroys pixel-scale noise, so a + round-trip assertion on a noisy scene measures interpolation loss, not the + property under test. + """ + rng = np.random.default_rng(seed) + yy, xx = np.mgrid[0:h, 0:w].astype(np.float64) + img = np.zeros((h, w), dtype=np.float64) + # A handful of gaussian blobs at irregular positions and widths. + for cy, cx, amp, sig in [ + (0.28 * h, 0.22 * w, 1.0, 5.0), + (0.61 * h, 0.44 * w, 0.7, 8.0), + (0.38 * h, 0.73 * w, 0.9, 4.0), + (0.79 * h, 0.66 * w, 0.5, 6.5), + (0.17 * h, 0.58 * w, 0.6, 3.5), + ]: + img += amp * np.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sig ** 2)) + if noise: + img += noise * rng.standard_normal((h, w)) + return img + + +def _shifted_stack(shifts, h=96, w=112, seed=3): + """Stack whose frame i is the scene translated by ``-shifts[i]``. + + So the CORRECTION needed for frame i is ``+shifts[i]`` — matching the + DriftModel sign convention. Built by Fourier phase ramp so sub-pixel truth is + exact rather than interpolated, which keeps the 0.1 px gate meaningful. + """ + base = _scene(h, w, seed) + fy = np.fft.fftfreq(h)[:, None] + fx = np.fft.fftfreq(w)[None, :] + F = np.fft.fft2(base) + frames = [] + for dy, dx in shifts: + # Applying -shift here means +shift is the correction. + ramp = np.exp(-2j * np.pi * (-dy * fy + -dx * fx)) + frames.append(np.real(np.fft.ifft2(F * ramp))) + return np.stack(frames).astype(np.float32) + + +class TestFrameSource: + def test_numpy_stack(self): + arr = np.zeros((5, 8, 9), dtype=np.uint16) + n, get, shape = frame_source(arr) + assert n == 5 and shape == (8, 9) + assert get(3).shape == (8, 9) + + def test_sequence_of_frames(self): + seq = [np.zeros((4, 6)) for _ in range(3)] + n, get, shape = frame_source(seq) + assert n == 3 and shape == (4, 6) + + def test_rejects_2d(self): + with pytest.raises(TypeError, match="3-D"): + frame_source(np.zeros((8, 9))) + + def test_rejects_unknown(self): + with pytest.raises(TypeError, match="cannot read frames"): + frame_source(object()) + + def test_hyperspy_signal_wrong_nav_dim_is_rejected(self): + """A 4D-STEM scan is not a movie; say so instead of solving nonsense.""" + class _AM: + navigation_dimension = 2 + signal_dimension = 2 + + class _Sig: + axes_manager = _AM() + data = np.zeros((3, 3, 4, 4)) + + with pytest.raises(TypeError, match="1-D navigation"): + frame_source(_Sig()) + + def test_dask_reads_one_frame_only(self): + """The Memory-Safety rule, enforced: never compute the whole array.""" + da = pytest.importorskip("dask.array") + arr = da.zeros((6, 8, 9), chunks=(1, 8, 9)) + n, get, shape = frame_source(arr) + assert n == 6 and shape == (8, 9) + called = {"full": 0} + real_compute = da.Array.compute + + def guard(self, *a, **k): + if self.shape == (6, 8, 9): + called["full"] += 1 + return real_compute(self, *a, **k) + + try: + da.Array.compute = guard + f = get(2) + finally: + da.Array.compute = real_compute + assert f.shape == (8, 9) + assert called["full"] == 0, "sliced a frame but computed the whole stack" + + +class TestSolveTranslationAccuracy: + def test_recovers_integer_shifts(self): + truth = np.array([[0, 0], [3, -4], [-6, 2], [1, 7]], dtype=float) + stack = _shifted_stack(truth) + model = solve_translation(stack, device="numpy", upsample=8) + assert np.allclose(model.shifts, truth, atol=GATE_PX), model.shifts + + def test_recovers_subpixel_shifts_inside_gate(self): + """The headline acceptance gate: < 0.1 px on sub-pixel ground truth. + + The truth values are deliberately **off** the ``1/upsample`` grid. Shifts + that happen to be multiples of 1/8 are recovered to 0.00000 px by an + upsample=8 solve — which looks like a spectacular result and actually + tests nothing, because the answer is exactly representable. Off-grid truth + is what makes the tolerance meaningful. + """ + truth = np.array( + [[0, 0], [1.37, -2.83], [-3.06, 0.61], [4.19, 5.44], [-0.72, -1.28]], + dtype=float, + ) + # None of these may land on the upsampled grid, or the test is vacuous. + assert not np.any(np.isclose(truth[1:] * 8, np.round(truth[1:] * 8))) + + stack = _shifted_stack(truth) + model = solve_translation(stack, device="numpy", upsample=8, + reference="first") + err = np.abs(model.shifts - truth) + assert err.max() < GATE_PX, f"max error {err.max():.4f} px\n{model.shifts}" + + def test_higher_upsample_reduces_error(self): + """Off-grid error should shrink as the upsampled grid gets finer. + + This is the test that would have caught the `_upsampled_dft` bug where + the frequency scaling was omitted: with that bug every result quantised to + 1/upsample regardless, so raising upsample changed the quantum but the + error stayed the same order. Here it must genuinely improve. + """ + truth = np.array([[0, 0], [2.31, -1.77], [-3.42, 4.09]], dtype=float) + stack = _shifted_stack(truth) + errs = {} + for u in (2, 8, 32): + m = solve_translation(stack, device="numpy", upsample=u, + reference="first") + errs[u] = float(np.abs(m.shifts - truth).max()) + assert errs[8] < errs[2], errs + assert errs[32] <= errs[8] + 1e-4, errs + + def test_frame_zero_is_the_origin(self): + stack = _shifted_stack(np.array([[0, 0], [2, 3]], dtype=float)) + model = solve_translation(stack, device="numpy") + assert tuple(model.shifts[0]) == (0.0, 0.0) + + def test_agrees_with_skimage_reference(self): + """Parity against the implementation we are replacing.""" + skreg = pytest.importorskip("skimage.registration") + truth = np.array([[0, 0], [2.25, -3.5], [-1.75, 4.125]], dtype=float) + stack = _shifted_stack(truth) + model = solve_translation(stack, device="numpy", upsample=8, + reference="first", apodize=False) + for i in range(1, len(truth)): + ref, _, _ = skreg.phase_cross_correlation( + stack[0], stack[i], upsample_factor=8, normalization="phase") + assert np.allclose(model.shifts[i], ref, atol=0.05), ( + f"frame {i}: ours={model.shifts[i]} skimage={ref}") + + def test_sequential_reference_accumulates(self): + """Sequential mode must return CUMULATIVE shifts, not per-pair deltas.""" + truth = np.array([[0, 0], [2, 0], [4, 0], [6, 0]], dtype=float) + stack = _shifted_stack(truth) + model = solve_translation(stack, device="numpy", reference="sequential", + upsample=4) + assert np.allclose(model.shifts, truth, atol=GATE_PX), model.shifts + + def test_running_reference_survives_one_corrupt_frame(self): + """Why 'running' is the default: a single bad frame must not poison it. + + The frames AFTER the corrupt one are what matters. The corrupt frame's own + shift is meaningless by construction and is not asserted on. + """ + truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3], [8, 4]], dtype=float) + stack = _shifted_stack(truth).copy() + rng = np.random.default_rng(0) + stack[2] = rng.standard_normal(stack.shape[1:]).astype(np.float32) # garbage + model = solve_translation(stack, device="numpy", upsample=8, max_shift=20) + good = [1, 3, 4] + err = np.abs(model.shifts[good] - truth[good]).max() + assert err < 0.5, f"good frames drifted after a corrupt frame: {model.shifts}" + assert model.params["rejected_from_reference"] >= 1, ( + "nothing was kept out of the reference, so this passed by luck rather " + "than by the outlier rejection it is meant to exercise") + + def test_outlier_rejection_can_be_disabled(self): + """And with it off, the corrupt frame really does poison the reference — + which is what makes the test above non-vacuous.""" + truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3], [8, 4]], dtype=float) + stack = _shifted_stack(truth).copy() + rng = np.random.default_rng(0) + stack[2] = rng.standard_normal(stack.shape[1:]).astype(np.float32) + model = solve_translation(stack, device="numpy", upsample=8, max_shift=20, + reject_outliers=False) + assert model.params["rejected_from_reference"] == 0 + good = [3, 4] + assert np.abs(model.shifts[good] - truth[good]).max() > 1.0, ( + "the corrupt frame no longer poisons an unprotected reference — if the " + "solver became robust some other way, check deliberately") + + def test_clean_stack_rejects_nothing(self): + """The rejection must not fire on ordinary frame-to-frame variation.""" + truth = np.array([[0, 0], [1.5, 0.5], [3, 1], [4.5, 1.5], [6, 2]], float) + model = solve_translation(_shifted_stack(truth), device="numpy", + upsample=8, max_shift=20) + assert model.params["rejected_from_reference"] == 0 + assert np.abs(model.shifts - truth).max() < GATE_PX + + +class TestSolveTranslationGuards: + def test_max_shift_rejects_far_peak(self): + """A shift beyond max_shift is clamped out of the search, not returned.""" + truth = np.array([[0, 0], [20, 0]], dtype=float) + stack = _shifted_stack(truth) + model = solve_translation(stack, device="numpy", max_shift=5, upsample=1) + assert abs(model.shifts[1][0]) <= 5.0 + 1e-6, model.shifts + + def test_impossible_bounds_raise(self): + stack = _shifted_stack(np.zeros((2, 2))) + with pytest.raises(ValueError, match="exclude every possible shift"): + solve_translation(stack, device="numpy", max_shift=1, min_shift=50) + + def test_bad_reference_name_raises(self): + stack = _shifted_stack(np.zeros((2, 2))) + with pytest.raises(ValueError, match="unknown reference"): + solve_translation(stack, device="numpy", reference="nonsense") + + def test_fixed_index_out_of_range_raises(self): + stack = _shifted_stack(np.zeros((2, 2))) + with pytest.raises(ValueError, match="outside"): + solve_translation(stack, device="numpy", reference="fixed:99") + + def test_progress_reports_every_frame(self): + stack = _shifted_stack(np.zeros((4, 2))) + seen = [] + solve_translation(stack, device="numpy", progress=lambda d, t: seen.append((d, t))) + assert seen[0] == (1, 4) and seen[-1] == (4, 4) + + def test_on_shift_streams_every_frame_as_it_solves(self): + """The drift caret draws its curve live; `progress` cannot carry that. + + `progress` is only a count, and the shift array is solver-local until the + return — so without this callback a UI can show a bar but not a trace. + """ + truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3]], dtype=float) + seen = [] + model = solve_translation(_shifted_stack(truth), device="numpy", + upsample=8, reference="first", + on_shift=lambda i, dy, dx, s: seen.append((i, dy, dx))) + assert [i for i, _, _ in seen] == list(range(len(truth))), ( + f"expected one callback per frame in order, got {seen}") + streamed = np.array([[dy, dx] for _, dy, dx in seen]) + assert np.allclose(streamed, model.shifts, equal_nan=True), ( + "the streamed values disagree with the returned array") + + def test_on_shift_is_optional(self): + stack = _shifted_stack(np.array([[0, 0], [1, 1]], float)) + assert solve_translation(stack, device="numpy").n_frames == 2 + + def test_cancel_leaves_nan_not_a_silent_partial(self): + stack = _shifted_stack(np.array([[0, 0], [1, 1], [2, 2], [3, 3]], float)) + calls = {"n": 0} + + def cancel(): + calls["n"] += 1 + return calls["n"] > 1 + + model = solve_translation(stack, device="numpy", cancel=cancel) + assert np.isnan(model.shifts[-1]).all(), ( + "a cancelled solve must be detectable, not quietly truncated") + + +class TestAlignmentROI: + """Correlating on a sub-region. Not a speed switch — often the RIGHT answer. + + Whole-frame correlation averages over everything that moved, so on a movie + where the sample itself evolves, the sample's motion contaminates the estimate + of the stage's. Restricting to a static landmark measures the stage alone. + """ + + def test_roi_recovers_the_same_shift_as_the_full_frame(self): + truth = np.array([[0, 0], [2.5, -1.75], [-3.25, 4.0]], dtype=float) + stack = _shifted_stack(truth, h=96, w=112) + full = solve_translation(stack, device="numpy", upsample=8, + reference="first") + roi = solve_translation(stack, device="numpy", upsample=8, + reference="first", roi=(20, 20, 56, 64)) + assert np.abs(roi.shifts - truth).max() < 0.3, roi.shifts + assert np.abs(roi.shifts - full.shifts).max() < 0.3, ( + "the ROI solve disagrees with the full-frame solve on the same data") + + def test_shifts_apply_to_the_whole_frame(self): + """A translation is a translation; the ROI only chooses where to measure.""" + truth = np.array([[0, 0], [3.0, -2.0]], dtype=float) + stack = _shifted_stack(truth, h=96, w=112) + model = solve_translation(stack, device="numpy", upsample=8, + reference="first", roi=(30, 30, 40, 48)) + from spyde.drift import shift_frame + aligned = shift_frame(stack[1], model.shifts[1], fill=0.0) + core = (slice(40, -40), slice(40, -40)) # far OUTSIDE the ROI + resid = np.abs(aligned[core] - stack[0][core]).max() + raw = np.abs(stack[1][core] - stack[0][core]).max() + assert resid < 0.25 * raw, ( + "correcting with an ROI-derived shift did not align the region " + "outside the ROI") + + def test_roi_is_recorded_in_params(self): + stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64) + m = solve_translation(stack, device="numpy", roi=(8, 8, 32, 32)) + assert m.params["roi"] == [8, 8, 32, 32] + assert m.params["frame_shape"] == [64, 64], ( + "frame_shape must stay the FULL frame — the shifts apply to it") + + def test_out_of_bounds_roi_raises_rather_than_clamping(self): + """A silently shrunk ROI would correlate on a region the user never chose.""" + stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64) + for bad in [(0, 0, 80, 32), (40, 40, 32, 32), (-4, 0, 32, 32)]: + with pytest.raises(ValueError, match="outside"): + solve_translation(stack, device="numpy", roi=bad) + + def test_tiny_roi_raises(self): + stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64) + with pytest.raises(ValueError, match="at least"): + solve_translation(stack, device="numpy", roi=(0, 0, 8, 8)) + + def test_malformed_roi_raises(self): + stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64) + with pytest.raises(ValueError, match=r"\(y0, x0, h, w\)"): + solve_translation(stack, device="numpy", roi=(1, 2, 3)) + + def test_roi_ignores_motion_outside_it(self): + """The point of the feature, on data built to punish whole-frame.""" + h, w, n = 96, 128, 5 + base = _scene(h, w, noise=0.0) + frames = [] + for t in range(n): + f = base.copy() + # A big bright square that moves the OTHER way, far from the ROI. + y, x = 60, 78 + 5 * t + f[y:y + 26, x:x + 26] += 3.0 + frames.append(f.astype(np.float32)) + stack = np.stack(frames) + roi = solve_translation(stack, device="numpy", upsample=8, + reference="first", roi=(4, 4, 44, 56)) + full = solve_translation(stack, device="numpy", upsample=8, + reference="first") + # The landmark region is STATIC, so the ROI answer should be ~zero. + assert np.abs(roi.shifts).max() < 0.6, ( + f"ROI solve drifted although its region never moved: {roi.shifts}") + assert np.abs(full.shifts).max() > np.abs(roi.shifts).max(), ( + "the whole-frame solve was not contaminated by the moving square, so " + "this fixture no longer demonstrates why the ROI exists") + + +class TestBackendParity: + """The GPU path has no independent reference, so it is pinned to numpy.""" + + def test_torch_matches_numpy(self): + torch = pytest.importorskip("torch") + truth = np.array([[0, 0], [2.5, -1.75], [-4.25, 3.5]], dtype=float) + stack = _shifted_stack(truth) + ref = solve_translation(stack, device="numpy", upsample=8) + got = solve_translation(stack, device="cpu", upsample=8) + assert got.params["backend"] == "torch" + assert np.allclose(got.shifts, ref.shifts, atol=1e-2), ( + f"torch={got.shifts}\nnumpy={ref.shifts}") + + +class TestWarp: + def test_integer_shift_is_exact_and_preserves_dtype(self): + f = np.arange(24, dtype=np.uint16).reshape(4, 6) + out = shift_frame(f, (1, 2), fill=0, preserve_dtype=True) + assert out.dtype == np.uint16 + # The interior must be bit-identical — no resampling on a whole-pixel move. + # Destination [1:, 2:] is fed by source [:-1, :-2]. + assert np.array_equal(out[1:, 2:], f[:-1, :-2]) + assert np.all(out[0, :] == 0) and np.all(out[:, :2] == 0) + + def test_nan_padding_marks_uncovered(self): + f = np.ones((5, 5), dtype=np.float32) + out = shift_frame(f, (2, 0)) + assert np.isnan(out[:2]).all() + assert np.allclose(out[2:], 1.0) + + def test_subpixel_shift_interpolates_and_pads(self): + f = _scene(32, 32).astype(np.float32) + out = shift_frame(f, (0.5, -0.5)) + assert out.dtype == np.float32 + assert np.isnan(out[0]).all(), "top row needs off-frame data" + assert np.isnan(out[:, -1]).all(), "right column needs off-frame data" + assert np.isfinite(out[3:-3, 3:-3]).all() + + def test_nan_does_not_bleed_into_real_data(self): + """Interpolating with NaN cval would smear it `order` px inward.""" + f = np.ones((16, 16), dtype=np.float32) + out = shift_frame(f, (2.5, 0.0)) + assert np.isfinite(out[4:, :]).all(), "NaN bled past the padded border" + assert np.allclose(out[5:-1, :], 1.0, atol=1e-5) + + def test_preserve_dtype_rejects_subpixel(self): + f = np.zeros((4, 4), dtype=np.uint16) + with pytest.raises(ValueError, match="whole-pixel"): + shift_frame(f, (0.5, 0), fill=0, preserve_dtype=True) + + def test_round_trip_is_sign_symmetric(self): + """Shifting out and back must land where it started. + + This pins the SIGN symmetry, not interpolation fidelity — so the scene is + noise-free (see :func:`_scene`) and cubic interpolation is used. Two + bilinear passes over pixel-scale noise would lose ~1% of amplitude for + entirely legitimate reasons and tell us nothing about the sign. + """ + f = _scene(64, 64, noise=0.0).astype(np.float32) + moved = shift_frame(f, (3.25, -2.5), fill=0.0, order=3) + back = shift_frame(moved, (-3.25, 2.5), fill=0.0, order=3) + core = (slice(10, -10), slice(10, -10)) + err = np.abs(back[core] - f[core]).max() + assert err < 0.01, f"round trip lost {err:.4f} — sign asymmetry?" + + def test_round_trip_beats_the_uncorrected_offset(self): + """Sanity: the corrected result is far closer than the shifted one.""" + f = _scene(64, 64, noise=0.0).astype(np.float32) + moved = shift_frame(f, (3.25, -2.5), fill=0.0) + back = shift_frame(moved, (-3.25, 2.5), fill=0.0) + core = (slice(10, -10), slice(10, -10)) + assert np.abs(back[core] - f[core]).max() < \ + 0.1 * np.abs(moved[core] - f[core]).max() + + def test_rejects_non_finite_shift(self): + with pytest.raises(ValueError, match="finite"): + shift_frame(np.zeros((4, 4)), (np.nan, 0)) + + def test_coverage_matches_finite_pixels(self): + f = np.ones((20, 20), dtype=np.float32) + for s in [(0, 0), (3, -2), (2.5, 1.25), (-4.75, 6.5)]: + out = shift_frame(f, s) + cov = coverage_mask((20, 20), s) + assert np.array_equal(np.isfinite(out), cov), f"mismatch at shift {s}" + + +class TestDriftModel: + def test_shape_validation(self): + with pytest.raises(ValueError, match=r"\(N, 2\)"): + DriftModel(shifts=np.zeros((4, 3))) + + def test_residual_length_validation(self): + with pytest.raises(ValueError, match="residuals must be"): + DriftModel(shifts=np.zeros((4, 2)), residuals=np.zeros(3)) + + def test_is_integer(self): + assert DriftModel(shifts=np.array([[0, 0], [2, -3]], float)).is_integer + assert not DriftModel(shifts=np.array([[0, 0], [2.5, 0]], float)).is_integer + + def test_max_abs_shift_ignores_nan(self): + m = DriftModel(shifts=np.array([[0, 0], [3, -7], [np.nan, np.nan]], float)) + assert m.max_abs_shift == 7.0 + + def test_frame_conversions_are_inverses(self): + m = DriftModel(shifts=np.array([[0, 0], [2.5, -1.5], [4, 3]], float)) + pos = np.array([[10.0, 12.0], [20.0, 22.0]]) + idx = np.array([1, 2]) + assert np.allclose(m.to_lab_frame(m.to_sample_frame(pos, idx), idx), pos) + + def test_to_sample_frame_removes_stage_motion(self): + """A particle that only *appears* to move because the stage drifted.""" + m = DriftModel(shifts=np.array([[0, 0], [-5, 0], [-10, 0]], float)) + # Same physical spot, drifting downward in the raw frames. + lab = np.array([[30.0, 40.0], [35.0, 40.0], [40.0, 40.0]]) + idx = np.array([0, 1, 2]) + sample = m.to_sample_frame(lab, idx) + assert np.allclose(sample[:, 0], 30.0), sample + + def test_save_load_round_trip(self, tmp_path): + m = DriftModel( + shifts=np.array([[0, 0], [1.25, -2.5]], float), + residuals=np.array([np.inf, 12.5], np.float32), + params={"upsample": 8}, provenance={"action": "drift"}, + reference="running", + ) + p = str(tmp_path / "d.npz") + m.save(p) + back = DriftModel.load(p) + assert np.array_equal(back.shifts, m.shifts) + assert back.params["upsample"] == 8 + assert back.provenance == {"action": "drift"} + assert back.reference == "running" + assert np.array_equal(back.residuals, m.residuals) + + def test_load_rejects_future_format(self, tmp_path): + import json + p = str(tmp_path / "bad.npz") + np.savez_compressed( + p, shifts=np.zeros((2, 2), np.float32), + meta=np.array(json.dumps({"format_version": 999}))) + with pytest.raises(ValueError, match="unsupported DriftModel format"): + DriftModel.load(p) + + +class TestEndToEnd: + def test_solve_then_warp_aligns_the_stack(self): + """The whole point: after correction, every frame agrees with frame 0.""" + truth = np.array( + [[0, 0], [2.5, -3.0], [-4.25, 1.75], [6.0, 4.5]], dtype=float) + stack = _shifted_stack(truth, h=80, w=80) + model = solve_translation(stack, device="numpy", upsample=8, + reference="first") + + core = (slice(12, -12), slice(12, -12)) + ref = stack[0][core] + for i in range(1, len(truth)): + aligned = shift_frame(stack[i], model.shifts[i], fill=0.0) + resid = np.abs(aligned[core] - ref).max() + raw = np.abs(stack[i][core] - ref).max() + assert resid < raw * 0.2, ( + f"frame {i}: correction barely helped (resid={resid:.4f} " + f"raw={raw:.4f}) — check the SIGN convention") diff --git a/spyde/tests/migrated/test_drift_wizard.py b/spyde/tests/migrated/test_drift_wizard.py new file mode 100644 index 00000000..91e0ed87 --- /dev/null +++ b/spyde/tests/migrated/test_drift_wizard.py @@ -0,0 +1,860 @@ +""" +The Drift Correction wizard backend (``drift_*`` staged handlers). + +Handlers are called directly as ``fn(session, plot, payload)`` and polled with +``_wait`` — the shape ``test_find_vectors_wizard.py`` establishes, because the +solve and the check sums both run on a worker thread. + +The four claims that matter: + +:class:`TestCheckWindow` + Plan A8 / README §6. The verification surface is a SEPARATE window, and a + bare ``figure`` is not a registered ``Plot`` — so it must be reachable + through ``session.controller_by_window_id`` and must disappear on close. A + check window that leaks is the exact bug README §6 documents. +:class:`TestSolve` + The solved shifts must match ``particle_movie``'s stamped ground truth, and + ``tree.drift`` must carry the model. Ground truth beats a golden number. +:class:`TestCommitIsLazy` + The CLAUDE.md memory-safety rule, guarded the way + ``test_find_vectors_memory.py`` guards it: a ``da.Array.compute`` spy that + counts calls on the full-dataset shape. And the corrected node has to be + genuinely better — an aligned stack sums SHARP, which is the whole claim the + check window makes to the user. +:class:`TestDoubleFire` + README §4 / StrictMode: open, close, open leaves exactly ONE controller and + exactly ONE check window. +""" +from __future__ import annotations + +import threading +import time + +import numpy as np +import pytest + +from spyde.actions import drift_action as dr + + +@pytest.fixture +def _capture_module_emit(window, monkeypatch): + """Route ``drift_action``'s own ``emit`` into the captured list. + + The module does ``from spyde.backend.ipc import emit`` at import, so + conftest's patch of ``ipc.emit`` never reaches that binding — the identical + hazard conftest already documents for ``session.py``, and the identical fix. + ``emit_status``/``emit_error`` need no patch: they resolve ``emit`` inside + ``ipc`` at call time. + + NOT autouse: requesting ``window`` (a real Session) is expensive, and the + pure-function classes below (``TestSharpnessNumber``, ``TestSchema``, + ``TestCoercion``) never dispatch a handler and so never need it. Classes + that DO dispatch handlers opt in via ``@pytest.mark.usefixtures``. + """ + monkeypatch.setattr(dr, "emit", window["messages"].append) + +# Small but enough for the drift curve to turn: `particle_movie`'s drift is a +# smooth excursion, and 8 frames already reach ~6 px, which is 5x the tolerance +# asserted below. +N_FRAMES = 8 + + +def _signal_plot(session): + return next((p for p in session._plots + if not p.is_navigator and p.plot_state is not None), None) + + +def _wait(pred, timeout=30.0): + end = time.time() + timeout + while time.time() < end: + if pred(): + return True + time.sleep(0.05) + return False + + +def _movie(window, frames: int = N_FRAMES): + session = window["window"] + session._load_test_data_particles({"frames": frames}) + plot = _wait(lambda: _signal_plot(session) is not None) and _signal_plot(session) + assert plot is not None, "the particle movie never produced a signal plot" + return session, plot, plot.signal_tree + + +def _opened(window, frames: int = N_FRAMES, **params): + session, plot, tree = _movie(window, frames) + dr.drift_open(session, plot, {"upsample": 8, "max_shift": 16, **params}) + assert _wait(lambda: getattr(tree, "_drift_wizard", None) is not None + and tree._drift_wizard.window_id is not None), \ + "the Drift Check window never opened" + return session, plot, tree, tree._drift_wizard + + +def _solved(window, frames: int = N_FRAMES): + session, plot, tree, wiz = _opened(window, frames) + dr.drift_run(session, plot, {"upsample": 8, "max_shift": 16}) + assert _wait(lambda: wiz.model is not None), "the solve never finished" + return session, plot, tree, wiz + + +def _of_type(messages, kind): + return [m for m in messages if isinstance(m, dict) and m.get("type") == kind] + + +def _sharpness(img) -> float: + """Mean squared gradient — an aligned sum has more of it than a blurred one.""" + a = np.nan_to_num(np.asarray(img, np.float64)) + gy, gx = np.gradient(a) + return float(np.mean(gy ** 2 + gx ** 2)) + + +class _FullComputeGuard: + """Count ``.compute()`` calls made on the whole movie.""" + + def __init__(self, shape): + self.shape = tuple(shape) + self.hits = 0 + + def __enter__(self): + import dask.array as da + self._real = da.Array.compute + guard = self + + def _spy(arr, *a, **k): + if tuple(arr.shape) == guard.shape: + guard.hits += 1 + return guard._real(arr, *a, **k) + + da.Array.compute = _spy + return self + + def __exit__(self, *exc): + import dask.array as da + da.Array.compute = self._real + return False + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestCheckWindow: + def test_open_registers_a_controller_for_the_bare_figure(self, window): + """README §6: a bare `figure` is not a Plot, so dispatch can only find + it through the window-controller registry.""" + session, _plot, _tree, wiz = _opened(window) + assert session.controller_by_window_id(wiz.window_id) is wiz + assert session._plot_by_window_id(wiz.window_id) is None, \ + "the check window is supposed to be a bare figure, not a Plot" + + def test_the_window_shows_the_uncorrected_sum(self, window): + session, _plot, _tree, wiz = _opened(window) + msgs = window["messages"] + figs = [m for m in _of_type(msgs, "figure") + if m.get("window_id") == wiz.window_id] + assert figs and figs[-1]["title"] == "Drift Check" + assert wiz._before_sum is not None and np.isfinite(wiz._before_sum).any() + + def test_open_solves_nothing(self, window): + """Plan A8: drift correction is explicit — nothing runs on load.""" + _s, _p, tree, wiz = _opened(window) + assert wiz.model is None + assert getattr(tree, "drift", None) is None + + def test_close_takes_the_window_with_it(self, window): + session, plot, tree, wiz = _opened(window) + wid = wiz.window_id + dr.drift_close(session, plot, {}) + assert getattr(tree, "_drift_wizard", None) is None + assert session.controller_by_window_id(wid) is None + from spyde.actions.figure_registry import _FIGS + assert wid not in _FIGS, "the check figure outlived its window" + + def test_the_summed_subset_is_bounded(self, window): + """A sum is a sharpness test, not a measurement — the cap is what keeps + the check window usable on a long movie.""" + _s, _p, _t, wiz = _opened(window) + wiz._sum_indices = None # as if the movie were long + idx = wiz.sum_indices(10_000) + assert idx.size <= dr._SUM_MAX_FRAMES + assert idx[0] == 0 and idx[-1] == 9_999 + assert wiz.sum_indices(10_000) is idx, ( + "the subset is memoised so the before and after sums cover the SAME " + "frames — comparing two different subsets means nothing") + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestMethodStubs: + def test_rigid_affine_says_so_too(self, window): + session, plot, _tree, wiz = _opened(window) + dr.drift_set_method(session, plot, {"method": "rigid_affine"}) + assert wiz.params["method"] == "rigid" + + def test_unknown_method_errors(self, window): + session, plot, _tree, _wiz = _opened(window) + msgs = window["messages"] + dr.drift_set_method(session, plot, {"method": "banana"}) + assert any("unknown model" in str(m.get("text", "")) + for m in _of_type(msgs, "error")) + + +class _Box: + """Stand-in for an anyplotlib RectangleWidget (x/y/w/h in IMAGE PIXELS). + + The headless session does build a real ``_plot2d``, so the wizard's own box + exists — but a test that wants a SPECIFIC region needs to place one, and + dragging a real widget means faking pointer events. Swapping this in is the + smaller lie, and it exercises the same ``roi_box()`` conversion. + """ + + def __init__(self, x, y, w, h): + self.x, self.y, self.w, self.h = float(x), float(y), float(w), float(h) + + def set(self, **kw): + for k, v in kw.items(): + setattr(self, k, float(v)) + + def hide(self): + pass + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestDiscoveryPreview: + """The centrepiece: a draggable box + a drift-corrected sum of it over ~20 + frames, so the user sees whether alignment works BEFORE paying for the whole + movie.""" + + def test_open_previews_the_default_box(self, window): + session, _plot, _tree, wiz = _opened(window) + msgs = window["messages"] + assert _wait(lambda: _of_type(msgs, "drift_preview")), \ + "opening the caret never produced a discovery preview" + prev = _of_type(msgs, "drift_preview")[-1] + assert prev["frames"] >= 2 + assert prev["gain"] > 1.0, ( + "the aligned sum of the default box is not sharper than the raw one " + "— the preview cannot discriminate anything if it never improves") + + def test_the_preview_never_computes_the_whole_movie(self, window): + session, plot, tree, _wiz = _opened(window) + msgs = window["messages"] + del msgs[:] + with _FullComputeGuard(tree.root.data.shape) as guard: + dr.drift_tune(session, plot, {"upsample": 4, "max_shift": 12}) + assert _wait(lambda: _of_type(msgs, "drift_preview")) + assert guard.hits == 0 + + def test_tune_stores_the_new_parameters(self, window): + session, plot, _tree, wiz = _opened(window) + dr.drift_tune(session, plot, {"upsample": 16, "max_shift": 9.0}) + assert wiz.params["upsample"] == 16 + assert wiz.params["max_shift"] == 9.0 + + def test_the_preview_uses_the_box_even_with_the_toggle_off(self, window): + """The toggle is the COMMITMENT (does the full solve restrict to the + box); the preview is the QUESTION and always asks it about the box.""" + session, plot, _tree, wiz = _opened(window) + msgs = window["messages"] + assert wiz.params["use_roi"] is False + wiz._roi_widget = _Box(10, 12, 60, 48) + del msgs[:] + dr.drift_tune(session, plot, {}) + assert _wait(lambda: _of_type(msgs, "drift_preview")) + assert _of_type(msgs, "drift_preview")[-1]["roi"] == [12, 10, 48, 60] + + def test_the_box_is_read_in_image_pixels_as_y0_x0_h_w(self, window): + """anyplotlib 2-D widgets report IMAGE PIXELS and solve_translation's + roi is in pixels — the two meet with NO scale conversion.""" + _s, _p, _t, wiz = _opened(window) + wiz._frame_shape = (96, 112) + wiz._roi_widget = _Box(x=20, y=8, w=40, h=32) + assert wiz.roi_box() == (8, 20, 32, 40) + + def test_the_box_is_clamped_into_the_frame(self, window): + _s, _p, _t, wiz = _opened(window) + wiz._frame_shape = (96, 112) + wiz._roi_widget = _Box(x=100, y=90, w=400, h=400) + y0, x0, h, w = wiz.roi_box() + assert 0 <= y0 and 0 <= x0 + assert y0 + h <= 96 and x0 + w <= 112 + + def test_a_box_below_the_solver_floor_is_refused(self, window): + """solve_translation REJECTS a too-small roi rather than clamping it, so + the caret must never hand it one.""" + from spyde.drift.translation import _MIN_ROI + assert dr._ROI_MIN_PX >= _MIN_ROI + _s, _p, _t, wiz = _opened(window) + wiz._frame_shape = (96, 112) + wiz._roi_widget = _Box(x=0, y=0, w=4, h=4) + y0, x0, h, w = wiz.roi_box() + assert h >= _MIN_ROI and w >= _MIN_ROI + + def test_a_superseded_preview_does_not_paint(self, window, monkeypatch): + """Latest-wins: a drag that outruns the solve must drop the stale + result, not paint it over the newer one. + + Drives a REAL ``drift_tune`` (worker thread → ``_run_preview`` → + ``is_current``-guarded ``_done``, drift_action.py) instead of poking the + generation counters directly — the old version bumped the generation + and asserted ``is_current`` without ever calling ``drift_tune``, so the + guarded ``_done`` path it claims to protect never ran and the assertion + held trivially. ``preview_alignment`` is gated on an Event so a "newer + drag" can be landed deterministically while the older one is still + computing, rather than racing real time. + """ + session, plot, tree, wiz = _opened(window) + msgs = window["messages"] + # Drain the automatic discovery preview `_opened` triggers so it can't + # land in the middle of the controlled run below. + assert _wait(lambda: _of_type(msgs, "drift_preview")) + + painted = [] + wiz.show_preview = lambda res: painted.append(res) + + started, release = threading.Event(), threading.Event() + real_preview_alignment = dr.preview_alignment + + def _blocking(*a, **k): + started.set() + release.wait(5.0) + return real_preview_alignment(*a, **k) + + monkeypatch.setattr(dr, "preview_alignment", _blocking) + + done = threading.Event() + real_is_current = dr.is_current + + def _is_current(owner, key, gen_): + result = real_is_current(owner, key, gen_) + if key == "_drift_preview_gen": + done.set() + return result + + monkeypatch.setattr(dr, "is_current", _is_current) + + dr.drift_tune(session, plot, {"upsample": 4, "max_shift": 12}) + assert started.wait(5.0), "the preview work never started" + dr.bump_generation(tree, "_drift_preview_gen") # a newer drag lands + release.set() + + assert _wait(lambda: done.is_set()), \ + "the superseded preview's _done never ran" + assert painted == [] + + def test_the_settle_timer_coalesces_a_drag(self, window): + """The widget's pointer_move fires at renderer frame rate; only the + RESTING geometry may solve. + + Coalescing is asserted by POLLING until the fired count stops + changing, not by racing the timer's own delay: asserting + ``fired == []`` immediately after scheduling is only true if the test + process outruns the 50ms timer, which a loaded box (other agents run + pytest concurrently — CLAUDE.md) does not guarantee. + """ + _s, _p, _t, wiz = _opened(window) + fired = [] + wiz._fire_preview = lambda: fired.append(1) + for _ in range(20): + wiz.schedule_preview(delay=0.05) + + assert _wait(lambda: len(fired) >= 1, timeout=3.0), \ + "the settle timer never fired" + last, stable_since, start = -1, time.time(), time.time() + while time.time() - start < 3.0: + n = len(fired) + if n != last: + last, stable_since = n, time.time() + elif time.time() - stable_since >= 0.3: + break + time.sleep(0.02) + assert len(fired) == 1, f"{len(fired)} solves for one drag" + + +class TestSharpnessNumber: + """The gain has to be an ANSWER, not decoration: a landmark and a + featureless patch must come out clearly different.""" + + @staticmethod + def _stack(n=16, size=140, pad=20): + """Textured on the left half, flat on the right, drifting rigidly.""" + from scipy.ndimage import gaussian_filter, map_coordinates + rng = np.random.default_rng(3) + canvas = np.zeros((size + 2 * pad, size + 2 * pad), np.float32) + 1.0 + tex = gaussian_filter(rng.standard_normal(canvas.shape), 1.5) * 0.6 + half = size // 2 + pad + canvas[:, :half] += tex[:, :half] + drift = np.stack([np.linspace(0, 8.0, n), np.linspace(0, -5.0, n)], 1) + yy, xx = np.mgrid[0:size, 0:size].astype(np.float64) + frames = np.empty((n, size, size), np.float32) + for t in range(n): + dy, dx = drift[t] + frames[t] = map_coordinates(canvas, [yy + pad - dy, xx + pad - dx], + order=1, mode="nearest") + frames += rng.normal(0, 0.02, frames.shape).astype(np.float32) + return frames + + def test_a_landmark_beats_a_featureless_patch(self): + frames = self._stack() + idx = np.arange(frames.shape[0]) + params = dict(dr.DEFAULTS) + good = dr.preview_alignment(frames.__getitem__, idx, (30, 5, 80, 55), + params=params) + bad = dr.preview_alignment(frames.__getitem__, idx, (30, 82, 80, 55), + params=params) + assert good["gain"] > 2.0, f"a real landmark only scored {good['gain']:.2f}" + assert bad["gain"] < 1.0, f"a featureless patch scored {bad['gain']:.2f}" + assert good["gain"] > 3 * bad["gain"] + + def test_the_nan_border_does_not_inflate_the_number(self): + """A shifted frame's uncovered edge is NaN (plan A7). Zero-filling it + manufactures a step whose gradient energy dwarfs the image's own — every + ROI would look brilliant.""" + a = np.ones((32, 32), np.float32) + a[:4, :] = np.nan + assert dr._gradient_energy(a) == 0.0 + + def test_the_two_sums_are_measured_on_the_same_pixels(self): + raw = np.ones((16, 16), np.float32) + aligned = raw.copy() + aligned[:3, :] = np.nan + both = np.isfinite(raw) & np.isfinite(aligned) + assert dr._gradient_energy(raw, both) == dr._gradient_energy(aligned, both) + + def test_the_preview_sample_spans_the_whole_movie(self): + """20 CONSECUTIVE frames of a long movie drift by almost nothing, so a + contiguous window would say "looks fine" for every box.""" + idx = dr._preview_indices(3000, 20, 64 * 64 * 4) + assert idx[0] == 0 and idx[-1] == 2999 and idx.size <= 20 + + def test_the_sample_is_thinned_to_fit_the_byte_cap(self): + """With no ROI the crop IS the frame — 20 × 4096² float32 is 1.3 GB.""" + idx = dr._preview_indices(3000, 20, 4096 * 4096 * 4) + assert 2 <= idx.size < 20 + + +def _ground_truth(tree): + import spyde.data.synthetic as sy + return np.asarray(sy.ground_truth(tree.root)["drift"], np.float64) + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestSolve: + def test_shifts_match_the_stamped_ground_truth(self, window): + _s, _p, tree, wiz = _solved(window) + truth = _ground_truth(tree)[:N_FRAMES] + err = np.abs(wiz.model.shifts - truth).max() + assert err < 0.25, f"worst per-axis drift error {err:.3f} px" + + def test_the_model_lands_on_the_tree(self, window): + _s, _p, tree, wiz = _solved(window) + assert tree.drift is wiz.model + assert wiz.model.kind == "rigid" + + def test_run_reports_progress_and_the_finished_trace(self, window): + session, plot, tree, wiz = _opened(window) + msgs = window["messages"] + dr.drift_run(session, plot, {}) + assert _wait(lambda: _of_type(msgs, "drift_result")) + prog = _of_type(msgs, "drift_progress") + assert prog and prog[-1]["done"] == prog[-1]["total"] == N_FRAMES + res = _of_type(msgs, "drift_result")[-1] + assert len(res["shifts"]) == N_FRAMES and not res["cancelled"] + assert res["max_abs_shift"] > 1.0 + + def test_run_never_computes_the_whole_movie(self, window): + session, plot, tree, wiz = _opened(window) + with _FullComputeGuard(tree.root.data.shape) as guard: + dr.drift_run(session, plot, {}) + assert _wait(lambda: wiz.model is not None) + assert guard.hits == 0, "the solve materialised the whole movie" + + def test_the_check_window_gets_a_sharper_corrected_sum(self, window): + """The claim the window makes to the user, asserted rather than drawn: + an aligned stack sums sharp, a misaligned one blurs.""" + _s, _p, _t, wiz = _solved(window) + n, get_frame, _shape = wiz.frames() + idx = wiz.sum_indices(n) + before = dr._stack_sum(get_frame, idx) + after = dr._stack_sum(get_frame, idx, wiz.model.shifts) + assert _sharpness(after) > 1.5 * _sharpness(before), ( + f"corrected sum {_sharpness(after):.5f} is not sharper than the raw " + f"{_sharpness(before):.5f} — check the sign convention in " + "spyde/drift/model.py") + + def test_closing_the_tree_cancels_the_solve(self, window, monkeypatch): + """Cancellation goes through BaseSignalTree.register_cancel, so closing + the tree has to stop it. + + Gated on an Event so ``tree.close()`` deterministically lands WHILE + the solve is genuinely mid-flight (right after frame 0), instead of + racing the solve's own speed — on this small movie an un-gated solve + can finish before ``close()`` ever gets a chance to cancel it, which + would make the "partial model" assertion below pass or fail by luck. + """ + session, plot, tree, wiz = _opened(window, frames=24) + import spyde.drift as drift_mod + real_solve = drift_mod.solve_translation + holding, proceed = threading.Event(), threading.Event() + + def _gated(data, *, progress=None, **kwargs): + def _progress(done, total): + if progress is not None: + progress(done, total) + if done == 1: + holding.set() + proceed.wait(10.0) + return real_solve(data, progress=_progress, **kwargs) + + monkeypatch.setattr(drift_mod, "solve_translation", _gated) + + dr.drift_run(session, plot, {}) + assert holding.wait(10.0), "the solve never reached its first frame" + tree.close() + proceed.set() + + assert _wait(lambda: wiz.model is not None or wiz._closed, timeout=60) + if wiz.model is not None: + # A cancelled solve leaves NaN for the frames it never reached, so a + # partial model is detectable rather than silently wrong. + assert not np.isfinite(wiz.model.shifts).all() + + def test_run_without_a_caret_errors(self, window): + session, plot, _tree = _movie(window) + msgs = window["messages"] + dr.drift_run(session, plot, {}) + assert any("caret is not open" in str(m.get("text", "")) + for m in _of_type(msgs, "error")) + + def test_use_roi_feeds_the_box_to_the_solver(self, window): + """The toggle's whole job: the same rectangle the preview tested is the + one the full solve correlates on.""" + session, plot, _tree, wiz = _opened(window) + msgs = window["messages"] + wiz._frame_shape = (96, 112) + wiz._roi_widget = _Box(x=16, y=12, w=64, h=64) + dr.drift_run(session, plot, {"use_roi": True}) + assert _wait(lambda: wiz.model is not None) + assert wiz.model.params["roi"] == [12, 16, 64, 64] + assert _of_type(msgs, "drift_result")[-1]["roi"] == [12, 16, 64, 64] + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestTraceWindow: + """The dy/dx curve is its OWN plot window, filled as the solve runs — not + caret furniture (plan §0.9a).""" + + def test_the_solve_opens_a_second_figure_window(self, window): + session, plot, _tree, wiz = _opened(window) + msgs = window["messages"] + dr.drift_run(session, plot, {}) + assert _wait(lambda: wiz.trace_window_id is not None) + figs = [m for m in _of_type(msgs, "figure") + if m.get("window_id") == wiz.trace_window_id] + assert figs and figs[-1]["title"] == "Drift dy/dx" + assert session.controller_by_window_id(wiz.trace_window_id) is wiz, \ + "a bare figure is only reachable through the controller registry" + assert session._plot_by_window_id(wiz.trace_window_id) is None + + def test_it_fills_from_the_on_shift_stream(self, window): + session, plot, _tree, wiz = _opened(window) + msgs = window["messages"] + # `_opened` fires the discovery preview, which streams `drift_trace` + # batches of its own (same window id, same leading indices). Under CI + # timing those interleave with the run's stream, so a raw point COUNT + # over-counts (seen: 10 for 8 frames on windows-py3.12). Snapshot the + # list and assert the run streamed every frame INDEX — the contract is + # "every solved frame went out", not "nobody else spoke". + start = len(msgs) + dr.drift_run(session, plot, {}) + assert _wait(lambda: wiz.model is not None) + assert _wait(lambda: int(wiz._trace.get("filled", 0)) == N_FRAMES) + streamed = {int(p[0]) for m in _of_type(msgs[start:], "drift_trace") + for p in m["points"]} + assert streamed == set(range(N_FRAMES)) + + def test_the_trace_matches_the_model(self, window): + _s, _p, _t, wiz = _solved(window) + assert _wait(lambda: int(wiz._trace.get("filled", 0)) == N_FRAMES) + np.testing.assert_allclose(wiz._trace["dy_data"], wiz.model.shifts[:, 0], + atol=1e-5) + np.testing.assert_allclose(wiz._trace["dx_data"], wiz.model.shifts[:, 1], + atol=1e-5) + + def test_only_the_solved_prefix_is_pushed(self, window): + """Pushing the NaN-padded whole array would leave anyplotlib's auto + y-range looking at one finite point.""" + _s, _p, _t, wiz = _opened(window) + wiz.trace_window_id = None + wiz.open_trace_window(50) + wiz.push_trace([(1, 3.0, -2.0), (2, 4.0, -3.0)]) + assert wiz._trace["filled"] == 3 + assert np.isnan(wiz._trace["dy_data"][3:]).all() + + def test_close_takes_both_windows(self, window): + session, plot, tree, wiz = _solved(window) + assert _wait(lambda: wiz.trace_window_id is not None) + check, trace = wiz.window_id, wiz.trace_window_id + dr.drift_close(session, plot, {}) + from spyde.actions.figure_registry import _FIGS + for wid in (check, trace): + assert session.controller_by_window_id(wid) is None + assert wid not in _FIGS, f"figure {wid} outlived its window" + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestDiscard: + def test_discard_drops_the_model_and_the_trace_window(self, window): + session, plot, tree, wiz = _solved(window) + assert _wait(lambda: wiz.trace_window_id is not None) + trace = wiz.trace_window_id + dr.drift_discard(session, plot, {}) + assert wiz.model is None + assert getattr(tree, "drift", None) is None + assert wiz.trace_window_id is None + assert session.controller_by_window_id(trace) is None + assert wiz.window_id is not None, "Discard must not close the caret" + + def test_discard_stops_a_solve_in_flight(self, window): + """Same user intent as Stop, so it is the same handler: bumping the run + generation FIRST means a solve that finishes anyway never installs. + + Waits on the wizard's own ``still()`` guard actually being evaluated + by the in-flight run's ``_done`` — the real signal that the race + between the solve and the discard has resolved — rather than a flat + sleep long enough to "probably" cover it. + """ + session, plot, tree, wiz = _opened(window, frames=24) + done = threading.Event() + real_still = wiz.still + + def _still(gen): + result = real_still(gen) + done.set() + return result + + wiz.still = _still + + dr.drift_run(session, plot, {}) + dr.drift_discard(session, plot, {}) + assert wiz._stop[0] is True + assert _wait(lambda: done.is_set(), timeout=30.0), \ + "the in-flight solve's _done never ran" + assert wiz.model is None + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestCommitIsLazy: + def test_commit_adds_a_lazy_node_without_computing(self, window): + session, plot, tree, wiz = _solved(window) + before = set(tree.root_node.children) + with _FullComputeGuard(tree.root.data.shape) as guard: + dr.drift_commit(session, plot, {}) + assert guard.hits == 0, "commit materialised the movie" + added = set(tree.root_node.children) - before + assert added == {"Drift corrected"} + node = tree.root_node.children["Drift corrected"] + assert node.signal._lazy + assert node.signal.data.shape == tree.root.data.shape + assert node.local is True, ( + "a per-frame shift IS local — without the tag the derived-view " + "reader falls back to the opaque path") + + def test_one_corrected_frame_costs_one_frame(self, window): + session, plot, tree, _wiz = _solved(window) + dr.drift_commit(session, plot, {}) + node = tree.root_node.children["Drift corrected"] + with _FullComputeGuard(tree.root.data.shape) as guard: + frame = np.asarray(node.signal.data[3].compute()) + assert guard.hits == 0 + assert frame.shape == tuple(tree.root.data.shape[1:]) + + def test_uncovered_pixels_are_nan_not_invented(self, window): + """Plan A7: nothing is cropped and nothing is filled with invented + data — segmentation would find 'particles' in a zero-filled border.""" + session, plot, tree, _wiz = _solved(window) + dr.drift_commit(session, plot, {}) + node = tree.root_node.children["Drift corrected"] + frame = np.asarray(node.signal.data[N_FRAMES - 1].compute()) + assert np.isnan(frame).any(), "no NaN padding on a shifted frame" + assert np.isfinite(frame).any(), "the whole frame is NaN" + + def test_the_corrected_node_is_actually_aligned(self, window): + session, plot, tree, wiz = _solved(window) + dr.drift_commit(session, plot, {}) + node = tree.root_node.children["Drift corrected"] + raw = np.nanmean(np.stack([np.asarray(tree.root.data[i].compute(), + np.float64) + for i in range(N_FRAMES)]), axis=0) + fixed = np.nanmean(np.stack([np.asarray(node.signal.data[i].compute(), + np.float64) + for i in range(N_FRAMES)]), axis=0) + assert _sharpness(fixed) > 1.5 * _sharpness(raw) + + def test_commit_stamps_provenance(self, window): + session, plot, tree, _wiz = _solved(window) + dr.drift_commit(session, plot, {}) + node = tree.root_node.children["Drift corrected"] + prov = node.signal.metadata.get_item("General.spyde_provenance") + assert prov["action"] == "Drift Correction" + assert prov["kind"] == "rigid" + + def test_commit_before_solving_errors(self, window): + session, plot, _tree, _wiz = _opened(window) + msgs = window["messages"] + dr.drift_commit(session, plot, {}) + assert any("solve first" in str(m.get("text", "")) + for m in _of_type(msgs, "error")) + + def test_a_model_of_the_wrong_length_is_refused(self, window): + """Re-solving after a crop must not silently pair frame 0's shift with + a different frame.""" + from spyde.drift import DriftModel + _s, _p, tree, _wiz = _opened(window) + bad = DriftModel(shifts=np.zeros((N_FRAMES + 3, 2), np.float32)) + with pytest.raises(ValueError, match="covers"): + dr.drift_corrected(tree.root, model=bad) + + +@pytest.mark.usefixtures("_capture_module_emit") +class TestDoubleFire: + def test_open_close_open_leaves_one_controller(self, window): + session, plot, tree = _movie(window) + msgs = window["messages"] + built: list = [] + still_done: dict[int, threading.Event] = {} + real_init = dr.DriftWizard.__init__ + + def _tracking(self, *a, **k): + real_init(self, *a, **k) + built.append(self) + evt = threading.Event() + still_done[id(self)] = evt + real_still = self.still + + def _still(gen, _real=real_still, _evt=evt): + result = _real(gen) + _evt.set() + return result + + self.still = _still + + dr.DriftWizard.__init__ = _tracking + try: + dr.drift_open(session, plot, {}) + dr.drift_close(session, plot, {}) + dr.drift_open(session, plot, {}) + finally: + dr.DriftWizard.__init__ = real_init + assert _wait(lambda: tree._drift_wizard is not None + and tree._drift_wizard.window_id is not None) + assert len(built) == 2, f"expected 2 wizards built, got {len(built)}" + # The FIRST open's deferred worker is still in flight when close() + # bumps its generation; wait for that worker's own `still()` guard to + # actually run rather than sleeping a guessed duration. + assert _wait(lambda: still_done[id(built[0])].is_set(), timeout=10.0), \ + "the superseded first open's worker never landed" + # The SURVIVING wizard's own open also kicks off the automatic + # discovery preview (drift_open's trailing `_run_preview`) on its own + # worker thread; drain it here too so it can't outlive the test and + # fire into a torn-down monkeypatch (dr.emit reverted, painting into + # real stdout). + assert _wait(lambda: _of_type(msgs, "drift_preview")), \ + "the surviving wizard's discovery preview never landed" + + alive = [w for w in built if not w._closed] + assert len(alive) == 1, \ + f"expected 1 live controller, got {len(alive)} of {len(built)} built" + assert tree._drift_wizard is alive[0] + # …and exactly one check window, because a superseded open's deferred + # build must be dropped rather than emitting a second figure. + wids = [w.window_id for w in built if w.window_id is not None] + assert len(wids) == 1, f"{len(wids)} check windows opened" + + def test_close_without_an_open_is_harmless(self, window): + session, plot, tree = _movie(window) + dr.drift_close(session, plot, {}) + assert getattr(tree, "_drift_wizard", None) is None + + +class TestSchema: + def test_schema_resolves_through_the_registry(self): + from spyde.actions import registry + schema = registry.wizard_parameters("drift") + assert schema and schema is not dr.DriftWizard.parameters + + def test_schema_defaults_match_the_handler_defaults(self): + from spyde.actions import registry + schema = registry.wizard_parameters("drift") + for key, spec in schema.items(): + assert key in dr.DEFAULTS, f"drift schema declares unknown param {key!r}" + assert spec["default"] == dr.DEFAULTS[key], \ + f"drift schema/{key} drifted from drift_action.DEFAULTS" + + def test_every_stage_is_registered(self): + from spyde.actions.registry import STAGED_HANDLERS, resolve_staged + for stage in ("drift_open", "drift_close", "drift_set_method", + "drift_tune", "drift_run", "drift_discard", + "drift_commit"): + assert stage in STAGED_HANDLERS + assert callable(resolve_staged(stage)) + + def test_the_default_face_is_two_toggles(self): + """§0.9a: everything that is not the task itself is tagged Advanced, so + any host renders the same small face. The caret is the enforcement; this + is the schema saying the same thing.""" + from spyde.actions import registry + schema = registry.wizard_parameters("drift") + face = [k for k, s in schema.items() if not s.get("tab")] + assert face == ["use_roi", "reject_outliers"], \ + f"the caret's default face grew to {face}" + + def test_toolbar_entry_gates_on_a_movie(self): + import spyde + meta = spyde.TOOLBAR_ACTIONS["functions"]["Drift Correction"] + assert meta["function"] == "spyde.actions.drift_action.drift_correction" + assert meta["signal_types"] == ["insitu"], ( + "the rigid solver needs a 1-D navigation axis; gating on `insitu` " + "is the same gate Play/Fast-Forward use for exactly that") + + @pytest.mark.parametrize("signal_type,offered", [ + ("insitu", True), + ("", False), # a single image has nothing to align + ("electron_diffraction", False), + ]) + def test_toolbar_gating(self, signal_type, offered): + import spyde + from spyde.drawing.toolbars.plot_control_toolbar import _action_matches_plot + + class _Sig: + _signal_type = signal_type + + class _Tree: + particles = None + diffraction_vectors = None + root = _Sig() + + class _Plot: + signal_tree = _Tree() + is_navigator = False + + class _State: + plot = _Plot() + current_signal = _Sig() + dimensions = 2 + navigation = False + + meta = spyde.TOOLBAR_ACTIONS["functions"]["Drift Correction"] + assert _action_matches_plot("Drift Correction", meta, _State()) is offered + + +class TestCoercion: + def test_unknown_method_falls_back(self): + assert dr._coerce({"method": "warp"})["method"] == dr.DEFAULTS["method"] + + def test_unknown_reference_falls_back(self): + assert dr._coerce({"reference": "later"})["reference"] == "running" + + def test_order_is_clamped(self): + assert dr._coerce({"order": 9})["order"] == 3 + assert dr._coerce({"order": -1})["order"] == 0 + + def test_a_junk_value_keeps_the_default(self): + assert dr._coerce({"upsample": "eight"})["upsample"] == \ + dr.DEFAULTS["upsample"] From c48d75c002a1ab6d527b47ee2bb6d6356141b386 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sat, 8 Aug 2026 23:56:29 -0500 Subject: [PATCH 2/6] feat(drift): rigid drift correction -- solver, model, warp, and the wizard Pulled file-level from feat/seg-fast-engine (PR #122) and trimmed to the rigid surface: - spyde/drift/{frames,translation,warp}.py verbatim (all rigid): streaming frame source, FFT phase correlation with running Fourier reference + Guizar-Sicairos sub-pixel refinement, per-frame NaN-padded warp. The measured contracts stay as docstrings (Tukey-not-Hann taper, the phase magnitude floor, outlier rejection thresholds). - model.py: DriftModel without the nonrigid 'extra' carrier and the uncalled shift_at accessor. - __init__.py: no nonrigid re-exports. - drift_action.py: METHODS drops 'nonrigid'; NONRIGID_MODELS, the decimated-stack reader, the nonrigid fit step, its DEFAULTS/schema/ _coerce keys and the drift_run nonrigid branch are all gone. Non-rigid arrives in its own PR. No numba anywhere in the package (FFT/scipy/optional torch only). --- spyde/actions/drift_action.py | 1343 +++++++++++++++++ .../toolbars/icons/drift_correction.svg | 44 + spyde/drift/__init__.py | 36 + spyde/drift/frames.py | 92 ++ spyde/drift/model.py | 165 ++ spyde/drift/translation.py | 699 +++++++++ spyde/drift/warp.py | 134 ++ 7 files changed, 2513 insertions(+) create mode 100644 spyde/actions/drift_action.py create mode 100644 spyde/drawing/toolbars/icons/drift_correction.svg create mode 100644 spyde/drift/__init__.py create mode 100644 spyde/drift/frames.py create mode 100644 spyde/drift/model.py create mode 100644 spyde/drift/translation.py create mode 100644 spyde/drift/warp.py diff --git a/spyde/actions/drift_action.py b/spyde/actions/drift_action.py new file mode 100644 index 00000000..67ac20e3 --- /dev/null +++ b/spyde/actions/drift_action.py @@ -0,0 +1,1343 @@ +""" +drift_action.py — the Drift Correction wizard (``drift_`` staged actions). + +Plan A8, rewritten under plan §0.9a (*"the caret shows ONE control; everything +else is Advanced"*) after the first review: **"way too complicated. Too many +options. Information overload."** + + drift_open caret mounted → Drift Check window, the alignment ROI on + the movie, and the first discovery preview + drift_close caret unmounted → tear all of it down + drift_set_method rigid | rigid+affine (lives in Advanced now) + drift_tune a toggle/parameter changed → re-run the discovery preview + drift_run solve the movie on a worker; opens the dy/dx window and + fills it progressively from the solver's ``on_shift`` + drift_discard drop the solved model (and stop a solve in flight) + drift_commit add the LAZY corrected node to the tree + +**The caret carries the TASK, not the algorithm.** Its default face is two +toggles and one button. Reference mode, sub-pixel factor, max shift, +interpolation order and the model tabs are all real and all still here — they +sit behind a collapsed *Advanced* in the caret, and the schema below (the one +source of truth, mirrored by ``registry._WIZARD_SCHEMAS``) tags them so any +host renders the same split. Nothing was deleted; provenance still records +every parameter. + +**Discovery comes before commitment.** The centrepiece is a draggable +rectangle on the movie plus a live drift-corrected sum of just that box over +~20 frames (:data:`_PREVIEW_FRAMES`). A good landmark sums sharp, a bad one +blurs, and the *gain* number (:func:`_gradient_energy` of the aligned sum over +the raw sum, measured on the SAME pixels) puts a figure on it. So the user sees +whether alignment works on a subset before paying for the whole movie — and the +"Use ROI for alignment" toggle then feeds that exact rectangle to +``solve_translation(roi=…)``, which is often the more CORRECT answer anyway: +whole-frame correlation is contaminated by the sample's own motion (see +``spyde/drift/translation.py``'s ``roi`` docs). + +**Geometry is in IMAGE PIXELS end to end.** anyplotlib's 2-D widgets report +``x/y/w/h`` in image pixels with no scale/offset applied, and +``solve_translation``'s ``roi=(y0, x0, h, w)`` is in pixels too, so the two meet +with no conversion. Do not add one "for consistency" — see +``spyde/actions/masks.py::_signal_k_grids`` for that bug class. + +**Two windows, each with one job.** The *Drift Check* window is the evidence: +the whole-movie raw/corrected sums on top, the discovery pair (ROI raw vs ROI +aligned) beneath. The *Drift dy/dx* window is the curve, opened when the solve +starts and filled progressively from ``on_shift`` — it is a normal figure +window, not caret furniture. Both are bare ``figure`` windows (NOT registered +``Plot``s), so each registers a controller via ``own_window`` and keeps its +figure referenced through ``figure_registry.keep_alive``, per +``actions/README.md`` §6. + +**Nothing here materialises the movie.** ``solve_translation`` streams one +frame at a time; the check sums stream over a bounded subset +(:data:`_SUM_MAX_FRAMES`); the preview reads one full frame at a time and keeps +only the small crop, under a byte cap (:data:`_PREVIEW_MAX_BYTES`); and +``drift_commit`` adds a ``map_blocks`` node so the corrected movie is a lazy +view, never a copy (plan §0.7). The corrected node is tagged ``local=True`` +because a rigid shift is exactly per-frame, which is what lets the existing +``LocalTransformReader`` scrub it. +""" +from __future__ import annotations + +import logging +import threading +import time +from typing import Any + +import numpy as np + +from spyde.actions.context import current_signal as _current_signal +from spyde.actions.context import src_plot_tree as _src_plot_tree +from spyde.actions.lifecycle import ( + bump_generation, is_current, run_on_worker, show_tree_node, +) +from spyde.actions.wizard import WizardController +from spyde.backend.ipc import emit, emit_error, emit_progress, emit_status + +log = logging.getLogger(__name__) + +#: Solver families. ``rigid`` is the only one ``spyde.drift`` implements today; +#: ``rigid_affine`` is declared so a host can render the choice, and it falls +#: back to ``rigid`` with an explicit status rather than silently doing +#: something the user did not ask for. It lives inside Advanced (§0.9a). +METHODS: tuple[str, ...] = ("rigid", "rigid_affine") + +_UNAVAILABLE = { + "rigid_affine": ("the affine drift search (plan A4) is not implemented in " + "spyde.drift yet"), +} + +#: Frames summed for the whole-movie before/after check images. A sum is a +#: SHARPNESS test, not a measurement — a few dozen frames already show the blur +#: unambiguously, and the cap is what keeps the check window responsive on a +#: movie whose full pass costs as much as the solve itself. Evenly spaced, and +#: the SAME indices for both sums, or the comparison means nothing. +_SUM_MAX_FRAMES = 64 + +# Frames per streamed drift-trace message / per dy-dx repaint. One message per +# frame would flood the PLOTAPP line protocol at the plan's target scale +# (thousands of frames) for a curve the eye cannot follow at that resolution; +# batching by 16 keeps the trace visibly live while cutting the message count +# by the same factor. +_TRACE_BATCH = 16 + +# …but a COUNT alone is not enough, and the first screenshot showed why: a +# 12-frame movie never reaches 16, so the curve stayed empty for the whole solve +# and appeared complete at the end — the exact opposite of "fills in as it is +# computed". Flush on whichever comes first, so the trace is live at any movie +# length and still capped at ~7 messages/s on a fast one. +_TRACE_MAX_INTERVAL = 0.15 + +#: Figure geometry for the two bare-figure windows. +#: +#: A bare figure never receives ``resize_figure`` (that path resolves a +#: registered ``Plot``), so its INITIAL px size is the one it keeps and anything +#: outside it is CLIPPED by the subwindow — which is what cut the check +#: window's bottom row in half. The renderer sizes a new window from the +#: ``aspect`` field as ``inner_h = clamp(460 / aspect, 130, 300)`` then +#: ``inner_w = inner_h * aspect`` (``MDIArea.windowSize``). At the height cap +#: the first clamp is active for any aspect below 460/300, so a figure exactly +#: :data:`_FIG_HEIGHT` tall lands pixel-for-pixel in its window at any width up +#: to 460 — pick the width, derive the aspect. +#: +#: The width is deliberately the renderer's OWN default (340). Widening the +#: check window to 460 made it no longer fit beside the movie, so the free-slot +#: packer wrapped it to the next row — straight on top of the caret, which is an +#: overlay the packer cannot see. Keeping the default width keeps the placement +#: the packer already gets right. +_FIG_WIDTH = 340 +_FIG_HEIGHT = 300 + + +def _figure_geometry(width: int = _FIG_WIDTH) -> tuple[tuple[int, int], float]: + """``(figsize, aspect)`` that opens a bare-figure window with no clipping.""" + w = int(min(460, max(190, width))) + return (w, _FIG_HEIGHT), w / float(_FIG_HEIGHT) + + +#: Frames the discovery preview aligns. ~20 is the brief's number and it is a +#: DEFAULT, not a law — ``preview_frames`` in Advanced moves it. +#: +#: Sampled EVENLY OVER THE WHOLE MOVIE, not the first 20 in a row. The question +#: a preview answers is "does this landmark survive the FULL excursion", and 20 +#: consecutive frames of a 3000-frame movie drift by almost nothing — a +#: contiguous window would answer "looks fine" for every box, including the +#: useless ones. The same reasoning (and the same spacing) as +#: :meth:`DriftWizard.sum_indices`. +_PREVIEW_FRAMES = 20 + +#: Byte ceiling on the preview's retained crop stack. The preview reads one +#: FULL frame at a time and keeps only the (usually small) ROI crop, so this +#: bounds the only thing that accumulates. With no ROI the crop IS the frame, +#: which is how 20 frames of a 4096² movie would otherwise become 1.3 GB; +#: over the cap the sampled frame count is thinned rather than the read being +#: abandoned. Never a reason to touch the full dataset (CLAUDE.md). +_PREVIEW_MAX_BYTES = 192 * 1024 * 1024 + +#: Settle delay for a preview re-solve driven by an ROI DRAG. The widget's +#: pointer_move fires at renderer frame rate; re-solving 20 frames on each one +#: would queue solves faster than they finish. ``drift_tune`` is NOT debounced +#: here — the renderer's ``useDebouncedAction`` already settles it, and +#: debouncing twice just adds latency. +_PREVIEW_SETTLE_S = 0.25 + +#: Smallest alignment box, in image pixels. MUST stay >= the solver's own +#: ``spyde.drift.translation._MIN_ROI``, which REJECTS a smaller box rather +#: than clamping it (a silently shrunk ROI would correlate somewhere the user +#: did not drag). Pinned by ``test_drift_wizard.py``. +_ROI_MIN_PX = 16 + +#: Default alignment box: this fraction of each frame dimension, centred. Half +#: the frame is deliberately generous — the ROI is FIXED in frame coordinates, +#: so the landmark drifts within it and the box wants to be comfortably larger +#: than the total excursion. +_ROI_DEFAULT_FRACTION = 0.5 + +_ROI_COLOR = "#94e2d5" + +#: **Off by default, and that is a measurement, not caution.** A guessed centre +#: box is NOT automatically the better correlation: on the ``particle_movie`` +#: fixture (96×112 frames, the default half-frame box = 48×56) the ROI solve +#: comes back 1.03 px from the stamped ground truth where the whole-frame solve +#: is 0.25 px — a quarter of the pixels is a quarter of the correlation signal, +#: and the Tukey taper eats a larger fraction of a small box. So the default +#: stays the answer we already know is right, and the ROI is what the user +#: reaches for when the whole frame is the problem (a moving sample, a mostly +#: featureless field). The preview runs on the box either way — that is the +#: discovery step, and it is what tells you the box is worth committing to. +DEFAULTS: dict[str, Any] = dict( + use_roi=False, + reject_outliers=True, + method="rigid", + upsample=8, + max_shift=32.0, + reference="running", + apodize=True, + normalize=True, + order=1, + preview_frames=_PREVIEW_FRAMES, +) + + +class DriftWizard(WizardController): + """Owns the drift caret's state: parameters, the alignment ROI and its live + preview, the solved model, and the two figure windows.""" + + key = "drift" + + #: One source of truth (mirrored by ``registry._WIZARD_SCHEMAS``). Entries + #: WITHOUT a ``tab`` are the caret's default face; everything tagged + #: ``"Advanced"`` renders behind the collapsed disclosure (§0.9a). + parameters = { + "use_roi": { + "name": "Use ROI for alignment", "type": "bool", + "default": DEFAULTS["use_roi"], + }, + "reject_outliers": { + "name": "Ignore bad frames", "type": "bool", + "default": DEFAULTS["reject_outliers"], + }, + "method": { + "name": "Model", "type": "enum", "default": DEFAULTS["method"], + "choices": list(METHODS), "tab": "Advanced", + }, + "reference": { + "name": "Reference", "type": "enum", "default": DEFAULTS["reference"], + "choices": ["running", "sequential", "first"], "tab": "Advanced", + }, + "upsample": { + "name": "Sub-pixel factor", "type": "int", "default": DEFAULTS["upsample"], + "min": 1, "max": 64, "tab": "Advanced", + }, + "max_shift": { + "name": "Max shift (px)", "type": "float", "default": DEFAULTS["max_shift"], + "min": 1.0, "max": 4096.0, "step": 1.0, "tab": "Advanced", + }, + "apodize": { + "name": "Edge taper", "type": "bool", "default": DEFAULTS["apodize"], + "tab": "Advanced", + }, + "normalize": { + "name": "Phase correlation", "type": "bool", + "default": DEFAULTS["normalize"], "tab": "Advanced", + }, + "order": { + "name": "Interpolation order", "type": "int", "default": DEFAULTS["order"], + "min": 0, "max": 3, "tab": "Advanced", + }, + "preview_frames": { + "name": "Preview frames", "type": "int", + "default": DEFAULTS["preview_frames"], "min": 4, "max": 200, + "tab": "Advanced", + }, + } + + def __init__(self, session, tree, src_plot): + super().__init__(session, tree) + self.src_plot = src_plot + self.src_window_id = getattr(src_plot, "window_id", None) + self.params: dict[str, Any] = dict(DEFAULTS) + self.model = None + #: The Drift Check window (a bare figure) and its four panels. + self.window_id: int | None = None + self._panels: dict[str, Any] = {} + self._sum_indices: np.ndarray | None = None + self._before_sum: np.ndarray | None = None + #: The dy/dx window — opened by the solve, filled from ``on_shift``. + self.trace_window_id: int | None = None + self._trace: dict[str, Any] = {} + #: The alignment ROI (discovery): widget + last preview result. + self._roi_widget = None + self._roi_handler = None + self._roi_clamping = False + self._frame_shape: tuple[int, int] | None = None + self._settle: threading.Timer | None = None + self.preview: dict[str, Any] | None = None + #: Cancel flag of the solve in flight (Discard/Stop flips it). + self._stop: list[bool] = [False] + + # ── the movie ──────────────────────────────────────────────────────────── + + def signal(self): + return _current_signal(self.src_plot) or self.tree.root + + def frames(self): + """``(n_frames, get_frame, (h, w))`` — one frame at a time.""" + from spyde.drift import frame_source + return frame_source(self.signal()) + + def sum_indices(self, n_frames: int) -> np.ndarray: + if self._sum_indices is None or self._sum_indices.size == 0: + k = min(int(n_frames), _SUM_MAX_FRAMES) + self._sum_indices = np.unique( + np.linspace(0, max(0, n_frames - 1), max(1, k)).round().astype(int)) + return self._sum_indices + + # ── the alignment ROI (the discovery feature) ──────────────────────────── + + def _plot2d(self): + return getattr(self.src_plot, "_plot2d", None) if self.src_plot else None + + def ensure_roi_widget(self, shape: tuple[int, int]) -> None: + """Draw the draggable alignment box on the source movie (idempotent). + + Geometry is IMAGE PIXELS — anyplotlib 2-D widgets report ``x/y/w/h`` + that way, and that is exactly what ``solve_translation(roi=…)`` wants. + A raw ``add_rectangle_widget`` rather than ``RectangleSelector``: the + selector caps itself at ``MAX_REGION_EXTENT_PER_DIM`` (16 px) because + it drives a nav-space region integrate, and a 16 px alignment box is + below the solver's own floor. + """ + h, w = int(shape[0]), int(shape[1]) + self._frame_shape = (h, w) + if self._roi_widget is not None: + return + plot2d = self._plot2d() + if plot2d is None: + return + if min(h, w) < 2 * _ROI_MIN_PX: + # Nothing sensible to drag; the whole frame IS the ROI. + return + bw = max(_ROI_MIN_PX, min(w, int(round(w * _ROI_DEFAULT_FRACTION)))) + bh = max(_ROI_MIN_PX, min(h, int(round(h * _ROI_DEFAULT_FRACTION)))) + try: + widget = plot2d.add_rectangle_widget( + x=float((w - bw) // 2), y=float((h - bh) // 2), + w=float(bw), h=float(bh), color=_ROI_COLOR, show_handles=True, + ) + from spyde.drawing.selectors.base_selector import event_handler_fn + handler = event_handler_fn(lambda event: self._on_roi_drag()) + widget.add_event_handler(handler, "pointer_move", "pointer_up") + self._roi_widget = widget + self._roi_handler = handler # keep a ref alive (weak callback) + except Exception as exc: + log.debug("[drift] alignment ROI widget failed: %s", exc) + + def _on_roi_drag(self) -> None: + """Clamp the box to the frame, then arm the settle timer. + + RE-ENTRANCY GUARD: anyplotlib ``Widget.set()`` fires ``pointer_move`` + UNCONDITIONALLY (even on a no-change write), so the clamp below + re-invokes this handler synchronously — unguarded, ONE JS drag frame + recursed ~2000 deep before RecursionError in the Crop box (see + ``actions/base.py``). A hard flag breaks the cycle; compare-before-set + is NOT sufficient. + """ + if self._roi_clamping or self._closed: + return + self._roi_clamping = True + try: + self._clamp_roi() + finally: + self._roi_clamping = False + self.schedule_preview() + + def _clamp_roi(self) -> None: + """Keep the box inside the frame and above the solver's floor. + + COMPARE BEFORE SET, with slack. ``Widget.set()`` pushes geometry back to + the renderer, and writing on every ``pointer_move`` echoes + python-sourced geometry into a live drag — the same failure the 1-D span + cap documents in CLAUDE.md (Live-Display §3). A box already resting on a + bound must be left alone. + """ + widget, shape = self._roi_widget, self._frame_shape + if widget is None or shape is None: + return + h, w = shape + try: + ww = min(max(float(widget.w), float(_ROI_MIN_PX)), float(w)) + hh = min(max(float(widget.h), float(_ROI_MIN_PX)), float(h)) + x = min(max(float(widget.x), 0.0), float(w) - ww) + y = min(max(float(widget.y), 0.0), float(h) - hh) + now = (float(widget.x), float(widget.y), + float(widget.w), float(widget.h)) + if max(abs(a - b) for a, b in zip(now, (x, y, ww, hh))) > 1e-6: + widget.set(x=x, y=y, w=ww, h=hh) + except Exception as exc: + log.debug("[drift] clamping the alignment ROI failed: %s", exc) + + def roi_box(self) -> tuple[int, int, int, int] | None: + """``(y0, x0, h, w)`` in IMAGE PIXELS, or None when there is no usable + box — the shape ``solve_translation``'s ``roi`` takes, with no scale or + offset applied because neither side has any.""" + widget, shape = self._roi_widget, self._frame_shape + if widget is None or shape is None: + return None + fh, fw = shape + try: + x0 = int(round(float(widget.x))) + y0 = int(round(float(widget.y))) + bw = int(round(float(widget.w))) + bh = int(round(float(widget.h))) + except Exception as exc: + log.debug("[drift] reading the alignment ROI failed: %s", exc) + return None + # SIZE first, then origin. The other order looks equivalent and is not: + # clamping the origin to the frame edge and only then applying the + # minimum size pushes the box back OUT past the edge, and + # solve_translation rejects an out-of-frame roi outright. + bw = max(_ROI_MIN_PX, min(bw, fw)) + bh = max(_ROI_MIN_PX, min(bh, fh)) + if bw > fw or bh > fh: + return None # frame smaller than the solver's floor + x0 = max(0, min(x0, fw - bw)) + y0 = max(0, min(y0, fh - bh)) + return (y0, x0, bh, bw) + + def active_roi(self) -> tuple[int, int, int, int] | None: + """The ROI the FULL SOLVE should use — None unless the toggle is on (and + None when no box is usable, which is a whole-frame correlation). + + The preview deliberately does NOT go through here: it always aligns the + box, toggle or not, because that is the question it exists to answer + ("is this landmark worth committing to?"). The toggle is the commitment. + """ + return self.roi_box() if self.params.get("use_roi") else None + + def remove_roi_widget(self) -> None: + widget, self._roi_widget = self._roi_widget, None + self._roi_handler = None + if widget is not None: + try: + widget.hide() # widgets have no remove(), only hide() + except Exception as exc: + log.debug("[drift] hiding the alignment ROI failed: %s", exc) + + # ── preview scheduling (latest-wins, cancellable) ──────────────────────── + + def schedule_preview(self, delay: float = _PREVIEW_SETTLE_S) -> None: + """(Re-)arm the settle timer for a drag-driven preview re-solve. + + Latest-wins in two places: the timer is restarted per pointer event so + only the RESTING geometry ever solves, and the solve that does run + carries a ``_drift_preview_gen`` generation so a superseded result is + dropped on arrival instead of painting over a newer one. + """ + self.cancel_preview() + if self._closed: + return + timer = threading.Timer(max(0.0, float(delay)), self._fire_preview) + timer.daemon = True + self._settle = timer + timer.start() + + def cancel_preview(self) -> None: + timer, self._settle = self._settle, None + if timer is not None: + try: + timer.cancel() + except Exception as exc: + log.debug("[drift] cancelling the preview timer failed: %s", exc) + + def _fire_preview(self) -> None: + """Timer thread → main thread → the worker. Reading widget geometry and + spawning the compute both belong on the main thread (thread marshal, + README §6); only the arithmetic runs on the worker.""" + self._settle = None + if self._closed: + return + dispatch = getattr(self.session, "_dispatch_to_main", None) + if dispatch is None: + _run_preview(self) + else: + dispatch(lambda: (None if self._closed else _run_preview(self))) + + # ── the check window ───────────────────────────────────────────────────── + + def open_check_window(self, before: np.ndarray, n_frames: int) -> None: + """Emit the bare-figure Drift Check window and register this controller + for it, so ✕ and ``Session._forget_window`` reach the wizard. + + Top row = the whole movie, raw and corrected (the solve's evidence). + Bottom row = the DISCOVERY pair, raw and aligned over ~20 frames of + whatever is being correlated (the ROI, or the whole frame when the + toggle is off). Side by side, because "is this landmark good" is + answered by comparing two sums, not by staring at one. + """ + import anyplotlib as apl + import anyplotlib._electron as _electron + from spyde.actions.figure_registry import keep_alive + from spyde.drawing.plots.plot import finalize_figure_html + + figsize, aspect = _figure_geometry() + fig, axes = apl.subplots(2, 2, figsize=figsize) + ax = np.array(axes, dtype=object).ravel() + before = np.asarray(before, np.float32) + zeros = np.zeros_like(before) + + panels = { + "before": ax[0].imshow(before, cmap="gray"), + "after": ax[1].imshow(zeros, cmap="gray"), + "roi_raw": ax[2].imshow(zeros, cmap="gray"), + "roi_aligned": ax[3].imshow(zeros, cmap="gray"), + } + titles = {"before": "Raw sum", "after": "Corrected sum", + "roi_raw": "ROI raw", "roi_aligned": "ROI aligned"} + self._panels = panels + for key, title in titles.items(): + self._set_panel_title(key, title) + + wid = self.session.next_window_id() + fig_id = _electron.register(fig) + html = finalize_figure_html(fig, fig_id) + keep_alive(int(wid), fig) + emit({"type": "figure", "fig_id": fig_id, "window_id": int(wid), + "html": html, "title": "Drift Check", "is_navigator": False, + "aspect": float(aspect)}) + self.window_id = int(wid) + self._before_sum = before + self.own_window(wid) + + def _set_panel_title(self, key: str, title: str) -> None: + panel = self._panels.get(key) + if panel is None: + return + try: + panel.set_title(title) + except Exception as exc: + log.debug("[drift] set_title(%s) failed: %s", key, exc) + + def update_check(self, *, after=None) -> None: + """Paint the whole-movie corrected sum (main thread only).""" + if after is None or not self._panels: + return + try: + self._panels["after"].set_data(np.asarray(after, np.float32)) + except Exception as exc: + log.debug("[drift] painting the corrected sum failed: %s", exc) + + def show_preview(self, result: dict) -> None: + """Paint the discovery pair + its titles (main thread only).""" + if not self._panels: + return + what = "ROI" if result.get("roi") is not None else "Whole frame" + n = int(result.get("frames", 0)) + gain = float(result.get("gain", float("nan"))) + for key, arr, title in ( + ("roi_raw", result.get("raw"), f"{what} raw · {n} frames"), + ("roi_aligned", result.get("aligned"), + f"{what} aligned · {gain:.1f}x sharper" if np.isfinite(gain) + else f"{what} aligned"), + ): + if arr is None: + continue + try: + self._panels[key].set_data(np.asarray(arr, np.float32)) + except Exception as exc: + log.debug("[drift] painting the %s panel failed: %s", key, exc) + self._set_panel_title(key, title) + + # ── the dy/dx window ───────────────────────────────────────────────────── + + def open_trace_window(self, n_frames: int) -> None: + """Open (or reset) the dy/dx figure window. + + Its OWN window, not caret furniture: the curve is the measurement, and + a 40 px inline sparkline could show that the stage crept but never + which frame jumped. One panel with two labelled lines rather than two + panels — drift is anisotropic, and a shared y-scale is what makes + "mostly x" readable at a glance. + """ + n = max(2, int(n_frames)) + if self.trace_window_id is not None and self._trace: + self.reset_trace(n) + return + import anyplotlib as apl + import anyplotlib._electron as _electron + from spyde.actions.figure_registry import keep_alive + from spyde.drawing.plots.plot import finalize_figure_html + + figsize, aspect = _figure_geometry() + fig, axes = apl.subplots(1, 1, figsize=figsize) + ax = np.array(axes, dtype=object).ravel()[0] + x0 = np.zeros(1, dtype=np.float64) + y0 = np.zeros(1, dtype=np.float64) + panel = ax.plot(y0, axes=[x0], units="frame", y_units="shift (px)", + color="#89b4fa", label="dy") + dx_line = panel.add_line(y0, x_axis=x0, color="#f38ba8", label="dx") + for setter, text in (("set_title", "Drift dy / dx"), + ("set_xlabel", "frame"), + ("set_ylabel", "shift (px)")): + try: + getattr(panel, setter)(text) + except Exception as exc: + log.debug("[drift] trace %s failed: %s", setter, exc) + + wid = self.session.next_window_id() + fig_id = _electron.register(fig) + html = finalize_figure_html(fig, fig_id) + keep_alive(int(wid), fig) + emit({"type": "figure", "fig_id": fig_id, "window_id": int(wid), + "html": html, "title": "Drift dy/dx", "is_navigator": False, + "aspect": float(aspect)}) + self.trace_window_id = int(wid) + self._trace = {"panel": panel, "dx": dx_line} + self.reset_trace(n) + self.own_window(wid) + + def reset_trace(self, n_frames: int) -> None: + n = max(2, int(n_frames)) + self._trace["dy_data"] = np.full(n, np.nan, np.float64) + self._trace["dx_data"] = np.full(n, np.nan, np.float64) + self._trace["filled"] = 0 + self.push_trace([(0, 0.0, 0.0)]) + + def push_trace(self, points) -> None: + """Append a batch of ``(index, dy, dx)`` and repaint (main thread only). + + Only the SOLVED PREFIX is pushed, so the curve grows left to right and + the y-scale tracks what has actually been measured — pushing the whole + NaN-padded array would make anyplotlib's auto-range see one point. + """ + if not self._trace: + return + dy = self._trace.get("dy_data") + dx = self._trace.get("dx_data") + if dy is None or dx is None: + return + hi = int(self._trace.get("filled", 0)) + for i, y, x in points: + i = int(i) + if 0 <= i < dy.size: + dy[i] = float(y) + dx[i] = float(x) + hi = max(hi, i + 1) + self._trace["filled"] = hi + if hi < 1: + return + xs = np.arange(hi, dtype=np.float64) + try: + self._trace["panel"].set_data(np.nan_to_num(dy[:hi]), x_axis=xs) + self._trace["dx"].set_data(np.nan_to_num(dx[:hi]), x_axis=xs) + except Exception as exc: + log.debug("[drift] painting the dy/dx trace failed: %s", exc) + + def close_trace_window(self) -> None: + self._trace = {} + wid, self.trace_window_id = self.trace_window_id, None + self._close_window(wid) + + # ── lifecycle ──────────────────────────────────────────────────────────── + + def _close_window(self, wid: int | None) -> None: + if wid is None: + return + forget = getattr(self.session, "_forget_window", None) + if forget is not None: + try: + forget(int(wid)) + except Exception as exc: + log.debug("[drift] forgetting window %s failed: %s", wid, exc) + return + # Bare / stub session: emit + unregister by hand. + try: + emit({"type": "window_closed", "window_id": int(wid)}) + except Exception as exc: + log.debug("[drift] closing window %s failed: %s", wid, exc) + reg = getattr(self.session, "_window_controllers", None) + if isinstance(reg, dict): + reg.pop(int(wid), None) + + def close(self) -> None: + """WindowController protocol — ``Session._forget_window`` calls this for + EITHER owned window, with no way to say which. + + Only the Drift Check window is the wizard's life; closing the dy/dx + window just drops the curve. ``_forget_window`` pops the controller for + the window that went away BEFORE calling here, so "is the check window + still registered?" identifies it exactly — and the programmatic path + (:meth:`close_trace_window`) clears ``trace_window_id`` first, so this + re-entry is a no-op rather than a recursion. + """ + reg = getattr(self.session, "_window_controllers", None) or {} + if self.window_id is not None and reg.get(int(self.window_id)) is self: + self._trace = {} + self.trace_window_id = None + return + self.remove() + + def remove(self) -> None: + if self._closed: + return + self._closed = True + self._stop[0] = True # stop a solve in flight + self.cancel_preview() + self.remove_roi_widget() + self._panels = {} + self._trace = {} + wid, self.window_id = self.window_id, None + twid, self.trace_window_id = self.trace_window_id, None + self._close_window(wid) + self._close_window(twid) + if getattr(self.tree, "_drift_wizard", None) is self: + self.tree._drift_wizard = None + + def commit(self): + """Add the lazy corrected node — see :func:`drift_commit`.""" + return _commit(self) + + +# ── parameters ─────────────────────────────────────────────────────────────── + +def _coerce(payload: dict | None) -> dict: + p = dict(DEFAULTS) + payload = payload or {} + for k, default in DEFAULTS.items(): + v = payload.get(k) + if v is None or v == "": + continue + try: + p[k] = bool(v) if isinstance(default, bool) else type(default)(v) + except (TypeError, ValueError) as exc: + log.debug("[drift] param %r=%r not coercible, keeping default: %s", + k, v, exc) + p["method"] = str(p["method"]).lower() + if p["method"] not in METHODS: + p["method"] = DEFAULTS["method"] + if p["reference"] not in ("running", "sequential", "first"): + p["reference"] = DEFAULTS["reference"] + p["upsample"] = max(1, int(p["upsample"])) + p["max_shift"] = max(1.0, float(p["max_shift"])) + p["order"] = int(min(3, max(0, p["order"]))) + p["preview_frames"] = int(min(200, max(4, p["preview_frames"]))) + return p + + +def _solver_kwargs(p: dict, roi=None) -> dict: + return dict(upsample=int(p["upsample"]), max_shift=float(p["max_shift"]), + reference=str(p["reference"]), apodize=bool(p["apodize"]), + normalize=bool(p["normalize"]), + reject_outliers=bool(p["reject_outliers"]), + roi=None if roi is None else tuple(int(v) for v in roi)) + + +def _wizard(session, plot) -> DriftWizard | None: + """Resolve the live wizard from any of its windows. + + The check and dy/dx windows are bare figures, so ``_plot_by_window_id`` + returns None for them and the plot-based lookup finds nothing — resolve by + window id through the controller registry first (README §6), then fall back + to the source tree's back-reference. + """ + wid = getattr(plot, "window_id", None) if plot is not None else None + lookup = getattr(session, "controller_by_window_id", None) + if wid is not None and lookup is not None: + ctrl = lookup(int(wid)) + if isinstance(ctrl, DriftWizard) and not ctrl._closed: + return ctrl + _src, tree = _src_plot_tree(session, plot) + wiz = getattr(tree, "_drift_wizard", None) if tree is not None else None + return wiz if (wiz is not None and not wiz._closed) else None + + +def _emit_state(wiz: DriftWizard, **extra) -> None: + roi = wiz.roi_box() + msg = {"type": "drift_state", + "window_id": wiz.src_window_id, + "check_window_id": wiz.window_id, + "trace_window_id": wiz.trace_window_id, + "method": wiz.params["method"], + "solved": wiz.model is not None, + "use_roi": bool(wiz.params["use_roi"]), + "roi": None if roi is None else [int(v) for v in roi], + "params": dict(wiz.params)} + msg.update(extra) + emit(msg) + + +# ── streaming sums + the sharpness number ──────────────────────────────────── + +def _stack_sum(get_frame, indices, shifts=None, *, order: int = 1) -> np.ndarray: + """Mean of the selected frames, optionally drift-corrected first. + + Streams: one frame resident at a time plus one float64 accumulator, so this + is safe at the plan's target scale however long the movie is. NaN padding + from :func:`spyde.drift.warp.shift_frame` is excluded per pixel rather than + zero-filled — a zero-filled border reads as a dark rim that looks like real + data and would be segmented as one. + """ + acc = None + hits = None + for i in indices: + frame = np.asarray(get_frame(int(i)), dtype=np.float32) + if shifts is not None: + s = shifts[int(i)] + if np.all(np.isfinite(s)): + from spyde.drift import shift_frame + frame = shift_frame(frame, s, order=order) + if acc is None: + acc = np.zeros(frame.shape, np.float64) + hits = np.zeros(frame.shape, np.int32) + good = np.isfinite(frame) + acc[good] += frame[good] + hits[good] += 1 + if acc is None: + return np.zeros((1, 1), np.float32) + with np.errstate(invalid="ignore", divide="ignore"): + out = np.where(hits > 0, acc / np.maximum(hits, 1), np.nan) + return out.astype(np.float32) + + +def _gradient_energy(img, mask=None) -> float: + """Mean squared forward-difference gradient over the valid pixels. + + The sharpness number, and it is NaN-aware by construction rather than by + ``nan_to_num``: an aligned sum's uncovered border is NaN (plan A7 — nothing + is cropped, nothing is invented), and zero-filling it manufactures a step + at the border whose gradient energy dwarfs the image's own, which would + make every ROI look brilliantly sharp. Differences touching a non-finite + (or masked-out) pixel are excluded from BOTH the sum and the count, so the + raw and aligned sums are measured over exactly the same pixels. + """ + a = np.asarray(img, np.float64) + ok = np.isfinite(a) + if mask is not None: + ok &= np.asarray(mask, bool) + a = np.where(ok, a, 0.0) + total = 0.0 + count = 0 + if a.shape[0] > 1: + m = ok[1:, :] & ok[:-1, :] + d = (a[1:, :] - a[:-1, :])[m] + total += float(np.sum(d * d)) + count += int(m.sum()) + if a.shape[1] > 1: + m = ok[:, 1:] & ok[:, :-1] + d = (a[:, 1:] - a[:, :-1])[m] + total += float(np.sum(d * d)) + count += int(m.sum()) + return total / count if count else float("nan") + + +def _preview_indices(n_frames: int, k: int, frame_bytes: int) -> np.ndarray: + """Evenly spaced sample of the movie for the preview, thinned to fit the + byte cap. See :data:`_PREVIEW_FRAMES` for why evenly spaced and not the + first *k* in a row.""" + n = max(1, int(n_frames)) + k = max(2, min(int(k), n)) + cap = max(2, int(_PREVIEW_MAX_BYTES // max(1, int(frame_bytes)))) + k = min(k, cap) + return np.unique(np.linspace(0, n - 1, k).round().astype(int)) + + +def preview_alignment(get_frame, indices, roi, *, params) -> dict: + """Align *indices* on *roi* alone and report how much sharper the sum got. + + Reads one FULL frame at a time and keeps only the crop, so the resident set + is one frame plus ``len(indices)`` crops (bounded by + :data:`_PREVIEW_MAX_BYTES` at the caller). The crops ARE the region to + correlate, so the solve runs with ``roi=None`` on them. + + Returns ``{roi, frames, raw, aligned, gain, raw_energy, aligned_energy, + max_abs_shift}``. *gain* is the whole point: > 1 means aligning this region + genuinely sharpened it, ~1 means alignment changes nothing here (a + featureless box), and it is measured on the pixels both sums cover. + """ + from spyde.drift import solve_translation + + crops: list[np.ndarray] = [] + for i in indices: + frame = np.asarray(get_frame(int(i)), np.float32) + if roi is not None: + y0, x0, h, w = (int(v) for v in roi) + frame = frame[y0:y0 + h, x0:x0 + w] + crops.append(np.ascontiguousarray(frame, dtype=np.float32)) + + model = solve_translation(crops, **_solver_kwargs(params)) + take = range(len(crops)) + raw = _stack_sum(crops.__getitem__, take) + aligned = _stack_sum(crops.__getitem__, take, model.shifts, + order=int(params["order"])) + both = np.isfinite(raw) & np.isfinite(aligned) + e_raw = _gradient_energy(raw, both) + e_aligned = _gradient_energy(aligned, both) + gain = (e_aligned / e_raw) if (np.isfinite(e_raw) and e_raw > 0) \ + else float("nan") + return {"roi": None if roi is None else tuple(int(v) for v in roi), + "frames": len(crops), "raw": raw, "aligned": aligned, + "gain": float(gain), "raw_energy": float(e_raw), + "aligned_energy": float(e_aligned), + "max_abs_shift": float(model.max_abs_shift)} + + +# ── staged handlers ────────────────────────────────────────────────────────── + +def drift_open(session, plot, payload) -> None: + """Caret mounted: build the controller, open the Drift Check window, draw + the alignment ROI, and run the first discovery preview. + + Nothing SOLVES here — plan A8 is explicit that drift correction is opt-in + and never runs on load. The compute is the bounded raw sum plus the + ~20-frame preview of the default box, which is what makes the caret's first + frame informative instead of an empty panel and a button. + """ + src, tree = _src_plot_tree(session, plot) + if src is None or tree is None: + emit_error("Drift Correction: no active dataset") + return + + existing = getattr(tree, "_drift_wizard", None) + if existing is not None and not existing._closed: + existing.params = _coerce({**existing.params, **(payload or {})}) + _emit_state(existing) + return + + wiz = DriftWizard(session, tree, src) + wiz.params = _coerce(payload) + try: + n_frames, get_frame, shape = wiz.frames() + except TypeError as exc: + emit_error(f"Drift Correction: {exc}") + return + # BEFORE the worker: StrictMode fires open/close/open synchronously and the + # close's bump has to be able to invalidate this open's deferred build. + gen = wiz.guard() + tree._drift_wizard = wiz + _emit_state(wiz, n_frames=int(n_frames)) + + def _work(): + return _stack_sum(get_frame, wiz.sum_indices(n_frames)) + + def _done(raw_sum): + if not wiz.still(gen) or wiz._closed: + return + wiz.open_check_window(raw_sum, int(n_frames)) + wiz.ensure_roi_widget(shape) + _emit_state(wiz, n_frames=int(n_frames)) + emit_status("Drift Correction: drag the box onto a landmark to test it, " + "then Correct Drift.") + _run_preview(wiz) + + def _fail(exc): + emit_error(f"Drift Correction: reading the movie failed — {exc}") + + run_on_worker(session, _work, name="drift-open", on_done=_done, on_error=_fail) + + +def drift_close(session, plot, payload=None) -> None: + """Caret unmounted: invalidate in-flight work FIRST, then tear down.""" + _src, tree = _src_plot_tree(session, plot) + wiz = _wizard(session, plot) + if tree is not None: + # The same `_drift_run_gen` key WizardController.cancel_inflight bumps, + # done on the TREE so it fires even when there is no controller yet: a + # StrictMode open whose worker has not landed must still be cancelled. + bump_generation(tree, "_drift_run_gen") + bump_generation(tree, "_drift_preview_gen") + if wiz is not None: + # Harmlessly re-bumps when the tree resolved above; the point is the + # case where it did not (the wizard was found through one of the + # figure windows' controller registry). + wiz.cancel_inflight() + wiz.remove() + + +def drift_set_method(session, plot, payload) -> None: + """Select the drift model (Advanced). + + Only ``rigid`` has a solver in ``spyde.drift`` today. Selecting the other + says so and stays on rigid — running a rigid solve while the caret claims + "rigid+affine" would put a wrong ``kind`` into the model's provenance, + which is worse than the missing feature. + """ + wiz = _wizard(session, plot) + if wiz is None: + return + method = str((payload or {}).get("method", "")).lower() + if method not in METHODS: + emit_error(f"Drift Correction: unknown model {method!r}") + return + reason = _UNAVAILABLE.get(method) + if reason: + emit_status(f"Drift Correction: {reason} — staying on the rigid solve.") + method = "rigid" + wiz.params["method"] = method + _emit_state(wiz) + + +def drift_tune(session, plot, payload) -> None: + """A toggle or Advanced parameter changed → re-run the discovery preview. + + NOT debounced here: the renderer's ``useDebouncedAction`` already settles + the send, and debouncing twice only adds latency. The drag path IS + debounced, on the backend, because widget pointer events arrive at renderer + frame rate (:meth:`DriftWizard.schedule_preview`). + """ + wiz = _wizard(session, plot) + if wiz is None: + return + wiz.params = _coerce({**wiz.params, **(payload or {})}) + _emit_state(wiz) + _run_preview(wiz) + + +def _run_preview(wiz: DriftWizard) -> None: + """Align ~20 sampled frames on the current box and report the gain. + + Latest-wins on ``_drift_preview_gen``: a drag that outruns the solve drops + the stale result rather than painting it over the newer one. The preview + never touches ``tree.drift`` or the caret's solved state — it is a question, + not an answer. + """ + if wiz._closed: + return + tree = wiz.tree + gen = bump_generation(tree, "_drift_preview_gen") + params = dict(wiz.params) + roi = wiz.roi_box() # the BOX, toggle or not — see active_roi() + try: + n_frames, get_frame, shape = wiz.frames() + except TypeError as exc: + log.debug("[drift] preview skipped: %s", exc) + return + if n_frames < 2: + return + h, w = (roi[2], roi[3]) if roi is not None else (int(shape[0]), int(shape[1])) + indices = _preview_indices(n_frames, params["preview_frames"], h * w * 4) + + def _work(): + return preview_alignment(get_frame, indices, roi, params=params) + + def _done(result): + if not is_current(tree, "_drift_preview_gen", gen) or wiz._closed: + return + wiz.preview = result + wiz.show_preview(result) + emit({"type": "drift_preview", "window_id": wiz.src_window_id, + "roi": None if result["roi"] is None else list(result["roi"]), + "frames": int(result["frames"]), + "gain": float(result["gain"]), + "max_abs_shift": float(result["max_abs_shift"]), + "params": dict(params)}) + + def _fail(exc): + if is_current(tree, "_drift_preview_gen", gen): + emit_error(f"Drift preview failed: {exc}") + + run_on_worker(wiz.session, _work, name="drift-preview", + on_done=_done, on_error=_fail) + + +def drift_run(session, plot, payload) -> None: + """Solve the whole movie on a worker: progress-reported and cancellable. + + Opens the dy/dx window FIRST and fills it from ``solve_translation``'s + ``on_shift`` stream, so the curve draws while it solves rather than + appearing whole at the end. + + Cancellation goes through ``BaseSignalTree.register_cancel`` so closing the + tree stops the solve, and ``solve_translation``'s own ``cancel()`` hook + polls the same flag — a cancelled solve leaves NaN shifts for the frames it + never reached, which is why a partial model is detectable rather than + silently wrong. Stop/Discard flips the same flag. + """ + src, tree = _src_plot_tree(session, plot) + if src is None or tree is None: + emit_error("Drift Correction: no active dataset") + return + wiz = _wizard(session, plot) + if wiz is None: + emit_error("Drift Correction: the caret is not open") + return + wiz.params = _coerce({**wiz.params, **(payload or {})}) + p = dict(wiz.params) + reason = _UNAVAILABLE.get(p["method"]) + if reason: + emit_status(f"Drift Correction: {reason} — solving rigid instead.") + p["method"] = wiz.params["method"] = "rigid" + + try: + n_frames, get_frame, _shape = wiz.frames() + except TypeError as exc: + emit_error(f"Drift Correction: {exc}") + return + if n_frames < 2: + emit_error("Drift Correction needs at least two frames") + return + + roi = wiz.active_roi() + if p["use_roi"] and roi is None: + emit_status("Drift Correction: no usable alignment box — correlating " + "the whole frame.") + + gen = wiz.guard() + stopped = [False] + wiz._stop = stopped + if hasattr(tree, "register_cancel"): + tree.register_cancel(flag=stopped) + wiz.open_trace_window(int(n_frames)) + _emit_state(wiz) + emit_status(f"Solving drift over {n_frames} frames…") + dispatch = getattr(session, "_dispatch_to_main", None) + + def _work(): + from spyde.drift import solve_translation + + def _progress(done, total): + emit_progress(int(done), int(total), "Drift") + emit({"type": "drift_progress", "window_id": wiz.src_window_id, + "done": int(done), "total": int(total)}) + + # Stream the curve as it solves. `progress` carries only a count and the + # shift array is solver-local until the return, so without this callback + # the caret could show a bar but not a trace. Batched rather than per + # frame: at thousands of frames one message each would flood the PLOTAPP + # line protocol for a curve the eye cannot follow that finely. The PAINT + # is marshalled — `on_shift` runs on the solver thread and figures are + # main-thread only (README §6). + pending: list[tuple[int, float, float]] = [] + last_flush = [time.monotonic()] + + def _flush(): + if not pending: + return + batch = pending[:] + pending.clear() + last_flush[0] = time.monotonic() + emit({"type": "drift_trace", "window_id": wiz.src_window_id, + "points": batch}) + if not wiz.still(gen): + return + if dispatch is None: + wiz.push_trace(batch) + else: + dispatch(lambda b=batch: (None if wiz._closed or not wiz.still(gen) + else wiz.push_trace(b))) + + def _on_shift(i, dy, dx, _sharp): + pending.append((int(i), float(dy), float(dx))) + if (len(pending) >= _TRACE_BATCH + or time.monotonic() - last_flush[0] >= _TRACE_MAX_INTERVAL): + _flush() + + model = solve_translation( + wiz.signal(), progress=_progress, on_shift=_on_shift, + cancel=lambda: stopped[0], + provenance={"action": "Drift Correction", "params": dict(p), + "roi": None if roi is None else [int(v) for v in roi]}, + **_solver_kwargs(p, roi)) + _flush() + if stopped[0]: + return model, None, float("nan") + # One extra streaming pass over the SAME bounded subset the raw sum + # used, so the two check images are comparable. + after = _stack_sum(get_frame, wiz.sum_indices(n_frames), model.shifts, + order=int(p["order"])) + # The same number the discovery preview reports, now for the whole + # movie — measured on the worker because a 4096² gradient energy is + # ~100 ms and the main thread is the navigator's. + before = wiz._before_sum + gain = float("nan") + if before is not None and before.shape == after.shape: + both = np.isfinite(before) & np.isfinite(after) + e_before = _gradient_energy(before, both) + if np.isfinite(e_before) and e_before > 0: + gain = _gradient_energy(after, both) / e_before + return model, after, gain + + def _done(res): + model, after, gain = res + try: + if not wiz.still(gen) or wiz._closed: + return + wiz.model = model + tree.drift = model + wiz.update_check(after=after) + emit({"type": "drift_result", "window_id": wiz.src_window_id, + "shifts": [[float(a), float(b)] for a, b in model.shifts], + "kind": model.kind, "reference": model.reference, + "roi": None if roi is None else [int(v) for v in roi], + "max_abs_shift": float(model.max_abs_shift), + "gain": float(gain), + "rejected": int(model.params.get("rejected_from_reference", 0)), + "cancelled": bool(stopped[0])}) + _emit_state(wiz) + solved = int(np.isfinite(model.shifts).all(axis=1).sum()) + if stopped[0]: + emit_status(f"Drift solve stopped after {solved} of " + f"{n_frames} frames") + else: + emit_status(f"Drift solved: max shift " + f"{model.max_abs_shift:.2f} px over {n_frames} frames") + finally: + if hasattr(tree, "unregister_cancel"): + try: + tree.unregister_cancel(flag=stopped) + except Exception as exc: + log.debug("[drift] unregister_cancel failed: %s", exc) + + def _fail(exc): + emit_error(f"Drift Correction failed: {exc}") + log.exception("drift solve failed") + if hasattr(tree, "unregister_cancel"): + try: + tree.unregister_cancel(flag=stopped) + except Exception as e2: + log.debug("[drift] unregister_cancel failed: %s", e2) + + run_on_worker(session, _work, name="drift-run", on_done=_done, on_error=_fail) + + +def drift_discard(session, plot, payload=None) -> None: + """Stop a solve in flight and/or throw the solved model away. + + One handler for both because they are the same user intent ("no, not + that"): the button reads *Stop* while the bar is moving and *Discard* + once there is a result. Bumping the run generation FIRST means a solve that + finishes anyway lands on a stale generation and never installs itself. + """ + wiz = _wizard(session, plot) + if wiz is None: + return + wiz._stop[0] = True + wiz.cancel_inflight() + wiz.model = None + if getattr(wiz.tree, "drift", None) is not None: + wiz.tree.drift = None + wiz.close_trace_window() + if wiz._panels and wiz._before_sum is not None: + try: + wiz._panels["after"].set_data(np.zeros_like(wiz._before_sum)) + except Exception as exc: + log.debug("[drift] clearing the corrected sum failed: %s", exc) + _emit_state(wiz) + emit_status("Drift result discarded.") + + +# ── the corrected node ─────────────────────────────────────────────────────── + +def drift_corrected(signal, *, model, order: int = 1, fill: float = float("nan")): + """A LAZY drift-corrected view of *signal*. Plan §0.7. + + Parameters + ---------- + signal + The source movie (1-D navigation, 2-D signal). + model + The :class:`~spyde.drift.model.DriftModel` to apply. ``shifts[i]`` is the + correction ADDED to frame *i* — go through the model rather than writing + the arithmetic out; the inverted sign doubles the drift and still looks + plausible (``spyde/drift/model.py``). + order + Interpolation order for sub-pixel shifts. A whole-pixel model takes an + exact slice-copy path inside :func:`~spyde.drift.warp.shift_frame`. + fill + Uncovered-pixel value. NaN by default, per the plan A7 edge policy — + nothing is cropped and nothing is filled with invented data. + + Notes + ----- + Built with ``map_blocks`` over the source's OWN chunking, deliberately: this + never calls ``.rechunk()`` and never computes anything, so a multi-GB movie + costs a graph and nothing else (CLAUDE.md memory-safety rule, and Live- + Display §1 on not reshuffling storage chunks). Each block warps its own + frames using ``block_info`` to recover their absolute indices, so a movie + stored several frames per chunk works unchanged. + """ + import dask.array as da + from spyde.drift import shift_frame + + data = signal.data + if getattr(data, "ndim", 0) != 3: + raise ValueError( + f"drift correction needs a (n, h, w) frame stack; got shape " + f"{getattr(data, 'shape', None)}") + shifts = np.asarray(model.shifts, dtype=np.float32) + if shifts.shape[0] != int(data.shape[0]): + raise ValueError( + f"the drift model covers {shifts.shape[0]} frames but the signal has " + f"{int(data.shape[0])} — solve again on this node") + + if not isinstance(data, da.Array): + # Already resident; wrapping it costs nothing and keeps the node lazy so + # the whole tree reads through one path. + data = da.from_array(data, chunks=(1,) + tuple(int(s) for s in data.shape[1:])) + + def _block(blk, block_info=None): + t0 = (0 if block_info is None + else int(block_info[0]["array-location"][0][0])) + out = np.empty(blk.shape, np.float32) + for k in range(blk.shape[0]): + s = shifts[t0 + k] + if not np.all(np.isfinite(s)): + # A frame the solve never reached (cancelled) keeps its raw + # pixels rather than becoming an all-NaN hole. + out[k] = np.asarray(blk[k], np.float32) + else: + out[k] = shift_frame(blk[k], s, order=int(order), fill=fill) + return out + + warped = da.map_blocks(_block, data, dtype=np.float32, + meta=np.zeros((0, 0, 0), np.float32)) + new = signal._deepcopy_with_new_data(warped) + if not new._lazy: + new._lazy = True + new._assign_subclass() + return new + + +def _commit(wiz: DriftWizard): + if wiz.model is None: + emit_error("Drift Correction: solve first, then Apply") + return None + parent = wiz.signal() + try: + new_signal = wiz.tree.add_transformation( + parent, function=drift_corrected, node_name="Drift corrected", + local=True, model=wiz.model, order=int(wiz.params["order"])) + except Exception as exc: + emit_error(f"Drift Correction: applying the model failed — {exc}") + log.exception("drift commit failed") + return None + if new_signal is None: + return None + wiz.tree.drift = wiz.model + try: + new_signal.metadata.set_item( + "General.spyde_provenance", + {"action": "Drift Correction", "params": dict(wiz.params), + "kind": wiz.model.kind, "reference": wiz.model.reference, + "roi": wiz.model.params.get("roi")}) + except Exception as exc: + log.debug("[drift] stamping provenance failed: %s", exc) + show_tree_node(wiz.src_plot, wiz.tree, new_signal) + emit_status(f"Drift corrected node added (max shift " + f"{wiz.model.max_abs_shift:.2f} px)") + return new_signal + + +def drift_commit(session, plot, payload=None) -> None: + """Add the lazy corrected node to the tree and show it.""" + wiz = _wizard(session, plot) + if wiz is None: + emit_error("Drift Correction: nothing to apply") + return + wiz.commit() + + +def drift_correction(ctx, action_name: str = "Drift Correction", **params): + """Toolbar entry — a no-op parent; the Electron toolbar opens the staged + caret, which drives the ``drift_*`` handlers (README §4).""" + return None diff --git a/spyde/drawing/toolbars/icons/drift_correction.svg b/spyde/drawing/toolbars/icons/drift_correction.svg new file mode 100644 index 00000000..b4354a4b --- /dev/null +++ b/spyde/drawing/toolbars/icons/drift_correction.svg @@ -0,0 +1,44 @@ + + + + + + + + + + + + + + + + + + diff --git a/spyde/drift/__init__.py b/spyde/drift/__init__.py new file mode 100644 index 00000000..d477f46d --- /dev/null +++ b/spyde/drift/__init__.py @@ -0,0 +1,36 @@ +""" +spyde.drift — drift correction for image stacks and in-situ movies. + +See ``DRIFT_AND_PARTICLES_PLAN.md`` (repo root) for the full design. The two +load-bearing constraints, both from CLAUDE.md: + +* **Nothing materialises the stack.** The target is thousands of frames at + 2048²–4096² (tens of GB). Every solver here STREAMS: read a frame, transform + it, discard it. The output of a solve is a small :class:`DriftModel` — an + ``(N, 2)`` shift array, not an aligned copy of the movie. +* **The corrected movie is a LAZY VIEW.** ``DriftModel`` describes the + correction; applying it is a per-frame numpy operation + (:func:`spyde.drift.warp.shift_frame`) wired into the signal tree as a lazy + transformation node. Nothing is ever written out. + +Public API:: + + from spyde.drift import DriftModel, solve_translation, shift_frame + + model = solve_translation(signal, upsample=8, max_shift=32) + aligned_frame = shift_frame(raw_frame, model.shifts[i]) +""" +from __future__ import annotations + +from spyde.drift.frames import frame_source +from spyde.drift.model import DriftModel +from spyde.drift.translation import solve_translation +from spyde.drift.warp import coverage_mask, shift_frame + +__all__ = [ + "DriftModel", + "solve_translation", + "shift_frame", + "coverage_mask", + "frame_source", +] diff --git a/spyde/drift/frames.py b/spyde/drift/frames.py new file mode 100644 index 00000000..23602b46 --- /dev/null +++ b/spyde/drift/frames.py @@ -0,0 +1,92 @@ +""" +frames.py — one streaming accessor for every kind of frame stack we accept. + +The whole point of this module is the Memory-Safety rule (CLAUDE.md): a drift +solve on a 3000 × 4096² movie must never hold more than a couple of frames at +once. So callers get a ``(n_frames, get_frame, frame_shape)`` triple and read +frames one at a time — they never touch ``.data`` directly and therefore cannot +accidentally ``.compute()`` the whole array. + +Accepted inputs, in the order they are tried: + +* a HyperSpy signal (lazy or eager) with 1-D navigation and 2-D signal axes +* a dask array of shape ``(n, h, w)`` — sliced and computed per frame +* a numpy array of shape ``(n, h, w)`` — already resident, just indexed +* any sequence of 2-D arrays +""" +from __future__ import annotations + +from typing import Callable + +import numpy as np + + +def _is_dask(obj) -> bool: + """True for a dask array, WITHOUT importing dask when it isn't already in.""" + return type(obj).__module__.startswith("dask.array") + + +def frame_source(data) -> tuple[int, Callable[[int], np.ndarray], tuple[int, int]]: + """Return ``(n_frames, get_frame, frame_shape)`` for *data*. + + ``get_frame(i)`` returns frame ``i`` as a **numpy** 2-D array. For a lazy + (dask) backing it computes exactly that one frame — never the whole stack. + + Raises + ------ + TypeError + If *data* is not a recognised stack, or is not 3-D / not a sequence of + 2-D frames. Failing loudly here is deliberate: a silently-wrong axis + order would produce a plausible but meaningless drift curve. + """ + # ── HyperSpy signal ────────────────────────────────────────────────────── + # Duck-typed on axes_manager so this module never imports hyperspy (import + # cost at backend startup) and so test doubles work. + if hasattr(data, "axes_manager") and hasattr(data, "data"): + am = data.axes_manager + nav = int(am.navigation_dimension) + sig = int(am.signal_dimension) + if sig != 2: + raise TypeError( + f"drift needs 2-D signal axes (images); got signal_dimension={sig}" + ) + if nav != 1: + raise TypeError( + "drift needs a 1-D navigation axis (a frame stack / movie); got " + f"navigation_dimension={nav}. Reduce a higher-dimensional " + "dataset to a movie first (e.g. a virtual image)." + ) + return frame_source(data.data) + + # ── dask / numpy 3-D array ─────────────────────────────────────────────── + if _is_dask(data) or isinstance(data, np.ndarray): + if data.ndim != 3: + raise TypeError(f"expected a 3-D (n, h, w) stack; got shape {data.shape}") + n, h, w = data.shape + if _is_dask(data): + def get_frame(i: int, _d=data) -> np.ndarray: + # ONE frame. Never `_d.compute()`. + return np.asarray(_d[int(i)].compute()) + else: + def get_frame(i: int, _d=data) -> np.ndarray: + return np.asarray(_d[int(i)]) + return int(n), get_frame, (int(h), int(w)) + + # ── sequence of 2-D frames ─────────────────────────────────────────────── + try: + n = len(data) + except TypeError as exc: + raise TypeError( + f"cannot read frames from {type(data).__name__}: expected a HyperSpy " + "signal, a 3-D array, or a sequence of 2-D arrays" + ) from exc + if n == 0: + raise TypeError("empty frame stack") + first = np.asarray(data[0]) + if first.ndim != 2: + raise TypeError(f"sequence elements must be 2-D frames; got ndim={first.ndim}") + + def get_frame(i: int, _d=data) -> np.ndarray: + return np.asarray(_d[int(i)]) + + return int(n), get_frame, (int(first.shape[0]), int(first.shape[1])) diff --git a/spyde/drift/model.py b/spyde/drift/model.py new file mode 100644 index 00000000..8e23cd45 --- /dev/null +++ b/spyde/drift/model.py @@ -0,0 +1,165 @@ +""" +model.py — :class:`DriftModel`, the small serialisable result of a drift solve. + +This is the ONLY thing a solve produces. It is deliberately tiny (an ``(N, 2)`` +float32 array plus metadata) because the alternative — an aligned copy of the +movie — is tens of GB (see the module docstring in ``spyde/drift/__init__.py``). + +Sign convention — read this before touching anything +----------------------------------------------------- +``shifts[i]`` is the **correction**: the ``(dy, dx)`` you ADD to frame *i* to +bring it into the reference frame. So:: + + aligned_i = scipy.ndimage.shift(frame_i, model.shifts[i]) + +This matches ``skimage.registration.phase_cross_correlation``, whose docstring +defines its return as "the shift vector required to register moving_image with +reference_image", and matches ``scipy.ndimage.shift``'s own sign. Keeping all +three identical is why the convention is stated here rather than inferred at each +call site — an inverted sign produces a drift curve that looks entirely plausible +and doubles the drift instead of removing it. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +# Bumped when the on-disk layout changes incompatibly. +FORMAT_VERSION = 1 + + +@dataclass +class DriftModel: + """The correction for a frame stack. + + Parameters + ---------- + shifts + ``(N, 2)`` float32, ``(dy, dx)`` per frame — the correction to ADD (see + the module docstring on sign convention). + kind + ``"rigid"``. Reserved for future model families (affine, non-rigid), + which will keep ``shifts`` as their rigid component. + reference + How the reference was formed: ``"running"`` (running Fourier average), + ``"sequential"`` (frame-to-frame, cumulative) or ``"fixed:"``. + residuals + Optional ``(N,)`` per-frame correlation peak sharpness — a weak but free + quality signal. NaN where not computed. + params + The solver arguments, for provenance and for re-running. + """ + + shifts: np.ndarray + kind: str = "rigid" + reference: str = "running" + residuals: np.ndarray | None = None + params: dict[str, Any] = field(default_factory=dict) + provenance: dict[str, Any] | None = None + + def __post_init__(self) -> None: + self.shifts = np.ascontiguousarray(self.shifts, dtype=np.float32) + if self.shifts.ndim != 2 or self.shifts.shape[1] != 2: + raise ValueError( + f"shifts must be (N, 2); got {self.shifts.shape}" + ) + if self.residuals is not None: + self.residuals = np.ascontiguousarray(self.residuals, dtype=np.float32) + if self.residuals.shape != (self.n_frames,): + raise ValueError( + f"residuals must be ({self.n_frames},); got {self.residuals.shape}" + ) + + # ── basic properties ───────────────────────────────────────────────────── + + @property + def n_frames(self) -> int: + return int(self.shifts.shape[0]) + + @property + def max_abs_shift(self) -> float: + """Largest single-axis correction, in pixels. Sizes the padded border.""" + if self.n_frames == 0: + return 0.0 + return float(np.nanmax(np.abs(self.shifts))) + + @property + def is_integer(self) -> bool: + """True when every shift is a whole number of pixels. + + An integer-only model can be applied by ``np.roll`` and so **preserves + the source dtype**; a sub-pixel model needs interpolation and therefore + float output. Callers use this to avoid promoting a uint16 movie to + float32 when they don't have to. + """ + if self.n_frames == 0: + return True + finite = self.shifts[np.isfinite(self.shifts)] + return bool(finite.size == 0 or np.all(finite == np.round(finite))) + + # ── frame-of-reference conversions ─────────────────────────────────────── + + def to_sample_frame(self, positions: np.ndarray, frame_index) -> np.ndarray: + """Map lab-frame ``(y, x)`` positions onto the corrected (sample) frame. + + *positions* is ``(M, 2)``; *frame_index* is a scalar or ``(M,)``. This is + how a trajectory measured on RAW frames is reported as if the stage had + never moved — the alternative to physically warping the movie first. + """ + pos = np.asarray(positions, dtype=np.float64) + if pos.ndim != 2 or pos.shape[1] != 2: + raise ValueError(f"positions must be (M, 2); got {pos.shape}") + idx = np.asarray(frame_index, dtype=np.intp) + return pos + self.shifts[idx].reshape(pos.shape) + + def to_lab_frame(self, positions: np.ndarray, frame_index) -> np.ndarray: + """Inverse of :meth:`to_sample_frame`.""" + pos = np.asarray(positions, dtype=np.float64) + if pos.ndim != 2 or pos.shape[1] != 2: + raise ValueError(f"positions must be (M, 2); got {pos.shape}") + idx = np.asarray(frame_index, dtype=np.intp) + return pos - self.shifts[idx].reshape(pos.shape) + + # ── serialisation ──────────────────────────────────────────────────────── + + def save(self, path: str) -> None: + """Write to a compressed ``.npz``. Small enough to sit beside the data.""" + meta = { + "format_version": FORMAT_VERSION, + "kind": self.kind, + "reference": self.reference, + "params": self.params, + "provenance": self.provenance, + } + arrays = {"shifts": self.shifts, "meta": np.array(json.dumps(meta))} + if self.residuals is not None: + arrays["residuals"] = self.residuals + np.savez_compressed(path, **arrays) + + @classmethod + def load(cls, path: str) -> "DriftModel": + with np.load(path, allow_pickle=False) as z: + meta = json.loads(str(z["meta"].item())) + ver = meta.get("format_version") + if ver != FORMAT_VERSION: + raise ValueError( + f"unsupported DriftModel format version {ver!r} " + f"(this build reads {FORMAT_VERSION})" + ) + return cls( + shifts=z["shifts"], + kind=meta.get("kind", "rigid"), + reference=meta.get("reference", "running"), + residuals=z["residuals"] if "residuals" in z.files else None, + params=meta.get("params") or {}, + provenance=meta.get("provenance"), + ) + + def __repr__(self) -> str: + return ( + f"DriftModel(kind={self.kind!r}, n_frames={self.n_frames}, " + f"reference={self.reference!r}, max_abs_shift={self.max_abs_shift:.2f} px)" + ) diff --git a/spyde/drift/translation.py b/spyde/drift/translation.py new file mode 100644 index 00000000..e0be84fa --- /dev/null +++ b/spyde/drift/translation.py @@ -0,0 +1,699 @@ +""" +translation.py — rigid (translation-only) drift solve. Plan step A1. + +Algorithm: FFT phase correlation with a **running Fourier average** reference and +Guizar-Sicairos matrix-multiply DFT upsampling for the sub-pixel peak. + +Three things here are deliberate and worth reading before changing them. + +**1. It streams.** One frame is resident at a time (plus the accumulated reference +FFT, which is frame-sized). A 3000 × 4096² movie is tens of GB; nothing here ever +holds more than a few hundred MB. This is the CLAUDE.md Memory-Safety rule. + +**2. The reference is accumulated in FOURIER space, aligned by a phase ramp.** +To add frame *i* to the running average *already aligned*, we multiply its FFT by +``exp(-2πi(dy·fy + dx·fx))`` rather than resampling the frame and re-transforming. +A translation is exactly a phase ramp in the Fourier domain, so this is not an +approximation — it is free, it is exact even for sub-pixel shifts, and it avoids +the interpolation blur that resample-then-average would accumulate over thousands +of frames. That blur is the reason a naive running average degrades as the stack +gets longer. + +**3. Sub-pixel refinement is a small matmul, not a padded inverse FFT.** Zero- +padding the cross-correlation to get ``1/upsample`` resolution costs an FFT of +``(H·u, W·u)`` — for u=8 on a 4096² frame that is a 32768² transform. The +matrix-multiply DFT evaluates the correlation only on the ~12×12 window around +the coarse peak, which is what the refinement actually needs. + +The torch and numpy paths run the *same* algorithm through a small operator +adapter, so ``test_drift_translation.py`` can assert they agree bit-closely — that +parity test is what protects the GPU path, since it has no independent reference. +""" +from __future__ import annotations + +import logging +import math +from typing import Any, Callable + +import numpy as np + +from spyde.drift.frames import frame_source +from spyde.drift.model import DriftModel + +log = logging.getLogger(__name__) + +# Guizar-Sicairos: the refinement window spans 1.5 upsampled pixels either side of +# the coarse peak. Matches skimage's `phase_cross_correlation` so the two agree. +_UPSAMPLED_REGION_FACTOR = 1.5 + +# Magnitude FLOOR for phase normalisation — NOT an additive epsilon. +# +# Pure phase correlation divides the cross-power spectrum by its own magnitude, +# which is only meaningful where there is signal. Bins whose magnitude is +# numerically zero must be left alone; dividing them by a tiny epsilon amplifies +# rounding noise to UNIT magnitude, and since there are far more empty bins than +# populated ones, that noise then dominates the inverse transform. +# +# This is not theoretical. With an additive `1e-12` the solver recovered the +# synthetic particle movie's drift to 25 px (worse than not correcting at all) +# as soon as apodisation was enabled — because windowing concentrates spectral +# energy and pushes many more bins down into the numerical floor. With the floor +# below it: 0.06 px. Matches skimage's `100 * finfo(float32).eps`. +_PHASE_FLOOR = 100.0 * float(np.finfo(np.float32).eps) # ~1.19e-5 + +# A frame whose correlation peak is weaker than this fraction of the running +# MEDIAN peak is kept out of the accumulated reference. +# +# The running-average reference exists to be robust to one bad frame, but folding +# every frame in unconditionally does the opposite: a dropped / blanked / saturated +# frame has a broadband spectrum, so after phase normalisation it contributes as +# much to the reference as a good frame and drags every subsequent registration +# with it. Measured on a 5-frame stack with one frame replaced by pure noise, the +# two frames AFTER the bad one came back ~3.9 px wrong; with this rejection they +# are correct and only the bad frame itself is wrong. +# +# Both constants come from measurement, not taste. Peak strength relative to the +# running median, measured across four stacks: +# +# worst NATURAL frame (clean sub-pixel stack, erratic peaks) 0.388 +# a frame replaced by pure noise 0.007 +# +# So there is a ~50x gap to put a threshold in, and 0.25 sits inside it with margin +# both ways. 0.5 was tried first and produced a FALSE rejection on the clean +# sub-pixel stack — which is why this is not simply "half". +# +# _REJECT_MIN_SAMPLES is 1, not 3, and that is deliberate: a short stack cannot +# afford a warm-up. On a 5-frame stack the bad frame arrives before three good ones +# have been seen, so a 3-sample warm-up let it into the reference and the rule never +# fired (measured: 0 rejections, and the two frames after it came back 3.9 px wrong). +# With a 1-sample warm-up those two frames are recovered EXACTLY. +# +# The asymmetry justifies being aggressive: keeping a good frame OUT of the +# reference only slows the averaging, while letting a bad frame IN corrupts every +# registration after it. A rejected frame still gets its own shift reported. +# +# Windowing the median over the last N accepted frames was tried and made NO +# difference at N=3, 5 or unbounded — the natural decay in peak strength as the +# reference averages more frames is not steep enough to matter. Don't add it back. +_REJECT_FRACTION = 0.25 +_REJECT_MIN_SAMPLES = 1 + +# Smallest alignment ROI worth correlating. Below roughly this the upsampled +# refinement window (1.5 x upsample, so 12 px at the default) approaches the box +# itself and the peak has nowhere to sit. +_MIN_ROI = 16 + + +# ── operator adapters ──────────────────────────────────────────────────────── +# The algorithm below is written once against this interface. `_TorchOps` is the +# production path; `_NumpyOps` is the reference the parity test pins it against. + +class _NumpyOps: + name = "numpy" + + def __init__(self, device=None): + self.device = None + + def to_backend(self, a): + return np.asarray(a, dtype=np.float32) + + def fft2(self, a): + return np.fft.fft2(a).astype(np.complex64) + + def ifft2(self, a): + return np.fft.ifft2(a) + + def conj(self, a): + return np.conj(a) + + def abs(self, a): + return np.abs(a) + + def clamp_min(self, a, floor): + return np.maximum(a, floor) + + def fftfreq(self, n): + return np.fft.fftfreq(n).astype(np.float32) + + def arange(self, n): + return np.arange(n, dtype=np.float32) + + def exp(self, a): + return np.exp(a) + + def masked_argmax(self, mag, mask): + m = np.where(mask, mag, -np.inf) + flat = int(np.argmax(m)) + return divmod(flat, mag.shape[1]) + + def argmax2d(self, mag): + flat = int(np.argmax(mag)) + return divmod(flat, mag.shape[1]) + + def tensordot_last(self, kernel, data): + """``kernel @ data`` contracting kernel's axis 1 with data's LAST axis.""" + return np.tensordot(kernel, data, axes=(1, -1)) + + def scalar(self, a) -> float: + return float(a) + + def mean_abs(self, a) -> float: + return float(np.mean(np.abs(a))) + + def to_numpy(self, a): + return np.asarray(a) + + +class _TorchOps: + name = "torch" + + def __init__(self, device): + import torch + self._torch = torch + self.device = device + + def to_backend(self, a): + t = self._torch + return t.as_tensor(np.ascontiguousarray(a, dtype=np.float32), device=self.device) + + def fft2(self, a): + return self._torch.fft.fft2(a).to(self._torch.complex64) + + def ifft2(self, a): + return self._torch.fft.ifft2(a) + + def conj(self, a): + return self._torch.conj(a) + + def abs(self, a): + return self._torch.abs(a) + + def clamp_min(self, a, floor): + return self._torch.clamp(a, min=float(floor)) + + def fftfreq(self, n): + return self._torch.fft.fftfreq(n, device=self.device, dtype=self._torch.float32) + + def arange(self, n): + return self._torch.arange(n, device=self.device, dtype=self._torch.float32) + + def exp(self, a): + return self._torch.exp(a) + + def masked_argmax(self, mag, mask): + t = self._torch + m = mag.masked_fill(~mask, float("-inf")) + flat = int(t.argmax(m.reshape(-1)).item()) + return divmod(flat, mag.shape[1]) + + def argmax2d(self, mag): + t = self._torch + flat = int(t.argmax(mag.reshape(-1)).item()) + return divmod(flat, mag.shape[1]) + + def tensordot_last(self, kernel, data): + return self._torch.tensordot(kernel, data, dims=([1], [data.ndim - 1])) + + def scalar(self, a) -> float: + return float(a.item()) if hasattr(a, "item") else float(a) + + def mean_abs(self, a) -> float: + return float(self._torch.mean(self._torch.abs(a)).item()) + + def to_numpy(self, a): + return a.detach().cpu().numpy() + + +def _resolve_ops(device: str | None): + """Pick the backend: ``cuda`` > ``mps`` > **torch CPU** > numpy. + + **torch CPU beats numpy even with no GPU**, and by a lot — measured on this + box, 120 × 512² frames at upsample=8:: + + numpy 6.57 s 18 frames/s + torch cpu 0.86 s 139 frames/s (7.7x) + torch cuda 0.42 s 284 frames/s (16x) + + The reason is mundane: ``np.fft.fft2`` is single-threaded, ``torch.fft.fft2`` + uses every core. A per-frame FFT is the entire cost of this solver, so that + one difference is the whole gap. An earlier version of this function preferred + numpy on CPU-only machines on the assumption that torch's dispatch overhead + would dominate at one frame at a time; that assumption was wrong by 7.7x. + + numpy is kept as an **explicitly** selectable reference path (``device= + "numpy"``), which is what the backend-parity test pins the torch path against. + """ + if device == "numpy": + return _NumpyOps(None) + try: + import torch + except Exception: + return _NumpyOps(None) + + if device is None: + if torch.cuda.is_available(): + device = "cuda" + elif getattr(torch.backends, "mps", None) is not None and \ + torch.backends.mps.is_available(): + device = "mps" + else: + device = "cpu" + try: + return _TorchOps(torch.device(device)) + except Exception as exc: # pragma: no cover — bad device string + log.warning("[drift] torch device %r unusable (%s); using numpy", device, exc) + return _NumpyOps(None) + + +# ── windows and masks (built once per solve) ────────────────────────────────── + +#: Default Tukey taper fraction. 0.25 tapers the outer ~12.5% at each edge and +#: leaves the middle 75% at unit weight. See :func:`_taper2d` for why this is +#: NOT 1.0 (a full Hann window). +DEFAULT_TAPER_ALPHA = 0.25 + + +def _tukey1d(n: int, alpha: float) -> np.ndarray: + """Tukey (cosine-tapered) window. ``alpha=0`` rectangular, ``alpha=1`` Hann.""" + if n < 2: + return np.ones(max(1, n), dtype=np.float32) + alpha = float(min(1.0, max(0.0, alpha))) + if alpha <= 0.0: + return np.ones(n, dtype=np.float32) + x = np.arange(n, dtype=np.float64) / (n - 1) + w = np.ones(n, dtype=np.float64) + lo = x < alpha / 2.0 + hi = x > 1.0 - alpha / 2.0 + w[lo] = 0.5 * (1.0 + np.cos(2.0 * np.pi / alpha * (x[lo] - alpha / 2.0))) + w[hi] = 0.5 * (1.0 + np.cos(2.0 * np.pi / alpha * (x[hi] - 1.0 + alpha / 2.0))) + return w.astype(np.float32) + + +def _taper2d(ops, h: int, w: int, alpha: float): + """Separable Tukey taper, built in numpy once per solve then moved on-device. + + **Why a Tukey taper and NOT a full Hann window.** Some apodisation is needed: + a feature entering or leaving at the frame edge otherwise correlates against + the *border discontinuity* rather than the sample. But a full Hann window + (``alpha=1``) reweights the entire frame, and once the drift is large the two + frames have different content under the taper — which manufactures a spurious + correlation peak that can outrank the true one. + + Measured on the synthetic particle movie, frame 23 (true drift ``(6.0, 2.9)``): + with a full Hann window the strongest peak sits at ``(-19, 19)`` scoring 0.121 + while the TRUE peak scores only 0.088, so the solve returns ``(-19.2, 18.6)`` + — a 25 px error, worse than not correcting at all. **skimage's + ``phase_cross_correlation`` returns the same wrong answer on the same windowed + input**, so this is a property of full-frame windowing, not of either + implementation. Tapering only the outer edge leaves the interior comparable + between frames and recovers 0.06 px. + + Do not "simplify" this back to a Hann window. + """ + win = _tukey1d(h, alpha)[:, None] * _tukey1d(w, alpha)[None, :] + return ops.to_backend(win) + + +def _shift_mask(ops, h: int, w: int, max_shift: float | None, + min_shift: float | None): + """Which cross-correlation bins are admissible shifts. + + The correlation is un-shifted, so bin ``k`` means shift ``k`` for + ``k <= n//2`` and ``k - n`` above that. Both bounds are separable, so this is + an outer product of two 1-D masks — built once and reused for every frame. + """ + def axis_mask(n: int): + k = np.arange(n) + s = np.where(k > n // 2, k - n, k).astype(np.float64) + ok = np.ones(n, dtype=bool) + if max_shift is not None: + ok &= np.abs(s) <= float(max_shift) + return ok, s + + oky, sy = axis_mask(h) + okx, sx = axis_mask(w) + mask = oky[:, None] & okx[None, :] + if min_shift is not None: + # Exclude the near-zero-shift core. Useful when the reference already + # contains this frame (a running average does), because the trivial + # self-correlation peak at the origin can then outrank the real one. + r = np.hypot(sy[:, None], sx[None, :]) + mask &= r >= float(min_shift) + if not mask.any(): + raise ValueError( + "max_shift/min_shift exclude every possible shift " + f"(max_shift={max_shift}, min_shift={min_shift}, frame={h}x{w})" + ) + if ops.name == "torch": + return ops._torch.as_tensor(mask, device=ops.device) + return mask + + +# ── the correlation ────────────────────────────────────────────────────────── + +def _upsampled_dft(ops, data, region_size: int, upsample: float, offsets): + """Evaluate the inverse DFT of *data* on a small upsampled window. + + Mirrors ``skimage.registration._masked_phase_cross_correlation._upsampled_dft``: + one kernel matmul per axis, contracting the last axis each time. + """ + out = data + shape = tuple(data.shape) + for axis in (1, 0): # last axis first + n = shape[axis] + off = offsets[axis] + # NOTE the /upsample: this is `np.fft.fftfreq(n, upsample)` — frequencies + # scaled to the UPSAMPLED grid. Without it the window still evaluates and + # still finds a peak, but at 1/upsample of the intended resolution, so + # every recovered shift lands on a multiple of 1/upsample and the + # refinement silently does nothing. + freq = ops.fftfreq(n) / float(upsample) + kern = (ops.arange(region_size).reshape(region_size, 1) - float(off)) * \ + freq.reshape(1, n) + if ops.name == "torch": + kern = ops.exp(ops._torch.complex( + ops._torch.zeros_like(kern), -2.0 * math.pi * kern)) + else: + kern = np.exp(-2j * math.pi * kern).astype(np.complex64) + out = ops.tensordot_last(kern, out) + return out + + +def _peak_shift(ops, ref_fft, mov_fft, mask, upsample: float, + normalize: bool) -> tuple[float, float, float]: + """Return ``(dy, dx, sharpness)`` registering *mov* onto *ref*. + + Sign matches ``skimage.registration.phase_cross_correlation`` and + ``scipy.ndimage.shift``: the result is the correction to ADD to the moving + frame (see :mod:`spyde.drift.model`). + """ + product = ref_fft * ops.conj(mov_fft) + if normalize: + # Phase correlation proper: discard magnitude, keep only phase. Gives a + # far sharper peak than plain cross-correlation on images whose spectra + # are dominated by low frequencies, which every real micrograph is. + # The divisor is FLOORED, not offset — see _PHASE_FLOOR. + product = product / ops.clamp_min(ops.abs(product), _PHASE_FLOOR) + + cc = ops.ifft2(product) + mag = ops.abs(cc) + py, px = ops.masked_argmax(mag, mask) + + h, w = int(mag.shape[0]), int(mag.shape[1]) + peak = ops.scalar(mag[py, px]) + baseline = ops.mean_abs(cc) + sharpness = float(peak / baseline) if baseline > 0 else float("nan") + + dy = float(py - h) if py > h // 2 else float(py) + dx = float(px - w) if px > w // 2 else float(px) + + if upsample and upsample > 1: + u = float(upsample) + dy = round(dy * u) / u + dx = round(dx * u) / u + region = int(math.ceil(u * _UPSAMPLED_REGION_FACTOR)) + dftshift = float(region // 2) + offsets = (dftshift - dy * u, dftshift - dx * u) + fine = _upsampled_dft(ops, ops.conj(product), region, u, offsets) + fmag = ops.abs(fine) + my, mx = ops.argmax2d(fmag) + dy += (my - dftshift) / u + dx += (mx - dftshift) / u + + return dy, dx, sharpness + + +def _validate_roi(roi, full_h: int, full_w: int): + """Normalise and bounds-check an ``(y0, x0, h, w)`` alignment ROI. + + Rejects rather than clamps. A silently shrunk ROI would correlate on a + different region than the one the user dragged, and the drift curve would be + wrong in a way nothing on screen could explain. + """ + if roi is None: + return None + try: + y0, x0, h, w = (int(v) for v in roi) + except (TypeError, ValueError): + raise ValueError( + f"roi must be (y0, x0, h, w) in pixels; got {roi!r}") from None + if h < _MIN_ROI or w < _MIN_ROI: + raise ValueError( + f"roi is {h}x{w} px; the correlation needs at least " + f"{_MIN_ROI}x{_MIN_ROI} to locate a peak at all") + if y0 < 0 or x0 < 0 or y0 + h > full_h or x0 + w > full_w: + raise ValueError( + f"roi (y0={y0}, x0={x0}, h={h}, w={w}) falls outside the " + f"{full_h}x{full_w} frame") + return (y0, x0, h, w) + + +def _accept_into_reference(sharpness: float, accepted: list[float], + enabled: bool) -> bool: + """Whether this frame is credible enough to join the running reference. + + Always True until there are :data:`_REJECT_MIN_SAMPLES` accepted frames to + take a median over — with nothing to compare against, rejecting would just be + guessing. See :data:`_REJECT_FRACTION`. + """ + if not enabled or len(accepted) < _REJECT_MIN_SAMPLES: + return True + if not math.isfinite(sharpness): + return False + return sharpness >= _REJECT_FRACTION * float(np.median(accepted)) + + +def _phase_ramp(ops, h: int, w: int, dy: float, dx: float): + """FFT multiplier that translates a frame by ``(dy, dx)`` exactly. + + ``F{f(y - dy, x - dx)}(k) = exp(-2πi(dy·fy + dx·fx))·F{f}(k)``. This is why + the running reference needs no resampling — see the module docstring. + """ + fy = ops.fftfreq(h).reshape(h, 1) + fx = ops.fftfreq(w).reshape(1, w) + ph = dy * fy + dx * fx + if ops.name == "torch": + return ops.exp(ops._torch.complex( + ops._torch.zeros_like(ph), -2.0 * math.pi * ph)) + return np.exp(-2j * math.pi * ph).astype(np.complex64) + + +# ── public solve ───────────────────────────────────────────────────────────── + +def solve_translation( + data, + *, + upsample: int = 8, + max_shift: float | None = 32.0, + min_shift: float | None = None, + reference: str = "running", + roi: tuple[int, int, int, int] | None = None, + apodize: bool | float = True, + normalize: bool = True, + reject_outliers: bool = True, + device: str | None = None, + progress: Callable[[int, int], None] | None = None, + on_shift: Callable[[int, float, float, float], None] | None = None, + cancel: Callable[[], bool] | None = None, + provenance: dict[str, Any] | None = None, +) -> DriftModel: + """Solve rigid drift for a frame stack. Returns a :class:`DriftModel`. + + Parameters + ---------- + data + A HyperSpy signal (1-D nav, 2-D signal), a 3-D numpy/dask array, or a + sequence of 2-D frames. Read one frame at a time — never materialised. + upsample + Sub-pixel factor. ``8`` resolves to 1/8 px, which is well past the + ~0.05 px accuracy floor set by noise on real data. + max_shift + Reject correlation peaks implying a larger per-frame shift, in pixels. + Guards against a spurious peak from a periodic lattice — the failure mode + where a crystalline sample locks onto the wrong lattice translation and + the drift curve jumps by exactly one lattice spacing. + min_shift + Exclude peaks *smaller* than this. Off by default; see + :func:`_shift_mask`. + reference + ``"running"`` — running Fourier average (default, robust to one bad + frame); ``"sequential"`` — register to the previous frame and accumulate + (handles large excursions, accumulates error); ``"first"`` or + ``"fixed:"`` — one fixed reference frame. + roi + ``(y0, x0, h, w)`` in pixels — correlate on this sub-region only. The + returned shifts still apply to the WHOLE frame; a translation is a + translation regardless of which window you measured it in. + + This is not merely a speed switch, it is often the more CORRECT answer. + Whole-frame correlation averages over everything that moved, so on an + in-situ movie where the sample is genuinely evolving — particles growing, + drifting, appearing — the sample's own motion contaminates the estimate of + the stage's. Restricting to a static, feature-rich landmark (a support + film edge, a fiducial, a stationary grain) measures the stage and nothing + else. It is also how a user can rescue a dataset where the field of view + is mostly featureless. + + **The ROI is FIXED in frame coordinates**, so the landmark drifts within + it. That is fine while the drift is small compared with the box, and it is + why the box wants to be comfortably larger than the total excursion — + the caret's preview exists so this is judged by eye rather than guessed. + A box smaller than the drift will lose the landmark and the solve will + wander. + apodize + Edge taper before transforming. ``True`` uses a Tukey window with + ``alpha=DEFAULT_TAPER_ALPHA``; a float sets alpha explicitly + (``1.0`` = full Hann, which is a trap — see :func:`_taper2d`); + ``False`` disables it. + normalize + True phase correlation (unit-magnitude spectrum). Sharper peak. + reject_outliers + ``running`` mode only: keep a frame out of the accumulated reference when + its correlation peak is not credible (see :data:`_REJECT_FRACTION`). Its + own shift is still reported — only the reference is protected. This is + what makes "robust to one bad frame" true rather than aspirational. + device + ``None`` auto-selects CUDA/MPS then falls back to numpy; ``"numpy"`` + forces the reference path; or an explicit torch device string. + progress, cancel + ``progress(done, total)`` is called as frames complete. + ``cancel()`` returning True aborts; frames not yet reached keep NaN + shifts, so a cancelled solve is detectable rather than silently partial. + on_shift + ``on_shift(index, dy, dx, sharpness)`` per frame, as each is solved. + + This exists so a UI can draw the drift curve **while** it solves, which + ``progress`` cannot support: it carries only a count, and the shift array + is solver-local until the return. Splitting the solve into chunks and + concatenating would not be equivalent either — the running Fourier + reference accumulates across the whole stack, so a restarted solve gives a + different (worse) answer. A callback is the only way to stream the trace + without changing the result. + + Called on the solver thread, so a UI implementation must marshal. + + Notes + ----- + Frame 0 is the origin by definition and always gets ``(0, 0)``. + """ + if reference not in ("running", "sequential", "first") and \ + not reference.startswith("fixed:"): + raise ValueError( + f"unknown reference {reference!r}; expected 'running', 'sequential', " + "'first' or 'fixed:'" + ) + if upsample < 1: + raise ValueError(f"upsample must be >= 1; got {upsample}") + + n_frames, get_frame, (full_h, full_w) = frame_source(data) + crop = _validate_roi(roi, full_h, full_w) + h, w = (full_h, full_w) if crop is None else (crop[2], crop[3]) + ops = _resolve_ops(device) + + shifts = np.full((n_frames, 2), np.nan, dtype=np.float32) + sharp = np.full((n_frames,), np.nan, dtype=np.float32) + + from spyde.device_lock import accelerator_lock + + # MPS is not thread-safe and every torch user in the process shares ONE lock + # (CLAUDE.md § GPU Computing). A null context off MPS, so CUDA keeps its + # stream concurrency. Held across the solve rather than per-frame: the solve + # runs on a worker thread and per-frame acquire/release would be pure + # overhead at thousands of frames. + with accelerator_lock(ops.device): + alpha = (DEFAULT_TAPER_ALPHA if apodize is True + else (0.0 if apodize is False else float(apodize))) + window = _taper2d(ops, h, w, alpha) if alpha > 0 else None + mask = _shift_mask(ops, h, w, max_shift, min_shift) + + def frame_fft(i: int): + raw = get_frame(i) + if crop is not None: + y0, x0, ch, cw = crop + raw = raw[y0:y0 + ch, x0:x0 + cw] + f = ops.to_backend(raw) + if window is not None: + f = f * window + return ops.fft2(f) + + fixed_index = 0 + if reference.startswith("fixed:"): + fixed_index = int(reference.split(":", 1)[1]) + if not 0 <= fixed_index < n_frames: + raise ValueError( + f"fixed reference index {fixed_index} outside 0..{n_frames - 1}" + ) + + first = frame_fft(fixed_index if reference.startswith("fixed:") else 0) + shifts[0] = (0.0, 0.0) + sharp[0] = np.inf if n_frames else np.nan + if on_shift is not None and n_frames: + on_shift(0, 0.0, 0.0, float("inf")) + + ref_fft = first # running accumulator / fixed reference + ref_count = 1 + prev_fft = first # sequential mode + cumulative = np.zeros(2, dtype=np.float64) + accepted_sharp: list[float] = [] # peak strengths folded into the reference + rejected = 0 + + if progress is not None: + progress(1, n_frames) + + for i in range(1, n_frames): + if cancel is not None and cancel(): + log.info("[drift] cancelled at frame %d/%d", i, n_frames) + break + + mov = frame_fft(i) + + if reference == "sequential": + dy, dx, s = _peak_shift(ops, prev_fft, mov, mask, upsample, normalize) + cumulative += (dy, dx) + shifts[i] = cumulative + prev_fft = mov + else: + dy, dx, s = _peak_shift(ops, ref_fft, mov, mask, upsample, normalize) + shifts[i] = (dy, dx) + if reference == "running" and _accept_into_reference( + s, accepted_sharp, reject_outliers): + # Fold the ALIGNED frame in via a phase ramp — exact, and no + # resampling blur accumulates over the stack. + aligned = mov * _phase_ramp(ops, h, w, dy, dx) + ref_fft = (ref_fft * ref_count + aligned) / (ref_count + 1) + ref_count += 1 + accepted_sharp.append(s) + elif reference == "running": + rejected += 1 + log.debug("[drift] frame %d kept out of the reference " + "(peak %.2f vs median %.2f)", i, s, + float(np.median(accepted_sharp))) + sharp[i] = s + + if on_shift is not None: + on_shift(i, float(shifts[i, 0]), float(shifts[i, 1]), float(s)) + if progress is not None: + progress(i + 1, n_frames) + + params = { + "upsample": int(upsample), + "max_shift": None if max_shift is None else float(max_shift), + "min_shift": None if min_shift is None else float(min_shift), + "reference": reference, + "apodize": float(alpha), + "normalize": bool(normalize), + "reject_outliers": bool(reject_outliers), + "rejected_from_reference": int(rejected), + "backend": ops.name, + "n_frames": int(n_frames), + "frame_shape": [int(full_h), int(full_w)], + "roi": None if crop is None else [int(v) for v in crop], + } + return DriftModel( + shifts=shifts, + kind="rigid", + reference=reference, + residuals=sharp, + params=params, + provenance=provenance, + ) diff --git a/spyde/drift/warp.py b/spyde/drift/warp.py new file mode 100644 index 00000000..03a46d64 --- /dev/null +++ b/spyde/drift/warp.py @@ -0,0 +1,134 @@ +""" +warp.py — apply a drift correction to ONE frame. + +Per-frame by design. This is the function a lazy signal-tree node calls, so the +aligned movie is never materialised (``spyde/drift/__init__.py``), and it is also +what the derived-view reader would call if the per-frame shift transform is ever +added to ``array_cache/readers/per_frame.py`` — which is a signal-tree read-path +change and therefore gated on review, NOT done here. + +Edge policy (locked, plan §A7): the frame keeps its full size and uncovered +pixels become **NaN**, with :func:`coverage_mask` giving the validity map. +Nothing is cropped and nothing is filled with invented data. + +**Downstream contract:** segmentation MUST respect the coverage mask. A NaN-padded +border is the single most likely integration bug in this feature — a threshold +applied to NaN, or a NaN-to-zero conversion, invents a large "particle" along the +edge that then nucleates a spurious track. +""" +from __future__ import annotations + +import numpy as np + + +def _split_shift(shift) -> tuple[np.ndarray, bool]: + s = np.asarray(shift, dtype=np.float64).reshape(-1) + if s.size != 2: + raise ValueError(f"shift must be (dy, dx); got {np.shape(shift)}") + if not np.all(np.isfinite(s)): + raise ValueError(f"shift must be finite; got {s!r}") + return s, bool(np.all(s == np.round(s))) + + +def coverage_mask(shape: tuple[int, int], shift) -> np.ndarray: + """Boolean map of which output pixels come from real source data. + + A pixel is covered when its source coordinate falls inside the source frame. + For a sub-pixel shift the border row/column that would need to interpolate + against off-frame data is treated as **uncovered** — bilinear interpolation + there would silently blend in the fill value. + """ + h, w = int(shape[0]), int(shape[1]) + s, integral = _split_shift(shift) + dy, dx = s + + mask = np.zeros((h, w), dtype=bool) + if integral: + y0, y1 = max(0, int(dy)), min(h, h + int(dy)) + x0, x1 = max(0, int(dx)), min(w, w + int(dx)) + else: + # Output pixel y draws from source y - dy; it needs floor and floor+1, + # so require 0 <= y - dy and y - dy + 1 <= h - 1. + y0 = int(np.ceil(max(0.0, dy))) + y1 = int(np.floor(min(float(h), h + dy - 1.0))) + 1 + x0 = int(np.ceil(max(0.0, dx))) + x1 = int(np.floor(min(float(w), w + dx - 1.0))) + 1 + if y1 > y0 and x1 > x0: + mask[y0:y1, x0:x1] = True + return mask + + +def shift_frame( + frame: np.ndarray, + shift, + *, + order: int = 1, + fill: float = np.nan, + preserve_dtype: bool = False, +) -> np.ndarray: + """Shift *frame* by ``(dy, dx)``, padding uncovered pixels with *fill*. + + Parameters + ---------- + frame + 2-D source frame, any dtype. + shift + ``(dy, dx)`` correction — the value ADDED to coordinates. See + :mod:`spyde.drift.model` for the sign convention. + order + Interpolation order for a sub-pixel shift (1 = bilinear, the default; + 3 = cubic). Ignored for a whole-pixel shift, which is done exactly. + fill + Value for uncovered pixels. NaN by default, which forces a float result. + preserve_dtype + Keep the source dtype. Only honoured for a whole-pixel shift with a + non-NaN *fill* — a uint16 frame cannot hold NaN, and interpolation + cannot be exact in an integer type. Raises otherwise rather than + silently returning something lossy. + + Notes + ----- + A whole-pixel shift takes an exact slice-copy path: no interpolation, no + float promotion, and bit-identical to the source pixels. This matters because + the common case for a well-behaved stage IS an integer shift, and running it + through ``scipy.ndimage.shift`` would resample (and blur) data that did not + need to move sub-pixel at all. + """ + src = np.asarray(frame) + if src.ndim != 2: + raise ValueError(f"frame must be 2-D; got shape {src.shape}") + s, integral = _split_shift(shift) + dy, dx = s + h, w = src.shape + + if preserve_dtype and not (integral and np.isfinite(fill)): + raise ValueError( + "preserve_dtype=True requires a whole-pixel shift and a finite fill " + f"(got shift={s.tolist()}, fill={fill!r}); a sub-pixel shift must " + "interpolate and NaN padding cannot be stored in an integer dtype" + ) + + out_dtype = src.dtype if preserve_dtype else np.float32 + + if integral: + out = np.full((h, w), fill, dtype=out_dtype) + iy, ix = int(dy), int(dx) + # Destination window, and the matching source window. + dy0, dy1 = max(0, iy), min(h, h + iy) + dx0, dx1 = max(0, ix), min(w, w + ix) + if dy1 > dy0 and dx1 > dx0: + out[dy0:dy1, dx0:dx1] = src[dy0 - iy:dy1 - iy, dx0 - ix:dx1 - ix] + return out + + from scipy.ndimage import shift as ndi_shift + + # ndimage cannot propagate NaN through its spline filter without smearing it, + # so interpolate with a finite sentinel and stamp the fill on afterwards using + # the analytic coverage map. This keeps the padded border crisp instead of + # letting a NaN bleed `order` pixels into real data. + work = src.astype(np.float32, copy=False) + out = ndi_shift(work, s, order=order, mode="constant", cval=0.0, prefilter=order > 1) + out = out.astype(out_dtype, copy=False) + cov = coverage_mask((h, w), s) + out[~cov] = fill + return out From 5fa279ce40b536e84009740a9f86e99c461ceff2 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sat, 8 Aug 2026 23:59:47 -0500 Subject: [PATCH 3/6] feat(drift): wire the wizard -- registry stages, toolbar entry, particle-movie fixture - registry: the seven drift_* staged handlers + the 'drift' wizard-schema entry (DriftWizard.parameters is the one source of truth). - toolbars.yaml: the Drift Correction toggle, gated signal_types [insitu] (the rigid solver needs a 1-D navigation axis -- the same gate Play/Fast-Forward use). - data/synthetic.py: the particle_movie generator with stamped per-frame drift ground truth (also carries the particle event table the later segmentation PR asserts against -- the fixture is shared, its identity must not fork between the two PRs). - test harness: load_test_data_particles, lazy at one frame per chunk like a real .mrc movie. The seg-only 'noise' payload knob stays behind. --- spyde/actions/registry.py | 9 + spyde/backend/_session_actions.py | 3 + spyde/backend/_session_testharness.py | 47 ++++++ spyde/data/__init__.py | 3 +- spyde/data/synthetic.py | 235 ++++++++++++++++++++++++++ spyde/toolbars.yaml | 10 ++ 6 files changed, 306 insertions(+), 1 deletion(-) diff --git a/spyde/actions/registry.py b/spyde/actions/registry.py index 34d92f87..d212018a 100644 --- a/spyde/actions/registry.py +++ b/spyde/actions/registry.py @@ -105,6 +105,14 @@ "crop_open": "spyde.actions.base.crop_open", "crop_close": "spyde.actions.base.crop_close", "crop_set_region": "spyde.actions.base.crop_set_region", + # Drift Correction (spyde/actions/drift_action.py) — plan A8. + "drift_open": "spyde.actions.drift_action.drift_open", + "drift_close": "spyde.actions.drift_action.drift_close", + "drift_set_method": "spyde.actions.drift_action.drift_set_method", + "drift_tune": "spyde.actions.drift_action.drift_tune", + "drift_run": "spyde.actions.drift_action.drift_run", + "drift_discard": "spyde.actions.drift_action.drift_discard", + "drift_commit": "spyde.actions.drift_action.drift_commit", "download_cancel": "spyde.backend.example_download.download_cancel", "compute_configure": "spyde.backend.compute_config.compute_configure", "set_log_level": "spyde.backend.log_stream.set_log_level", @@ -241,6 +249,7 @@ def register_staged(name: str, dotted_path: str) -> None: "vom": ("spyde.actions.vector_orientation_om", "VomWizard"), "ebsd": ("spyde.actions.ebsd_action", "EbsdWizard"), "czb": ("spyde.actions.center_zero_beam", "PARAMETERS"), + "drift": ("spyde.actions.drift_action", "DriftWizard"), # YAML-declared (resolved from spyde.TOOLBAR_ACTIONS): "fv": ("__yaml__", "Find Diffraction Vectors"), "om": ("__yaml__", "Orientation Mapping"), diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py index c1012705..4a26a395 100644 --- a/spyde/backend/_session_actions.py +++ b/spyde/backend/_session_actions.py @@ -38,6 +38,7 @@ "load_test_data_si_grains", "load_test_data_sped_ag", "load_test_data_eels", "load_test_data_eds", "load_test_data_ebsd", "load_test_data_line", "load_test_data_movie", "load_test_data_5d", + "load_test_data_particles", "test_nav_drag", "test_region_scrub", "test_add_second_navigator", "load_test_vectors", "run_test_orientation", "dump_dask_state", @@ -148,6 +149,8 @@ def dispatch_action(self, msg: dict) -> None: self._load_test_data_movie(payload) elif action == "load_test_data_5d": self._load_test_data_5d(payload) + elif action == "load_test_data_particles": + self._load_test_data_particles(payload) elif action == "test_add_second_navigator": self._test_add_second_navigator() elif action == "test_nav_drag": diff --git a/spyde/backend/_session_testharness.py b/spyde/backend/_session_testharness.py index 3acce329..7d9afc20 100644 --- a/spyde/backend/_session_testharness.py +++ b/spyde/backend/_session_testharness.py @@ -439,6 +439,53 @@ def _load_test_data_movie(self, payload: dict | None = None) -> None: s.set_signal_type("insitu") # gates the Play/Fast Forward toolbar buttons self._add_signal(s, source_path="test_data_movie") + def _load_test_data_particles(self, payload: dict | None = None) -> None: + """The synthetic in-situ PARTICLE movie — the drift-correction fixture. + + Wraps :func:`spyde.data.synthetic.particle_movie`, whose ground truth + (per-frame drift, particle radii, and the nucleation / dissolution / + merge frames) is stamped into ``metadata.Spyde.synthetic`` — so a test + can assert against the numbers the data was built from instead of a + golden screenshot. + + Loaded **lazy at one frame per chunk**, like ``load_test_data_movie`` + and like a real ``.mrc`` in-situ movie: each nav move is then a small + cold read of just that frame, which is the path a user actually drags. + + Payload: ``{"frames": int, "size": [ny, nx], "eager": bool}``. + ``eager`` keeps it in RAM, which skips the whole array-cache path — + useful only for a test that wants to isolate something else (an eager + fixture is how ``si_grains`` silently made a cache benchmark measure + nothing). + """ + import dask.array as da + + from spyde.backend.heavy_imports import ensure_heavy_imports + from spyde.data.synthetic import particle_movie + + # BEFORE set_signal_type: racing the startup prewarm's pyxem import + # partially initialises the module and the cast then silently fails, + # taking every gated toolbar action with it. + ensure_heavy_imports() + + payload = payload or {} + n_frames = int(payload.get("frames", 24)) + shape = tuple(payload.get("size", (96, 112))) + + s = particle_movie(n_frames=n_frames, shape=shape) + if not payload.get("eager"): + eager = s.data + ny, nx = eager.shape[1:] + # `as_lazy()` carries the axes, the signal type AND the stamped + # ground truth across (metadata has no setter, so copying it by + # hand is not an option), then swap in the correctly-chunked array. + # Assigning `.data` rather than calling `.rechunk()` keeps this a + # graph construction over an array already in RAM — never a + # scheduler shuffle. + s = s.as_lazy() + s.data = da.from_array(eager, chunks=(1, ny, nx)) + self._add_signal(s, source_path="test_data_particles") + @staticmethod def _write_movie_mrc(frames) -> str: """Write ``frames`` as a minimal MRC2014 stack (mode 6 = uint16) into the diff --git a/spyde/data/__init__.py b/spyde/data/__init__.py index 44896f38..6872664b 100644 --- a/spyde/data/__init__.py +++ b/spyde/data/__init__.py @@ -26,7 +26,8 @@ eds_si, eels_si, ground_truth, + particle_movie, ) __all__ = ["eels_si", "eds_si", "ebsd_patterns", "atom_lattice", - "ground_truth"] + "particle_movie", "ground_truth"] diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py index 8d5315e7..d813a18b 100644 --- a/spyde/data/synthetic.py +++ b/spyde/data/synthetic.py @@ -481,3 +481,238 @@ def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55), pc=np.asarray(pc, float), n_bands=int(len(normals)), grain2_mask=grain2) return s + + +# --------------------------------------------------------------------------- +# In-situ particle movie +# --------------------------------------------------------------------------- + +# The particle table. Hand-written rather than randomised so every number in +# the ground truth is exact and readable: a test that says "particle 4 +# dissolves at frame 16" can be checked against this table by eye. +# +# Columns: y0, x0, radius, amp, birth, death, vy, vx, faint +# y0/x0 position at t=0 in the SAMPLE frame, pixels +# birth first frame the particle exists (0 = present from the start) +# death first frame it is GONE (-1 = never dissolves) +# vy/vx sample-frame velocity, px/frame +# faint amplitude is scaled down to `faint_amplitude` -- the Section 0.9 probes +_PARTICLES: tuple[dict, ...] = ( + # two bright anchors, static and persistent + dict(y0=24.0, x0=22.0, radius=7.0, amp=1.00, birth=0, death=-1, vy=0.0, vx=0.0, faint=False), + dict(y0=70.0, x0=30.0, radius=9.0, amp=0.90, birth=0, death=-1, vy=0.0, vx=0.0, faint=False), + # a mover, for trails and displacement + dict(y0=30.0, x0=80.0, radius=6.0, amp=0.85, birth=0, death=-1, vy=1.10, vx=-0.70, faint=False), + # nucleation + dict(y0=60.0, x0=94.0, radius=5.0, amp=0.80, birth=8, death=-1, vy=0.0, vx=0.0, faint=False), + # dissolution + dict(y0=16.0, x0=58.0, radius=6.0, amp=0.75, birth=0, death=16, vy=0.0, vx=0.0, faint=False), + # the merge pair: converge along x until they overlap + dict(y0=84.0, x0=60.0, radius=5.0, amp=0.80, birth=0, death=-1, vy=0.0, vx=0.45, faint=False), + dict(y0=84.0, x0=82.0, radius=5.0, amp=0.80, birth=0, death=-1, vy=0.0, vx=-0.45, faint=False), + # faint, low-contrast -- these are what plan Section 0.9 is about + dict(y0=46.0, x0=102.0, radius=4.0, amp=1.00, birth=0, death=-1, vy=0.0, vx=0.0, faint=True), + dict(y0=88.0, x0=14.0, radius=3.0, amp=0.85, birth=0, death=-1, vy=0.0, vx=0.0, faint=True), +) + +#: The two particles that merge, plus the nucleating and dissolving ones. +MERGE_PAIR = (5, 6) +NUCLEATION_INDEX = 3 +DISSOLUTION_INDEX = 4 + + +def _drift_curve(n_frames: int, amplitude: float) -> np.ndarray: + """``(n_frames, 2)`` per-frame CORRECTION, matching ``DriftModel``'s sign. + + ``drift[t]`` is what you ADD to frame *t* to align it, so the sample appears + displaced by ``-drift[t]`` in the raw frame. + + The two axes get deliberately DIFFERENT shapes -- y grows monotonically with + a slight curve, x swings negative and comes back. A swapped or negated axis + therefore shows up as a wrong-shaped curve rather than merely a wrong + number, which is the whole point of an asymmetric fixture. + """ + if n_frames < 2: + return np.zeros((max(1, n_frames), 2), np.float64) + t = np.arange(n_frames, dtype=np.float64) / (n_frames - 1) + dy = amplitude * t ** 1.3 + dx = -0.6 * amplitude * np.sin(1.7 * np.pi * t) + return np.stack([dy, dx], axis=-1) + + +def _soft_disc(yy, xx, cy, cx, radius, edge=0.9): + """A disc with a smooth ~1 px edge, so sub-pixel centroids are meaningful. + + A hard-edged disc quantises its own centroid to the pixel grid, which would + make any sub-pixel tracking assertion against this fixture untestable. + """ + r = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2) + return 0.5 * (1.0 - np.tanh((r - radius) / edge)) + + +def _particle_arrays(drift: np.ndarray, faint_amplitude: float) -> dict: + """The per-particle arrays, in ONE place so the renderer and the stamped + ground truth cannot disagree about the motion model.""" + return dict( + drift=drift, + p_y0=np.array([p["y0"] for p in _PARTICLES]), + p_x0=np.array([p["x0"] for p in _PARTICLES]), + p_radius=np.array([p["radius"] for p in _PARTICLES]), + p_amp=np.array([(faint_amplitude * p["amp"]) if p["faint"] else p["amp"] + for p in _PARTICLES]), + p_birth=np.array([p["birth"] for p in _PARTICLES]), + p_death=np.array([p["death"] for p in _PARTICLES]), + p_vy=np.array([p["vy"] for p in _PARTICLES]), + p_vx=np.array([p["vx"] for p in _PARTICLES]), + p_faint=np.array([p["faint"] for p in _PARTICLES]), + n_particles=len(_PARTICLES), + ) + + +def particle_truth_at(truth: dict, t: int): + """Expected ``(positions, radii, present)`` at frame *t*. + + ``positions`` is ``(N, 2)`` ``(y, x)`` in **pixels, in the LAB frame** -- + what a segmenter looking at the RAW frame should find, drift included. + ``present`` is a boolean mask of which particles exist in that frame. + + Every consumer computes its expectations through here rather than + re-deriving the motion model, so a test cannot pass by repeating the same + mistake the generator made. + """ + t = int(t) + y0 = np.asarray(truth["p_y0"], float) + x0 = np.asarray(truth["p_x0"], float) + vy = np.asarray(truth["p_vy"], float) + vx = np.asarray(truth["p_vx"], float) + birth = np.asarray(truth["p_birth"], int) + death = np.asarray(truth["p_death"], int) + drift = np.asarray(truth["drift"], float) + + sample = np.stack([y0 + vy * t, x0 + vx * t], axis=-1) + lab = sample - drift[t] # see _drift_curve on the sign + present = (t >= birth) & ((death < 0) | (t < death)) + return lab, np.asarray(truth["p_radius"], float), present + + +def _merge_frame(arrays: dict, n_frames: int) -> int: + """First frame where the merge pair's discs overlap. + + Derived from the same motion model that draws them rather than hard-coded, + so the stamped truth stays correct if the table or the frame count changes. + Returns -1 if they never touch within the movie. + """ + a, b = MERGE_PAIR + ra, rb = _PARTICLES[a]["radius"], _PARTICLES[b]["radius"] + for t in range(n_frames): + pos, _, present = particle_truth_at(arrays, t) + if present[a] and present[b] and np.hypot(*(pos[a] - pos[b])) <= ra + rb: + return t + return -1 + + +def particle_movie(n_frames: int = 24, shape=(96, 112), *, + drift_amplitude: float = 6.0, + faint_amplitude: float = 0.11, + scale: float = 0.5, seed: int = 0, noise: float = 0.015): + """An in-situ particle movie whose every event is known exactly. + + A 1-D navigation (time) axis over 2-D images: nine particles on a drifting + support film, with one nucleation, one dissolution, one merge, one mover and + two deliberately faint low-contrast probes. + + Everything a downstream test needs to assert is stamped as ground truth -- + read it with :func:`ground_truth` and evaluate the motion model with + :func:`particle_truth_at`. + + Parameters + ---------- + n_frames + Number of time points. The nucleation (frame 8) and dissolution + (frame 16) frames are fixed, so keep this above ~20 for both to occur. + shape + ``(ny, nx)``. **Deliberately non-square** so a transposed frame is + obvious at a glance. + drift_amplitude + Peak per-frame drift correction, pixels. The support film drifts + rigidly; particles move relative to it. + faint_amplitude + Peak amplitude of the two faint probes. The default puts them around + 7x the noise sigma -- findable, but not by a threshold tuned for the + bright particles. This is the plan's Section 0.9 gate. + scale + Pixel size in nm, written onto both signal axes. + seed + Everything random here (film speckle, noise) comes from this. + noise + Gaussian sigma added last. 0 disables. + + Notes + ----- + **Why there is a speckled support film.** Drift is only recoverable if + something static dominates the correlation. Particles move, appear and + vanish, so they cannot serve; a smooth gradient has too little + high-frequency content to correlate sharply. A rigidly-drifting speckle + field is both physically right and what lets + :func:`spyde.drift.solve_translation` recover ``drift`` from this movie. + + The film is generated once on a padded canvas and sampled per frame with + bilinear interpolation, so there are no edge artifacts at any drift. + Particles are drawn **analytically** at their lab positions, so their + centroids stay exact regardless of how the film was resampled. + """ + import hyperspy.api as hs + from scipy.ndimage import gaussian_filter, map_coordinates + + ny, nx = int(shape[0]), int(shape[1]) + n_frames = int(n_frames) + rng = np.random.default_rng(seed) + drift = _drift_curve(n_frames, float(drift_amplitude)) + arrays = _particle_arrays(drift, float(faint_amplitude)) + + # The support film, on a canvas padded to cover every drift. + pad = int(np.ceil(np.abs(drift).max())) + 3 + fy, fx = ny + 2 * pad, nx + 2 * pad + yy_p, _ = np.mgrid[0:fy, 0:fx] + film = 0.10 + 0.12 * (yy_p / fy) # ramp along y ONLY (asymmetric) + film += 0.09 * gaussian_filter(rng.standard_normal((fy, fx)), 1.6) + + yy, xx = np.mgrid[0:ny, 0:nx].astype(np.float64) + frames = np.empty((n_frames, ny, nx), dtype=np.float32) + + for t in range(n_frames): + dy, dx = drift[t] + # Sample the film at the lab-frame offset: content sits at -drift. + frame = map_coordinates(film, [yy + pad - dy, xx + pad - dx], + order=1, mode="nearest") + pos, radii, present = particle_truth_at(arrays, t) + for i, p in enumerate(_PARTICLES): + if not present[i]: + continue + amp = (faint_amplitude * p["amp"]) if p["faint"] else p["amp"] + frame = frame + amp * _soft_disc(yy, xx, pos[i, 0], pos[i, 1], radii[i]) + frames[t] = frame + + if noise: + frames += (noise * rng.standard_normal(frames.shape)).astype(np.float32) + + s = hs.signals.Signal2D(frames) + for ax in s.axes_manager.signal_axes: + ax.scale, ax.units = float(scale), "nm" + tax = s.axes_manager.navigation_axes[0] + tax.name, tax.units, tax.scale = "time", "s", 0.05 + s.metadata.General.title = "Synthetic particle movie" + # `insitu` is registered by SpyDE's OWN hyperspy extension, so unlike the + # EELS/EBSD generators this cannot silently fail for a missing extra. + _try_signal_type(s, "insitu") + _stamp(s, kind="particle_movie", **arrays, + n_frames=n_frames, frame_shape=np.asarray([ny, nx]), + scale=float(scale), noise=float(noise), + faint_amplitude=float(faint_amplitude), + merge_pair=np.asarray(MERGE_PAIR), + nucleation_index=NUCLEATION_INDEX, + dissolution_index=DISSOLUTION_INDEX, + nucleation_frame=int(_PARTICLES[NUCLEATION_INDEX]["birth"]), + dissolution_frame=int(_PARTICLES[DISSOLUTION_INDEX]["death"]), + merge_frame=_merge_frame(arrays, n_frames)) + return s diff --git a/spyde/toolbars.yaml b/spyde/toolbars.yaml index 2794dfe0..b6276ea6 100644 --- a/spyde/toolbars.yaml +++ b/spyde/toolbars.yaml @@ -364,6 +364,16 @@ functions: toolbar_side: bottom navigation: False + Drift Correction: + description: Solve and remove sample drift across an image stack. Check the result in a separate window — an aligned stack sums sharp, a misaligned one blurs — then Apply to add a lazy corrected node (nothing is copied). + icon: drawing/toolbars/icons/drift_correction.svg + function: spyde.actions.drift_action.drift_correction + signal_types: [insitu] + plot_dim: [2] + toolbar_side: bottom + navigation: False + toggle: True + Fit: description: Fit a model to every spectrum. Add components line by line, watch the model curve on the spectrum, then Run to fit the whole scan on the GPU. Commit turns each component's integrated area into a map. icon: drawing/toolbars/icons/fit.svg From 8ed53a55721211baa044883209f388f978be2313 Mon Sep 17 00:00:00 2001 From: Carter Francis Date: Sun, 9 Aug 2026 00:04:08 -0500 Subject: [PATCH 4/6] test(drift): e2e -- the caret, discovery preview, solve and apply, on real pixels drift_wizard.spec.ts pulled from the quarry and pruned of its non-rigid stages (the 'nonrigid selectable/field controls' and 'nonrigid solve runs the fit' tests). DriftWizard.tsx pruned to the rigid tabs -- METHODS, DEFAULTS and the params() payload mirror drift_action 1:1. Wiring: FloatingToolbar mounts the caret for the Drift Correction toggle; SpyDEContext re-broadcasts the five drift_* messages as CustomEvents; protocol.ts types them. --- .../renderer/src/components/DriftWizard.tsx | 395 ++++++++++++++++++ .../src/components/FloatingToolbar.tsx | 8 + .../src/renderer/src/kernel/SpyDEContext.tsx | 10 + electron/src/renderer/src/kernel/protocol.ts | 81 ++++ electron/tests/drift_wizard.spec.ts | 161 +++++++ 5 files changed, 655 insertions(+) create mode 100644 electron/src/renderer/src/components/DriftWizard.tsx create mode 100644 electron/tests/drift_wizard.spec.ts diff --git a/electron/src/renderer/src/components/DriftWizard.tsx b/electron/src/renderer/src/components/DriftWizard.tsx new file mode 100644 index 00000000..62bc15d6 --- /dev/null +++ b/electron/src/renderer/src/components/DriftWizard.tsx @@ -0,0 +1,395 @@ +/** + * DriftWizard.tsx — the Drift Correction caret (`drift_` staged actions, + * backend: spyde/actions/drift_action.py; plan §A8 + §0.9a). + * + * **Two toggles and a button.** The first version of this caret had thirteen + * controls on its face — three model tabs, four numeric fields, three + * checkboxes, Solve/Apply/Cancel — and the review was "way too complicated. + * Too many options. Information overload." Plan §0.9a is the rule that came out + * of it: the default face carries the TASK, not the algorithm. Reference mode, + * sub-pixel factor, max shift, interpolation order and the model tabs all still + * exist, all still reach the backend, and all still land in provenance — they + * live behind the collapsed `Advanced` disclosure, because drift's parameters + * have one right answer we already know. + * + * **What the caret does NOT show.** The dy/dx curve used to be a 40 px inline + * SVG here; it is now its own figure window (`Drift dy/dx`), opened by + * `drift_run` and filled progressively from the solver's `on_shift` stream. A + * sparkline could show that the stage crept 30 px; only a real plot shows WHICH + * frame jumped. The before/after sums stay in the `Drift Check` window, whose + * bottom row is the discovery pair. + * + * **Discovery, not configuration.** The backend puts a draggable box on the + * movie the moment this mounts, aligns ~20 frames sampled across the whole + * movie on that box alone, and reports how much sharper the sum got. That + * number (`drift_preview.gain`) is what the readout under the toggles shows — + * drag the box onto a landmark and watch it rise, drag it onto empty film and + * watch it fall below 1. `Use ROI for alignment` is then the commitment: the + * full solve correlates on that same rectangle. It is OFF by default because + * a guessed box is not automatically better than the whole frame (measured: + * 1.03 px vs 0.25 px against ground truth on the test movie) — the preview is + * how you find out whether yours is. + * + * Only `rigid` has a solver. `rigid+affine` is shown LOCKED inside Advanced + * with the backend's own reason rather than silently falling back — a rigid + * solve under a caret claiming "rigid+affine" puts a wrong `kind` into the + * model's provenance, which is worse than the missing feature. + */ +import React from 'react' +import { WizardShell, TabRow, Field, NumInput, Select, Check, S } from './WizardShell' +import { useWizardLifecycle, useDebouncedAction, useWizardEvent, CommitButton } from './wizardHooks' +import type { SendAction } from './wizardHooks' + +interface Props { + caretPos: React.CSSProperties + windowId: number + sendAction: SendAction + onClose: () => void +} + +/** `drift_action.METHODS`. */ +type Method = 'rigid' | 'rigid_affine' +type TabLabel = 'Rigid' | 'Rigid+Affine' +const TABS: readonly TabLabel[] = ['Rigid', 'Rigid+Affine'] +const METHOD_OF: Record = { + 'Rigid': 'rigid', 'Rigid+Affine': 'rigid_affine', +} +const TAB_OF: Record = { + rigid: 'Rigid', rigid_affine: 'Rigid+Affine', +} +/** Verbatim from `drift_action._UNAVAILABLE` — the reason the backend gives. + * + * This list is DUPLICATED from the backend, which is a trap worth naming: + * implementing a model there while leaving its tab locked here makes the + * finished feature unreachable — with every headless test still green, + * because none of them can see a disabled tab. If a model is added or + * implemented, both ends move. */ +const UNAVAILABLE: Partial> = { + rigid_affine: 'the affine drift search (plan A4) is not implemented in spyde.drift yet', +} + +type Reference = 'running' | 'sequential' | 'first' +const REFERENCES: readonly { value: Reference; label: string }[] = [ + { value: 'running', label: 'Running average' }, + { value: 'sequential', label: 'Previous frame' }, + { value: 'first', label: 'First frame' }, +] + +/** Mirrors `drift_action.DEFAULTS`. */ +interface DriftSaved { + useRoi: boolean + rejectOutliers: boolean + method: Method + reference: Reference + upsample: number + maxShift: number + apodize: boolean + normalize: boolean + order: number + previewFrames: number +} +const DEFAULTS: DriftSaved = { + useRoi: false, rejectOutliers: true, method: 'rigid', reference: 'running', + upsample: 8, maxShift: 32, apodize: true, normalize: true, order: 1, + previewFrames: 20, +} +const _driftStore = new Map() + +interface Preview { roi: number[] | null; frames: number; gain: number } +interface Result { maxShift: number; gain: number; rejected: number; cancelled: boolean } + +export function DriftWizard({ caretPos, windowId, sendAction, onClose }: Props) { + const saved = _driftStore.get(windowId) ?? DEFAULTS + const [useRoi, setUseRoi] = React.useState(saved.useRoi) + const [rejectOutliers, setRejectOutliers] = React.useState(saved.rejectOutliers) + const [method, setMethod] = React.useState(saved.method) + const [reference, setReference] = React.useState(saved.reference) + const [upsample, setUpsample] = React.useState(saved.upsample) + const [maxShift, setMaxShift] = React.useState(saved.maxShift) + const [apodize, setApodize] = React.useState(saved.apodize) + const [normalize, setNormalize] = React.useState(saved.normalize) + const [order, setOrder] = React.useState(saved.order) + const [previewFrames, setPreviewFrames] = React.useState(saved.previewFrames) + + const [advanced, setAdvanced] = React.useState(false) + const [nFrames, setNFrames] = React.useState(0) + const [solved, setSolved] = React.useState(false) + const [running, setRunning] = React.useState(false) + const [progress, setProgress] = React.useState<{ done: number; total: number } | null>(null) + const [preview, setPreview] = React.useState(null) + const [result, setResult] = React.useState(null) + const [status, setStatus] = React.useState('Drag the box onto a landmark to test it.') + + const vals = React.useRef(saved) + vals.current = { + useRoi, rejectOutliers, method, reference, + upsample, maxShift, apodize, normalize, order, previewFrames, + } + React.useEffect(() => { _driftStore.set(windowId, vals.current) }) + + /** The backend's parameter names (`drift_action.DEFAULTS` keys). */ + const params = (): Record => { + const v = vals.current + return { + use_roi: v.useRoi, reject_outliers: v.rejectOutliers, method: v.method, + reference: v.reference, upsample: v.upsample, max_shift: v.maxShift, + apodize: v.apodize, normalize: v.normalize, order: v.order, + preview_frames: v.previewFrames, + } + } + + // Mount → drift_open (Drift Check window + the alignment box + the first + // discovery preview; nothing SOLVES — plan A8 is explicit that drift + // correction never runs on load). Unmount → drift_close. StrictMode-safe. + useWizardLifecycle({ + windowId, sendAction, + openAction: 'drift_open', openPayload: params, closeAction: 'drift_close', + }) + + // A toggle/parameter change re-runs the ~20-frame discovery preview. Only + // debounced HERE — the backend deliberately doesn't debounce drift_tune + // again (it debounces the ROI DRAG, whose events arrive at frame rate). + const sendTune = useDebouncedAction(sendAction, 'drift_tune', windowId) + const tune = () => sendTune(params) + const live = (set: (v: T) => void) => (v: T) => { set(v); tune() } + + useWizardEvent('spyde:drift_state', windowId, (d) => { + if (typeof d.n_frames === 'number') setNFrames(d.n_frames) + if (typeof d.solved === 'boolean') { + setSolved(d.solved) + if (!d.solved) setResult(null) + } + // The backend refuses an unimplemented model and stays on rigid, so the + // tab follows what it actually selected — never what was clicked. + const m = String(d.method ?? '') as Method + if (m in TAB_OF) setMethod(m) + }) + + useWizardEvent('spyde:drift_preview', windowId, (d) => { + const gain = Number(d.gain) + setPreview({ + roi: Array.isArray(d.roi) ? (d.roi as number[]).map(Number) : null, + frames: Number(d.frames ?? 0), + gain: Number.isFinite(gain) ? gain : NaN, + }) + }) + + useWizardEvent('spyde:drift_progress', windowId, (d) => { + const done = Number(d.done ?? 0), total = Number(d.total ?? 0) + const live = total > 0 && done < total + setProgress(live ? { done, total } : null) + if (live) setRunning(true) + }) + + useWizardEvent('spyde:drift_result', windowId, (d) => { + const gain = Number(d.gain) + setResult({ + maxShift: Number(d.max_abs_shift ?? 0), + gain: Number.isFinite(gain) ? gain : NaN, + rejected: Number(d.rejected ?? 0), + cancelled: Boolean(d.cancelled), + }) + setProgress(null) + setRunning(false) + setSolved(true) + setStatus(d.cancelled ? 'Stopped — partial model' : 'Solved.') + }) + + const onMethod = (t: TabLabel) => { + const m = METHOD_OF[t] + setMethod(m) + vals.current = { ...vals.current, method: m } + sendAction('drift_set_method', { method: m }, windowId) + } + + const solve = () => { + setResult(null) + setRunning(true) + setStatus(`Correcting drift over ${nFrames || '…'} frames`) + sendAction('drift_run', params(), windowId) + } + + const discard = () => { + setRunning(false) + setProgress(null) + setResult(null) + setSolved(false) + setStatus('Discarded.') + sendAction('drift_discard', {}, windowId) + } + + const locked = UNAVAILABLE[method] + const pct = progress ? Math.round((progress.done / progress.total) * 100) : 0 + + return ( + + {/* The whole default face: two toggles, one number, one button. */} + + + + + + + + {progress && ( +
+
+ {progress.done}/{progress.total} +
+ )} + + {result && ( + <> +
+ {result.cancelled ? '◐' : '✓'}{' '} + {Number.isFinite(result.gain) ? `${result.gain.toFixed(1)}x sharper · ` : ''} + {result.maxShift.toFixed(1)} px drift + {result.rejected ? ` · ${result.rejected} bad frames` : ''} +
+
+ {/* Apply adds the LAZY corrected node (map_blocks over the source's + own chunking) — nothing is copied, so this is cheap even on a + multi-GB movie. */} + + +
+ + )} + + setAdvanced(v => !v)}> + Boolean(UNAVAILABLE[METHOD_OF[t]])} + testid={(t) => `drift-tab-${METHOD_OF[t]}`} + /> + {/* The stub is locked, so this names it rather than waiting for a + click that cannot happen. Text is the backend's own wording. */} +
+ {locked ?? 'Rigid+Affine is not implemented in spyde.drift yet.'} +
+ +