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:
+34
-15
@@ -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, \
|
||||
|
||||
Reference in New Issue
Block a user