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

Merge the Documents/sras-viewer working copy into this repo

Browse Source 
The two clones had diverged from origin/main (191d1b8) and both had
grown uncommitted work. This merges the other clone's three commits and
reconciles four conflicting files.

Both clones independently grew a feature called "batch export", but they
are different sweeps and both are kept:
- BatchExportWorker (this clone) — every angle x channel of the one open
  file, fixed colorbar per channel, under a new Export menu.
- BatchExportImagesWorker (other clone) — the current view across many
  selected files, process-pooled, under Convert.

Conflict resolutions of note:
- sras_average.py: the other clone's memory-bounded rewrite had silently
  reverted this clone's float32 -> int16 accumulator fix, which exists to
  keep averaged output from landing 1 ADC count off the original tool.
  Kept the rewrite, re-applied the fix, and corrected the two docstrings
  that still claimed float32.
- tests/test_compute.py: kept the other clone's meta["waveforms"] key
  (meta["data"] no longer exists for this fixture and would KeyError)
  and its laser_freq_hz assertions, plus this clone's int16 expectation
  and its dropped subprocess returncode assertions.
- ComputeWorker: row_avg_n and min_freq_mhz were added independently on
  either side; both are now threaded through.
- compute_rf_image's `exact` parameter is gone, removed by the other
  clone's PyFFTW peak-search rewrite. It had no remaining callers.

Verified: 149 passed with a complete dependency set. The failures seen
with this clone's own .venv are a missing scikit-image, which predates
this merge and reproduces identically on the pre-merge commit.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales [M S E]
2026-08-12 10:36:25 -05:00
parent efecc154fe a8b962e317
commit 2937c017a7
19 changed files with 2130 additions and 617 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+171 -244
View File
@@ -1,9 +1,9 @@
#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing
child can import it without loading Qt or matplotlib — which matters
because Python 3.14 on macOS spawns rather than forks.
"""
import atexit
@@ -15,61 +15,34 @@ from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pyfftw
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
adc_to_mv)
# ---------------------------------------------------------------------------
# FFT backend
# ---------------------------------------------------------------------------
try:
import pyfftw
PYFFTW_AVAILABLE = True
except ImportError:
PYFFTW_AVAILABLE = False
try:
from threadpoolctl import threadpool_limits
except ImportError:
threadpool_limits = None
_fft_backend = "scipy" # "scipy" or "pyfftw"; set via set_fft_backend()
def set_fft_backend(name: str):
"""Select the rfft implementation. Module-level state, so it must be set
explicitly inside each multiprocessing child — it does not survive a
spawn. "numpy" is accepted as a legacy alias for "scipy"."""
global _fft_backend
if name == "numpy":
name = "scipy"
_fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy"
def get_fft_backend() -> str:
return _fft_backend
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
workers: int = 1) -> np.ndarray:
"""Dispatch rfft to the selected backend with optional multithreading."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
# ---------------------------------------------------------------------------
# FFT worker pool and per-thread pyFFTW plans
# ---------------------------------------------------------------------------
_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine:
# smaller blocks serialise on GIL-held numpy dispatch,
# larger ones lose cache residency and task granularity.
_ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf
_FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0))
_FFT_BLOCK_MAX = 512 # ceiling on waveforms/task: the knee measured on a
# 16-core machine at natural resolution (n_len == spf).
# Smaller blocks serialise on GIL-held numpy dispatch,
# larger ones lose cache residency and task granularity.
_FFT_BLOCK_MIN = 32 # floor, so task granularity never collapses at
# extreme pad factors — at the cost of exceeding
# _FFT_PLAN_BYTES_BUDGET there (see _fft_block_for).
_FFT_PLAN_BYTES_BUDGET = int(
os.environ.get("SRAS_FFT_PLAN_BUDGET_MB", 16)) * 1024 * 1024
# Per-thread ceiling on one cached pyFFTW plan's resident input+output
# buffers (see _fft_block_for). This is PERMANENT memory — the plan
# cache is never evicted — so the worst case across the whole pool for
# one distinct (spf, n_len) is _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET.
# Deliberately separate from _TOTAL_BYTES_BUDGET/SRAS_MEM_BUDGET_MB,
# which bounds transient concurrently-live chunk buffers, freed at the
# end of each chunk — this one bounds a permanent per-thread tax that
# compounds with worker count, not chunk concurrency.
_pool_lock = threading.Lock()
_pool: ThreadPoolExecutor | None = None
@@ -115,23 +88,38 @@ def _save_wisdom():
pass
def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
def _fft_block_for(spf: int, n_len: int) -> int:
"""Waveforms per FFT task/plan at transform length n_len. Bounds a
cached pyFFTW plan's resident input+output buffers (permanent, one per
pool thread, never evicted) to _FFT_PLAN_BYTES_BUDGET regardless of pad
factor. Reproduces _FFT_BLOCK_MAX exactly at natural resolution
(n_len == spf) — see docs/design.md for the worked numbers at high pad."""
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
block = _FFT_PLAN_BYTES_BUDGET // max(1, bytes_per_wf)
return int(min(_FFT_BLOCK_MAX, max(_FFT_BLOCK_MIN, block)))
def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
"""rfft of a (B, spf) float32 block via a cached per-thread FFTW plan.
Plans have a fixed (_FFT_BLOCK, spf) input shape so each worker thread
plans once per transform length; a remainder block runs through the same
plan with its tail rows ignored. The returned array is the plan's output
buffer — consume it before the next call on the same thread.
Plans have a fixed (block, spf) input shape, block = _fft_block_for(spf,
n), so each worker thread plans once per (block size, transform length)
pair; a remainder block runs through the same plan with its tail rows
ignored. Single-threaded (threads=1): outer parallelism comes from the
pool, so the transform itself must not spin up threads. The returned
array is the plan's output buffer — consume it before the next call on
the same thread.
"""
n_wf, spf = waves.shape
block = _fft_block_for(spf, n)
plans = getattr(_fftw_local, "plans", None)
if plans is None:
plans = _fftw_local.plans = {}
key = (_FFT_BLOCK, spf, n)
key = (block, spf, n)
plan = plans.get(key)
if plan is None:
_load_wisdom_once()
buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32")
buf = pyfftw.empty_aligned((block, spf), dtype="float32")
# No overwrite_input: FFTW must not scribble on input_array, whose
# zero-padded tail (columns spf..n) is zeroed exactly once here.
plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1,
@@ -144,132 +132,20 @@ def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
return plan()[:n_wf]
def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
"""Single-threaded rfft of one block; outer parallelism comes from the
pool, so the transform itself must not spin up threads."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return _fftw_block_rfft(waves, n)
return scipy_fft.rfft(waves, n=n, axis=-1, workers=1)
# ---------------------------------------------------------------------------
# Zoom peak search: coarse rfft + local fine DFT around the winning bin
# Peak search: full transform, block-parallel
# ---------------------------------------------------------------------------
_ZOOM_HALFWIDTH = 0.75 # refinement window half-width, in coarse spacings.
# Every fine bin lies within 0.5 spacings of its
# nearest coarse bin, and that bin is guaranteed to
# be a candidate (see _ZOOM_CAND_RATIO), so 0.5
# suffices; 0.75 adds rounding margin.
_ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of
# its row's coarse maximum. Quarter-natural-bin
# scalloping at the 2x-oversampled coarse grid can
# understate a peak by at most ~19% in power, so 0.7
# keeps a wide margin — near-contenders are resolved
# on the fine grid, never ranked from coarse samples.
@dataclass
class _ZoomPlan:
"""Constants for the coarse+refine peak search, built once per
compute_rf_image call and shared across worker threads. *phases* is a
lazily-filled window-start -> phase-vector cache; a benign duplicate
compute under concurrency is cheaper than locking."""
n_fft: int
n_coarse: int
m: int # fine bins per refinement window
n_bins_fine: int
E: np.ndarray # (spf, m) complex64 fine-DFT matrix, relative bins
j: np.ndarray # (spf,) float64 sample indices
phases: dict
def _zoom_plan(spf: int, n_fft: int) -> _ZoomPlan:
# 2x-oversampled coarse grid: the padded power spectrum is a trig
# polynomial of degree spf-1, so at 2x sampling its global max cannot
# hide between coarse bins.
n_coarse = scipy_fft.next_fast_len(2 * spf, real=True)
n_bins_fine = n_fft // 2 + 1
m = int(np.ceil(2 * _ZOOM_HALFWIDTH * n_fft / n_coarse)) + 1
m = min(m, n_bins_fine - 1)
j = np.arange(spf, dtype=np.float64)
E = np.exp((-2j * np.pi / n_fft) * np.outer(j, np.arange(m))).astype(np.complex64)
return _ZoomPlan(n_fft, n_coarse, m, n_bins_fine, E, j, {})
def _window_start(k_c: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
"""First fine bin of the refinement window around each coarse bin.
Clipped to [1, ...] so fine bin 0 stays excluded (DC suppression)."""
k0 = np.floor((k_c - _ZOOM_HALFWIDTH) * (zp.n_fft / zp.n_coarse)).astype(np.int64)
return np.clip(k0, 1, max(1, zp.n_bins_fine - zp.m))
def _refine_window(waves: np.ndarray, rows: np.ndarray, k0: int, zp: _ZoomPlan,
best_pow: np.ndarray, best_bin: np.ndarray):
"""Evaluate fine bins k0..k0+m-1 for the given rows and fold the result
into the per-row best (power, bin), preserving np.argmax's lowest-bin
tie-break."""
ph = zp.phases.get(k0)
if ph is None:
ph = np.exp((-2j * np.pi * k0 / zp.n_fft) * zp.j).astype(np.complex64)
zp.phases[k0] = ph
F = (waves[rows] * ph) @ zp.E
q = F.real ** 2
q += F.imag ** 2
i = np.argmax(q, axis=1)
p = q[np.arange(len(rows)), i]
b = k0 + i
upd = (p > best_pow[rows]) | ((p == best_pow[rows]) & (b < best_bin[rows]))
ridx = rows[upd]
best_pow[ridx] = p[upd]
best_bin[ridx] = b[upd]
def _peak_bins_zoom(waves: np.ndarray, zp: _ZoomPlan) -> np.ndarray:
"""Peak fine-bin per waveform without materialising the padded spectrum:
coarse rfft, then a small fine DFT (one gemm per shared window) on the
exact n_fft bin grid. Identity with the full padded argmax is enforced
by tests/test_compute.py::test_zoom_identity and the golden-hash sweep."""
n_wf = waves.shape[0]
S = _block_rfft(waves, zp.n_coarse)
P = S.real ** 2
P += S.imag ** 2
P[:, 0] = 0.0
p_c = np.max(P, axis=1)
best_pow = np.full(n_wf, -1.0, dtype=np.float32)
best_bin = np.full(n_wf, np.iinfo(np.int64).max, dtype=np.int64)
# For a clean signal this yields one or two windows; a noise spectrum
# (many near-equal peaks) yields a dozen or so — still a tiny fraction
# of the padded grid.
thr = np.where(p_c > 0, np.float32(_ZOOM_CAND_RATIO) * p_c,
np.float32(np.inf))
rows_c, bins_c = np.nonzero(P >= thr[:, None])
k0_c = _window_start(bins_c, zp)
# Zeroing the coarse DC bin (suppression) blinds the candidate scan to
# fine bins closer to DC than the first coarse sample — where the
# DC-leakage skirt of an un-subtracted offset peaks. Always refine the
# DC-adjacent window too.
rows_c = np.concatenate([rows_c, np.arange(n_wf)])
k0_c = np.concatenate([k0_c, np.ones(n_wf, dtype=np.int64)])
pair = np.unique(np.stack([rows_c, k0_c], axis=1), axis=0)
for k0 in np.unique(pair[:, 1]):
_refine_window(waves, pair[pair[:, 1] == k0, 0], int(k0), zp,
best_pow, best_bin)
# An all-zero spectrum must reproduce argmax-of-zeros = bin 0.
best_bin[p_c == 0.0] = 0
return best_bin
def _peak_bins_direct(waves: np.ndarray, n_len: int) -> np.ndarray:
"""Peak bin per waveform via the full transform — for pad factors too
small for the zoom search to pay."""
def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray:
"""Peak bin per waveform via the full transform. *min_bin* zeroes every
bin below it (always >= 1, so true DC stays suppressed) before the
argmax, so the peak search never returns a bin below the caller's
frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks,
fanned out across _fft_pool() — see docs/design.md."""
S = _block_rfft(waves, n_len)
power = S.real ** 2
power += S.imag ** 2
power[:, 0] = 0.0
power[:, :min_bin] = 0.0
return np.argmax(power, axis=1)
@@ -300,9 +176,11 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int,
def _plan_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int:
"""Rows per outer chunk for the block-FFT path. Only the float32
waveform buffer scales with the chunk (in-flight block spectra total a
few MB across the whole pool), so budget it with 2x slack and let the
block fan-out saturate the pool regardless of pad factor."""
waveform read buffer scales with the chunk; spectrum memory is bounded
independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET
regardless of chunk size or pad factor (see docs/design.md), so this
formula budgets the read buffer alone, with 2x slack, and lets the block
fan-out saturate the pool."""
bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
@@ -508,7 +386,8 @@ def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int:
def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
apply_bg_sub: bool, row_avg_n: int) -> list[str]:
apply_bg_sub: bool, row_avg_n: int,
min_freq_mhz: float = 0.0) -> list[str]:
"""Why this file's stored FFT images can't answer a request, as human-
readable phrases; empty means they can.
@@ -518,8 +397,23 @@ def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
serves a stored image iff this returns nothing, and callers that want to
*explain* the miss (rather than silently recompute) format these same
strings, so the two can't drift apart.
The min peak frequency floor is the asymmetric case: a stored floor
*above* the request is a genuine mismatch (bins below the stored floor
were never searched, so the stored numbers can't answer a request that
wants them considered), but a stored floor at or *below* the request is
servable — the difference is re-applied as a display-time mask by
cached_rf_image (pixels whose stored peak falls below the requested
floor are marked invalid rather than re-resolved). Compared in whole
kHz, the header field's own fixed-point grid, so a value that
round-trips through the file can never miscompare against itself.
"""
reasons = []
if round(sras.precomputed_min_freq_mhz * 1000) > round(min_freq_mhz * 1000):
reasons.append(
f"stored with a {sras.precomputed_min_freq_mhz:g} MHz "
f"min-peak-freq floor, requested {min_freq_mhz:g} MHz "
f"(bins below the stored floor were never searched)")
pad = pad_factor_for(sras, n_fft)
if pad != sras.precomputed_pad_factor:
want = f"pad {pad}x" if pad else "a ragged n_fft"
@@ -543,7 +437,8 @@ def cached_rf_image(sras: SrasFile, angle_idx: int,
n_fft: int | None = None,
dc4_mv: np.ndarray | None = None,
row_avg_n: int = 0,
allow_dc_recompute: bool = True) -> np.ndarray | None:
allow_dc_recompute: bool = True,
min_freq_mhz: float = 0.0) -> np.ndarray | None:
"""The precomputed-cache fast path for compute_rf_image: a ready-to-
display peak-frequency image if this file already has one matching
every setting the caller cares about, else None (caller must run a
@@ -571,10 +466,20 @@ def cached_rf_image(sras: SrasFile, angle_idx: int,
channel on the caller's behalf; this returns None instead so a caller
that wants to stay off the I/O path (e.g. a GUI thread) can choose to
fall through to a real compute rather than block.
*min_freq_mhz* above the stored floor is re-imposed here as a mask:
pixels whose stored peak is below it are set to 0.0 — the same sentinel
the DC mask uses, and unambiguous, since a genuine peak can never be
0.0 (bin 0 is always excluded from the search). Masked-below-floor
pixels are marked invalid, not re-resolved; only a real recompute can
recover the strongest peak *above* the floor for them. A request whose
floor is *below* the stored one can't be answered at all (see
cache_mismatch_reasons) and returns None like any other mismatch.
"""
cached_freq = sras.precomputed_freq_mhz[angle_idx]
if cached_freq is None or cache_mismatch_reasons(
sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n):
sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n,
min_freq_mhz=min_freq_mhz):
return None
dc4_img = None
if dc_threshold_mv is not None:
@@ -594,6 +499,10 @@ def cached_rf_image(sras: SrasFile, angle_idx: int,
freq_img = cached_freq.copy()
if dc4_img is not None:
freq_img[dc4_img < dc_threshold_mv] = 0.0
if min_freq_mhz > 0.0:
# Strict <, matching the compute path: the first bin at or above
# the floor survives searchsorted there, so it must survive here.
freq_img[freq_img < min_freq_mhz] = 0.0
return freq_img
@@ -605,8 +514,9 @@ 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,
use_stored: bool = True) -> 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,36 +534,59 @@ 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. A *min_freq_mhz* at or above the
stored floor is part of that service: pixels whose stored peak falls
below it come back masked to 0.0 rather than re-resolved. A request
*below* the stored floor is a genuine mismatch and falls through to
the real FFT here.
*use_stored=False* skips the fast path entirely and always runs the
real FFT. Batch recompute (cache_file) needs this: served-then-masked
pixels written back into the store would permanently replace peaks a
real recompute re-resolves above the floor.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub,
n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n)
if fast is not None:
return fast
if use_stored:
fast = cached_rf_image(sras, angle_idx, dc_threshold_mv,
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
dc4_mv=dc4_mv, row_avg_n=row_avg_n,
min_freq_mhz=min_freq_mhz)
if fast is not None:
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)
@@ -665,26 +598,15 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
dc4_full = dc4_mv
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):
@@ -748,23 +670,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.
@@ -773,18 +686,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
@@ -793,16 +698,17 @@ 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:
row_avg_n: int = 0,
min_freq_mhz: float = 0.0) -> str:
"""Compute and store DC or FFT images for every angle of one file,
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
@@ -818,6 +724,16 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
using, since a pad-1 cache is dead weight to a padded view and vice
versa (cached_rf_image refuses the mismatch rather than showing peaks
resolved at the wrong resolution).
*min_freq_mhz* is the peak-search floor (see compute_rf_image), also
recorded in the SFFT block; it is quantized up front to the header
field's 0.001 MHz grid so the stored images and the recorded floor
can never disagree. Both FFT modes force use_stored=False on
compute_rf_image: a batch recompute must always run the real FFT, never
serve this file's own existing cache back to itself — with a floor (or
a changed DC threshold in fft_rowavg mode) that would silently bake a
masked copy of the old image into the store in place of a genuine
recompute.
"""
global _MAX_WORKERS
try:
@@ -828,7 +744,12 @@ 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 not (np.isfinite(min_freq_mhz) and min_freq_mhz >= 0):
return f"min_freq_mhz must be a finite value >= 0, got {min_freq_mhz}"
# Quantize to the SFFT header's fixed-point grid (whole kHz) before
# computing, so the floor the FFT actually ran with is exactly the
# floor the header records.
min_freq_mhz = round(min_freq_mhz * 1000) / 1000.0
if max_workers:
_MAX_WORKERS = max_workers
@@ -860,13 +781,16 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
# internally over blocks, so angles run one at a time with the
# full budget.
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
apply_bg_sub=effective_bg, n_fft=n_fft)
apply_bg_sub=effective_bg, n_fft=n_fft,
min_freq_mhz=min_freq_mhz,
use_stored=False)
for a in range(n)]
# new_row_avg_n=0 explicitly: these images are raw per-pixel FFTs,
# and carrying forward a row_avg_n left by an earlier fft_rowavg
# write would label them as something they are not.
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
new_row_avg_n=0, new_pad_factor=pad_factor)
new_row_avg_n=0, new_pad_factor=pad_factor,
new_min_freq_mhz=min_freq_mhz)
else: # "fft_rowavg"
if row_avg_n <= 0:
return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode"
@@ -876,11 +800,14 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
effective_bg = apply_bg_sub and sras.background is not None
freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv,
apply_bg_sub=effective_bg, n_fft=n_fft,
row_avg_n=row_avg_n)
row_avg_n=row_avg_n,
min_freq_mhz=min_freq_mhz,
use_stored=False)
for a in range(n)]
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
new_row_avg_n=row_avg_n,
new_pad_factor=pad_factor)
new_pad_factor=pad_factor,
new_min_freq_mhz=min_freq_mhz)
return ""
except Exception as exc:
return str(exc)