tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rewrite the FFT peak search: block-parallel zoom refinement

Browse Source 
At pad 40 the old path materialised a ~9 GB padded spectrum per row,
which collapsed the chunk planner to one worker and one rfft call with
workers=1 — synthesis ran single-threaded, ~1 hour per angle on real
files.

The padded spectrum is never materialised now. Each block of 512
waveforms gets a coarse rfft at next_fast_len(2*spf); every coarse bin
within 0.7 of its row's max (plus the DC-adjacent window, which coarse
DC suppression would otherwise blind) is refined onto the exact n_fft
grid by a small complex gemm. The selected bin is bit-identical to the
full padded argmax — enforced by test_zoom_identity, a 25-seed fuzz
test over adversarial spectra, and a clean golden-hash diff against the
pre-rewrite baseline across pads {1,2,4,8,40}, masked/unmasked, bg
on/off, int8/int16, and both backends.

Blocks fan out over a persistent thread pool; pyFFTW runs through
per-thread FFTW_MEASURE builder plans with wisdom persisted to
~/.cache/sras-viewer, and threadpoolctl clamps BLAS under the pool.
compute_rf_image(exact=True) (or SRAS_FFT_EXACT=1) keeps the reference
padded path for audits.

tools/bench_fft.py measures: pad 40, 16 cores, 8192x2500 synthetic —
exact serial 717 wf/s -> zoom pool 25100 wf/s (35x, pyFFTW backend;
19x scipy), every variant verified equal to the reference.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-06 10:20:14 -05:00
parent 007089dd48
commit 11ff3b62e2
4 changed files with 522 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+322 -41
View File
@@ -6,8 +6,10 @@ multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
"""
import atexit
import json
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
@@ -24,12 +26,16 @@ from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
try:
import pyfftw
pyfftw.interfaces.cache.enable()
PYFFTW_AVAILABLE = True
except ImportError:
PYFFTW_AVAILABLE = False
_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend()
try:
from threadpoolctl import threadpool_limits
except ImportError:
threadpool_limits = None
_fft_backend = "numpy" # "numpy" (scipy.fft) or "pyfftw"; set via set_fft_backend()
def set_fft_backend(name: str):
@@ -52,6 +58,218 @@ def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
# ---------------------------------------------------------------------------
# FFT worker pool and per-thread pyFFTW plans
# ---------------------------------------------------------------------------
_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine:
# smaller blocks serialise on GIL-held numpy dispatch,
# larger ones lose cache residency and task granularity.
_ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf
_FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0))
_pool_lock = threading.Lock()
_pool: ThreadPoolExecutor | None = None
def _fft_pool() -> ThreadPoolExecutor:
"""The persistent process-wide pool for FFT block tasks."""
global _pool
with _pool_lock:
if _pool is None:
_pool = ThreadPoolExecutor(max_workers=_MAX_WORKERS,
thread_name_prefix="sras-fft")
atexit.register(_pool.shutdown, wait=False, cancel_futures=True)
return _pool
_WISDOM_PATH = Path.home() / ".cache" / "sras-viewer" / "fftw_wisdom"
_wisdom_lock = threading.Lock()
_wisdom_loaded = False
_fftw_local = threading.local()
def _load_wisdom_once():
"""Import saved FFTW wisdom so FFTW_MEASURE planning is a one-time cost
per machine. Purely an optimisation: failures are ignored."""
global _wisdom_loaded
with _wisdom_lock:
if _wisdom_loaded:
return
_wisdom_loaded = True
try:
pyfftw.import_wisdom(_WISDOM_PATH.read_bytes().split(b"\x00\n"))
except Exception:
pass
def _save_wisdom():
with _wisdom_lock:
try:
_WISDOM_PATH.parent.mkdir(parents=True, exist_ok=True)
_WISDOM_PATH.write_bytes(b"\x00\n".join(pyfftw.export_wisdom()))
except Exception:
pass
def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
"""rfft of a (B, spf) float32 block via a cached per-thread FFTW plan.
Plans have a fixed (_FFT_BLOCK, spf) input shape so each worker thread
plans once per transform length; a remainder block runs through the same
plan with its tail rows ignored. The returned array is the plan's output
buffer — consume it before the next call on the same thread.
"""
n_wf, spf = waves.shape
plans = getattr(_fftw_local, "plans", None)
if plans is None:
plans = _fftw_local.plans = {}
key = (_FFT_BLOCK, spf, n)
plan = plans.get(key)
if plan is None:
_load_wisdom_once()
buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32")
# No overwrite_input: FFTW must not scribble on input_array, whose
# zero-padded tail (columns spf..n) is zeroed exactly once here.
plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1,
planner_effort="FFTW_MEASURE")
plan.input_array[:] = 0.0
plans[key] = plan
_save_wisdom()
inp = plan.input_array
inp[:n_wf, :spf] = waves
return plan()[:n_wf]
def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
"""Single-threaded rfft of one block; outer parallelism comes from the
pool, so the transform itself must not spin up threads."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return _fftw_block_rfft(waves, n)
return scipy_fft.rfft(waves, n=n, axis=-1, workers=1)
# ---------------------------------------------------------------------------
# Zoom peak search: coarse rfft + local fine DFT around the winning bin
# ---------------------------------------------------------------------------
_ZOOM_HALFWIDTH = 0.75 # refinement window half-width, in coarse spacings.
# Every fine bin lies within 0.5 spacings of its
# nearest coarse bin, and that bin is guaranteed to
# be a candidate (see _ZOOM_CAND_RATIO), so 0.5
# suffices; 0.75 adds rounding margin.
_ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of
# its row's coarse maximum. Quarter-natural-bin
# scalloping at the 2x-oversampled coarse grid can
# understate a peak by at most ~19% in power, so 0.7
# keeps a wide margin — near-contenders are resolved
# on the fine grid, never ranked from coarse samples.
@dataclass
class _ZoomPlan:
"""Constants for the coarse+refine peak search, built once per
compute_rf_image call and shared across worker threads. *phases* is a
lazily-filled window-start -> phase-vector cache; a benign duplicate
compute under concurrency is cheaper than locking."""
n_fft: int
n_coarse: int
m: int # fine bins per refinement window
n_bins_fine: int
E: np.ndarray # (spf, m) complex64 fine-DFT matrix, relative bins
j: np.ndarray # (spf,) float64 sample indices
phases: dict
def _zoom_plan(spf: int, n_fft: int) -> _ZoomPlan:
# 2x-oversampled coarse grid: the padded power spectrum is a trig
# polynomial of degree spf-1, so at 2x sampling its global max cannot
# hide between coarse bins.
n_coarse = scipy_fft.next_fast_len(2 * spf, real=True)
n_bins_fine = n_fft // 2 + 1
m = int(np.ceil(2 * _ZOOM_HALFWIDTH * n_fft / n_coarse)) + 1
m = min(m, n_bins_fine - 1)
j = np.arange(spf, dtype=np.float64)
E = np.exp((-2j * np.pi / n_fft) * np.outer(j, np.arange(m))).astype(np.complex64)
return _ZoomPlan(n_fft, n_coarse, m, n_bins_fine, E, j, {})
def _window_start(k_c: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
"""First fine bin of the refinement window around each coarse bin.
Clipped to [1, ...] so fine bin 0 stays excluded (DC suppression)."""
k0 = np.floor((k_c - _ZOOM_HALFWIDTH) * (zp.n_fft / zp.n_coarse)).astype(np.int64)
return np.clip(k0, 1, max(1, zp.n_bins_fine - zp.m))
def _refine_window(waves: np.ndarray, rows: np.ndarray, k0: int, zp: _ZoomPlan,
best_pow: np.ndarray, best_bin: np.ndarray):
"""Evaluate fine bins k0..k0+m-1 for the given rows and fold the result
into the per-row best (power, bin), preserving np.argmax's lowest-bin
tie-break."""
ph = zp.phases.get(k0)
if ph is None:
ph = np.exp((-2j * np.pi * k0 / zp.n_fft) * zp.j).astype(np.complex64)
zp.phases[k0] = ph
F = (waves[rows] * ph) @ zp.E
q = F.real ** 2
q += F.imag ** 2
i = np.argmax(q, axis=1)
p = q[np.arange(len(rows)), i]
b = k0 + i
upd = (p > best_pow[rows]) | ((p == best_pow[rows]) & (b < best_bin[rows]))
ridx = rows[upd]
best_pow[ridx] = p[upd]
best_bin[ridx] = b[upd]
def _peak_bins_zoom(waves: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
"""Peak fine-bin per waveform without materialising the padded spectrum:
coarse rfft, then a small fine DFT (one gemm per shared window) on the
exact n_fft bin grid. Identity with the full padded argmax is enforced
by tests/test_compute.py::test_zoom_identity and the golden-hash sweep."""
n_wf = waves.shape[0]
S = _block_rfft(waves, zp.n_coarse)
P = S.real ** 2
P += S.imag ** 2
P[:, 0] = 0.0
p_c = np.max(P, axis=1)
best_pow = np.full(n_wf, -1.0, dtype=np.float32)
best_bin = np.full(n_wf, np.iinfo(np.int64).max, dtype=np.int64)
# For a clean signal this yields one or two windows; a noise spectrum
# (many near-equal peaks) yields a dozen or so — still a tiny fraction
# of the padded grid.
thr = np.where(p_c > 0, np.float32(_ZOOM_CAND_RATIO) * p_c,
np.float32(np.inf))
rows_c, bins_c = np.nonzero(P >= thr[:, None])
k0_c = _window_start(bins_c, zp)
# Zeroing the coarse DC bin (suppression) blinds the candidate scan to
# fine bins closer to DC than the first coarse sample — where the
# DC-leakage skirt of an un-subtracted offset peaks. Always refine the
# DC-adjacent window too.
rows_c = np.concatenate([rows_c, np.arange(n_wf)])
k0_c = np.concatenate([k0_c, np.ones(n_wf, dtype=np.int64)])
pair = np.unique(np.stack([rows_c, k0_c], axis=1), axis=0)
for k0 in np.unique(pair[:, 1]):
_refine_window(waves, pair[pair[:, 1] == k0, 0], int(k0), zp,
best_pow, best_bin)
# An all-zero spectrum must reproduce argmax-of-zeros = bin 0.
best_bin[p_c == 0.0] = 0
return best_bin
def _peak_bins_direct(waves: np.ndarray, n_len: int) -> np.ndarray:
"""Peak bin per waveform via the full transform — for pad factors too
small for the zoom search to pay."""
S = _block_rfft(waves, n_len)
power = S.real ** 2
power += S.imag ** 2
power[:, 0] = 0.0
return np.argmax(power, axis=1)
# ---------------------------------------------------------------------------
# Chunking / parallel budget
# ---------------------------------------------------------------------------
@@ -76,17 +294,21 @@ _TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
# An rfft chunk holds, live at once: the float32 input, the complex64
# transform, and the float32 power spectrum — roughly 3x the input buffer.
_FFT_LIVE_MULTIPLIER = 3
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
budget: int = _TOTAL_BYTES_BUDGET) -> int:
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
def _plan_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int:
"""Rows per outer chunk for the block-FFT path. Only the float32
waveform buffer scales with the chunk (in-flight block spectra total a
few MB across the whole pool), so budget it with 2x slack and let the
block fan-out saturate the pool regardless of pad factor."""
bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
live_multiplier: int = 1,
max_workers: int | None = None,
@@ -215,7 +437,8 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
dc4_mv: np.ndarray | None = None,
max_workers: int | None = None,
budget: int | None = None,
should_stop=None) -> np.ndarray:
should_stop=None,
exact: bool = False) -> np.ndarray:
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
@@ -233,6 +456,13 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
At pad factors >= _ZOOM_MIN_PAD the padded spectrum is never
materialised: a coarse rfft finds each peak and a local fine DFT
resolves it on the exact n_fft bin grid (_peak_bins_zoom). *exact*
forces the reference full-padded transform instead — it exists for
tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and
memory-hungry at high pad.
Fast path: if the file has a precomputed peak-frequency image for this
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
matches, the stored image is used directly — no FFT is run.
@@ -257,17 +487,31 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
return freq_img
# ---- Chunked FFT path --------------------------------------------------
freq_axis = sras.freq_axis_mhz(n_fft)
exact = exact or _FFT_EXACT_ENV
spf = sras.samples_per_frame
n_len = n_fft if n_fft is not None else spf
freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
img = np.zeros((n_rows, n_frames), dtype=np.float32)
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
chunk_rows, n_workers = _plan_chunks(
n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
live_multiplier=_FFT_LIVE_MULTIPLIER, max_workers=max_workers,
budget=budget)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
def chunk(r0: int, r1: int):
zp = (_zoom_plan(spf, n_fft)
if not exact and n_fft is not None and n_fft >= _ZOOM_MIN_PAD * spf
else None)
total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
cap = max_workers if max_workers is not None else _MAX_WORKERS
if exact:
# The reference path materialises the full padded spectrum, so rows
# are budgeted against it (complex64 + power + temp per bin) and the
# chunk runs as one serial transform.
bytes_per_row = max(1, n_frames * (4 * spf + 16 * (n_len // 2 + 1)))
chunk_rows = int(max(1, min(_CHUNK_ROWS_MAX, total // bytes_per_row)))
cap = 1
else:
chunk_rows = _plan_fft_rows(n_frames, spf, total)
pool = _fft_pool() if cap > 1 else None
def process(r0: int, r1: int):
if dc_threshold_mv is None:
valid = None
else:
@@ -280,36 +524,73 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
if not valid.any():
return
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so only the
# selected elements are actually read from disk; masked-out pixels'
# pages are never paged in at all.
raw = data[r0:r1, CH1_IDX, :, :]
waves = (raw[valid] if valid is not None else raw).astype(np.float32)
counts = (valid.sum(axis=1) if valid is not None
else np.full(r1 - r0, n_frames, dtype=np.int64))
offs = np.concatenate(([0], np.cumsum(counts)))
n_wf = int(offs[-1])
waves = np.empty((n_wf, spf), dtype=np.float32)
if background is not None:
waves -= background # background is 1-D (spf,)
def read_row(i: int):
if should_stop is not None and should_stop():
return
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so only
# the selected elements are actually read from disk; masked-out
# pixels' pages are never paged in at all.
raw = data[r0 + i, CH1_IDX]
dst = waves[offs[i]:offs[i + 1]]
dst[:] = raw[valid[i]] if valid is not None else raw
if background is not None:
dst -= background # background is 1-D (spf,)
# Parallelism comes from the outer chunk loop, so keep the inner
# transform single-threaded to avoid oversubscribing the machine.
spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
del waves
# |z|^2 without np.abs()'s extra full-size temporary.
power = spectrum.real ** 2
power += spectrum.imag ** 2
del spectrum
# [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames,
# bins) shape, where [:, 0] would blank a whole frame.
power[..., 0] = 0.0 # suppress DC bin
peak_bins = np.argmax(power, axis=-1)
del power
if valid is not None:
img[r0:r1][valid] = freq_axis[peak_bins]
if pool is None:
for i in range(r1 - r0):
read_row(i)
else:
img[r0:r1] = freq_axis[peak_bins]
list(pool.map(read_row, range(r1 - r0)))
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
out = np.empty(n_wf, dtype=np.float32)
def fft_block(b0: int):
if should_stop is not None and should_stop():
return
b1 = min(b0 + _FFT_BLOCK, n_wf)
w = waves[b0:b1]
bins = (_peak_bins_zoom(w, zp) if zp is not None
else _peak_bins_direct(w, n_len))
out[b0:b1] = freq32[bins]
if exact:
spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
power = spectrum.real ** 2
power += spectrum.imag ** 2
power[:, 0] = 0.0 # suppress DC bin
out[:] = freq32[np.argmax(power, axis=1)]
elif pool is None:
for b0 in range(0, n_wf, _FFT_BLOCK):
fft_block(b0)
else:
list(pool.map(fft_block, range(0, n_wf, _FFT_BLOCK)))
# Boolean scatter is row-major, matching the read pass's
# concatenation order.
if valid is not None:
img[r0:r1][valid] = out
else:
img[r0:r1] = out.reshape(r1 - r0, n_frames)
# BLAS must not thread under the pool (the fine-DFT gemm would multiply
# against the pool's own workers).
limiter = (threadpool_limits(limits=1)
if pool is not None and threadpool_limits is not None else None)
try:
for r0 in range(0, n_rows, chunk_rows):
if should_stop is not None and should_stop():
break
process(r0, min(r0 + chunk_rows, n_rows))
finally:
if limiter is not None:
limiter.unregister()
return img