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
+48 -35
View File
@@ -32,45 +32,58 @@ share of the budget. Capping the workers alone is not enough: the chunk would
still be sized against the whole budget, and N concurrent callers would each
allocate all of it.
## FFT peak search: block-parallel zoom refinement (`sras_compute.py`)
## FFT peak search: block-parallel direct transform (`sras_compute.py`)
The displayed RF value per pixel is the argmax of the zero-padded power
spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for
mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan
row, which is what used to collapse the old row-chunk planner to one worker
and make synthesis single-threaded.
spectrum of that pixel's CH1 waveform. This used to run through `_peak_bins_zoom`,
a coarse-rfft-plus-local-fine-DFT refinement that avoided ever materialising
a padded spectrum — at the pad factor of 40 needed for mapping resolution, a
full padded spectrum is ~9 GB per scan row, which used to collapse the old
row-chunk planner to one worker and make synthesis single-threaded. That
refinement was removed once pyFFTW became the sole, mandatory FFT backend
(SciPy dropped as a compute backend entirely): `_peak_bins` now always runs
the real full transform, and the memory problem zoom dodged is instead
solved by bounding *per-block* spectrum memory rather than avoiding full
spectra altogether.
`_peak_bins_zoom` never materialises the padded spectrum:
`_peak_bins` runs the full transform (`_block_rfft`, a cached pyFFTW
`builders.rfft` plan, FFTW_MEASURE, wisdom persisted under
`~/.cache/sras-viewer/`) and argmaxes the power spectrum, in blocks fanned
out task-parallel over a persistent thread pool (`_fft_pool()`) — each pool
thread runs one single-threaded transform at a time, so aggregate
parallelism equals the pool's worker count. This block-streaming structure
is not just about parallelism: it is also what keeps memory bounded, by
never materialising more than one block's worth of full padded spectrum at
a time, regardless of how many waveforms a chunk holds.
1. a coarse rfft at `next_fast_len(2*spf)` — 2× oversampled, so the padded
power spectrum (a trig polynomial of degree spf−1) cannot hide its global
max between coarse samples;
2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max
becomes a refinement candidate. Quarter-natural-bin scalloping at the 2×
grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide
margin. The DC-adjacent window is always refined too: the coarse DC bin
is zeroed for suppression, which would otherwise blind the scan to fine
bins closer to DC than the first coarse sample (where the leakage skirt
of an un-subtracted offset peaks);
3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every
fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated
on the exact `n_fft` grid by one small complex gemm, with np.argmax's
lowest-bin tie-break preserved across windows.
The block size is the part that has to adapt to pad factor.
`_fft_block_for(spf, n_len)` derives waveforms-per-task from a fixed
per-thread byte budget (`_FFT_PLAN_BYTES_BUDGET`, `SRAS_FFT_PLAN_BUDGET_MB`,
default 16 MB) rather than a fixed constant, because a cached pyFFTW plan's
input+output buffers are *permanent* per-thread memory (the plan cache is
never evicted) — with a fixed 512-waveform block, pad 40 at a 2500-sample
frame costs ~210 MB per pool thread (~3.3 GB total across 16 threads);
`_fft_block_for` bounds that to ~16 MB per thread (~260 MB across 16
threads) at the same pad factor, while still reproducing the old tuned 512
exactly at natural resolution (pad 1), where it cost nothing to begin with.
`_FFT_BLOCK_MAX` (512) and `_FFT_BLOCK_MIN` (32) cap and floor the result:
the ceiling is the measured knee on a 16-core machine at natural resolution
(smaller blocks serialise on GIL-held numpy dispatch, larger ones lose cache
residency and task granularity); the floor keeps task granularity from
collapsing at extreme pad factors, at the cost of exceeding the byte budget
there.
The selected bin is bit-identical to the full padded argmax — enforced by
`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial
spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose
baseline was captured on the old full-padded path.
The outer row-chunk sizing (`_plan_fft_rows`) needed no companion change.
It only ever budgets the raw float32 waveform *read* buffer, which this
change doesn't touch — spectrum memory is bounded independently by
`_fft_block_for`, and since the pool only ever runs as many blocks
concurrently as it has workers, peak transient spectrum memory during a
chunk's FFT phase is `_MAX_WORKERS * block * bytes_per_wf`, the same bound
whether the chunk holds 10 rows or 10,000. Queuing more rows into one chunk
to keep the read-row pool busy therefore can't blow up spectrum memory.
Work fans out over a persistent thread pool in `_FFT_BLOCK` = 512-waveform
tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose
cache residency and task granularity (measured on a 16-core machine, where
this path runs ~35× faster than the old serial padded transform at pad 40).
pyFFTW runs through per-thread `builders` plans (FFTW_MEASURE, wisdom
persisted under `~/.cache/sras-viewer/`), and `threadpoolctl` clamps BLAS to
one thread under the pool so the refinement gemm cannot oversubscribe.
`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
full-padded path for audits.
`tools/check_equivalence.py`'s golden-hash harness remains the end-to-end
regression baseline for this path, unaffected by this change.
## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
@@ -121,8 +134,8 @@ same edge case the masked convolution already handles.
The averaging step doubles the live per-row scratch memory (a full-width
`(n_frames, spf)` buffer on top of the existing compacted `waves` buffer),
so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before
`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row
chunking" above. On the largest real scans `_plan_chunks` is already
`_plan_fft_rows` runs — see "Memory budget and row chunking" above. On the
largest real scans `_plan_chunks` is already
clamped to its floor of one row regardless, so this costs no concurrency
where it matters most; it mainly protects moderate-sized scans from an
unexpected regression.