Files
sras-viewer/tools/bench_fft.py
Thomas Ales 9a9a2557d6 Persist FFT settings, fix backend naming, parallelise batch DC caching
- FFT backend and pad factor persist across sessions via QSettings
  (IniFormat; tests redirect the settings path for hermeticity).
- The default backend was labelled "NumPy FFT" but always dispatched to
  scipy.fft — rename the canonical value to "scipy" ("numpy" stays as a
  legacy alias) and fix the dialog label.
- cache_file: DC caching fans out over angles via _parallel_map with
  per-angle budgets (the DcPrecomputeWorker pattern); FFT caching stays
  serial per angle because compute_rf_image now parallelises internally
  over blocks. Documented that the v7 FFT cache is natural-resolution
  (pad 1) by design.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-08-06 10:26:50 -05:00

107 lines
4.1 KiB
Python

#!/usr/bin/env python3
"""Benchmark the FFT peak-search path: exact vs zoom, 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.
Usage:
python tools/bench_fft.py # synthetic, pads 1/8/40
python tools/bench_fft.py --pads 40 --spf 2500 --rows 8 --frames 1024
python tools/bench_fft.py --real /path/big.sras --real-rows 32 --pads 40
"""
import argparse
import resource
import sys
import tempfile
import time
from pathlib import Path
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
def _timed(fn):
r0 = resource.getrusage(resource.RUSAGE_SELF)
t0 = time.perf_counter()
out = fn()
wall = time.perf_counter() - t0
r1 = resource.getrusage(resource.RUSAGE_SELF)
cpu = (r1.ru_utime - r0.ru_utime) + (r1.ru_stime - r0.ru_stime)
return out, wall, cpu / max(wall, 1e-9)
def bench(sras, pads, backends):
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")
for pad in pads:
n_fft = spf * pad if pad > 1 else None
for backend in backends:
set_fft_backend(backend)
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'}")
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--pads", default="1,8,40",
help="comma-separated pad factors (default 1,8,40)")
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)
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)
if __name__ == "__main__":
main()