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
tools/bench_fft.py
+38 -35
View File
@@ -1,8 +1,10 @@
#!/usr/bin/env python3
"""Benchmark the FFT peak-search path: exact vs zoom, serial vs pooled.
"""Benchmark the FFT peak-search path: serial vs pooled.
Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
cores), and verifies every variant against the exact reference image.
cores), the estimated per-thread resident pyFFTW plan footprint at each pad
factor (see sras_compute._fft_block_for), and verifies pooled output against
the serial reference.
Usage:
python tools/bench_fft.py # synthetic, pads 1/8/40
@@ -22,10 +24,10 @@ import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import sras_compute as compute # noqa: E402
from sras_compute import compute_rf_image, set_fft_backend # noqa: E402
from sras_format import SrasFile # noqa: E402
import tools.make_test_sras as gen # noqa: E402
from tools.check_equivalence import row_slice # noqa: E402
from sras_compute import compute_rf_image # noqa: E402
from sras_format import SrasFile # noqa: E402
import tools.make_test_sras as gen # noqa: E402
from tools.check_equivalence import row_slice # noqa: E402
def _timed(fn):
@@ -38,35 +40,42 @@ def _timed(fn):
return out, wall, cpu / max(wall, 1e-9)
def bench(sras, pads, backends):
def _plan_mb(spf: int, n_len: int, n_workers: int) -> float:
"""Estimated resident pyFFTW plan-buffer footprint across the whole
pool at this transform length (see sras_compute._fft_block_for)."""
block = compute._fft_block_for(spf, n_len)
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
return block * bytes_per_wf * n_workers / (1024 * 1024)
def bench(sras, pads):
n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
for a in range(sras.n_angles))
spf = sras.samples_per_frame
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)")
print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} "
f"{'wf/s':>10} {'util':>6} match")
n_workers = compute._MAX_WORKERS
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), "
f"{n_workers} workers")
print(f"{'pad':>4} {'variant':>8} {'wall':>9} {'wf/s':>10} {'util':>6} "
f"{'plan MB':>9} match")
for pad in pads:
n_fft = spf * pad if pad > 1 else None
for backend in backends:
set_fft_backend(backend)
n_len = n_fft if n_fft is not None else spf
plan_mb = _plan_mb(spf, n_len, n_workers)
def run(**kw):
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
apply_bg_sub=True, n_fft=n_fft, **kw)
for a in range(sras.n_angles)]
return np.concatenate([i.ravel() for i in imgs])
def run(**kw):
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
apply_bg_sub=True, n_fft=n_fft, **kw)
for a in range(sras.n_angles)]
return np.concatenate([i.ravel() for i in imgs])
ref, wall, util = _timed(lambda: run(exact=True))
rows = [("exact(serial)", ref, wall, util, True)]
for label, kw in (("zoom(serial)", dict(max_workers=1)),
("zoom(pool)", {})):
img, wall, util = _timed(lambda: run(**kw))
rows.append((label, img, wall, util, bool(np.array_equal(img, ref))))
for label, img, wall, util, ok in rows:
print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s "
f"{n_wf / wall:>10.0f} {util:>5.1f}x "
f"{'OK' if ok else 'MISMATCH'}")
ref, wall, util = _timed(lambda: run(max_workers=1))
rows = [("serial", ref, wall, util, True)]
img, wall, util = _timed(lambda: run())
rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref))))
for label, img, wall, util, ok in rows:
print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} "
f"{util:>5.1f}x {plan_mb:>8.1f} {'OK' if ok else 'MISMATCH'}")
def main():
@@ -76,30 +85,24 @@ def main():
p.add_argument("--spf", type=int, default=2500)
p.add_argument("--rows", type=int, default=8)
p.add_argument("--frames", type=int, default=1024)
p.add_argument("--backends", default=None,
help="comma-separated (default: scipy,pyfftw if available)")
p.add_argument("--real", help="path to a real .sras file")
p.add_argument("--real-rows", type=int, default=32,
help="rows of angle 0 to use from the real file")
args = p.parse_args()
pads = [int(x) for x in args.pads.split(",")]
if args.backends:
backends = args.backends.split(",")
else:
backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
if args.real:
sras = row_slice(SrasFile(args.real), 0, args.real_rows)
sras.data = [sras.data[0]]
sras.n_angles = 1
bench(sras, pads, backends)
bench(sras, pads)
else:
with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
path = Path(tmp) / "bench.sras"
gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
geometry=[(args.rows, args.frames)])
bench(SrasFile(str(path)), pads, backends)
bench(SrasFile(str(path)), pads)
if __name__ == "__main__":