Replace zoom FFT peak search with a budget-bounded PyFFTW direct transform

Drop the coarse+fine zoom refinement, the SciPy FFT backend, and the
exact= audit path in favor of a single always-on full-transform peak
search (_peak_bins). Block size is now derived from a per-thread memory
budget (_fft_block_for/SRAS_FFT_PLAN_BUDGET_MB) instead of a fixed
constant, so the existing block-parallel PyFFTW pool stays memory-safe
at high pad factors without the zoom algorithm's bookkeeping. Also
removes the now-unused threadpoolctl dependency and the FFT backend
selector from the UI.

Also includes a pre-existing min_freq_mhz peak-search floor (excludes
bins below a caller-supplied frequency from the argmax) that was
already implemented and tested in the working tree.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_compute.py

@@ -122,7 +122,7 @@ def test_parallel_identity(tmp_path, monkeypatch):
     # the block size so every chunk splits into many FFT tasks — the worst
     # case for boundary bugs.
     monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
-    monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
+    monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
     fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
     assert fft_rows < n_rows, \
         f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
@@ -155,43 +155,10 @@ def test_parallel_identity(tmp_path, monkeypatch):
         "rf image identical (masked, padded/zoom)"
 
 
-@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)])
-def test_zoom_identity(tmp_path, monkeypatch, spf, bps):
-    """The zoom peak search must reproduce the full padded-rfft argmax
-    bit-for-bit, across pad factors, masking, bg-sub, dtype, and backend."""
-    path = tmp_path / f"zoom_{spf}.sras"
-    gen.write(path, n_angles=2, seed=6, samples_per_frame=spf, bps=bps)
-    sras = SrasFile(str(path))
-    dc4 = dc_image_mv(sras, 0, CH4_IDX)
-    thr = float(np.median(dc4))
-
-    backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
-    for backend in backends:
-        monkeypatch.setattr(compute, "_fft_backend", backend)
-        for pad in (4, 8, 40):
-            n_fft = spf * pad
-            for thr_v in (None, thr):
-                for bg in (False, True):
-                    ref = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
-                                           apply_bg_sub=bg, n_fft=n_fft,
-                                           exact=True)
-                    zoom = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
-                                            apply_bg_sub=bg, n_fft=n_fft)
-                    diff = int((ref != zoom).sum())
-                    assert diff == 0, \
-                        (f"{diff} px differ: backend={backend} pad={pad} "
-                         f"thr={thr_v} bg={bg} spf={spf}")
-
-    # A threshold above every pixel masks everything: both paths must agree
-    # on an all-zero image.
-    all_masked = compute_rf_image(sras, 0, dc_threshold_mv=1e9, n_fft=spf * 8)
-    assert not all_masked.any()
-
-
-def test_zoom_identity_fuzz():
-    """Hammer _peak_bins_zoom directly with adversarial spectra: noise,
+def test_peak_bins_fuzz():
+    """Hammer _peak_bins directly with adversarial spectra: noise,
     un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs,
-    and all-zero rows."""
+    and all-zero rows — against an independent scipy.fft reference."""
     import scipy.fft as scipy_fft
 
     rng = np.random.default_rng(42)
@@ -222,14 +189,102 @@ def test_zoom_identity_fuzz():
         P[:, 0] = 0.0
         ref = np.argmax(P, axis=1)
 
-        zp = compute._zoom_plan(spf, n_fft)
-        got = compute._peak_bins_zoom(w, zp)
+        got = compute._peak_bins(w, n_fft)
         bad = np.nonzero(ref != got)[0]
         assert not len(bad), \
             (f"spf={spf} pad={pad}: rows {bad.tolist()} picked "
              f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
 
 
+def test_peak_bins_fuzz_min_freq():
+    """Same adversarial-spectra fuzz as test_peak_bins_fuzz, but with a swept
+    min_bin floor: bins below the floor must be excluded from the argmax
+    exactly as the independent scipy.fft reference is, when zeroed the same
+    way before argmax."""
+    import scipy.fft as scipy_fft
+
+    rng = np.random.default_rng(43)
+    for _ in range(25):
+        spf = int(rng.integers(16, 220))
+        pad = int(rng.choice([4, 5, 8, 16, 40]))
+        n_fft = spf * pad
+        n_wf = 24
+        w = rng.normal(scale=20.0, size=(n_wf, spf))
+        t = np.arange(spf)
+        for r in range(6):
+            f = rng.uniform(1.0, spf / 2 - 1)
+            w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
+        f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
+        w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+            + 49.9 * np.sin(2 * np.pi * f2 * t / spf)
+        w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+            + 50 * np.cos(2 * np.pi * f2 * t / spf)
+        w[8] = 90 + rng.normal(scale=5.0, size=spf)
+        w[9] = 0.0
+        w = w.astype(np.float32)
+
+        S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
+        P = S.real ** 2
+        P += S.imag ** 2
+        n_bins_fine = n_fft // 2 + 1
+        min_bin = int(rng.integers(1, max(2, n_bins_fine // 3)))
+        P[:, :min_bin] = 0.0
+        ref = np.argmax(P, axis=1)
+
+        got = compute._peak_bins(w, n_fft, min_bin)
+        bad = np.nonzero(ref != got)[0]
+        assert not len(bad), \
+            (f"spf={spf} pad={pad} min_bin={min_bin}: rows {bad.tolist()} picked "
+             f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
+
+
+@pytest.mark.parametrize("spf", [64, 500, 2500])
+def test_fft_block_for(spf):
+    """Block size == _FFT_BLOCK_MAX at natural resolution, shrinks and stays
+    >= _FFT_BLOCK_MIN as n_len grows, and the implied per-thread byte
+    estimate respects _FFT_PLAN_BYTES_BUDGET except when the floor is
+    engaged."""
+    block_natural = compute._fft_block_for(spf, spf)
+    assert block_natural == compute._FFT_BLOCK_MAX
+
+    prev = compute._FFT_BLOCK_MAX
+    for pad in (2, 4, 8, 40, 500):
+        n_len = spf * pad
+        block = compute._fft_block_for(spf, n_len)
+        assert compute._FFT_BLOCK_MIN <= block <= prev
+        bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
+        if block > compute._FFT_BLOCK_MIN:
+            assert block * bytes_per_wf <= compute._FFT_PLAN_BYTES_BUDGET
+        prev = block
+
+
+def test_compute_rf_image_min_freq_mhz(tmp_path):
+    """min_freq_mhz threads through compute_rf_image end-to-end, for both
+    the natural-resolution and padded paths: 0.0 (default) must reproduce
+    the pre-existing image exactly, and a floor above every real peak must
+    collapse the image to bin 0 (0 MHz) — the same fallback the low-level
+    search uses when nothing survives the floor."""
+    path = tmp_path / "floor_e2e.sras"
+    gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
+    sras = SrasFile(str(path))
+    huge_floor = float(sras.freq_axis_mhz(None)[-1]) + 1.0   # above Nyquist
+
+    for n_fft in (None, 64 * 8):
+        unfiltered = compute_rf_image(sras, 0, dc_threshold_mv=None,
+                                      apply_bg_sub=False, n_fft=n_fft)
+        same = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
+                                n_fft=n_fft, min_freq_mhz=0.0)
+        assert np.array_equal(unfiltered, same), \
+            f"n_fft={n_fft}: min_freq_mhz=0.0 changed the output"
+        assert unfiltered.any(), \
+            f"n_fft={n_fft}: fixture should have real signal"
+
+        collapsed = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
+                                     n_fft=n_fft, min_freq_mhz=huge_floor)
+        assert not collapsed.any(), \
+            f"n_fft={n_fft}: floor above Nyquist should collapse to 0 MHz"
+
+
 def test_nomask_equals_low_threshold(tmp_path):
     """dc_threshold_mv=None must equal a threshold below every pixel, while
     skipping the CH4 read."""

tests/test_gui.py

@@ -223,6 +223,46 @@ def test_threshold_change_recomputes(ctx):
         f"masking zeroed some pixels ({n_zero} of {win._current_image.size})"
 
 
+def test_highlight_masked_pixels(ctx):
+    """Masked (below-threshold) pixels are drawn as NaN - filled with a
+    highlight color, separate from the normal colormap - so they can't be
+    mistaken for a real, possibly-low-frequency pixel; unchecking restores
+    the old behavior where both blend into the same plain 0."""
+    win = ctx.win
+    assert win.chk_highlight_masked.isChecked(), "on by default"
+    dc4 = win._dc_cache[(0, CH4_IDX)]
+    expect_masked = dc4 < win.spin_threshold_mv.value()
+    assert expect_masked.any() and not expect_masked.all(), \
+        "fixture threshold should mask some but not all pixels"
+
+    calls = []
+    orig = win.image_canvas.show_image
+
+    def spy(img, *a, **kw):
+        calls.append((np.array(img, copy=True), kw.get("bad_color")))
+        return orig(img, *a, **kw)
+
+    with patch.object(win.image_canvas, "show_image", side_effect=spy):
+        win._redraw_image(win._current_image)
+    shown, bad_color = calls[-1]
+    assert bad_color is not None, "highlight color set while checkbox is on"
+    assert np.array_equal(np.isnan(shown), expect_masked), \
+        "NaN exactly where DC4 is below threshold, nowhere else"
+
+    win.chk_highlight_masked.setChecked(False)
+    calls.clear()
+    with patch.object(win.image_canvas, "show_image", side_effect=spy):
+        win._redraw_image(win._current_image)
+    shown2, bad_color2 = calls[-1]
+    assert bad_color2 is None, "no highlight color once unchecked"
+    assert not np.isnan(shown2).any(), "unchecked: no pixel pulled out to NaN"
+    assert np.array_equal(shown2, win._current_image), \
+        "unchecked: displayed array is the raw, unmodified image"
+
+    win.chk_highlight_masked.setChecked(True)
+    pump(60)
+
+
 def test_bg_sub_toggle(ctx):
     """bg-sub no longer gates the display: it only affects a future live
     compute for an angle with nothing cached yet, or an explicit batch

tests/test_row_average.py

@@ -14,7 +14,7 @@ import pytest
 
 import sras_compute as compute
 from sras_compute import compute_rf_image, dc_image_mv
-from sras_format import CH4_IDX, SrasFile
+from sras_format import CH1_IDX, CH4_IDX, SrasFile
 import tools.make_test_sras as gen
 
 
@@ -193,26 +193,45 @@ def test_row_average_respects_own_center_mask(tmp_path):
 
 def test_row_average_composes_with_padding(tmp_path):
     """row_avg_n and n_fft (zero-padding) are independent knobs: using them
-    together must not raise, and must still agree with the exact (non-zoom)
-    reference path at that pad factor -- i.e. row-averaging composes with
-    the zoom peak search correctly, not just with the direct one."""
+    together must not raise, and must agree bit-for-bit with an independent
+    reference that row-averages the raw waveforms and background-subtracts
+    them, then runs a plain scipy rfft + argmax at the same pad factor --
+    i.e. row-averaging composes correctly with the padded peak search."""
+    import scipy.fft as scipy_fft
+
     path = tmp_path / "padded_rowavg.sras"
     gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
     sras = SrasFile(str(path))
     spf = sras.samples_per_frame
+    n_rows, n_frames = sras.image_shape(0)
+    n_fft = spf * 40
 
     raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
     avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
                                    row_avg_n=3)
-    avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
-                                       row_avg_n=3, n_fft=spf * 40)
-    avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
-                                        row_avg_n=3, n_fft=spf * 40, exact=True)
+    avg_padded = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
+                                  row_avg_n=3, n_fft=n_fft)
 
-    assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape
-    assert np.all(np.isfinite(avg_padded_zoom))
-    assert np.array_equal(avg_padded_zoom, avg_padded_exact), \
-        "row-averaged waveforms feed the zoom and exact FFT paths identically"
+    assert avg_natural.shape == raw_natural.shape == avg_padded.shape
+    assert np.all(np.isfinite(avg_padded))
+
+    weights = compute._row_average_weights(3)
+    freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
+    data = sras.data[0]
+    expected = np.zeros((n_rows, n_frames), dtype=np.float32)
+    for r in range(n_rows):
+        v = np.ones(n_frames, dtype=bool)
+        avg = compute._row_average_waveforms(
+            data[r, CH1_IDX].astype(np.float32), v, weights)
+        avg = avg - sras.background
+        S = scipy_fft.rfft(avg, n=n_fft, axis=-1, workers=1)
+        power = S.real ** 2 + S.imag ** 2
+        power[:, 0] = 0.0
+        expected[r] = freq32[np.argmax(power, axis=1)]
+
+    assert np.array_equal(avg_padded, expected), \
+        "row-averaged waveforms feed the padded peak search identically " \
+        "to an independent reference"
 
 
 def test_row_average_improves_snr_recovery():
@@ -233,8 +252,8 @@ def test_row_average_improves_snr_recovery():
     weights = compute._row_average_weights(8)   # wide window: lots of averaging
     averaged = compute._row_average_waveforms(raw, valid, weights)
 
-    raw_bins = compute._peak_bins_direct(raw, spf)
-    avg_bins = compute._peak_bins_direct(averaged, spf)
+    raw_bins = compute._peak_bins(raw, spf)
+    avg_bins = compute._peak_bins(averaged, spf)
 
     raw_hits = int(np.sum(raw_bins == true_bin))
     avg_hits = int(np.sum(avg_bins == true_bin))
@@ -257,7 +276,7 @@ def test_row_average_parallel_identity(tmp_path, monkeypatch):
     sras = SrasFile(str(path))
 
     monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
-    monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
+    monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
     # row_avg_n > 0 halves the effective budget before chunk planning.
     fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
     assert fft_rows < n_rows, \

tests/test_stored_cache.py

@@ -90,7 +90,7 @@ def no_fft(monkeypatch):
     """Make any real FFT work loud: returns a list that stays empty unless a
     peak search actually runs."""
     calls = []
-    for name in ("_peak_bins_direct", "_peak_bins_zoom"):
+    for name in ("_peak_bins",):
         original = getattr(compute, name)
 
         def spy(*args, _f=original, **kwargs):
@@ -147,7 +147,7 @@ def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
     gen.write(path, n_angles=2, seed=13, samples_per_frame=256)
     spf = SrasFile(str(path)).samples_per_frame
 
-    assert cache_file(str(path), "fft", True, "scipy", 0, pad) == ""
+    assert cache_file(str(path), "fft", True, 0, pad) == ""
     sras = SrasFile(str(path))
     assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"