tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Thomas Ales [M S E]
2026-08-10 14:14:38 -05:00
parent 191d1b8946
commit 1caf6373cb
13 changed files with 489 additions and 418 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_compute.py
+94 -39
View File
@@ -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."""