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>
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
+34 -15
View File
@@ -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, \