Thomas Ales [M S E]
1caf6373cb
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>
110 lines
4.1 KiB
Python
110 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Benchmark the FFT peak-search path: serial vs pooled.
|
|
|
|
Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
|
|
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
|
|
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 # 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 _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
|
|
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
|
|
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])
|
|
|
|
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():
|
|
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("--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.real:
|
|
sras = row_slice(SrasFile(args.real), 0, args.real_rows)
|
|
sras.data = [sras.data[0]]
|
|
sras.n_angles = 1
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|