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

Implement row-averaged FFT feature for same-row SNR cleanup

Browse Source 
Add row-averaged FFT feature with configurable window size (row_avg_n
parameter) for improved signal-to-noise ratio on noisy scans. Includes:

- Gaussian-weighted same-row neighbor averaging (never crosses rows)
- Masked/renormalized convolution handling for edge cases and masked samples
- Cache format v2 with row_avg_n tracking to prevent silent cache mismatches
- GUI dialog option for row-average window configuration
- Comprehensive tests validating kernel properties, background subtraction
  invariance, and cache dispatch

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-08-07 21:50:05 -05:00
parent eecb0e3a82
commit 3989c2a1b8
10 changed files with 1363 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_compute.py
+181 -27
View File
@@ -420,6 +420,118 @@ def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
*sras.cal(ch_idx))
# ---------------------------------------------------------------------------
# Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing
#
# Averages a pixel's CH1 waveform with its same-row neighbors *before* the
# FFT peak search, to improve SNR on noisy scans. Never crosses rows: pixel
# pitch is strongly anisotropic and varies by scan, but the X pitch within
# one row is a single file-wide constant (SrasFile.pixel_x_mm), so a
# Gaussian in pixel-index distance along a row and one in true physical mm
# distance are the same function up to that constant scale factor — the
# kernel itself needs no pitch, only the GUI's physical-width hint label
# does. Rationale and the masked-average/background-subtraction proofs:
# docs/design.md ("Row-averaged FFT").
# ---------------------------------------------------------------------------
_ROW_AVG_SIGMA_FRAC = 0.5 # sigma = n * this; edge weight (distance n) is
# exp(-1/(2*frac**2)) ~= 0.135 of the center tap
def _row_average_weights(row_avg_n: int) -> np.ndarray:
"""(2n+1,) float32 Gaussian weights for distance-weighted row averaging,
symmetric around the center tap. Not pre-normalized to sum to 1 —
_row_average_waveforms renormalizes per output pixel by the actual sum
of included, valid, in-window neighbor weights, not a fixed total."""
n = max(0, int(row_avg_n))
if n == 0:
return np.ones(1, dtype=np.float32)
d = np.arange(-n, n + 1, dtype=np.float64)
sigma = n * _ROW_AVG_SIGMA_FRAC
return np.exp(-0.5 * (d / sigma) ** 2).astype(np.float32)
def _row_average_waveforms(masked_waves: np.ndarray, valid: np.ndarray,
weights: np.ndarray) -> np.ndarray:
"""Distance-weighted mean of each row position's CH1 waveform with its
same-row neighbors, counting only neighbors where *valid* is True.
masked_waves: (n_frames, spf) float32 — CH1 samples at valid[f]
positions; must already be 0.0 elsewhere (the caller must never
have read the raw memmap at an invalid position).
valid: (n_frames,) bool.
weights: (2n+1,) float32 from _row_average_weights.
Returns (n_frames, spf) float32, meaningful only where valid is True
(the caller compacts by that same mask right after, matching
compute_rf_image's existing masked-compaction contract).
Two 1-D correlations along the frame axis — the numerator against the
raw masked-zero waveforms, the denominator against the validity mask
itself — so a masked neighbor contributes *zero weight* rather than a
zero-amplitude sample at full weight, and a window truncated at a row's
edge renormalizes correctly with no separate edge case: mode="constant",
cval=0.0 pads both convolutions with zero beyond the row's own ends.
"""
num = scipy_ndimage.correlate1d(masked_waves, weights, axis=0,
mode="constant", cval=0.0)
den = scipy_ndimage.correlate1d(valid.astype(np.float32), weights, axis=0,
mode="constant", cval=0.0)
den_safe = np.where(valid, den, np.float32(1.0))
return num / den_safe[:, None]
def cached_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
n_fft: int | None = None,
dc4_mv: np.ndarray | None = None,
row_avg_n: int = 0,
allow_dc_recompute: bool = True) -> 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
real FFT).
A stored image stands in only when *all* hold: this angle actually has
a stored image (v5 PREC or v7 CACH); no custom zero-padding is
requested (n_fft is None) — the store is always natural-resolution;
the stored bg-sub flag matches what the caller wants; and
sras.precomputed_row_avg_n == row_avg_n exactly (0 == "raw") — this
last check is what stops a raw request from ever being silently served
a row-averaged image, or vice versa, or a request at one window size
being served a cache stored at a different one.
If *allow_dc_recompute* is False and no DC4 image is already cached or
supplied via *dc4_mv*, applying the mask would mean reading a whole
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.
"""
cached_freq = sras.precomputed_freq_mhz[angle_idx]
if (cached_freq is None
or n_fft is not None
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)
or sras.precomputed_row_avg_n != row_avg_n):
return None
freq_img = cached_freq.copy()
if dc_threshold_mv is not None:
# DC4 mask, in priority order: already-cached DC block, caller-
# supplied image, or a fresh (cheap — no FFT) recompute.
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
if dc4_img is None:
if dc4_mv is not None:
dc4_img = dc4_mv
elif allow_dc_recompute:
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
*sras.cal(CH4_IDX))
else:
return None
freq_img[dc4_img < dc_threshold_mv] = 0.0
return freq_img
def compute_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
@@ -428,7 +540,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) -> np.ndarray:
exact: bool = False,
row_avg_n: int = 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
@@ -453,28 +566,23 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and
memory-hungry at high pad.
*row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n
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.
Fast path: if the file has a precomputed peak-frequency image for this
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
matches, the stored image is used directly — no FFT is run.
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.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
cached_freq = sras.precomputed_freq_mhz[angle_idx]
if (cached_freq is not None
and n_fft is None # no custom zero-padding
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
freq_img = cached_freq.copy()
if dc_threshold_mv is not None:
# DC4 mask, in priority order: already-cached DC block, caller-
# supplied image, or a fresh (cheap — no FFT) recompute.
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
if dc4_img is None:
dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
freq_img[dc4_img < dc_threshold_mv] = 0.0
return freq_img
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
# ---- Chunked FFT path --------------------------------------------------
exact = exact or _FFT_EXACT_ENV
@@ -484,11 +592,18 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
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.
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
@@ -523,13 +638,31 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
def read_row(i: int):
if should_stop is not None and should_stop():
return
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so only
# the selected elements are actually read from disk; masked-out
# pixels' pages are never paged in at all.
raw = data[r0 + i, CH1_IDX]
dst = waves[offs[i]:offs[i + 1]]
dst[:] = raw[valid[i]] if valid is not None else raw
row_valid = valid[i] if valid is not None else None
if row_avg_weights is not None:
v = row_valid if row_valid is not None else np.ones(n_frames, dtype=bool)
# A plain cast + broadcast multiply, not a boolean-indexed
# scatter into a zeroed buffer: numpy's fancy indexing holds
# the GIL for its whole duration (measured zero speedup
# across threads, and net negative once threads outnumber
# physical cores), while a cast and a multiply are ordinary
# ufuncs that release it — this is what lets read_row actually
# parallelize across the pool instead of serializing on the
# scatter/gather. Trade-off: every sample in the row is read,
# valid or not (row averaging needs broad neighbor context
# regardless, unlike the plain path below, which still skips
# masked-out pixels entirely).
full = raw.astype(np.float32) * v[:, None]
avg = _row_average_waveforms(full, v, row_avg_weights)
dst[:] = avg[v] if row_valid is not None else avg
else:
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so
# only the selected elements are actually read from disk;
# masked-out pixels' pages are never paged in at all.
dst[:] = raw[row_valid] if row_valid is not None else raw
if background is not None:
dst -= background # background is 1-D (spf,)
@@ -589,7 +722,9 @@ 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) -> str:
fft_backend: str = "scipy", max_workers: int = 0,
dc_threshold_mv: float | None = None,
row_avg_n: int = 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.
@@ -597,14 +732,21 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
FFT backend and worker cap are passed explicitly because module globals
do not survive a spawn.
mode is "dc", "fft" (raw per-pixel FFT, natural-resolution, unmasked —
masking applied at display time), or "fft_rowavg" (same-row,
distance-weighted CH1 averaging before the FFT — see compute_rf_image's
row_avg_n). fft_rowavg needs *dc_threshold_mv* up front, unlike plain
"fft": neighbor validity is baked into the stored numbers, so it can't
be deferred to display time the way plain masking can.
The stored FFT cache is always natural-resolution (pad 1): the v7 SFFT
block records no pad factor, and padded views compute live fast enough
(see _peak_bins_zoom) that caching them is not worth a format change.
"""
global _MAX_WORKERS
try:
if mode not in ("dc", "fft"):
return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
if mode not in ("dc", "fft", "fft_rowavg"):
return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')"
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers
@@ -628,7 +770,7 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
budget=angle_budget), *sras.cal(CH4_IDX)),
range(n), n_workers)
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
else:
elif mode == "fft":
effective_bg = apply_bg_sub and sras.background is not None
# dc_threshold_mv=None: store unmasked images and mask at display
# time (same convention as v5's PREC block). Skipping the mask
@@ -639,6 +781,18 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
apply_bg_sub=effective_bg)
for a in range(n)]
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
else: # "fft_rowavg"
if row_avg_n <= 0:
return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode"
if dc_threshold_mv is None:
return ("fft_rowavg mode requires a DC threshold "
"(neighbor validity depends on it)")
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, row_avg_n=row_avg_n)
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)
return ""
except Exception as exc:
return str(exc)