Replace zoom FFT peak search with a budget-bounded PyFFTW direct transform

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
+48 -35
View File
@@ -32,45 +32,58 @@ share of the budget. Capping the workers alone is not enough: the chunk would
still be sized against the whole budget, and N concurrent callers would each still be sized against the whole budget, and N concurrent callers would each
allocate all of it. allocate all of it.
## FFT peak search: block-parallel zoom refinement (`sras_compute.py`) ## FFT peak search: block-parallel direct transform (`sras_compute.py`)
The displayed RF value per pixel is the argmax of the zero-padded power The displayed RF value per pixel is the argmax of the zero-padded power
spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for spectrum of that pixel's CH1 waveform. This used to run through `_peak_bins_zoom`,
mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan a coarse-rfft-plus-local-fine-DFT refinement that avoided ever materialising
row, which is what used to collapse the old row-chunk planner to one worker a padded spectrum — at the pad factor of 40 needed for mapping resolution, a
and make synthesis single-threaded. full padded spectrum is ~9 GB per scan row, which used to collapse the old
row-chunk planner to one worker and make synthesis single-threaded. That
refinement was removed once pyFFTW became the sole, mandatory FFT backend
(SciPy dropped as a compute backend entirely): `_peak_bins` now always runs
the real full transform, and the memory problem zoom dodged is instead
solved by bounding *per-block* spectrum memory rather than avoiding full
spectra altogether.
`_peak_bins_zoom` never materialises the padded spectrum: `_peak_bins` runs the full transform (`_block_rfft`, a cached pyFFTW
`builders.rfft` plan, FFTW_MEASURE, wisdom persisted under
`~/.cache/sras-viewer/`) and argmaxes the power spectrum, in blocks fanned
out task-parallel over a persistent thread pool (`_fft_pool()`) — each pool
thread runs one single-threaded transform at a time, so aggregate
parallelism equals the pool's worker count. This block-streaming structure
is not just about parallelism: it is also what keeps memory bounded, by
never materialising more than one block's worth of full padded spectrum at
a time, regardless of how many waveforms a chunk holds.
1. a coarse rfft at `next_fast_len(2*spf)` — 2× oversampled, so the padded The block size is the part that has to adapt to pad factor.
power spectrum (a trig polynomial of degree spf−1) cannot hide its global `_fft_block_for(spf, n_len)` derives waveforms-per-task from a fixed
max between coarse samples; per-thread byte budget (`_FFT_PLAN_BYTES_BUDGET`, `SRAS_FFT_PLAN_BUDGET_MB`,
2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max default 16 MB) rather than a fixed constant, because a cached pyFFTW plan's
becomes a refinement candidate. Quarter-natural-bin scalloping at the 2× input+output buffers are *permanent* per-thread memory (the plan cache is
grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide never evicted) — with a fixed 512-waveform block, pad 40 at a 2500-sample
margin. The DC-adjacent window is always refined too: the coarse DC bin frame costs ~210 MB per pool thread (~3.3 GB total across 16 threads);
is zeroed for suppression, which would otherwise blind the scan to fine `_fft_block_for` bounds that to ~16 MB per thread (~260 MB across 16
bins closer to DC than the first coarse sample (where the leakage skirt threads) at the same pad factor, while still reproducing the old tuned 512
of an un-subtracted offset peaks); exactly at natural resolution (pad 1), where it cost nothing to begin with.
3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every `_FFT_BLOCK_MAX` (512) and `_FFT_BLOCK_MIN` (32) cap and floor the result:
fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated the ceiling is the measured knee on a 16-core machine at natural resolution
on the exact `n_fft` grid by one small complex gemm, with np.argmax's (smaller blocks serialise on GIL-held numpy dispatch, larger ones lose cache
lowest-bin tie-break preserved across windows. residency and task granularity); the floor keeps task granularity from
collapsing at extreme pad factors, at the cost of exceeding the byte budget
there.
The selected bin is bit-identical to the full padded argmax — enforced by The outer row-chunk sizing (`_plan_fft_rows`) needed no companion change.
`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial It only ever budgets the raw float32 waveform *read* buffer, which this
spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose change doesn't touch — spectrum memory is bounded independently by
baseline was captured on the old full-padded path. `_fft_block_for`, and since the pool only ever runs as many blocks
concurrently as it has workers, peak transient spectrum memory during a
chunk's FFT phase is `_MAX_WORKERS * block * bytes_per_wf`, the same bound
whether the chunk holds 10 rows or 10,000. Queuing more rows into one chunk
to keep the read-row pool busy therefore can't blow up spectrum memory.
Work fans out over a persistent thread pool in `_FFT_BLOCK` = 512-waveform `tools/check_equivalence.py`'s golden-hash harness remains the end-to-end
tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose regression baseline for this path, unaffected by this change.
cache residency and task granularity (measured on a 16-core machine, where
this path runs ~35× faster than the old serial padded transform at pad 40).
pyFFTW runs through per-thread `builders` plans (FFTW_MEASURE, wisdom
persisted under `~/.cache/sras-viewer/`), and `threadpoolctl` clamps BLAS to
one thread under the pool so the refinement gemm cannot oversubscribe.
`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
full-padded path for audits.
## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`) ## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
@@ -121,8 +134,8 @@ same edge case the masked convolution already handles.
The averaging step doubles the live per-row scratch memory (a full-width The averaging step doubles the live per-row scratch memory (a full-width
`(n_frames, spf)` buffer on top of the existing compacted `waves` buffer), `(n_frames, spf)` buffer on top of the existing compacted `waves` buffer),
so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before
`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row `_plan_fft_rows` runs — see "Memory budget and row chunking" above. On the
chunking" above. On the largest real scans `_plan_chunks` is already largest real scans `_plan_chunks` is already
clamped to its floor of one row regardless, so this costs no concurrency clamped to its floor of one row regardless, so this costs no concurrency
where it matters most; it mainly protects moderate-sized scans from an where it matters most; it mainly protects moderate-sized scans from an
unexpected regression. unexpected regression.
+1 -3
View File
@@ -14,10 +14,8 @@ dependencies = [
"scipy==1.18.0", "scipy==1.18.0",
# Angle alignment only: masked FFT phase correlation (skimage.registration). # Angle alignment only: masked FFT phase correlation (skimage.registration).
"scikit-image==0.26.0", "scikit-image==0.26.0",
# Faster rfft backend; the viewer falls back to scipy.fft without it. # Mandatory rfft backend for the peak search (no SciPy fallback).
"pyFFTW==0.15.1", "pyFFTW==0.15.1",
# Clamps BLAS threading under the FFT worker pool.
"threadpoolctl==3.6.0",
] ]
[project.optional-dependencies] [project.optional-dependencies]
+88 -226
View File
@@ -1,9 +1,9 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans. """Image computation and angle alignment for .sras scans.
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing
multiprocessing child can import it without loading Qt or matplotlib — child can import it without loading Qt or matplotlib — which matters
which matters because Python 3.14 on macOS spawns rather than forks. because Python 3.14 on macOS spawns rather than forks.
""" """
import atexit import atexit
@@ -15,61 +15,34 @@ from dataclasses import dataclass
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
import pyfftw
import scipy.fft as scipy_fft import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage import scipy.ndimage as scipy_ndimage
from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile, from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
adc_to_mv) 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 worker pool and per-thread pyFFTW plans
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine: _FFT_BLOCK_MAX = 512 # ceiling on waveforms/task: the knee measured on a
# smaller blocks serialise on GIL-held numpy dispatch, # 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. # larger ones lose cache residency and task granularity.
_ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf _FFT_BLOCK_MIN = 32 # floor, so task granularity never collapses at
_FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0)) # 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_lock = threading.Lock()
_pool: ThreadPoolExecutor | None = None _pool: ThreadPoolExecutor | None = None
@@ -115,23 +88,38 @@ def _save_wisdom():
pass 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. """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 have a fixed (block, spf) input shape, block = _fft_block_for(spf,
plans once per transform length; a remainder block runs through the same n), so each worker thread plans once per (block size, transform length)
plan with its tail rows ignored. The returned array is the plan's output pair; a remainder block runs through the same plan with its tail rows
buffer — consume it before the next call on the same thread. 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 n_wf, spf = waves.shape
block = _fft_block_for(spf, n)
plans = getattr(_fftw_local, "plans", None) plans = getattr(_fftw_local, "plans", None)
if plans is None: if plans is None:
plans = _fftw_local.plans = {} plans = _fftw_local.plans = {}
key = (_FFT_BLOCK, spf, n) key = (block, spf, n)
plan = plans.get(key) plan = plans.get(key)
if plan is None: if plan is None:
_load_wisdom_once() _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 # No overwrite_input: FFTW must not scribble on input_array, whose
# zero-padded tail (columns spf..n) is zeroed exactly once here. # zero-padded tail (columns spf..n) is zeroed exactly once here.
plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1, 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] 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. def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray:
# Every fine bin lies within 0.5 spacings of its """Peak bin per waveform via the full transform. *min_bin* zeroes every
# nearest coarse bin, and that bin is guaranteed to bin below it (always >= 1, so true DC stays suppressed) before the
# be a candidate (see _ZOOM_CAND_RATIO), so 0.5 argmax, so the peak search never returns a bin below the caller's
# suffices; 0.75 adds rounding margin. frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks,
_ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of fanned out across _fft_pool() — see docs/design.md."""
# 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) S = _block_rfft(waves, n_len)
power = S.real ** 2 power = S.real ** 2
power += S.imag ** 2 power += S.imag ** 2
power[:, 0] = 0.0 power[:, :min_bin] = 0.0
return np.argmax(power, axis=1) 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: 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 """Rows per outer chunk for the block-FFT path. Only the float32
waveform buffer scales with the chunk (in-flight block spectra total a waveform read buffer scales with the chunk; spectrum memory is bounded
few MB across the whole pool), so budget it with 2x slack and let the independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET
block fan-out saturate the pool regardless of pad factor.""" 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) bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row))) 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, max_workers: int | None = None,
budget: int | None = None, budget: int | None = None,
should_stop=None, should_stop=None,
exact: bool = False, row_avg_n: int = 0,
row_avg_n: int = 0) -> np.ndarray: min_freq_mhz: float = 0.0) -> np.ndarray:
"""FFT of each CH1 waveform; pixel = peak frequency in MHz. """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 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 DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here. (n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
At pad factors >= _ZOOM_MIN_PAD the padded spectrum is never The peak search runs the full transform (_peak_bins) in
materialised: a coarse rfft finds each peak and a local fine DFT _fft_block_for(spf, n_len)-sized blocks, fanned out across the
resolves it on the exact n_fft bin grid (_peak_bins_zoom). *exact* pyFFTW-plan-caching thread pool — see docs/design.md for how the block
forces the reference full-padded transform instead — it exists for size is bounded so this stays memory-safe at high pad factors.
tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and
memory-hungry at high pad.
*row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n *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-row neighbors (distance-weighted, valid-neighbors-only per the
same dc_threshold_mv mask) before the FFT runs — see same dc_threshold_mv mask) before the FFT runs — see
_row_average_waveforms. 0 (default) is the raw, unaveraged behavior. _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 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 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 — 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) n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx] data = sras.data[angle_idx]
@@ -650,34 +538,26 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
return fast return fast
# ---- Chunked FFT path -------------------------------------------------- # ---- Chunked FFT path --------------------------------------------------
exact = exact or _FFT_EXACT_ENV
spf = sras.samples_per_frame spf = sras.samples_per_frame
n_len = n_fft if n_fft is not None else spf 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) 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) 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 background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX) cal4 = sras.cal(CH4_IDX)
row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None 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) total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
if row_avg_n > 0: if row_avg_n > 0:
# One extra same-sized transient buffer (the pre-averaging full-row # One extra same-sized transient buffer (the pre-averaging full-row
# scratch array) is live per in-flight row; halve the budget so # scratch array) is live per in-flight row; halve the budget so
# _plan_fft_rows/the exact-path sizing accounts for it rather than # _plan_fft_rows accounts for it rather than relying on its existing
# relying on _plan_fft_rows's existing 2x slack to happen to cover it. # 2x slack to happen to cover it.
total = max(1, total // 2) total = max(1, total // 2)
cap = max_workers if max_workers is not None else _MAX_WORKERS 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 pool = _fft_pool() if cap > 1 else None
@@ -742,23 +622,14 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
def fft_block(b0: int): def fft_block(b0: int):
if should_stop is not None and should_stop(): if should_stop is not None and should_stop():
return return
b1 = min(b0 + _FFT_BLOCK, n_wf) b1 = min(b0 + block, n_wf)
w = waves[b0:b1] out[b0:b1] = freq32[_peak_bins(waves[b0:b1], n_len, min_bin)]
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: if pool is None:
spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) for b0 in range(0, n_wf, block):
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) fft_block(b0)
else: 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 # Boolean scatter is row-major, matching the read pass's
# concatenation order. # concatenation order.
@@ -767,18 +638,10 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
else: else:
img[r0:r1] = out.reshape(r1 - r0, n_frames) 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): for r0 in range(0, n_rows, chunk_rows):
if should_stop is not None and should_stop(): if should_stop is not None and should_stop():
break break
process(r0, min(r0 + chunk_rows, n_rows)) process(r0, min(r0 + chunk_rows, n_rows))
finally:
if limiter is not None:
limiter.unregister()
return img 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, 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, pad_factor: int = 1,
dc_threshold_mv: float | None = None, dc_threshold_mv: float | None = None,
row_avg_n: int = 0) -> str: 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. converting v6 → v7 in place. Returns "" on success or an error message.
Module-level and picklable so it can run in a ProcessPoolExecutor. The Module-level and picklable so it can run in a ProcessPoolExecutor. The
FFT backend and worker cap are passed explicitly because module globals worker cap is passed explicitly because module globals do not survive a
do not survive a spawn. spawn.
mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at
display time), or "fft_rowavg" (same-row, distance-weighted CH1 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. # between a bad argument costing nothing and costing the whole run.
if not (1 <= pad_factor <= MAX_PAD_FACTOR): if not (1 <= pad_factor <= MAX_PAD_FACTOR):
return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}" return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}"
set_fft_backend(fft_backend)
if max_workers: if max_workers:
_MAX_WORKERS = max_workers _MAX_WORKERS = max_workers
+9 -2
View File
@@ -166,11 +166,14 @@ class ImageCanvas(FigureCanvasQTAgg):
def show_image(self, img: np.ndarray, extent: list[float], cmap, def show_image(self, img: np.ndarray, extent: list[float], cmap,
vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
colorbar_label: str = "", cb_ticks=None, norm=None): colorbar_label: str = "", cb_ticks=None, norm=None,
bad_color=None):
"""*cmap* may be a name or a Colormap instance. *norm* (which overrides """*cmap* may be a name or a Colormap instance. *norm* (which overrides
vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image — vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image —
the wizard's overlap-count view — with whole-number colorbar bands the wizard's overlap-count view — with whole-number colorbar bands
instead of a continuous shade.""" instead of a continuous shade. *bad_color*, if given, is the fill for
NaN pixels — a copy of *cmap* is made so the shared, registered
instance is never mutated."""
self.figure.clf() self.figure.clf()
self.ax = self.figure.add_subplot(111) self.ax = self.figure.add_subplot(111)
# Patches and lines are destroyed by figure.clf(); drop stale refs. # Patches and lines are destroyed by figure.clf(); drop stale refs.
@@ -179,6 +182,10 @@ class ImageCanvas(FigureCanvasQTAgg):
self._extent = extent self._extent = extent
self._img_shape = img.shape self._img_shape = img.shape
if bad_color is not None:
cmap = (cmap if hasattr(cmap, "with_extremes")
else mpl.colormaps[cmap]).with_extremes(bad=bad_color)
kw = ({"norm": norm} if norm is not None kw = ({"norm": norm} if norm is not None
else {"vmin": vmin, "vmax": vmax}) else {"vmin": vmin, "vmax": vmax})
im = self.ax.imshow( im = self.ax.imshow(
+4
View File
@@ -26,6 +26,10 @@ CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX)
CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"]
# Fill color for masked (below-DC-threshold) pixels when "Highlight masked
# pixels" is on, chosen to stand out against every colormap in CMAPS above.
_MASKED_HIGHLIGHT_COLOR = "magenta"
# (mode_str, status-bar unit, colorbar label) per channel index # (mode_str, status-bar unit, colorbar label) per channel index
_CHANNEL_DISPLAY = { _CHANNEL_DISPLAY = {
CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"),
+1 -29
View File
@@ -13,7 +13,6 @@ from PyQt6.QtWidgets import (
QScrollArea, QSpinBox, QVBoxLayout, QWidget, QScrollArea, QSpinBox, QVBoxLayout, QWidget,
) )
from sras_compute import PYFFTW_AVAILABLE
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES
from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label
@@ -24,7 +23,7 @@ from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group,
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class FftOptionsDialog(QDialog): class FftOptionsDialog(QDialog):
"""Configure FFT backend and zero-padding. """Configure FFT zero-padding.
Changes take effect only when the user clicks Apply. Cancel discards Changes take effect only when the user clicks Apply. Cancel discards
all pending edits. The live 'frequency resolution' label updates as all pending edits. The live 'frequency resolution' label updates as
@@ -33,7 +32,6 @@ class FftOptionsDialog(QDialog):
""" """
def __init__(self, parent=None, *, def __init__(self, parent=None, *,
current_backend: str,
current_pad_factor: int, current_pad_factor: int,
samples_per_frame: int | None, samples_per_frame: int | None,
sample_rate_hz: float | None, sample_rate_hz: float | None,
@@ -49,29 +47,6 @@ class FftOptionsDialog(QDialog):
layout = QVBoxLayout(self) layout = QVBoxLayout(self)
# ---- Backend ---------------------------------------------------
grp_backend = QGroupBox("FFT Backend")
bl = QVBoxLayout(grp_backend)
self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)")
self._btn_pyfftw = QRadioButton(
"pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE
else "pyFFTW (not installed — run: pip install pyfftw)")
self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE)
self._backend_group = QButtonGroup(self)
self._backend_group.addButton(self._btn_scipy, id=0)
self._backend_group.addButton(self._btn_pyfftw, id=1)
if current_backend == "pyfftw" and PYFFTW_AVAILABLE:
self._btn_pyfftw.setChecked(True)
else:
self._btn_scipy.setChecked(True)
bl.addWidget(self._btn_scipy)
bl.addWidget(self._btn_pyfftw)
layout.addWidget(grp_backend)
# ---- Zero-padding ---------------------------------------------- # ---- Zero-padding ----------------------------------------------
grp_zp = QGroupBox("Zero-Padding") grp_zp = QGroupBox("Zero-Padding")
zl = QVBoxLayout(grp_zp) zl = QVBoxLayout(grp_zp)
@@ -134,9 +109,6 @@ class FftOptionsDialog(QDialog):
f"Velocity bin: {vel_res_ms:.3f} m/s " f"Velocity bin: {vel_res_ms:.3f} m/s "
f"(at grating = {self._grating_um:.2f} µm)") f"(at grating = {self._grating_um:.2f} µm)")
def get_backend(self) -> str:
return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy"
def get_pad_factor(self) -> int: def get_pad_factor(self) -> int:
return max(1, self._spin_pad.value()) return max(1, self._spin_pad.value())
+118 -22
View File
@@ -30,8 +30,8 @@ from .canvases import ImageCanvas, WaveformCanvas
from .common import ( from .common import (
CH1_DERIVED_MODES, CH_LABELS, CMAPS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY, CH1_DERIVED_MODES, CH_LABELS, CMAPS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY,
_CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W, _CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W,
_RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, _scroll_panel, _MASKED_HIGHLIGHT_COLOR, _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin,
_wrap_label, _scroll_panel, _wrap_label,
) )
from .align_wizard import AlignmentWizard from .align_wizard import AlignmentWizard
from .dialogs import ( from .dialogs import (
@@ -58,6 +58,7 @@ class SrasViewerWindow(QMainWindow):
self._pending_ch: int = 0 self._pending_ch: int = 0
self._pending_bg_sub: bool = True self._pending_bg_sub: bool = True
self._pending_threshold: float = 50.0 # mV self._pending_threshold: float = 50.0 # mV
self._pending_min_freq_mhz: float = 0.0
self._pending_fft_pad_factor: int = 1 self._pending_fft_pad_factor: int = 1
# Live background jobs, keyed by role — see _run_worker. # Live background jobs, keyed by role — see _run_worker.
@@ -70,7 +71,6 @@ class SrasViewerWindow(QMainWindow):
self._settings = QSettings(QSettings.Format.IniFormat, self._settings = QSettings(QSettings.Format.IniFormat,
QSettings.Scope.UserScope, QSettings.Scope.UserScope,
"sras-viewer", "sras-viewer") "sras-viewer", "sras-viewer")
compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy")))
try: try:
pad = int(self._settings.value("fft/pad_factor", 1)) pad = int(self._settings.value("fft/pad_factor", 1))
except (TypeError, ValueError): except (TypeError, ValueError):
@@ -272,6 +272,28 @@ class SrasViewerWindow(QMainWindow):
self.lbl_threshold_adc = _wrap_label( self.lbl_threshold_adc = _wrap_label(
f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED)
tl.addWidget(self.lbl_threshold_adc) tl.addWidget(self.lbl_threshold_adc)
min_freq_form = _form()
self.spin_min_freq_mhz = _make_dspin(0.0, 10000.0, 3, suffix=" MHz",
value=0.0, step=1.0)
self.spin_min_freq_mhz.setEnabled(False)
self.spin_min_freq_mhz.setToolTip(
"Excludes every FFT bin below this frequency from the peak\n"
"search (0 = off, only true DC is excluded). A pixel can pass\n"
"the DC threshold above yet still carry only weak real signal —\n"
"when that happens, the un-subtracted background's DC-leakage\n"
"skirt can have more power than the genuine signal, so the peak\n"
"search resolves to a near-zero frequency even though the pixel\n"
"is real, valid data. Raising this floor above that skirt forces\n"
"the search to report the strongest peak that is plausibly real\n"
"signal instead.\n"
"Only affects a fresh/live compute — it cannot change an image\n"
"already shown, or one already stored in this file's own cache\n"
"(use Batch Compute to regenerate those)."
)
self.spin_min_freq_mhz.editingFinished.connect(self._on_min_freq_changed)
min_freq_form.addRow("Min peak freq:", self.spin_min_freq_mhz)
tl.addLayout(min_freq_form)
vl.addWidget(self.grp_threshold) vl.addWidget(self.grp_threshold)
# Background subtraction (v4+ files only) # Background subtraction (v4+ files only)
@@ -417,6 +439,22 @@ class SrasViewerWindow(QMainWindow):
self.chk_auto.toggled.connect(self._on_autoscale_toggled) self.chk_auto.toggled.connect(self._on_autoscale_toggled)
dl.addWidget(self.chk_auto) dl.addWidget(self.chk_auto)
self.chk_highlight_masked = QCheckBox("Highlight masked (no-data) pixels")
self.chk_highlight_masked.setChecked(True)
self.chk_highlight_masked.setEnabled(False)
self.chk_highlight_masked.setToolTip(
"RF/Velocity only. A pixel below the DC threshold has no FFT\n"
"result at all and is otherwise shown as 0 on the same grayscale\n"
"ramp as a real, valid pixel whose peak just happens to be a low\n"
"frequency — the two look identical (both near-black). When\n"
"checked, masked pixels are drawn in a distinct highlight color\n"
"instead, excluded from the colormap's data range, so real\n"
"low-frequency pixels keep their own true shade. Uncheck to\n"
"restore the old behavior where both blend together."
)
self.chk_highlight_masked.toggled.connect(self._on_highlight_masked_toggled)
dl.addWidget(self.chk_highlight_masked)
range_form = _form() range_form = _form()
for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")):
spin = _make_dspin(-1e9, 1e9, 4) spin = _make_dspin(-1e9, 1e9, 4)
@@ -440,7 +478,7 @@ class SrasViewerWindow(QMainWindow):
fft_menu = menubar.addMenu("&FFT") fft_menu = menubar.addMenu("&FFT")
fft_act = QAction("FFT &Options…", self) fft_act = QAction("FFT &Options…", self)
fft_act.setStatusTip("Configure FFT backend and zero-padding") fft_act.setStatusTip("Configure FFT zero-padding")
fft_act.triggered.connect(self._on_fft_options) fft_act.triggered.connect(self._on_fft_options)
fft_menu.addAction(fft_act) fft_menu.addAction(fft_act)
@@ -684,7 +722,9 @@ class SrasViewerWindow(QMainWindow):
# Threshold and bg-sub apply to all CH1 modes # Threshold and bg-sub apply to all CH1 modes
self.spin_threshold_mv.setEnabled(is_ch1) self.spin_threshold_mv.setEnabled(is_ch1)
self.spin_min_freq_mhz.setEnabled(is_ch1)
self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1)
self.chk_highlight_masked.setEnabled(is_ch1)
self.spin_grating_um.setEnabled(is_vel) self.spin_grating_um.setEnabled(is_vel)
self.grp_velocity.setVisible(is_vel) self.grp_velocity.setVisible(is_vel)
@@ -733,6 +773,14 @@ class SrasViewerWindow(QMainWindow):
if self._is_fft_mode(): if self._is_fft_mode():
self._refresh_display() self._refresh_display()
def _on_min_freq_changed(self):
# Like the DC threshold, this changes what the FFT itself produces —
# a genuine cache-key change — but unlike the threshold it can't be
# re-applied to an already-computed image; it only takes effect on a
# fresh compute (see _fft_cache_key / compute_rf_image's docstring).
if self._is_fft_mode():
self._refresh_display()
def _on_autoscale_toggled(self, checked: bool): def _on_autoscale_toggled(self, checked: bool):
manual = not checked manual = not checked
self.spin_vmin.setEnabled(manual and self._sras is not None) self.spin_vmin.setEnabled(manual and self._sras is not None)
@@ -740,6 +788,10 @@ class SrasViewerWindow(QMainWindow):
if self._sras is not None and self._current_image is not None: if self._sras is not None and self._current_image is not None:
self._redraw_image(self._current_image) self._redraw_image(self._current_image)
def _on_highlight_masked_toggled(self, checked: bool):
if self._sras is not None and self._current_image is not None:
self._redraw_image(self._current_image)
def _on_manual_range_changed(self): def _on_manual_range_changed(self):
if not self.chk_auto.isChecked() and self._current_image is not None: if not self.chk_auto.isChecked() and self._current_image is not None:
self._redraw_image(self._current_image) self._redraw_image(self._current_image)
@@ -1014,15 +1066,19 @@ class SrasViewerWindow(QMainWindow):
return freq_mhz return freq_mhz
def _fft_cache_key(self, angle_idx: int) -> tuple: def _fft_cache_key(self, angle_idx: int) -> tuple:
"""Keyed by angle and DC threshold only. Once any FFT image exists """Keyed by angle, DC threshold, and min peak frequency only. Once
for an angle this session — live-computed or pulled from the file's any FFT image exists for an angle this session — live-computed or
own stored cache — it stays the displayed image for that angle pulled from the file's own stored cache — it stays the displayed
regardless of later bg-sub/pad toggles; those only affect a future image for that angle regardless of later bg-sub/pad toggles; those
live compute for an angle with nothing cached yet, or an explicit only affect a future live compute for an angle with nothing cached
batch recompute (see _stored_fft_image). Threshold stays in the key yet, or an explicit batch recompute (see _stored_fft_image).
because re-masking against it is free and meant to stay interactive Threshold and min-freq stay in the key so revisiting a combination
(see _on_threshold_changed).""" already computed this session is instant — even though only
return (angle_idx, self.spin_threshold_mv.value()) threshold can be cheaply re-applied to a stored image; a min-freq
change against a stored image still falls through to
_stored_fft_image, which ignores it (see _on_min_freq_changed)."""
return (angle_idx, self.spin_threshold_mv.value(),
self.spin_min_freq_mhz.value())
def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple:
"""Mirrors _fft_cache's key granularity so a stale aligned image is """Mirrors _fft_cache's key granularity so a stale aligned image is
@@ -1100,8 +1156,9 @@ class SrasViewerWindow(QMainWindow):
produces. See _cache_mismatch_notes for the informational (non- produces. See _cache_mismatch_notes for the informational (non-
blocking) note when the live controls diverge from what's shown. blocking) note when the live controls diverge from what's shown.
Only the DC threshold is taken live: re-masking a stored image Only the DC threshold is taken live: re-masking a stored image
against it is free, unlike bg-sub/pad/row-averaging which are baked against it is free, unlike bg-sub/pad/row-averaging — or the min
irreversibly into the stored numbers. peak frequency floor — which are baked irreversibly into the stored
numbers.
allow_dc_recompute=False keeps this off the I/O path: if the mask allow_dc_recompute=False keeps this off the I/O path: if the mask
would mean reading a whole CH4 channel, this declines and the caller would mean reading a whole CH4 channel, this declines and the caller
@@ -1194,6 +1251,24 @@ class SrasViewerWindow(QMainWindow):
self._redraw_image(img) self._redraw_image(img)
self._update_roi_ui() self._update_roi_ui()
def _dc_validity_mask(self, angle_idx: int, aligned: bool) -> np.ndarray | None:
"""True where a CH1/Velocity pixel passed the DC threshold and so
actually has a real FFT result, on the same grid as display_img.
None if the CH4 DC image for this angle isn't cached yet — the
background DC precompute hasn't reached it, and this is deliberately
not worth a synchronous recompute just to redraw."""
dc4 = self._dc_cache.get((angle_idx, CH4_IDX))
if dc4 is None:
return None
valid = dc4 >= self.spin_threshold_mv.value()
if aligned:
# apply_alignment defaults to nearest-neighbor (order=0), so a
# 0.0/1.0 float warp stays exactly 0 or 1 - no blending at mask
# edges to second-guess with a >= 0.5 cutoff.
valid = apply_alignment(self._alignment_result, angle_idx,
valid.astype(np.float32)) >= 0.5
return valid
def _redraw_image(self, img: np.ndarray): def _redraw_image(self, img: np.ndarray):
s = self._sras s = self._sras
angle_idx = self._current_angle angle_idx = self._current_angle
@@ -1214,8 +1289,27 @@ class SrasViewerWindow(QMainWindow):
dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
extent = _axes_extent(x_axis, y_axis, dx, dy) extent = _axes_extent(x_axis, y_axis, dx, dy)
# Masked-out (below-threshold) pixels are stored as a plain 0, the
# same value a real but low-frequency pixel can legitimately have -
# the two are indistinguishable once both land near the bottom of a
# linear colormap. Pull masked pixels out to NaN (drawn in a
# distinct highlight color, excluded from the auto-scale range) so a
# real low-frequency pixel keeps its own true shade instead of
# disappearing into the same black as "no data".
highlight_masked = (ch_idx in CH1_DERIVED_MODES
and self.chk_highlight_masked.isChecked())
mask_valid = (self._dc_validity_mask(angle_idx, aligned)
if highlight_masked else None)
if mask_valid is not None and mask_valid.shape == display_img.shape:
display_img = display_img.astype(np.float32, copy=True)
display_img[~mask_valid] = np.nan
else:
mask_valid = None
if self.chk_auto.isChecked(): if self.chk_auto.isChecked():
vmin, vmax = float(display_img.min()), float(display_img.max()) vmin, vmax = float(np.nanmin(display_img)), float(np.nanmax(display_img))
if not np.isfinite(vmin):
vmin, vmax = 0.0, 0.0 # every pixel masked out
for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)):
with QSignalBlocker(spin): with QSignalBlocker(spin):
spin.setValue(val) spin.setValue(val)
@@ -1239,6 +1333,7 @@ class SrasViewerWindow(QMainWindow):
vmin=vmin, vmax=vmax, vmin=vmin, vmax=vmax,
xlabel="X (mm)", ylabel="Y (mm)", xlabel="X (mm)", ylabel="Y (mm)",
title=title, colorbar_label=colorbar_label, title=title, colorbar_label=colorbar_label,
bad_color=_MASKED_HIGHLIGHT_COLOR if mask_valid is not None else None,
) )
self.statusBar().showMessage( self.statusBar().showMessage(
f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° "
@@ -1262,6 +1357,7 @@ class SrasViewerWindow(QMainWindow):
self._pending_ch = ch_idx self._pending_ch = ch_idx
self._pending_bg_sub = self.chk_bg_sub.isChecked() self._pending_bg_sub = self.chk_bg_sub.isChecked()
self._pending_threshold = self.spin_threshold_mv.value() self._pending_threshold = self.spin_threshold_mv.value()
self._pending_min_freq_mhz = self.spin_min_freq_mhz.value()
self._pending_fft_pad_factor = self._fft_pad_factor self._pending_fft_pad_factor = self._fft_pad_factor
worker = ComputeWorker( worker = ComputeWorker(
@@ -1274,6 +1370,7 @@ class SrasViewerWindow(QMainWindow):
# need to re-read the CH4 channel from disk. # need to re-read the CH4 channel from disk.
dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
is_fft_mode=is_fft, is_fft_mode=is_fft,
min_freq_mhz=self._pending_min_freq_mhz,
) )
if not self._run_worker( if not self._run_worker(
Jobs.COMPUTE, worker, Jobs.COMPUTE, worker,
@@ -1302,9 +1399,10 @@ class SrasViewerWindow(QMainWindow):
already be cached.""" already be cached."""
if (self.spin_angle.value(), self.combo_channel.currentIndex(), if (self.spin_angle.value(), self.combo_channel.currentIndex(),
self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(),
self._fft_pad_factor) != ( self.spin_min_freq_mhz.value(), self._fft_pad_factor) != (
self._pending_angle, self._pending_ch, self._pending_bg_sub, self._pending_angle, self._pending_ch, self._pending_bg_sub,
self._pending_threshold, self._pending_fft_pad_factor): self._pending_threshold, self._pending_min_freq_mhz,
self._pending_fft_pad_factor):
self._refresh_display() self._refresh_display()
def _on_compute_done(self, result): def _on_compute_done(self, result):
@@ -1315,7 +1413,8 @@ class SrasViewerWindow(QMainWindow):
ch_idx = self._pending_ch ch_idx = self._pending_ch
if ch_idx in CH1_DERIVED_MODES: if ch_idx in CH1_DERIVED_MODES:
self._fft_cache[(angle_idx, self._pending_threshold)] = result key = (angle_idx, self._pending_threshold, self._pending_min_freq_mhz)
self._fft_cache[key] = result
img = self._scale_for_display(result, ch_idx) img = self._scale_for_display(result, ch_idx)
else: else:
img = result img = result
@@ -1620,7 +1719,6 @@ class SrasViewerWindow(QMainWindow):
def _on_fft_options(self): def _on_fft_options(self):
dlg = FftOptionsDialog( dlg = FftOptionsDialog(
self, self,
current_backend=compute.get_fft_backend(),
current_pad_factor=self._fft_pad_factor, current_pad_factor=self._fft_pad_factor,
samples_per_frame=self._sras.samples_per_frame if self._sras else None, samples_per_frame=self._sras.samples_per_frame if self._sras else None,
sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, sample_rate_hz=self._sras.sample_rate_hz if self._sras else None,
@@ -1628,9 +1726,7 @@ class SrasViewerWindow(QMainWindow):
) )
if dlg.exec() != QDialog.DialogCode.Accepted: if dlg.exec() != QDialog.DialogCode.Accepted:
return return
compute.set_fft_backend(dlg.get_backend())
self._fft_pad_factor = dlg.get_pad_factor() self._fft_pad_factor = dlg.get_pad_factor()
self._settings.setValue("fft/backend", compute.get_fft_backend())
self._settings.setValue("fft/pad_factor", self._fft_pad_factor) self._settings.setValue("fft/pad_factor", self._fft_pad_factor)
# Pad factor no longer gates the display: it only affects a future # Pad factor no longer gates the display: it only affects a future
# live compute for an angle with nothing cached yet, or an explicit # live compute for an angle with nothing cached yet, or an explicit
+6 -4
View File
@@ -115,7 +115,8 @@ class ComputeWorker(CancellableWorker):
apply_bg_sub: bool = True, n_fft: int | None = None, apply_bg_sub: bool = True, n_fft: int | None = None,
dc_threshold_mv: float = 0.0, dc_threshold_mv: float = 0.0,
dc4_mv: np.ndarray | None = None, dc4_mv: np.ndarray | None = None,
is_fft_mode: bool = False): is_fft_mode: bool = False,
min_freq_mhz: float = 0.0):
super().__init__() super().__init__()
self._sras = sras self._sras = sras
self._angle = angle_idx self._angle = angle_idx
@@ -125,6 +126,7 @@ class ComputeWorker(CancellableWorker):
self._dc_threshold = dc_threshold_mv self._dc_threshold = dc_threshold_mv
self._dc4_mv = dc4_mv self._dc4_mv = dc4_mv
self._is_fft_mode = is_fft_mode self._is_fft_mode = is_fft_mode
self._min_freq_mhz = min_freq_mhz
def run(self): def run(self):
try: try:
@@ -132,7 +134,8 @@ class ComputeWorker(CancellableWorker):
img = compute_rf_image( img = compute_rf_image(
self._sras, self._angle, dc_threshold_mv=self._dc_threshold, self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
dc4_mv=self._dc4_mv, should_stop=self._stopped) dc4_mv=self._dc4_mv, should_stop=self._stopped,
min_freq_mhz=self._min_freq_mhz)
else: else:
img = dc_image_mv(self._sras, self._angle, self._ch, img = dc_image_mv(self._sras, self._angle, self._ch,
should_stop=self._stopped) should_stop=self._stopped)
@@ -240,7 +243,7 @@ class BatchCacheWorker(QObject):
with ProcessPoolExecutor(max_workers=n_procs) as executor: with ProcessPoolExecutor(max_workers=n_procs) as executor:
futures = { futures = {
executor.submit(cache_file, p, self._mode, self._apply_bg_sub, executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), per_proc_workers, per_proc_workers,
pad_factor=self._pad_factor, pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold, dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n): p row_avg_n=self._row_avg_n): p
@@ -266,7 +269,6 @@ class BatchCacheWorker(QObject):
for path in paths: for path in paths:
try: try:
err = cache_file(path, self._mode, self._apply_bg_sub, err = cache_file(path, self._mode, self._apply_bg_sub,
compute.get_fft_backend(),
compute.default_max_workers(), compute.default_max_workers(),
pad_factor=self._pad_factor, pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold, dc_threshold_mv=self._dc_threshold,
+94 -39
View File
@@ -122,7 +122,7 @@ def test_parallel_identity(tmp_path, monkeypatch):
# the block size so every chunk splits into many FFT tasks — the worst # the block size so every chunk splits into many FFT tasks — the worst
# case for boundary bugs. # case for boundary bugs.
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4) monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
monkeypatch.setattr(compute, "_FFT_BLOCK", 4) monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET) fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
assert fft_rows < n_rows, \ assert fft_rows < n_rows, \
f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})" f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
@@ -155,43 +155,10 @@ def test_parallel_identity(tmp_path, monkeypatch):
"rf image identical (masked, padded/zoom)" "rf image identical (masked, padded/zoom)"
@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)]) def test_peak_bins_fuzz():
def test_zoom_identity(tmp_path, monkeypatch, spf, bps): """Hammer _peak_bins directly with adversarial spectra: noise,
"""The zoom peak search must reproduce the full padded-rfft argmax
bit-for-bit, across pad factors, masking, bg-sub, dtype, and backend."""
path = tmp_path / f"zoom_{spf}.sras"
gen.write(path, n_angles=2, seed=6, samples_per_frame=spf, bps=bps)
sras = SrasFile(str(path))
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.median(dc4))
backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
for backend in backends:
monkeypatch.setattr(compute, "_fft_backend", backend)
for pad in (4, 8, 40):
n_fft = spf * pad
for thr_v in (None, thr):
for bg in (False, True):
ref = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
apply_bg_sub=bg, n_fft=n_fft,
exact=True)
zoom = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
apply_bg_sub=bg, n_fft=n_fft)
diff = int((ref != zoom).sum())
assert diff == 0, \
(f"{diff} px differ: backend={backend} pad={pad} "
f"thr={thr_v} bg={bg} spf={spf}")
# A threshold above every pixel masks everything: both paths must agree
# on an all-zero image.
all_masked = compute_rf_image(sras, 0, dc_threshold_mv=1e9, n_fft=spf * 8)
assert not all_masked.any()
def test_zoom_identity_fuzz():
"""Hammer _peak_bins_zoom directly with adversarial spectra: noise,
un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs, un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs,
and all-zero rows.""" and all-zero rows — against an independent scipy.fft reference."""
import scipy.fft as scipy_fft import scipy.fft as scipy_fft
rng = np.random.default_rng(42) rng = np.random.default_rng(42)
@@ -222,14 +189,102 @@ def test_zoom_identity_fuzz():
P[:, 0] = 0.0 P[:, 0] = 0.0
ref = np.argmax(P, axis=1) ref = np.argmax(P, axis=1)
zp = compute._zoom_plan(spf, n_fft) got = compute._peak_bins(w, n_fft)
got = compute._peak_bins_zoom(w, zp)
bad = np.nonzero(ref != got)[0] bad = np.nonzero(ref != got)[0]
assert not len(bad), \ assert not len(bad), \
(f"spf={spf} pad={pad}: rows {bad.tolist()} picked " (f"spf={spf} pad={pad}: rows {bad.tolist()} picked "
f"{got[bad].tolist()} instead of {ref[bad].tolist()}") f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
def test_peak_bins_fuzz_min_freq():
"""Same adversarial-spectra fuzz as test_peak_bins_fuzz, but with a swept
min_bin floor: bins below the floor must be excluded from the argmax
exactly as the independent scipy.fft reference is, when zeroed the same
way before argmax."""
import scipy.fft as scipy_fft
rng = np.random.default_rng(43)
for _ in range(25):
spf = int(rng.integers(16, 220))
pad = int(rng.choice([4, 5, 8, 16, 40]))
n_fft = spf * pad
n_wf = 24
w = rng.normal(scale=20.0, size=(n_wf, spf))
t = np.arange(spf)
for r in range(6):
f = rng.uniform(1.0, spf / 2 - 1)
w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+ 49.9 * np.sin(2 * np.pi * f2 * t / spf)
w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
+ 50 * np.cos(2 * np.pi * f2 * t / spf)
w[8] = 90 + rng.normal(scale=5.0, size=spf)
w[9] = 0.0
w = w.astype(np.float32)
S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
P = S.real ** 2
P += S.imag ** 2
n_bins_fine = n_fft // 2 + 1
min_bin = int(rng.integers(1, max(2, n_bins_fine // 3)))
P[:, :min_bin] = 0.0
ref = np.argmax(P, axis=1)
got = compute._peak_bins(w, n_fft, min_bin)
bad = np.nonzero(ref != got)[0]
assert not len(bad), \
(f"spf={spf} pad={pad} min_bin={min_bin}: rows {bad.tolist()} picked "
f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
@pytest.mark.parametrize("spf", [64, 500, 2500])
def test_fft_block_for(spf):
"""Block size == _FFT_BLOCK_MAX at natural resolution, shrinks and stays
>= _FFT_BLOCK_MIN as n_len grows, and the implied per-thread byte
estimate respects _FFT_PLAN_BYTES_BUDGET except when the floor is
engaged."""
block_natural = compute._fft_block_for(spf, spf)
assert block_natural == compute._FFT_BLOCK_MAX
prev = compute._FFT_BLOCK_MAX
for pad in (2, 4, 8, 40, 500):
n_len = spf * pad
block = compute._fft_block_for(spf, n_len)
assert compute._FFT_BLOCK_MIN <= block <= prev
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
if block > compute._FFT_BLOCK_MIN:
assert block * bytes_per_wf <= compute._FFT_PLAN_BYTES_BUDGET
prev = block
def test_compute_rf_image_min_freq_mhz(tmp_path):
"""min_freq_mhz threads through compute_rf_image end-to-end, for both
the natural-resolution and padded paths: 0.0 (default) must reproduce
the pre-existing image exactly, and a floor above every real peak must
collapse the image to bin 0 (0 MHz) — the same fallback the low-level
search uses when nothing survives the floor."""
path = tmp_path / "floor_e2e.sras"
gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
sras = SrasFile(str(path))
huge_floor = float(sras.freq_axis_mhz(None)[-1]) + 1.0 # above Nyquist
for n_fft in (None, 64 * 8):
unfiltered = compute_rf_image(sras, 0, dc_threshold_mv=None,
apply_bg_sub=False, n_fft=n_fft)
same = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
n_fft=n_fft, min_freq_mhz=0.0)
assert np.array_equal(unfiltered, same), \
f"n_fft={n_fft}: min_freq_mhz=0.0 changed the output"
assert unfiltered.any(), \
f"n_fft={n_fft}: fixture should have real signal"
collapsed = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
n_fft=n_fft, min_freq_mhz=huge_floor)
assert not collapsed.any(), \
f"n_fft={n_fft}: floor above Nyquist should collapse to 0 MHz"
def test_nomask_equals_low_threshold(tmp_path): def test_nomask_equals_low_threshold(tmp_path):
"""dc_threshold_mv=None must equal a threshold below every pixel, while """dc_threshold_mv=None must equal a threshold below every pixel, while
skipping the CH4 read.""" skipping the CH4 read."""
+40
View File
@@ -223,6 +223,46 @@ def test_threshold_change_recomputes(ctx):
f"masking zeroed some pixels ({n_zero} of {win._current_image.size})" f"masking zeroed some pixels ({n_zero} of {win._current_image.size})"
def test_highlight_masked_pixels(ctx):
"""Masked (below-threshold) pixels are drawn as NaN - filled with a
highlight color, separate from the normal colormap - so they can't be
mistaken for a real, possibly-low-frequency pixel; unchecking restores
the old behavior where both blend into the same plain 0."""
win = ctx.win
assert win.chk_highlight_masked.isChecked(), "on by default"
dc4 = win._dc_cache[(0, CH4_IDX)]
expect_masked = dc4 < win.spin_threshold_mv.value()
assert expect_masked.any() and not expect_masked.all(), \
"fixture threshold should mask some but not all pixels"
calls = []
orig = win.image_canvas.show_image
def spy(img, *a, **kw):
calls.append((np.array(img, copy=True), kw.get("bad_color")))
return orig(img, *a, **kw)
with patch.object(win.image_canvas, "show_image", side_effect=spy):
win._redraw_image(win._current_image)
shown, bad_color = calls[-1]
assert bad_color is not None, "highlight color set while checkbox is on"
assert np.array_equal(np.isnan(shown), expect_masked), \
"NaN exactly where DC4 is below threshold, nowhere else"
win.chk_highlight_masked.setChecked(False)
calls.clear()
with patch.object(win.image_canvas, "show_image", side_effect=spy):
win._redraw_image(win._current_image)
shown2, bad_color2 = calls[-1]
assert bad_color2 is None, "no highlight color once unchecked"
assert not np.isnan(shown2).any(), "unchecked: no pixel pulled out to NaN"
assert np.array_equal(shown2, win._current_image), \
"unchecked: displayed array is the raw, unmodified image"
win.chk_highlight_masked.setChecked(True)
pump(60)
def test_bg_sub_toggle(ctx): def test_bg_sub_toggle(ctx):
"""bg-sub no longer gates the display: it only affects a future live """bg-sub no longer gates the display: it only affects a future live
compute for an angle with nothing cached yet, or an explicit batch compute for an angle with nothing cached yet, or an explicit batch
+34 -15
View File
@@ -14,7 +14,7 @@ import pytest
import sras_compute as compute import sras_compute as compute
from sras_compute import compute_rf_image, dc_image_mv from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH4_IDX, SrasFile from sras_format import CH1_IDX, CH4_IDX, SrasFile
import tools.make_test_sras as gen import tools.make_test_sras as gen
@@ -193,26 +193,45 @@ def test_row_average_respects_own_center_mask(tmp_path):
def test_row_average_composes_with_padding(tmp_path): def test_row_average_composes_with_padding(tmp_path):
"""row_avg_n and n_fft (zero-padding) are independent knobs: using them """row_avg_n and n_fft (zero-padding) are independent knobs: using them
together must not raise, and must still agree with the exact (non-zoom) together must not raise, and must agree bit-for-bit with an independent
reference path at that pad factor -- i.e. row-averaging composes with reference that row-averages the raw waveforms and background-subtracts
the zoom peak search correctly, not just with the direct one.""" them, then runs a plain scipy rfft + argmax at the same pad factor --
i.e. row-averaging composes correctly with the padded peak search."""
import scipy.fft as scipy_fft
path = tmp_path / "padded_rowavg.sras" path = tmp_path / "padded_rowavg.sras"
gen.write(path, n_angles=1, seed=12, samples_per_frame=64) gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
sras = SrasFile(str(path)) sras = SrasFile(str(path))
spf = sras.samples_per_frame spf = sras.samples_per_frame
n_rows, n_frames = sras.image_shape(0)
n_fft = spf * 40
raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True) raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3) row_avg_n=3)
avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, avg_padded = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40) row_avg_n=3, n_fft=n_fft)
avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40, exact=True)
assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape assert avg_natural.shape == raw_natural.shape == avg_padded.shape
assert np.all(np.isfinite(avg_padded_zoom)) assert np.all(np.isfinite(avg_padded))
assert np.array_equal(avg_padded_zoom, avg_padded_exact), \
"row-averaged waveforms feed the zoom and exact FFT paths identically" weights = compute._row_average_weights(3)
freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
data = sras.data[0]
expected = np.zeros((n_rows, n_frames), dtype=np.float32)
for r in range(n_rows):
v = np.ones(n_frames, dtype=bool)
avg = compute._row_average_waveforms(
data[r, CH1_IDX].astype(np.float32), v, weights)
avg = avg - sras.background
S = scipy_fft.rfft(avg, n=n_fft, axis=-1, workers=1)
power = S.real ** 2 + S.imag ** 2
power[:, 0] = 0.0
expected[r] = freq32[np.argmax(power, axis=1)]
assert np.array_equal(avg_padded, expected), \
"row-averaged waveforms feed the padded peak search identically " \
"to an independent reference"
def test_row_average_improves_snr_recovery(): def test_row_average_improves_snr_recovery():
@@ -233,8 +252,8 @@ def test_row_average_improves_snr_recovery():
weights = compute._row_average_weights(8) # wide window: lots of averaging weights = compute._row_average_weights(8) # wide window: lots of averaging
averaged = compute._row_average_waveforms(raw, valid, weights) averaged = compute._row_average_waveforms(raw, valid, weights)
raw_bins = compute._peak_bins_direct(raw, spf) raw_bins = compute._peak_bins(raw, spf)
avg_bins = compute._peak_bins_direct(averaged, spf) avg_bins = compute._peak_bins(averaged, spf)
raw_hits = int(np.sum(raw_bins == true_bin)) raw_hits = int(np.sum(raw_bins == true_bin))
avg_hits = int(np.sum(avg_bins == true_bin)) avg_hits = int(np.sum(avg_bins == true_bin))
@@ -257,7 +276,7 @@ def test_row_average_parallel_identity(tmp_path, monkeypatch):
sras = SrasFile(str(path)) sras = SrasFile(str(path))
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4) monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
monkeypatch.setattr(compute, "_FFT_BLOCK", 4) monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
# row_avg_n > 0 halves the effective budget before chunk planning. # row_avg_n > 0 halves the effective budget before chunk planning.
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2) fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
assert fft_rows < n_rows, \ assert fft_rows < n_rows, \
+2 -2
View File
@@ -90,7 +90,7 @@ def no_fft(monkeypatch):
"""Make any real FFT work loud: returns a list that stays empty unless a """Make any real FFT work loud: returns a list that stays empty unless a
peak search actually runs.""" peak search actually runs."""
calls = [] calls = []
for name in ("_peak_bins_direct", "_peak_bins_zoom"): for name in ("_peak_bins",):
original = getattr(compute, name) original = getattr(compute, name)
def spy(*args, _f=original, **kwargs): def spy(*args, _f=original, **kwargs):
@@ -147,7 +147,7 @@ def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
gen.write(path, n_angles=2, seed=13, samples_per_frame=256) gen.write(path, n_angles=2, seed=13, samples_per_frame=256)
spf = SrasFile(str(path)).samples_per_frame spf = SrasFile(str(path)).samples_per_frame
assert cache_file(str(path), "fft", True, "scipy", 0, pad) == "" assert cache_file(str(path), "fft", True, 0, pad) == ""
sras = SrasFile(str(path)) sras = SrasFile(str(path))
assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip" assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"
+29 -26
View File
@@ -1,8 +1,10 @@
#!/usr/bin/env python3 #!/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 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: Usage:
python tools/bench_fft.py # synthetic, pads 1/8/40 python tools/bench_fft.py # synthetic, pads 1/8/40
@@ -22,7 +24,7 @@ import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import sras_compute as compute # noqa: E402 import sras_compute as compute # noqa: E402
from sras_compute import compute_rf_image, set_fft_backend # noqa: E402 from sras_compute import compute_rf_image # noqa: E402
from sras_format import SrasFile # noqa: E402 from sras_format import SrasFile # noqa: E402
import tools.make_test_sras as gen # noqa: E402 import tools.make_test_sras as gen # noqa: E402
from tools.check_equivalence import row_slice # noqa: E402 from tools.check_equivalence import row_slice # noqa: E402
@@ -38,18 +40,28 @@ def _timed(fn):
return out, wall, cpu / max(wall, 1e-9) 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]) n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
for a in range(sras.n_angles)) for a in range(sras.n_angles))
spf = sras.samples_per_frame spf = sras.samples_per_frame
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)") n_workers = compute._MAX_WORKERS
print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} " print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), "
f"{'wf/s':>10} {'util':>6} match") 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: for pad in pads:
n_fft = spf * pad if pad > 1 else None n_fft = spf * pad if pad > 1 else None
for backend in backends: n_len = n_fft if n_fft is not None else spf
set_fft_backend(backend) plan_mb = _plan_mb(spf, n_len, n_workers)
def run(**kw): def run(**kw):
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None, imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
@@ -57,16 +69,13 @@ def bench(sras, pads, backends):
for a in range(sras.n_angles)] for a in range(sras.n_angles)]
return np.concatenate([i.ravel() for i in imgs]) return np.concatenate([i.ravel() for i in imgs])
ref, wall, util = _timed(lambda: run(exact=True)) ref, wall, util = _timed(lambda: run(max_workers=1))
rows = [("exact(serial)", ref, wall, util, True)] rows = [("serial", ref, wall, util, True)]
for label, kw in (("zoom(serial)", dict(max_workers=1)), img, wall, util = _timed(lambda: run())
("zoom(pool)", {})): rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref))))
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: for label, img, wall, util, ok in rows:
print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s " print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} "
f"{n_wf / wall:>10.0f} {util:>5.1f}x " f"{util:>5.1f}x {plan_mb:>8.1f} {'OK' if ok else 'MISMATCH'}")
f"{'OK' if ok else 'MISMATCH'}")
def main(): def main():
@@ -76,30 +85,24 @@ def main():
p.add_argument("--spf", type=int, default=2500) p.add_argument("--spf", type=int, default=2500)
p.add_argument("--rows", type=int, default=8) p.add_argument("--rows", type=int, default=8)
p.add_argument("--frames", type=int, default=1024) 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", help="path to a real .sras file")
p.add_argument("--real-rows", type=int, default=32, p.add_argument("--real-rows", type=int, default=32,
help="rows of angle 0 to use from the real file") help="rows of angle 0 to use from the real file")
args = p.parse_args() args = p.parse_args()
pads = [int(x) for x in args.pads.split(",")] 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: if args.real:
sras = row_slice(SrasFile(args.real), 0, args.real_rows) sras = row_slice(SrasFile(args.real), 0, args.real_rows)
sras.data = [sras.data[0]] sras.data = [sras.data[0]]
sras.n_angles = 1 sras.n_angles = 1
bench(sras, pads, backends) bench(sras, pads)
else: else:
with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp: with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
path = Path(tmp) / "bench.sras" path = Path(tmp) / "bench.sras"
gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf, gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
geometry=[(args.rows, args.frames)]) geometry=[(args.rows, args.frames)])
bench(SrasFile(str(path)), pads, backends) bench(SrasFile(str(path)), pads)
if __name__ == "__main__": if __name__ == "__main__":