diff --git a/docs/design.md b/docs/design.md index 7b0c583..1543828 100644 --- a/docs/design.md +++ b/docs/design.md @@ -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 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 -spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for -mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan -row, which is what used to collapse the old row-chunk planner to one worker -and make synthesis single-threaded. +spectrum of that pixel's CH1 waveform. This used to run through `_peak_bins_zoom`, +a coarse-rfft-plus-local-fine-DFT refinement that avoided ever materialising +a padded spectrum — at the pad factor of 40 needed for mapping resolution, a +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 - power spectrum (a trig polynomial of degree spf−1) cannot hide its global - max between coarse samples; -2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max - becomes a refinement candidate. Quarter-natural-bin scalloping at the 2× - grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide - margin. The DC-adjacent window is always refined too: the coarse DC bin - is zeroed for suppression, which would otherwise blind the scan to fine - bins closer to DC than the first coarse sample (where the leakage skirt - of an un-subtracted offset peaks); -3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every - fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated - on the exact `n_fft` grid by one small complex gemm, with np.argmax's - lowest-bin tie-break preserved across windows. +The block size is the part that has to adapt to pad factor. +`_fft_block_for(spf, n_len)` derives waveforms-per-task from a fixed +per-thread byte budget (`_FFT_PLAN_BYTES_BUDGET`, `SRAS_FFT_PLAN_BUDGET_MB`, +default 16 MB) rather than a fixed constant, because a cached pyFFTW plan's +input+output buffers are *permanent* per-thread memory (the plan cache is +never evicted) — with a fixed 512-waveform block, pad 40 at a 2500-sample +frame costs ~210 MB per pool thread (~3.3 GB total across 16 threads); +`_fft_block_for` bounds that to ~16 MB per thread (~260 MB across 16 +threads) at the same pad factor, while still reproducing the old tuned 512 +exactly at natural resolution (pad 1), where it cost nothing to begin with. +`_FFT_BLOCK_MAX` (512) and `_FFT_BLOCK_MIN` (32) cap and floor the result: +the ceiling is the measured knee on a 16-core machine at natural resolution +(smaller blocks serialise on GIL-held numpy dispatch, larger ones lose cache +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 -`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial -spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose -baseline was captured on the old full-padded path. +The outer row-chunk sizing (`_plan_fft_rows`) needed no companion change. +It only ever budgets the raw float32 waveform *read* buffer, which this +change doesn't touch — spectrum memory is bounded independently by +`_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 -tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose -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. +`tools/check_equivalence.py`'s golden-hash harness remains the end-to-end +regression baseline for this path, unaffected by this change. ## 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 `(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 -`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row -chunking" above. On the largest real scans `_plan_chunks` is already +`_plan_fft_rows` runs — see "Memory budget and row chunking" above. On the +largest real scans `_plan_chunks` is already 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 unexpected regression. diff --git a/pyproject.toml b/pyproject.toml index 9f7942e..8617fab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,10 +14,8 @@ dependencies = [ "scipy==1.18.0", # Angle alignment only: masked FFT phase correlation (skimage.registration). "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", - # Clamps BLAS threading under the FFT worker pool. - "threadpoolctl==3.6.0", ] [project.optional-dependencies] diff --git a/sras_compute.py b/sras_compute.py index 02b0ae2..28af21e 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -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 diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py index 8c0a1e4..a4c8373 100644 --- a/sras_viewer/canvases.py +++ b/sras_viewer/canvases.py @@ -166,11 +166,14 @@ class ImageCanvas(FigureCanvasQTAgg): def show_image(self, img: np.ndarray, extent: list[float], cmap, 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 vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image — 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.ax = self.figure.add_subplot(111) # Patches and lines are destroyed by figure.clf(); drop stale refs. @@ -179,6 +182,10 @@ class ImageCanvas(FigureCanvasQTAgg): self._extent = extent 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 else {"vmin": vmin, "vmax": vmax}) im = self.ax.imshow( diff --git a/sras_viewer/common.py b/sras_viewer/common.py index ad587e9..7c9b050 100644 --- a/sras_viewer/common.py +++ b/sras_viewer/common.py @@ -26,6 +26,10 @@ CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) 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 _CHANNEL_DISPLAY = { CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index 87379eb..dfd6da7 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -13,7 +13,6 @@ from PyQt6.QtWidgets import ( QScrollArea, QSpinBox, QVBoxLayout, QWidget, ) -from sras_compute import PYFFTW_AVAILABLE 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 @@ -24,7 +23,7 @@ from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, # --------------------------------------------------------------------------- class FftOptionsDialog(QDialog): - """Configure FFT backend and zero-padding. + """Configure FFT zero-padding. Changes take effect only when the user clicks Apply. Cancel discards all pending edits. The live 'frequency resolution' label updates as @@ -33,7 +32,6 @@ class FftOptionsDialog(QDialog): """ def __init__(self, parent=None, *, - current_backend: str, current_pad_factor: int, samples_per_frame: int | None, sample_rate_hz: float | None, @@ -49,29 +47,6 @@ class FftOptionsDialog(QDialog): 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 ---------------------------------------------- grp_zp = QGroupBox("Zero-Padding") zl = QVBoxLayout(grp_zp) @@ -134,9 +109,6 @@ class FftOptionsDialog(QDialog): f"Velocity bin: {vel_res_ms:.3f} m/s " 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: return max(1, self._spin_pad.value()) diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index b02d6e7..481ea63 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -30,8 +30,8 @@ from .canvases import ImageCanvas, WaveformCanvas from .common import ( 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, - _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, _scroll_panel, - _wrap_label, + _MASKED_HIGHLIGHT_COLOR, _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, + _scroll_panel, _wrap_label, ) from .align_wizard import AlignmentWizard from .dialogs import ( @@ -58,6 +58,7 @@ class SrasViewerWindow(QMainWindow): self._pending_ch: int = 0 self._pending_bg_sub: bool = True self._pending_threshold: float = 50.0 # mV + self._pending_min_freq_mhz: float = 0.0 self._pending_fft_pad_factor: int = 1 # Live background jobs, keyed by role — see _run_worker. @@ -70,7 +71,6 @@ class SrasViewerWindow(QMainWindow): self._settings = QSettings(QSettings.Format.IniFormat, QSettings.Scope.UserScope, "sras-viewer", "sras-viewer") - compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy"))) try: pad = int(self._settings.value("fft/pad_factor", 1)) except (TypeError, ValueError): @@ -272,6 +272,28 @@ class SrasViewerWindow(QMainWindow): self.lbl_threshold_adc = _wrap_label( f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) 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) # Background subtraction (v4+ files only) @@ -417,6 +439,22 @@ class SrasViewerWindow(QMainWindow): self.chk_auto.toggled.connect(self._on_autoscale_toggled) 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() for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): spin = _make_dspin(-1e9, 1e9, 4) @@ -440,7 +478,7 @@ class SrasViewerWindow(QMainWindow): fft_menu = menubar.addMenu("&FFT") 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_menu.addAction(fft_act) @@ -684,7 +722,9 @@ class SrasViewerWindow(QMainWindow): # Threshold and bg-sub apply to all CH1 modes 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_highlight_masked.setEnabled(is_ch1) self.spin_grating_um.setEnabled(is_vel) self.grp_velocity.setVisible(is_vel) @@ -733,6 +773,14 @@ class SrasViewerWindow(QMainWindow): if self._is_fft_mode(): 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): manual = not checked 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: 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): if not self.chk_auto.isChecked() and self._current_image is not None: self._redraw_image(self._current_image) @@ -1014,15 +1066,19 @@ class SrasViewerWindow(QMainWindow): return freq_mhz def _fft_cache_key(self, angle_idx: int) -> tuple: - """Keyed by angle and DC threshold only. Once any FFT image exists - for an angle this session — live-computed or pulled from the file's - own stored cache — it stays the displayed image for that angle - regardless of later bg-sub/pad toggles; those only affect a future - live compute for an angle with nothing cached yet, or an explicit - batch recompute (see _stored_fft_image). Threshold stays in the key - because re-masking against it is free and meant to stay interactive - (see _on_threshold_changed).""" - return (angle_idx, self.spin_threshold_mv.value()) + """Keyed by angle, DC threshold, and min peak frequency only. Once + any FFT image exists for an angle this session — live-computed or + pulled from the file's own stored cache — it stays the displayed + image for that angle regardless of later bg-sub/pad toggles; those + only affect a future live compute for an angle with nothing cached + yet, or an explicit batch recompute (see _stored_fft_image). + Threshold and min-freq stay in the key so revisiting a combination + already computed this session is instant — even though only + 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: """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- blocking) note when the live controls diverge from what's shown. 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 - irreversibly into the stored numbers. + against it is free, unlike bg-sub/pad/row-averaging — or the min + 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 would mean reading a whole CH4 channel, this declines and the caller @@ -1194,6 +1251,24 @@ class SrasViewerWindow(QMainWindow): self._redraw_image(img) 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): s = self._sras 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 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(): - 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)): with QSignalBlocker(spin): spin.setValue(val) @@ -1239,6 +1333,7 @@ class SrasViewerWindow(QMainWindow): vmin=vmin, vmax=vmax, xlabel="X (mm)", ylabel="Y (mm)", title=title, colorbar_label=colorbar_label, + bad_color=_MASKED_HIGHLIGHT_COLOR if mask_valid is not None else None, ) self.statusBar().showMessage( f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " @@ -1262,6 +1357,7 @@ class SrasViewerWindow(QMainWindow): self._pending_ch = ch_idx self._pending_bg_sub = self.chk_bg_sub.isChecked() 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 worker = ComputeWorker( @@ -1274,6 +1370,7 @@ class SrasViewerWindow(QMainWindow): # need to re-read the CH4 channel from disk. dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), is_fft_mode=is_fft, + min_freq_mhz=self._pending_min_freq_mhz, ) if not self._run_worker( Jobs.COMPUTE, worker, @@ -1302,9 +1399,10 @@ class SrasViewerWindow(QMainWindow): already be cached.""" if (self.spin_angle.value(), self.combo_channel.currentIndex(), 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_threshold, self._pending_fft_pad_factor): + self._pending_threshold, self._pending_min_freq_mhz, + self._pending_fft_pad_factor): self._refresh_display() def _on_compute_done(self, result): @@ -1315,7 +1413,8 @@ class SrasViewerWindow(QMainWindow): ch_idx = self._pending_ch 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) else: img = result @@ -1620,7 +1719,6 @@ class SrasViewerWindow(QMainWindow): def _on_fft_options(self): dlg = FftOptionsDialog( self, - current_backend=compute.get_fft_backend(), current_pad_factor=self._fft_pad_factor, 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, @@ -1628,9 +1726,7 @@ class SrasViewerWindow(QMainWindow): ) if dlg.exec() != QDialog.DialogCode.Accepted: return - compute.set_fft_backend(dlg.get_backend()) 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) # Pad factor no longer gates the display: it only affects a future # live compute for an angle with nothing cached yet, or an explicit diff --git a/sras_workers.py b/sras_workers.py index d3761f7..0d8f05d 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -115,7 +115,8 @@ class ComputeWorker(CancellableWorker): apply_bg_sub: bool = True, n_fft: int | None = None, dc_threshold_mv: float = 0.0, dc4_mv: np.ndarray | None = None, - is_fft_mode: bool = False): + is_fft_mode: bool = False, + min_freq_mhz: float = 0.0): super().__init__() self._sras = sras self._angle = angle_idx @@ -125,6 +126,7 @@ class ComputeWorker(CancellableWorker): self._dc_threshold = dc_threshold_mv self._dc4_mv = dc4_mv self._is_fft_mode = is_fft_mode + self._min_freq_mhz = min_freq_mhz def run(self): try: @@ -132,7 +134,8 @@ class ComputeWorker(CancellableWorker): img = compute_rf_image( self._sras, self._angle, dc_threshold_mv=self._dc_threshold, 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: img = dc_image_mv(self._sras, self._angle, self._ch, should_stop=self._stopped) @@ -240,7 +243,7 @@ class BatchCacheWorker(QObject): with ProcessPoolExecutor(max_workers=n_procs) as executor: futures = { 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, dc_threshold_mv=self._dc_threshold, row_avg_n=self._row_avg_n): p @@ -266,7 +269,6 @@ class BatchCacheWorker(QObject): for path in paths: try: err = cache_file(path, self._mode, self._apply_bg_sub, - compute.get_fft_backend(), compute.default_max_workers(), pad_factor=self._pad_factor, dc_threshold_mv=self._dc_threshold, diff --git a/tests/test_compute.py b/tests/test_compute.py index 3c6e400..f49b78b 100644 --- a/tests/test_compute.py +++ b/tests/test_compute.py @@ -122,7 +122,7 @@ def test_parallel_identity(tmp_path, monkeypatch): # the block size so every chunk splits into many FFT tasks — the worst # case for boundary bugs. 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) assert fft_rows < 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)" -@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)]) -def test_zoom_identity(tmp_path, monkeypatch, spf, bps): - """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, +def test_peak_bins_fuzz(): + """Hammer _peak_bins directly with adversarial spectra: noise, 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 rng = np.random.default_rng(42) @@ -222,14 +189,102 @@ def test_zoom_identity_fuzz(): P[:, 0] = 0.0 ref = np.argmax(P, axis=1) - zp = compute._zoom_plan(spf, n_fft) - got = compute._peak_bins_zoom(w, zp) + got = compute._peak_bins(w, n_fft) bad = np.nonzero(ref != got)[0] assert not len(bad), \ (f"spf={spf} pad={pad}: rows {bad.tolist()} picked " 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): """dc_threshold_mv=None must equal a threshold below every pixel, while skipping the CH4 read.""" diff --git a/tests/test_gui.py b/tests/test_gui.py index 4b184ff..b27a6ca 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -223,6 +223,46 @@ def test_threshold_change_recomputes(ctx): 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): """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 diff --git a/tests/test_row_average.py b/tests/test_row_average.py index 448cee4..dbb00d2 100644 --- a/tests/test_row_average.py +++ b/tests/test_row_average.py @@ -14,7 +14,7 @@ import pytest import sras_compute as compute 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 @@ -193,26 +193,45 @@ def test_row_average_respects_own_center_mask(tmp_path): def test_row_average_composes_with_padding(tmp_path): """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) - reference path at that pad factor -- i.e. row-averaging composes with - the zoom peak search correctly, not just with the direct one.""" + together must not raise, and must agree bit-for-bit with an independent + reference that row-averages the raw waveforms and background-subtracts + 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" gen.write(path, n_angles=1, seed=12, samples_per_frame=64) sras = SrasFile(str(path)) 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) avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, row_avg_n=3) - avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, - row_avg_n=3, n_fft=spf * 40) - 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) + avg_padded = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, + row_avg_n=3, n_fft=n_fft) - assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape - assert np.all(np.isfinite(avg_padded_zoom)) - assert np.array_equal(avg_padded_zoom, avg_padded_exact), \ - "row-averaged waveforms feed the zoom and exact FFT paths identically" + assert avg_natural.shape == raw_natural.shape == avg_padded.shape + assert np.all(np.isfinite(avg_padded)) + + 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(): @@ -233,8 +252,8 @@ def test_row_average_improves_snr_recovery(): weights = compute._row_average_weights(8) # wide window: lots of averaging averaged = compute._row_average_waveforms(raw, valid, weights) - raw_bins = compute._peak_bins_direct(raw, spf) - avg_bins = compute._peak_bins_direct(averaged, spf) + raw_bins = compute._peak_bins(raw, spf) + avg_bins = compute._peak_bins(averaged, spf) raw_hits = int(np.sum(raw_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)) 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. fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2) assert fft_rows < n_rows, \ diff --git a/tests/test_stored_cache.py b/tests/test_stored_cache.py index f802e53..b73adc9 100644 --- a/tests/test_stored_cache.py +++ b/tests/test_stored_cache.py @@ -90,7 +90,7 @@ def no_fft(monkeypatch): """Make any real FFT work loud: returns a list that stays empty unless a peak search actually runs.""" calls = [] - for name in ("_peak_bins_direct", "_peak_bins_zoom"): + for name in ("_peak_bins",): original = getattr(compute, name) 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) 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)) assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip" diff --git a/tools/bench_fft.py b/tools/bench_fft.py index 68a1cc3..a509b19 100644 --- a/tools/bench_fft.py +++ b/tools/bench_fft.py @@ -1,8 +1,10 @@ #!/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 -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: python tools/bench_fft.py # synthetic, pads 1/8/40 @@ -22,10 +24,10 @@ import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import sras_compute as compute # noqa: E402 -from sras_compute import compute_rf_image, set_fft_backend # noqa: E402 -from sras_format import SrasFile # noqa: E402 -import tools.make_test_sras as gen # noqa: E402 -from tools.check_equivalence import row_slice # noqa: E402 +from sras_compute import compute_rf_image # noqa: E402 +from sras_format import SrasFile # noqa: E402 +import tools.make_test_sras as gen # noqa: E402 +from tools.check_equivalence import row_slice # noqa: E402 def _timed(fn): @@ -38,35 +40,42 @@ def _timed(fn): 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]) for a in range(sras.n_angles)) spf = sras.samples_per_frame - print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)") - print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} " - f"{'wf/s':>10} {'util':>6} match") + n_workers = compute._MAX_WORKERS + print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), " + f"{n_workers} workers") + print(f"{'pad':>4} {'variant':>8} {'wall':>9} {'wf/s':>10} {'util':>6} " + f"{'plan MB':>9} match") for pad in pads: n_fft = spf * pad if pad > 1 else None - for backend in backends: - set_fft_backend(backend) + n_len = n_fft if n_fft is not None else spf + plan_mb = _plan_mb(spf, n_len, n_workers) - def run(**kw): - imgs = [compute_rf_image(sras, a, dc_threshold_mv=None, - apply_bg_sub=True, n_fft=n_fft, **kw) - for a in range(sras.n_angles)] - return np.concatenate([i.ravel() for i in imgs]) + def run(**kw): + imgs = [compute_rf_image(sras, a, dc_threshold_mv=None, + apply_bg_sub=True, n_fft=n_fft, **kw) + for a in range(sras.n_angles)] + return np.concatenate([i.ravel() for i in imgs]) - ref, wall, util = _timed(lambda: run(exact=True)) - rows = [("exact(serial)", ref, wall, util, True)] - for label, kw in (("zoom(serial)", dict(max_workers=1)), - ("zoom(pool)", {})): - img, wall, util = _timed(lambda: run(**kw)) - rows.append((label, img, wall, util, bool(np.array_equal(img, ref)))) - for label, img, wall, util, ok in rows: - print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s " - f"{n_wf / wall:>10.0f} {util:>5.1f}x " - f"{'OK' if ok else 'MISMATCH'}") + ref, wall, util = _timed(lambda: run(max_workers=1)) + rows = [("serial", ref, wall, util, True)] + img, wall, util = _timed(lambda: run()) + rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref)))) + for label, img, wall, util, ok in rows: + print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} " + f"{util:>5.1f}x {plan_mb:>8.1f} {'OK' if ok else 'MISMATCH'}") def main(): @@ -76,30 +85,24 @@ def main(): p.add_argument("--spf", type=int, default=2500) p.add_argument("--rows", type=int, default=8) p.add_argument("--frames", type=int, default=1024) - p.add_argument("--backends", default=None, - help="comma-separated (default: scipy,pyfftw if available)") p.add_argument("--real", help="path to a real .sras file") p.add_argument("--real-rows", type=int, default=32, help="rows of angle 0 to use from the real file") args = p.parse_args() pads = [int(x) for x in args.pads.split(",")] - if args.backends: - backends = args.backends.split(",") - else: - backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else []) if args.real: sras = row_slice(SrasFile(args.real), 0, args.real_rows) sras.data = [sras.data[0]] sras.n_angles = 1 - bench(sras, pads, backends) + bench(sras, pads) else: with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp: path = Path(tmp) / "bench.sras" gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf, geometry=[(args.rows, args.frames)]) - bench(SrasFile(str(path)), pads, backends) + bench(SrasFile(str(path)), pads) if __name__ == "__main__":