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
sras_compute.py
+94 -232
View File
@@ -1,9 +1,9 @@
#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing
child can import it without loading Qt or matplotlib — which matters
because Python 3.14 on macOS spawns rather than forks.
"""
import atexit
@@ -15,61 +15,34 @@ from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pyfftw
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
adc_to_mv)
# ---------------------------------------------------------------------------
# FFT backend
# ---------------------------------------------------------------------------
try:
import pyfftw
PYFFTW_AVAILABLE = True
except ImportError:
PYFFTW_AVAILABLE = False
try:
from threadpoolctl import threadpool_limits
except ImportError:
threadpool_limits = None
_fft_backend = "scipy" # "scipy" or "pyfftw"; set via set_fft_backend()
def set_fft_backend(name: str):
"""Select the rfft implementation. Module-level state, so it must be set
explicitly inside each multiprocessing child — it does not survive a
spawn. "numpy" is accepted as a legacy alias for "scipy"."""
global _fft_backend
if name == "numpy":
name = "scipy"
_fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy"
def get_fft_backend() -> str:
return _fft_backend
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
workers: int = 1) -> np.ndarray:
"""Dispatch rfft to the selected backend with optional multithreading."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
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))
_FFT_BLOCK_MAX = 512 # ceiling on waveforms/task: the knee measured on a
# 16-core machine at natural resolution (n_len == spf).
# Smaller blocks serialise on GIL-held numpy dispatch,
# larger ones lose cache residency and task granularity.
_FFT_BLOCK_MIN = 32 # floor, so task granularity never collapses at
# extreme pad factors — at the cost of exceeding
# _FFT_PLAN_BYTES_BUDGET there (see _fft_block_for).
_FFT_PLAN_BYTES_BUDGET = int(
os.environ.get("SRAS_FFT_PLAN_BUDGET_MB", 16)) * 1024 * 1024
# Per-thread ceiling on one cached pyFFTW plan's resident input+output
# buffers (see _fft_block_for). This is PERMANENT memory — the plan
# cache is never evicted — so the worst case across the whole pool for
# one distinct (spf, n_len) is _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET.
# Deliberately separate from _TOTAL_BYTES_BUDGET/SRAS_MEM_BUDGET_MB,
# which bounds transient concurrently-live chunk buffers, freed at the
# end of each chunk — this one bounds a permanent per-thread tax that
# compounds with worker count, not chunk concurrency.
_pool_lock = threading.Lock()
_pool: ThreadPoolExecutor | None = None
@@ -115,23 +88,38 @@ def _save_wisdom():
pass
def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
def _fft_block_for(spf: int, n_len: int) -> int:
"""Waveforms per FFT task/plan at transform length n_len. Bounds a
cached pyFFTW plan's resident input+output buffers (permanent, one per
pool thread, never evicted) to _FFT_PLAN_BYTES_BUDGET regardless of pad
factor. Reproduces _FFT_BLOCK_MAX exactly at natural resolution
(n_len == spf) — see docs/design.md for the worked numbers at high pad."""
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
block = _FFT_PLAN_BYTES_BUDGET // max(1, bytes_per_wf)
return int(min(_FFT_BLOCK_MAX, max(_FFT_BLOCK_MIN, block)))
def _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.
Plans have a fixed (block, spf) input shape, block = _fft_block_for(spf,
n), so each worker thread plans once per (block size, transform length)
pair; a remainder block runs through the same plan with its tail rows
ignored. Single-threaded (threads=1): outer parallelism comes from the
pool, so the transform itself must not spin up threads. The returned
array is the plan's output buffer — consume it before the next call on
the same thread.
"""
n_wf, spf = waves.shape
block = _fft_block_for(spf, n)
plans = getattr(_fftw_local, "plans", None)
if plans is None:
plans = _fftw_local.plans = {}
key = (_FFT_BLOCK, spf, n)
key = (block, spf, n)
plan = plans.get(key)
if plan is None:
_load_wisdom_once()
buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32")
buf = pyfftw.empty_aligned((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,
@@ -144,132 +132,20 @@ def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
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
# Peak search: full transform, block-parallel
# ---------------------------------------------------------------------------
_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."""
def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray:
"""Peak bin per waveform via the full transform. *min_bin* zeroes every
bin below it (always >= 1, so true DC stays suppressed) before the
argmax, so the peak search never returns a bin below the caller's
frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks,
fanned out across _fft_pool() — see docs/design.md."""
S = _block_rfft(waves, n_len)
power = S.real ** 2
power += S.imag ** 2
power[:, 0] = 0.0
power[:, :min_bin] = 0.0
return np.argmax(power, axis=1)
@@ -300,9 +176,11 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int,
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."""
waveform read buffer scales with the chunk; spectrum memory is bounded
independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET
regardless of chunk size or pad factor (see docs/design.md), so this
formula budgets the read buffer alone, with 2x slack, and lets the block
fan-out saturate the pool."""
bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
@@ -605,8 +483,8 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
max_workers: int | None = None,
budget: int | None = None,
should_stop=None,
exact: bool = False,
row_avg_n: int = 0) -> np.ndarray:
row_avg_n: int = 0,
min_freq_mhz: float = 0.0) -> 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
@@ -624,22 +502,32 @@ 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.
The peak search runs the full transform (_peak_bins) in
_fft_block_for(spf, n_len)-sized blocks, fanned out across the
pyFFTW-plan-caching thread pool — see docs/design.md for how the block
size is bounded so this stays memory-safe at high pad factors.
*row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n
same-row neighbors (distance-weighted, valid-neighbors-only per the
same dc_threshold_mv mask) before the FFT runs — see
_row_average_waveforms. 0 (default) is the raw, unaveraged behavior.
*min_freq_mhz* excludes every bin below it (true DC, bin 0, is always
excluded regardless) from the peak search, on top of — not instead of —
the dc_threshold_mv mask. Without it, a pixel that passes the DC-bias
threshold but carries only weak real signal can still resolve to a
near-zero frequency: the un-subtracted background's DC-leakage skirt
then has more power than the genuine (but weak) signal peak. Raising
the floor above that skirt forces the search to report the strongest
peak that is plausibly real signal instead. 0.0 (default) disables it
(only true DC is excluded, the long-standing behavior).
Fast path: if the file has a precomputed peak-frequency image for this
angle (v5 PREC or v7 CACH) matching every one of the caller's settings
— including row_avg_n exactly — the stored image is used directly, no
FFT is run. See cached_rf_image.
FFT is run. See cached_rf_image. *min_freq_mhz* plays no part in that
match — a stored image was baked without any floor, so it is returned
as-is; the floor only ever affects a real (re)compute.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
@@ -650,35 +538,27 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
return fast
# ---- Chunked FFT path --------------------------------------------------
exact = exact or _FFT_EXACT_ENV
spf = sras.samples_per_frame
n_len = n_fft if n_fft is not None else spf
block = _fft_block_for(spf, n_len)
freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
# First fine bin at or above the floor; searchsorted is exact here since
# freq32 is the very axis the floor is expressed against.
min_bin = max(1, int(np.searchsorted(freq32, min_freq_mhz)))
img = np.zeros((n_rows, n_frames), dtype=np.float32)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None
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)
if row_avg_n > 0:
# One extra same-sized transient buffer (the pre-averaging full-row
# scratch array) is live per in-flight row; halve the budget so
# _plan_fft_rows/the exact-path sizing accounts for it rather than
# relying on _plan_fft_rows's existing 2x slack to happen to cover it.
# _plan_fft_rows accounts for it rather than relying on its existing
# 2x slack to happen to cover it.
total = max(1, total // 2)
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)
chunk_rows = _plan_fft_rows(n_frames, spf, total)
pool = _fft_pool() if cap > 1 else None
def process(r0: int, r1: int):
@@ -742,23 +622,14 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
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]
b1 = min(b0 + block, n_wf)
out[b0:b1] = freq32[_peak_bins(waves[b0:b1], n_len, min_bin)]
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):
if pool is None:
for b0 in range(0, n_wf, block):
fft_block(b0)
else:
list(pool.map(fft_block, range(0, n_wf, _FFT_BLOCK)))
list(pool.map(fft_block, range(0, n_wf, block)))
# Boolean scatter is row-major, matching the read pass's
# concatenation order.
@@ -767,18 +638,10 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
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()
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))
return img
@@ -787,7 +650,7 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
# ---------------------------------------------------------------------------
def cache_file(path: str, mode: str, apply_bg_sub: bool,
fft_backend: str = "scipy", max_workers: int = 0,
max_workers: int = 0,
pad_factor: int = 1,
dc_threshold_mv: float | None = None,
row_avg_n: int = 0) -> str:
@@ -795,8 +658,8 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
converting v6 → v7 in place. Returns "" on success or an error message.
Module-level and picklable so it can run in a ProcessPoolExecutor. The
FFT backend and worker cap are passed explicitly because module globals
do not survive a spawn.
worker cap is passed explicitly because module globals do not survive a
spawn.
mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at
display time), or "fft_rowavg" (same-row, distance-weighted CH1
@@ -822,7 +685,6 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
# between a bad argument costing nothing and costing the whole run.
if not (1 <= pad_factor <= MAX_PAD_FACTOR):
return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}"
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers