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:
+48
-35
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user