3 Commits

Author SHA1 Message Date
Thomas Ales 3989c2a1b8 Implement row-averaged FFT feature for same-row SNR cleanup
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>
2026-08-07 21:50:05 -05:00
Thomas Ales eecb0e3a82 Merge remote-tracking branch 'origin/refactor/parallel-and-dedup' 2026-08-06 21:20:58 -05:00
tka b0bfe6e8c6 Merge pull request 'Refactor/parallel and dedup' (#1) from refactor/parallel-and-dedup into main
Reviewed-on: #1
2026-07-31 15:20:33 -05:00
11 changed files with 1367 additions and 294 deletions
+62
View File
@@ -72,6 +72,68 @@ one thread under the pool so the refinement gemm cannot oversubscribe.
`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
full-padded path for audits.
## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
`compute_rf_image`'s `row_avg_n` parameter averages each pixel's CH1
waveform with its up-to-n 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 (5 µm × 50 µm on a typical scan, but as
stretched as 5 µm × 1 mm on others), so a physically meaningful "neighbor"
set can't be a fixed-shape 2-D window — but the X pitch *within one row* is
a single file-wide constant (`SrasFile.pixel_x_mm`), so restricting to the
row axis sidesteps the anisotropy question entirely rather than solving it
with an elliptical or physically-scaled 2-D kernel.
`_row_average_weights` is a Gaussian in pixel-index distance, not physical
mm distance — deliberately: within one row those are the same function up
to a fixed scale factor (`pixel_x_mm` is constant along a row), so the
kernel itself needs no pitch at all. `pixel_x_mm` is used for real exactly
once, in the GUI's options dialog, to show the window's physical width —
not in the kernel math, where it would only ever cancel out.
`_row_average_waveforms` is a masked/renormalized convolution (two
`correlate1d` calls, numerator and denominator, divided) rather than a
single fixed-normalized convolution, because a masked neighbor must
contribute *zero weight*, not a zero-amplitude sample at full weight — the
latter would bias every average near a masked run or a row's own edge
toward zero. The same two-correlation trick handles row-edge truncation for
free: `mode="constant", cval=0.0` zero-pads both the numerator and the
denominator beyond a row's own ends, so the output renormalizes by whatever
weight sum actually landed inside the row, no separate edge case.
Background subtraction stays exactly where it already was (subtracted once
from the fully-assembled `waves` buffer) rather than being threaded into the
per-neighbor gather. This is exact, not an approximation: because
`_row_average_waveforms`'s denominator is always the *actual* sum of
included, valid weights (never a fixed total), `Σwᵢ·(rawᵢ−bg) / Σwᵢ`
distributes to `avg − bg·(Σwᵢ/Σwᵢ) = avg − bg` regardless of which or how
many neighbors were included — subtracting background from the averaged
waveform is identical to subtracting it from every neighbor first, for any
window, at any row edge, with any number of masked-out neighbors.
No cross-row halo is needed: `compute_rf_image`'s chunk loop already splits
on rows only, and `read_row` already reads one row's complete
`(n_frames, spf)` slice at a time — averaging happens entirely inside that
one row's own frame axis, so a chunk boundary (which falls between rows)
can never truncate a window. Only a row's own start/end can, and that's the
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
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.
Persistence: `cached_rf_image` (the extracted fast-path check) requires
`sras.precomputed_row_avg_n == row_avg_n` exactly, so a raw request can
never be silently served a row-averaged cache or vice versa, and a request
at one window size can never be served a cache at another — see
`scan_format.md`'s Cache Tail / CACH tail version history sections for the
on-disk `row_avg_n` field this depends on.
## Angle alignment coordinate frames (`sras_compute.py`)
Alignment puts every angle's images onto one shared, zero-padded pixel grid
+39 -8
View File
@@ -220,7 +220,7 @@ actions.
| Offset | Size | Type | Field | Description |
|--------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `cach_magic` | `CACH` (ASCII). Missing/wrong magic → treat file as having no cache. |
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `1`. Readers must treat the file as uncached if this is not a version they understand (unlike v5's `PREC` section, which read but never validated its version byte). |
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `2`; readers also accept `1` (a `1` tail predates row-averaged FFT caching — see the `SFFT` block below and [CACH tail version history](#cach-tail-version-history)). Any other value → treat the file as uncached (unlike v5's `PREC` section, which read but never validated its version byte). |
| 5 | 1 | `u8` | `block_flags` | Bit 0 = DC block (`SDCB`) follows. Bit 1 = FFT block (`SFFT`) follows, immediately after the DC block if both are present. Bits 2–7 reserved, must be zero on write. |
### DC block `SDCB` (present iff `block_flags & 0x01`)
@@ -249,13 +249,20 @@ value, same as v5's `PREC` section.
### FFT block `SFFT` (present iff `block_flags & 0x02`)
7-byte block header, format `">4sBH"`:
Block header layout depends on `cach_version`:
- **`cach_version` 1**: 7 bytes, format `">4sBH"` — magic, flags, n_stored.
- **`cach_version` 2**: 8 bytes, format `">4sBHB"` — magic, flags, n_stored,
`row_avg_n`. Always written by current code; a `cach_version` 1 tail (no
trailing byte) is still read, with `row_avg_n` taken as `0` for every
entry it stores.
| Offset (rel) | Size | Type | Field | Description |
|--------------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `magic` | `SFFT` |
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bits 1–7 reserved. |
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bit 1 = `row_averaged` — `peak_freq_mhz` came from same-row, distance-weighted averaged CH1 waveforms rather than raw per-pixel ones; `row_avg_n` (below) is the neighbor half-width used. Bits 2–7 reserved. |
| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow |
| 7 | 1 | `u8` | `row_avg_n` | *`cach_version` 2 only.* Same-row neighbor half-width, in pixels, that `peak_freq_mhz` was averaged over before its FFT; `0` = raw (unaveraged). Meaningful only when `flags` bit 1 is set — a `cach_version` 1 tail has no such byte and is always `row_avg_n = 0`. |
followed by `n_stored` entries, each:
@@ -264,9 +271,10 @@ u16 angle_idx — index into the angle table (0-ba
f32[n_rows[angle_idx] × n_frames[angle_idx]] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
```
**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the
FFT is run on every pixel unconditionally, same as v5's `PREC` convention).
Readers apply the DC4 threshold at display time:
**`peak_freq_mhz`** for a raw store (`row_avg_n == 0`) is computed without
any DC-threshold masking (i.e. the FFT is run on every pixel
unconditionally, same as v5's `PREC` convention). Readers apply the DC4
threshold at display time:
```
pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
@@ -276,10 +284,22 @@ display_value = peak_freq_mhz[r][f] if valid, else 0
using the DC4 image from the DC block if that angle is also cached there,
else computed on demand.
For a row-averaged store (`row_avg_n > 0`), the DC4 threshold is applied
*during* the store — a pixel below threshold is left at `0` and never
contributes to any neighbor's average — since neighbor validity can't be
deferred to display time the way plain masking can. The threshold value
itself is not recorded, only that averaging happened and at what window
size. Readers still apply their own live DC4 threshold at display time
exactly as for a raw store, using whatever mask they currently have.
Readers must fall back to real-time FFT computation (ignoring stored
`peak_freq_mhz`) under the same conditions as v5's PREC fast path: time-domain
gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested, or
the reader's background-subtraction setting doesn't match `flags.bg_sub_applied`.
gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested,
the reader's background-subtraction setting doesn't match
`flags.bg_sub_applied`, or the reader's requested `row_avg_n` doesn't match
the stored value exactly — a raw request must never be served a
row-averaged store, or vice versa, and a request at one window size must
never be served a store at another.
### In-place write ordering
@@ -294,6 +314,17 @@ interrupted write leaves harmless trailing bytes rather than a corrupt file,
and the next successful write overwrites them via the same deterministic
`cache_offset`.
### CACH tail version history
Distinct from the outer `.sras` file `version` byte (top of this document),
which has stayed `7` since the Cache Tail was introduced — this is the inner
`cach_version` byte inside the `CACH` header itself.
| cach_version | Change |
|--------------|--------|
| 1 | Initial Cache Tail: `SDCB` (DC) and `SFFT` (FFT, 7-byte header) blocks. |
| 2 | `SFFT` header grows one byte, `row_avg_n` — the same-row neighbor half-width the stored `peak_freq_mhz` was averaged over before its FFT, `0` = raw. Readers still accept a `cach_version` 1 tail, treated as `row_avg_n = 0` for every angle it stores, so files cached before this change keep working without a recompute. |
---
## Acquisition Settings (fixed by sc3_aui_app.py)
+170 -65
View File
@@ -18,9 +18,7 @@ import numpy as np
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
from sras_format import (
CH1_IDX, CH3_IDX, CH4_IDX, PAD_FACTOR_MAX, SrasFile, adc_to_mv,
)
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
# ---------------------------------------------------------------------------
# FFT backend
@@ -422,52 +420,114 @@ 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 stored peak-frequency image for one angle (v5 PREC or v7 CACH),
masked and ready to display — or None if no stored image can serve these
settings and a real FFT is needed.
"""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).
Stored images are unmasked, so they serve any threshold — but they carry
one bg-sub state and one zero-padding, and a view asking for either of the
others is asking for different numbers, not a different rendering of the
same ones. Padding is the sharp edge: a pad-10 view resolves peaks on a
10x finer bin grid, so handing back a pad-1 image would quietly undo the
setting. v5 PREC and CACH v1 tails record no pad and are natural
resolution by construction.
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.
Cheap enough to call on the GUI thread — a copy plus a comparison —
*except* when the CH4 mask has to be read from disk to build it. Pass
*allow_dc_recompute* False to return None in that case instead, and leave
the whole-channel read to a worker.
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]
# n_fft None means "no padding", i.e. exactly samples_per_frame points.
want_n_fft = n_fft if n_fft is not None else sras.samples_per_frame
if (cached_freq is None
or want_n_fft != sras.precomputed_pad_factor * sras.samples_per_frame
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not 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
if dc_threshold_mv is None:
return cached_freq.copy()
# DC4 mask, in priority order: stored DC block, caller-supplied image,
# or a fresh (cheap — no FFT) recompute.
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
if dc4_img is None:
if not allow_dc_recompute:
return None
elif allow_dc_recompute:
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
*sras.cal(CH4_IDX))
freq_img = cached_freq.copy()
else:
return None
freq_img[dc4_img < dc_threshold_mv] = 0.0
return freq_img
@@ -480,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
@@ -505,22 +566,25 @@ 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.
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. That test is
cached_rf_image, which the viewer also applies before it ever dispatches
a compute at all.
"""
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
cached = cached_rf_image(sras, angle_idx, dc_threshold_mv,
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
dc4_mv=dc4_mv)
if cached is not None:
return cached
*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.
# ---- Chunked FFT path --------------------------------------------------
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.
"""
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
# ---- 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
@@ -528,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
@@ -567,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,)
@@ -634,23 +723,30 @@ 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,
pad_factor: int = 1) -> str:
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.
Module-level and picklable so it can run in a ProcessPoolExecutor. The
FFT backend, worker cap and pad factor are passed explicitly because
module globals do not survive a spawn.
FFT backend and worker cap are passed explicitly because module globals
do not survive a spawn.
*pad_factor* is stored alongside the images: a cache is only usable by a
view asking for the same padding, so caching at the viewer's own setting
is the difference between a batch that pays off and one that can never be
read back (see cached_rf_image).
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
@@ -674,20 +770,29 @@ 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
pad = max(1, min(PAD_FACTOR_MAX, int(pad_factor)))
n_fft = None if pad <= 1 else sras.samples_per_frame * pad
# dc_threshold_mv=None: store unmasked images and mask at display
# time (same convention as v5's PREC block). Skipping the mask
# also skips reading CH4 entirely. The FFT path parallelises
# 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)
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_pad_factor=pad)
new_row_avg_n=row_avg_n)
return ""
except Exception as exc:
return str(exc)
+66 -40
View File
@@ -60,12 +60,11 @@ PREC_FLAG_BG_SUB = 0x01
CACH_MAGIC = b"CACH"
CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags
CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT)
# v2 added the SFFT pad factor. Writers always emit v2; readers still accept
# v1, whose FFT block is natural-resolution by construction. An older reader
# meeting a v2 tail rejects the whole section on the version check and
# recomputes — no cache, but no misreading either.
CACH_VERSION = 2
CACH_VERSION_MIN = 1
CACH_VERSION = 2 # written on every fresh write
CACH_VERSIONS_READABLE = (1, 2) # accepted on read — see
# _read_sfft_block: a v1 tail
# predates row-averaged FFT
# caching and reads as row_avg_n=0
CACH_FLAG_DC = 0x01
CACH_FLAG_FFT = 0x02
@@ -74,12 +73,16 @@ SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT)
SFFT_MAGIC = b"SFFT"
SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored
SFFT_HDR_FMT = ">4sBHH" # magic, flags, n_stored, pad_factor
SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored (cach_version 1)
SFFT_HDR_FMT = ">4sBHB" # + row_avg_n (cach_version 2)
SFFT_HDR_SIZE_V1 = struct.calcsize(SFFT_HDR_FMT_V1)
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
SFFT_FLAG_BG_SUB = 0x01
# The viewer clamps its pad factor to this, and the field is a uint16.
PAD_FACTOR_MAX = 256
SFFT_FLAG_ROW_AVG = 0x02 # peak_freq_mhz came from same-row,
# distance-weighted averaged CH1
# waveforms, not raw per-pixel ones;
# row_avg_n is the neighbor half-width
# (pixels) used. Bits 2-7 reserved.
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
@@ -177,9 +180,7 @@ class SrasFile:
``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``,
always as ragged per-angle lists (``list[np.ndarray | None]``, one entry
per angle, ``None`` where that angle was never stored) regardless of
source version. The settings the FFT images were computed under travel
with them as ``precomputed_bg_sub`` / ``precomputed_pad_factor``, since a
stored image is only usable by a view asking for the same two.
source version.
"""
def __init__(self, path: str):
@@ -236,9 +237,7 @@ class SrasFile:
self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_bg_sub: bool = False
# Zero-padding the stored FFT images were computed at; 1 = natural
# resolution, which is all a v5 PREC or CACH v1 tail can hold.
self.precomputed_pad_factor: int = 1
self.precomputed_row_avg_n: int = 0
def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""A stored DC image (already in mV) for (angle, channel), or None."""
@@ -509,18 +508,16 @@ class SrasFile:
return end
def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
stores: list[list]) -> tuple | None:
stores: list[list]) -> int | None:
"""Read one CACH sub-block header, then its per-angle image entries
into *stores* (one list per image the block stores per angle).
Every sub-block header is (magic, flags, n_stored, *extras). Returns
everything after the magic, so a caller that has extras knows how to
read them, or None if the block is malformed.
Returns the header's flags byte, or None if the block is malformed.
"""
raw = f.read(struct.calcsize(hdr_fmt))
if len(raw) < struct.calcsize(hdr_fmt):
return None
block_magic, flags, n_stored, *extras = struct.unpack(hdr_fmt, raw)
block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
if block_magic != magic:
return None
for _ in range(n_stored):
@@ -530,7 +527,35 @@ class SrasFile:
shape = self.image_shape(angle_idx)
for store in stores:
store[angle_idx] = _read_f32_image(f, shape)
return (flags, *extras)
return flags
def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int] | None:
"""Read the SFFT block header — its layout depends on cach_version,
since v2 appended a trailing row_avg_n byte — then n_stored per-angle
peak_freq_mhz entries (unchanged across versions).
Returns (flags, row_avg_n), or None if the block is malformed.
row_avg_n is always 0 for a v1 tail, which predates row-averaged FFT
caching entirely.
"""
hdr_fmt = SFFT_HDR_FMT_V1 if cach_version == 1 else SFFT_HDR_FMT
raw = f.read(struct.calcsize(hdr_fmt))
if len(raw) < struct.calcsize(hdr_fmt):
return None
if cach_version == 1:
magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
row_avg_n = 0
else:
magic, flags, n_stored, row_avg_n = struct.unpack(hdr_fmt, raw)
if magic != SFFT_MAGIC:
return None
for _ in range(n_stored):
(angle_idx,) = _read_struct(f, ">H")
if angle_idx >= self.n_angles:
break
self.precomputed_freq_mhz[angle_idx] = _read_f32_image(
f, self.image_shape(angle_idx))
return flags, row_avg_n
def _parse_cach_section(self, offset: int):
"""Parse the v7 CACH tail that holds precomputed DC/FFT images."""
@@ -540,8 +565,7 @@ class SrasFile:
if len(header_raw) < CACH_HDR_SIZE:
return
magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
if (magic != CACH_MAGIC
or not CACH_VERSION_MIN <= cach_version <= CACH_VERSION):
if magic != CACH_MAGIC or cach_version not in CACH_VERSIONS_READABLE:
return
if block_flags & CACH_FLAG_DC:
@@ -551,23 +575,19 @@ class SrasFile:
return
if block_flags & CACH_FLAG_FFT:
# v1 has no pad field: those images are natural-resolution.
hdr_fmt = SFFT_HDR_FMT if cach_version >= 2 else SFFT_HDR_FMT_V1
header = self._read_cache_block(
f, hdr_fmt, SFFT_MAGIC, [self.precomputed_freq_mhz])
if header is None:
result = self._read_sfft_block(f, cach_version)
if result is None:
return
flags, *extras = header
flags, row_avg_n = result
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
self.precomputed_pad_factor = (
min(max(1, extras[0]), PAD_FACTOR_MAX) if extras else 1)
self.precomputed_row_avg_n = row_avg_n if (flags & SFFT_FLAG_ROW_AVG) else 0
def write_v7_cache(self, *,
new_dc3_mv: list[np.ndarray | None] | None = None,
new_dc4_mv: list[np.ndarray | None] | None = None,
new_freq_mhz: list[np.ndarray | None] | None = None,
new_bg_sub: bool | None = None,
new_pad_factor: int | None = None):
new_row_avg_n: int | None = None):
"""Store computed DC and/or FFT images into this file's CACH tail,
in place, converting a v6 source to v7 (or updating an existing v7
file). Only the block(s) passed in are recomputed; whichever block
@@ -575,6 +595,12 @@ class SrasFile:
``SrasFile`` already has in memory (from parsing, or a prior write
in this same session) — its bytes are never re-read from disk.
*new_row_avg_n* is the same-row neighbor half-width (pixels) the
passed *new_freq_mhz* was averaged over before its FFT, 0 for a raw
(unaveraged) compute — carried forward like *new_bg_sub* when None.
It describes the whole stored FFT block, not per-angle, mirroring
how bg-sub has never been tracked per-angle either.
The waveform data itself is never touched: the cache tail always
starts at ``_cache_tail_offset()``, a fixed offset derived from the
header and geometry table alone.
@@ -587,11 +613,10 @@ class SrasFile:
final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv
final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz
final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub
final_pad = (new_pad_factor if new_pad_factor is not None
else self.precomputed_pad_factor)
if not 1 <= final_pad <= PAD_FACTOR_MAX:
raise ValueError(
f"pad factor {final_pad} outside 1..{PAD_FACTOR_MAX}")
final_row_avg_n = (new_row_avg_n if new_row_avg_n is not None
else self.precomputed_row_avg_n)
if not (0 <= final_row_avg_n <= 255):
raise ValueError(f"row_avg_n must fit in a byte (0-255), got {final_row_avg_n}")
dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]
@@ -611,8 +636,9 @@ class SrasFile:
if fft_entries:
fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
fft_flags |= SFFT_FLAG_ROW_AVG if final_row_avg_n else 0
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags,
len(fft_entries), final_pad)
len(fft_entries), final_row_avg_n)
for a in fft_entries:
payload += struct.pack(">H", a)
payload += final_freq[a].astype(">f4").tobytes()
@@ -640,7 +666,7 @@ class SrasFile:
self.precomputed_dc4_mv = final_dc4
self.precomputed_freq_mhz = final_freq
self.precomputed_bg_sub = final_bg_sub
self.precomputed_pad_factor = final_pad
self.precomputed_row_avg_n = final_row_avg_n
# ------------------------------------------------------------------
# Axes helpers
+3 -1
View File
@@ -20,5 +20,7 @@ faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
from .canvases import ImageCanvas, RoiQuad, WaveformCanvas # noqa: E402,F401
from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401
from .dialogs import FftOptionsDialog, ManualAlignmentDialog # noqa: E402,F401
from .dialogs import ( # noqa: E402,F401
FftOptionsDialog, ManualAlignmentDialog, RowAverageFftOptionsDialog,
)
from .main_window import SrasViewerWindow, main # noqa: E402,F401
+100 -18
View File
@@ -48,8 +48,7 @@ class FftOptionsDialog(QDialog):
current_pad_factor: int,
samples_per_frame: int | None,
sample_rate_hz: float | None,
grating_um: float,
cached_pad_factor: int | None = None):
grating_um: float):
super().__init__(parent)
self.setWindowTitle("FFT Options")
self.setModal(True)
@@ -58,10 +57,6 @@ class FftOptionsDialog(QDialog):
self._samples_per_frame = samples_per_frame
self._sample_rate_hz = sample_rate_hz
self._grating_um = grating_um
# The pad the open file's stored FFT cache was computed at, if it has
# one — picking anything else here makes that cache unreadable, which
# is worth saying before Apply rather than after.
self._cached_pad_factor = cached_pad_factor
layout = QVBoxLayout(self)
@@ -114,11 +109,6 @@ class FftOptionsDialog(QDialog):
lbl.setStyleSheet(_CSS_HINT)
zl.addWidget(lbl)
self._lbl_cache = QLabel()
self._lbl_cache.setWordWrap(True)
self._lbl_cache.setStyleSheet(_CSS_WARN)
zl.addWidget(self._lbl_cache)
layout.addWidget(grp_zp)
# ---- Buttons ---------------------------------------------------
@@ -136,13 +126,6 @@ class FftOptionsDialog(QDialog):
sr = self._sample_rate_hz
pad = self._spin_pad.value()
self._lbl_cache.setText(
"" if self._cached_pad_factor in (None, pad) else
f"! This file's stored FFT cache was computed at pad "
f"{self._cached_pad_factor}x, so at {pad}x it can't be used and "
f"CH1/Velocity will recompute. Re-run Convert -> Batch Compute FFT "
f"to cache at {pad}x.")
if spf is None or sr is None:
self._lbl_nfft.setText("Load a file to preview FFT parameters.")
self._lbl_freq_res.setText("")
@@ -169,6 +152,105 @@ class FftOptionsDialog(QDialog):
return max(1, self._spin_pad.value())
# ---------------------------------------------------------------------------
# Row-Averaged FFT Options dialog
# ---------------------------------------------------------------------------
class RowAverageFftOptionsDialog(QDialog):
"""Configure the same-row, distance-weighted neighbor averaging applied
to each pixel's CH1 waveform before 'Batch Compute Row-Averaged FFT and
Store' re-runs the FFT peak search — a same-row SNR cleanup pass, never
mixing across rows/Y (see sras_compute._row_average_waveforms).
Unlike the plain FFT batch action (which stores unmasked and defers
masking to display time), the DC threshold here is required up front:
it decides which same-row neighbors are eligible to contribute to a
pixel's average, so it can't be deferred.
Changes take effect only when the user clicks Apply. Cancel discards
all pending edits.
"""
def __init__(self, parent=None, *,
current_n: int,
current_threshold_mv: float,
pixel_x_mm: float | None):
super().__init__(parent)
self.setWindowTitle("Row-Averaged FFT Options")
self.setModal(True)
self.setMinimumWidth(380)
self._pixel_x_mm = pixel_x_mm
layout = QVBoxLayout(self)
# ---- Neighbor window ---------------------------------------------
grp_window = QGroupBox("Same-Row Neighbor Window")
wl = QVBoxLayout(grp_window)
n_row = QHBoxLayout()
n_row.addWidget(QLabel("Neighbor half-width (n):"))
self._spin_n = QSpinBox()
self._spin_n.setRange(1, 50)
self._spin_n.setValue(max(1, current_n))
self._spin_n.setToolTip(
"Each pixel's CH1 waveform is averaged with up to n same-row\n"
"neighbors on each side, distance-weighted (Gaussian) and\n"
"counting only neighbors that already pass the DC threshold\n"
"below. Never mixes across rows/Y.")
self._spin_n.valueChanged.connect(self._update_info)
n_row.addWidget(self._spin_n)
wl.addLayout(n_row)
self._lbl_width = QLabel()
self._lbl_width.setStyleSheet(_CSS_HINT)
wl.addWidget(self._lbl_width)
layout.addWidget(grp_window)
# ---- DC threshold ------------------------------------------------
grp_thr = QGroupBox("Neighbor Validity")
tl = QVBoxLayout(grp_thr)
thr_row = QHBoxLayout()
thr_row.addWidget(QLabel("DC threshold:"))
self._spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
value=current_threshold_mv, step=0.025)
self._spin_threshold.setToolTip(
"A same-row neighbor only contributes to a pixel's average if\n"
"its own CH4 signal is at or above this threshold -- the same\n"
"test used for RF mask display. A pixel below threshold stays\n"
"masked, exactly as today; it is never rescued by its neighbors.")
thr_row.addWidget(self._spin_threshold)
tl.addLayout(thr_row)
layout.addWidget(grp_thr)
# ---- Buttons -----------------------------------------------------
buttons = QDialogButtonBox()
buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
).clicked.connect(self.accept)
buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
).clicked.connect(self.reject)
layout.addWidget(buttons)
self._update_info()
def _update_info(self):
n = self._spin_n.value()
if self._pixel_x_mm is None:
self._lbl_width.setText("Load a file to preview the window's physical width.")
return
width_um = 2 * n * self._pixel_x_mm * 1e3
self._lbl_width.setText(
f"Window: ±{n} px = {width_um:.2f} µm full width "
f"(pixel pitch {self._pixel_x_mm * 1e3:.3g} µm)")
def get_half_width(self) -> int:
return self._spin_n.value()
def get_threshold_mv(self) -> float:
return self._spin_threshold.value()
class ManualAlignmentDialog(QDialog):
"""Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...).
+74 -102
View File
@@ -33,7 +33,7 @@ from .common import (
_CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W,
_RIGHT_PANEL_W, Jobs, _form, _group, _make_dspin, _scroll_panel, _wrap_label,
)
from .dialogs import FftOptionsDialog, ManualAlignmentDialog
from .dialogs import FftOptionsDialog, ManualAlignmentDialog, RowAverageFftOptionsDialog
# ---------------------------------------------------------------------------
# Main window
@@ -73,21 +73,22 @@ class SrasViewerWindow(QMainWindow):
except (TypeError, ValueError):
pad = 1
self._fft_pad_factor: int = max(1, min(256, pad)) # 1 = no padding
try:
row_avg_n = int(self._settings.value("fft/row_avg_n", 3))
except (TypeError, ValueError):
row_avg_n = 3
self._pending_row_avg_n: int = max(1, min(50, row_avg_n))
# Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
self._batch_errors: list[str] = []
# Display-only settings (colormap, grating) never trigger a
# recompute — they're applied to cached data on redraw. DC images
# (CH3/CH4) are cheap: they come straight from the file's stored
# cache when Batch Compute DC has been run for it (see
# _stored_dc_image), and are otherwise precomputed for every angle
# in the background right after load. CH1/Velocity FFT images come
# from the file's own stored cache when Batch Compute FFT has been
# run for it (see _stored_fft_image), and are otherwise computed
# lazily (with a progress popup) the first time an angle/threshold
# combination is viewed — using the cached DC4 image to skip the
# FFT entirely for masked-out pixels. Either way the result is
# (CH3/CH4) are cheap and precomputed for every angle in the
# background right after load. CH1/Velocity FFT images are
# computed lazily (with a progress popup) the first time an
# angle/threshold combination is viewed — using the cached DC4
# image to skip the FFT entirely for masked-out pixels — and
# cached per (angle, bg_sub, n_fft, threshold) so revisiting the
# same combination is free.
self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
@@ -223,14 +224,7 @@ class SrasViewerWindow(QMainWindow):
self.spin_angle.setRange(0, 0)
self.spin_angle.setEnabled(False)
self.spin_angle.setMinimumWidth(64)
# Keyboard tracking off is what makes valueChanged safe to act on: it
# stops the signal firing per keystroke, so typing "12" is one change
# to angle 12 (on Return or focus-out) rather than a trip to angle 1
# first. Arrow clicks and the wheel still emit immediately, which
# editingFinished alone did not — that is why stepping the angle used
# to leave the plot on the previous one until the box lost focus.
self.spin_angle.setKeyboardTracking(False)
self.spin_angle.valueChanged.connect(self._on_view_changed)
self.spin_angle.editingFinished.connect(self._on_view_changed)
self.lbl_angle_deg = QLabel("—")
angle_field = QWidget()
ar = QHBoxLayout(angle_field)
@@ -470,6 +464,17 @@ class SrasViewerWindow(QMainWindow):
self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
convert_menu.addAction(self._batch_fft_act)
self._batch_fft_rowavg_act = QAction(
"Batch Compute Row-A&veraged FFT and Store…", self)
self._batch_fft_rowavg_act.setStatusTip(
"Select .sras files and compute+store a same-row, distance-weighted "
"smoothed FFT peak-frequency image for every angle — improves SNR "
"on noisy regions. DC images and raw waveform data are never "
"touched; the raw (unsmoothed) FFT is always a recompute away. "
"Converts v6 files to v7 in place.")
self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg)
convert_menu.addAction(self._batch_fft_rowavg_act)
# ------------------------------------------------------------------
# Drag-and-drop
# ------------------------------------------------------------------
@@ -609,32 +614,15 @@ class SrasViewerWindow(QMainWindow):
n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None)
if n_dc or n_fft:
bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
pad_note = (f", pad {s.precomputed_pad_factor}x"
if n_fft and s.precomputed_pad_factor > 1 else "")
avg_note = (f", row-averaged n={s.precomputed_row_avg_n}"
if s.precomputed_row_avg_n else "")
notes.append(
f"Cached images: DC {n_dc}/{s.n_angles} angles, "
f"FFT {n_fft}/{s.n_angles} angles"
f"{bg_note + pad_note if n_fft else ''} "
f"FFT {n_fft}/{s.n_angles} angles{bg_note if n_fft else ''}"
f"{avg_note if n_fft else ''} "
"— display is instant for cached angles")
elif s.version == 7:
notes.append("v7 format: no cache blocks stored yet")
# A stored FFT cache the current settings can't read back is the one
# failure here with no other symptom: the images are right there in
# the file and every angle still recomputes. Say why.
if n_fft:
reasons = []
if s.precomputed_pad_factor != self._fft_pad_factor:
reasons.append(f"stored at pad {s.precomputed_pad_factor}x, "
f"FFT Options is set to {self._fft_pad_factor}x")
if s.precomputed_bg_sub != (self.chk_bg_sub.isChecked()
and s.background is not None):
reasons.append("stored with bg-sub "
f"{'on' if s.precomputed_bg_sub else 'off'}")
if reasons:
notes.append(f"! Cached FFT unusable as configured ({'; '.join(reasons)})"
" — CH1/Velocity will recompute. Re-run Convert -> "
"Batch Compute FFT to cache at the current settings.")
self.lbl_frame_warn.setText("\n".join(notes))
# ------------------------------------------------------------------
@@ -671,6 +659,7 @@ class SrasViewerWindow(QMainWindow):
can_batch = not self._job_running(Jobs.BATCH)
self._batch_dc_act.setEnabled(can_batch)
self._batch_fft_act.setEnabled(can_batch)
self._batch_fft_rowavg_act.setEnabled(can_batch)
self._alignment_act.setEnabled(
has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN))
@@ -687,9 +676,6 @@ class SrasViewerWindow(QMainWindow):
# Background subtraction changes the FFT input, so it genuinely
# invalidates the cached raw FFT (the cache key includes it) —
# _refresh_display() recomputes only on a miss for the new state.
# It can also decide whether the file's stored cache is readable, so
# the info panel's verdict on that has to be redrawn with it.
self._update_scan_info_labels()
if self._is_fft_mode():
self._refresh_display()
@@ -872,14 +858,6 @@ class SrasViewerWindow(QMainWindow):
return None
return self._sras.samples_per_frame * self._fft_pad_factor
def _cached_fft_pad_factor(self) -> int | None:
"""The pad factor the open file's stored FFT images were computed at,
or None if it has none stored."""
if self._sras is None or not any(x is not None
for x in self._sras.precomputed_freq_mhz):
return None
return self._sras.precomputed_pad_factor
def _is_fft_mode(self) -> bool:
"""Is the selected channel an FFT-derived (CH1/Velocity) mode?"""
return (self._sras is not None
@@ -919,45 +897,6 @@ class SrasViewerWindow(QMainWindow):
self._aligned_cache[key] = cached
return cached
def _stored_fft_image(self, angle_idx: int) -> np.ndarray | None:
"""The file's own batch-computed FFT image for this angle, masked for
the current threshold — or None if it can't serve the current settings.
Batch Compute FFT stores a peak-frequency image per angle in the file
itself, and the whole point of paying for that once is that display is
then instant. Without this check the viewer only ever found the stored
image deep inside ComputeWorker, so every angle change after a batch
still dispatched a background job behind a "Computing FFT…" popup for
an image that was already sitting on disk.
"""
img = compute.cached_rf_image(
self._sras, angle_idx,
dc_threshold_mv=self.spin_threshold_mv.value(),
apply_bg_sub=self.chk_bg_sub.isChecked(),
n_fft=self._current_n_fft(),
dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
# Masking is a copy plus a comparison, but *reading* CH4 to build
# the mask is a whole-channel read — that case belongs on a worker,
# so it returns None here and falls through to _start_compute.
allow_dc_recompute=False)
if img is not None:
self._fft_cache[self._fft_cache_key(angle_idx)] = img
return img
def _stored_dc_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""The file's own batch-computed DC image for (angle, channel), if it
has one — the CH3/CH4 half of _stored_fft_image.
The worker this saves would have handed back the very same stored
array (dc_image_mv prefers it over recomputing), so all it ever cost
was a thread and a progress popup — but that is exactly what made
"Batch Compute DC and Store" look like it had done nothing.
"""
img = self._sras.cached_dc_mv(angle_idx, ch_idx)
if img is not None:
self._dc_cache[(angle_idx, ch_idx)] = img
return img
def _refresh_display(self):
"""Show the image for the current angle/channel/threshold, using
cached data whenever possible and only falling back to a background
@@ -969,16 +908,12 @@ class SrasViewerWindow(QMainWindow):
if ch_idx in CH1_DERIVED_MODES:
raw = self._fft_cache.get(self._fft_cache_key(angle_idx))
if raw is None:
raw = self._stored_fft_image(angle_idx)
if raw is not None:
self._show_image_now(self._scale_for_display(raw, ch_idx),
angle_idx, ch_idx)
return
else:
cached = self._dc_cache.get((angle_idx, ch_idx))
if cached is None:
cached = self._stored_dc_image(angle_idx, ch_idx)
if cached is not None:
self._show_image_now(cached, angle_idx, ch_idx)
return
@@ -1248,11 +1183,7 @@ class SrasViewerWindow(QMainWindow):
return
self._batch_errors = []
# Cache at the pad factor currently configured, not a fixed pad 1 —
# a cache stored at any other padding is one this viewer can never
# read back (see _stored_fft_image).
worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked(),
self._fft_pad_factor)
worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked())
started = self._run_worker(
Jobs.BATCH, worker,
connect=(
@@ -1267,10 +1198,53 @@ class SrasViewerWindow(QMainWindow):
self._batch_dc_act.setEnabled(False)
self._batch_fft_act.setEnabled(False)
self._batch_fft_rowavg_act.setEnabled(False)
self._show_progress(
Jobs.BATCH, f"Batch computing {label} for {len(paths)} file(s)…",
maximum=100)
def _on_batch_compute_row_avg(self):
if self._job_running(Jobs.BATCH):
return
dlg = RowAverageFftOptionsDialog(
self, current_n=self._pending_row_avg_n,
current_threshold_mv=self.spin_threshold_mv.value(),
pixel_x_mm=self._sras.pixel_x_mm if self._sras is not None else None)
if dlg.exec() != QDialog.DialogCode.Accepted:
return
n, threshold_mv = dlg.get_half_width(), dlg.get_threshold_mv()
self._pending_row_avg_n = n
self._settings.setValue("fft/row_avg_n", n)
paths, _ = QFileDialog.getOpenFileNames(
self, "Select .sras files to batch-compute Row-Averaged FFT", "",
"SRAS files (*.sras);;All files (*)")
if not paths:
return
self._batch_errors = []
worker = BatchCacheWorker(paths, "fft_rowavg", self.chk_bg_sub.isChecked(),
dc_threshold_mv=threshold_mv, row_avg_n=n)
started = self._run_worker(
Jobs.BATCH, worker,
connect=(
("progress", lambda pct: self._set_progress(Jobs.BATCH, pct)),
("file_done", self._on_batch_file_done),
("finished", lambda p=paths: self._on_batch_finished(p)),
),
on_done=self._after_batch,
)
if not started:
return # a second trigger snuck in while a dialog was open
self._batch_dc_act.setEnabled(False)
self._batch_fft_act.setEnabled(False)
self._batch_fft_rowavg_act.setEnabled(False)
self._show_progress(
Jobs.BATCH,
f"Batch computing row-averaged FFT (n={n}, threshold={threshold_mv:.1f} mV) "
f"for {len(paths)} file(s)…", maximum=100)
def _on_batch_file_done(self, path: str, err: str):
if err:
self._batch_errors.append(f"{Path(path).name} — {err}")
@@ -1298,6 +1272,7 @@ class SrasViewerWindow(QMainWindow):
def _after_batch(self):
self._batch_dc_act.setEnabled(True)
self._batch_fft_act.setEnabled(True)
self._batch_fft_rowavg_act.setEnabled(True)
# ------------------------------------------------------------------
# Fusion: angle alignment
@@ -1427,7 +1402,6 @@ class SrasViewerWindow(QMainWindow):
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,
grating_um=self.spin_grating_um.value(),
cached_pad_factor=self._cached_fft_pad_factor(),
)
if dlg.exec() != QDialog.DialogCode.Accepted:
return
@@ -1437,9 +1411,7 @@ class SrasViewerWindow(QMainWindow):
self._settings.setValue("fft/pad_factor", self._fft_pad_factor)
# Pad factor changes the FFT bin count, so it genuinely invalidates
# the cached raw FFT (part of the cache key) — _refresh_display()
# recomputes only on a cache miss. It also decides whether the file's
# stored cache can be read back at all, which the info panel reports.
self._update_scan_info_labels()
# recomputes only on a cache miss.
if self._is_fft_mode():
self._refresh_display()
+11 -9
View File
@@ -187,12 +187,11 @@ class BatchCacheWorker(QObject):
an existing v7 file's cache blocks without disturbing whatever the other
block already holds.
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
*mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
images, unmasked — masking is applied at display time, same as v5's PREC
convention). FFT images are computed and stored at *pad_factor*, the
viewer's own zero-padding setting: a stored image only serves a view
asking for the same padding, so caching at any other one produces a file
the viewer will never read back.
convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
averaging before the FFT — needs *dc_threshold_mv* and a positive
*row_avg_n*; see ``sras_compute.cache_file``).
Files are processed one per subprocess: they are fully independent, each
opens its own memmap and writes only its own bytes, and only path strings
@@ -205,12 +204,13 @@ class BatchCacheWorker(QObject):
finished = pyqtSignal()
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
pad_factor: int = 1):
dc_threshold_mv: float | None = None, row_avg_n: int = 0):
super().__init__()
self._paths = paths
self._mode = mode
self._apply_bg_sub = apply_bg_sub
self._pad_factor = pad_factor
self._dc_threshold = dc_threshold_mv
self._row_avg_n = row_avg_n
def _report(self, path: str, err: str, done: int, total: int):
self.file_done.emit(path, err)
@@ -236,7 +236,8 @@ class BatchCacheWorker(QObject):
futures = {
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), per_proc_workers,
self._pad_factor): p
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n): p
for p in paths
}
for fut in as_completed(futures):
@@ -261,7 +262,8 @@ class BatchCacheWorker(QObject):
err = cache_file(path, self._mode, self._apply_bg_sub,
compute.get_fft_backend(),
compute.default_max_workers(),
self._pad_factor)
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n)
except Exception as exc:
err = str(exc)
done += 1
+1 -46
View File
@@ -94,18 +94,11 @@ def test_dc_precompute_all_angles(ctx):
def test_angle_switching_from_cache(ctx):
"""Setting the angle must redraw on its own — no nudge from the test.
The spinbox used to be wired to editingFinished alone, so stepping it with
the arrows changed the number and left the plot on the previous angle until
the box happened to lose focus.
"""
win, s = ctx.win, ctx.s
for a in range(s.n_angles):
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
assert win._current_angle == a, \
f"angle {a} redrawn from the spinbox alone, showing {win._current_angle}"
expected = win._sras.image_shape(a)
assert win._current_image.shape == expected, \
f"angle {a} shows its own geometry {expected}, got {win._current_image.shape}"
@@ -113,44 +106,6 @@ def test_angle_switching_from_cache(ctx):
"no compute job needed for cached DC angles"
def test_angle_steps_and_typing(ctx):
"""The arrows step and redraw one angle at a time; typing commits once,
without a stop at every intermediate value."""
win, s = ctx.win, ctx.s
assert s.n_angles > 2, "fixture needs several angles"
win.spin_angle.setValue(0)
pump(60)
seen = []
win.spin_angle.valueChanged.connect(seen.append)
try:
win.spin_angle.stepUp()
pump(60)
assert win._current_angle == 1, f"arrow stepped to 1, showing {win._current_angle}"
win.spin_angle.stepDown()
pump(60)
assert win._current_angle == 0, f"arrow stepped to 0, showing {win._current_angle}"
# Typing must not redraw per keystroke — that is what editingFinished
# bought and what keyboard tracking has to keep buying.
seen.clear()
last = s.n_angles - 1
win.spin_angle.lineEdit().selectAll()
QTest.keyClicks(win.spin_angle.lineEdit(), str(last))
pump(30)
assert seen == [] and win._current_angle == 0, \
f"half-typed text must not redraw ({seen}, showing {win._current_angle})"
QTest.keyClick(win.spin_angle, Qt.Key.Key_Return)
pump(60)
assert seen == [last], f"committed exactly once on Return: {seen}"
assert win._current_angle == last
finally:
win.spin_angle.valueChanged.disconnect(seen.append)
win.spin_angle.clearFocus()
win.spin_angle.setValue(0)
pump(60)
def test_channel_switching(ctx):
win = ctx.win
win.spin_angle.setValue(0)
+278
View File
@@ -0,0 +1,278 @@
"""Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing.
Covers the properties the design depends on: the kernel is symmetric and
n=0 is a true no-op; the masked/renormalized convolution matches an
independent brute-force reference and gives masked neighbors exactly zero
weight regardless of their content; background subtraction after averaging
is algebraically identical to subtracting before; chunking/worker count
never changes the result; and a pixel that's itself masked is never
"rescued" by averaging.
"""
import numpy as np
import pytest
import sras_compute as compute
from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH4_IDX, SrasFile
import tools.make_test_sras as gen
def _reference_row_average(masked_waves: np.ndarray, valid: np.ndarray,
weights: np.ndarray) -> np.ndarray:
"""Independent, unvectorized reference for _row_average_waveforms: for
each row position, sum weighted valid neighbors within the kernel's
radius and normalize by the actual included weight sum. Same
definition, computed by brute-force nested loops instead of
correlate1d, so it can't share a bug with the implementation."""
n_frames, spf = masked_waves.shape
n = len(weights) // 2
out = np.zeros_like(masked_waves)
for i in range(n_frames):
num = np.zeros(spf, dtype=np.float64)
den = 0.0
for d in range(-n, n + 1):
j = i + d
if 0 <= j < n_frames and valid[j]:
w = float(weights[d + n])
num += w * masked_waves[j].astype(np.float64)
den += w
out[i] = num / den if den > 0 else 0.0
return out
# ---------------------------------------------------------------------------
# _row_average_weights
# ---------------------------------------------------------------------------
def test_row_average_weights_shape_and_symmetry():
w0 = compute._row_average_weights(0)
assert w0.shape == (1,) and w0[0] == 1.0
for n in (1, 2, 5):
w = compute._row_average_weights(n)
assert w.shape == (2 * n + 1,)
assert w[n] == pytest.approx(1.0), "center tap is the peak weight"
assert np.allclose(w, w[::-1]), "symmetric about the center"
half = w[n:]
assert np.all(np.diff(half) < 0), "strictly decreasing away from center"
# ---------------------------------------------------------------------------
# _row_average_waveforms
# ---------------------------------------------------------------------------
def test_row_average_matches_hand_rolled_reference():
rng = np.random.default_rng(0)
n_frames, spf = 15, 6
raw = rng.integers(-50, 51, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[2, 3, 9]] = False # a run of two invalid, plus a lone invalid
masked = raw.copy()
masked[~valid] = 0.0
weights = compute._row_average_weights(3)
got = compute._row_average_waveforms(masked, valid, weights)
ref = _reference_row_average(masked, valid, weights)
assert np.allclose(got[valid], ref[valid], atol=1e-4)
def test_row_average_edge_of_row():
"""A window wider than the row itself must still renormalize correctly
at both ends -- mode='constant', cval=0.0 zero-pads both the numerator
and denominator, so this is not a special case, but it's the one most
likely to break if that padding were ever mismatched between the two."""
n_frames, spf = 6, 3
raw = np.arange(n_frames * spf, dtype=np.float32).reshape(n_frames, spf)
valid = np.ones(n_frames, dtype=bool)
weights = compute._row_average_weights(4) # window (9 taps) > n_frames (6)
got = compute._row_average_waveforms(raw, valid, weights)
ref = _reference_row_average(raw, valid, weights)
assert np.allclose(got, ref, atol=1e-4)
def test_row_average_excludes_masked_neighbor_from_normalization():
"""A masked neighbor must contribute zero *weight* to the normalization,
not participate as a legitimate zero-valued sample at full weight --
the two give different answers, and only the former is correct. (Note:
masked_waves must already be 0 at invalid positions per
_row_average_waveforms's contract -- that's what read_row's zero-filled
scratch buffer guarantees in production -- so the only way to vary
"what a masked position looks like" while respecting that contract is
whether its weight is excluded from the denominator at all.)"""
n_frames, spf = 9, 4
weights = compute._row_average_weights(2)
# A: position 4 is masked -- excluded from the weight sum entirely.
valid_a = np.ones(n_frames, dtype=bool)
valid_a[4] = False
masked_a = np.zeros((n_frames, spf), dtype=np.float32)
masked_a[valid_a] = 1.0
got_a = compute._row_average_waveforms(masked_a, valid_a, weights)
# B: position 4 is valid but genuinely zero-valued -- included in the
# weight sum, diluting neighbors' averages.
valid_b = np.ones(n_frames, dtype=bool)
masked_b = np.ones((n_frames, spf), dtype=np.float32)
masked_b[4] = 0.0
got_b = compute._row_average_waveforms(masked_b, valid_b, weights)
# Every position whose window reaches index 4 must average *higher* in
# A (excluded from the denominator) than in B (included as a real zero).
affected = [2, 3, 5, 6]
assert np.all(got_a[affected] > got_b[affected]), \
"masking must exclude a neighbor from normalization, not just zero its value"
# Positions outside the window (radius 2) are unaffected either way.
assert np.allclose(got_a[[0, 1, 7, 8]], got_b[[0, 1, 7, 8]])
def test_background_subtracted_once_equals_subtract_then_average():
"""Algebraic identity the implementation relies on: subtracting a fixed
background from the already-averaged waveform equals subtracting it
from every valid neighbor first, because the denominator is always the
*actual* included weight sum (never a fixed total)."""
rng = np.random.default_rng(1)
n_frames, spf = 11, 8
raw = rng.integers(-40, 41, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[1, 7]] = False
masked = raw.copy()
masked[~valid] = 0.0
background = rng.integers(-5, 6, size=spf).astype(np.float32)
weights = compute._row_average_weights(3)
# Order A (what the code does): average first, subtract background once.
order_a = compute._row_average_waveforms(masked, valid, weights) - background
# Order B: subtract background from every valid neighbor first (restoring
# the "0 at invalid positions" contract afterward), then average.
bg_subbed = masked - background
bg_subbed[~valid] = 0.0
order_b = compute._row_average_waveforms(bg_subbed, valid, weights)
assert np.allclose(order_a[valid], order_b[valid], atol=1e-3)
# ---------------------------------------------------------------------------
# compute_rf_image(row_avg_n=...) integration
# ---------------------------------------------------------------------------
def test_row_average_zero_is_identity(tmp_path):
"""row_avg_n=0 must take the exact same code path as before this
feature existed (row_avg_weights stays None), not a single-tap kernel
that merely computes to the same answer."""
path = tmp_path / "zero.sras"
gen.write(path, n_angles=1, seed=10, samples_per_frame=64)
sras = SrasFile(str(path))
plain = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
explicit_zero = compute_rf_image(sras, 0, dc_threshold_mv=None,
apply_bg_sub=True, row_avg_n=0)
assert np.array_equal(plain, explicit_zero)
def test_row_average_respects_own_center_mask(tmp_path):
"""A pixel that's itself below threshold stays masked (0) after row
averaging -- averaging never rescues a masked pixel, matching the
'valid neighbors only' design (masked pixels are excluded from other
pixels' averages, and are never themselves smoothed)."""
path = tmp_path / "center_mask.sras"
gen.write(path, n_angles=1, seed=13, samples_per_frame=64)
sras = SrasFile(str(path))
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.percentile(dc4, 50))
mask = dc4 >= thr
assert mask.any() and not mask.all(), "threshold actually splits the image"
img = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(img == 0, ~mask), \
"masked pixels stay exactly 0 after row averaging; valid ones don't"
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."""
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
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)
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"
def test_row_average_improves_snr_recovery():
"""The actual point of the feature: averaging same-row waveforms that
share a true underlying tone but carry independent noise recovers that
tone far more reliably than any single raw (unaveraged) waveform does."""
rng = np.random.default_rng(42)
n_frames, spf = 21, 128
true_bin = 9
t = np.arange(spf)
tone = 15.0 * np.sin(2 * np.pi * true_bin * t / spf) # same true signal
# at every position
noise_sigma = 40.0 # much larger than the tone -- deliberately poor SNR
raw = (tone[None, :] + rng.normal(scale=noise_sigma, size=(n_frames, spf))
).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
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_hits = int(np.sum(raw_bins == true_bin))
avg_hits = int(np.sum(avg_bins == true_bin))
assert avg_hits > raw_hits, (
f"row averaging should recover the true bin ({true_bin}) more often "
f"than raw per-pixel estimates: raw {raw_hits}/{n_frames}, "
f"averaged {avg_hits}/{n_frames}")
assert avg_hits >= n_frames * 0.7, \
f"averaged recovery should be reliable, not just barely better: {avg_hits}/{n_frames}"
def test_row_average_parallel_identity(tmp_path, monkeypatch):
"""Forcing 1 worker vs many must give an identical row-averaged image --
catches chunk-boundary bugs (there should be none, since averaging never
crosses rows, but this is the empirical proof, not just inspection)."""
path = tmp_path / "parallel_rowavg.sras"
n_rows, n_frames, spf = 40, 13, 128
gen.write(path, n_angles=1, seed=11, samples_per_frame=spf,
geometry=[(n_rows, n_frames)])
sras = SrasFile(str(path))
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
monkeypatch.setattr(compute, "_FFT_BLOCK", 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, \
f"row-averaged FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.median(dc4))
monkeypatch.setattr(compute, "_MAX_WORKERS", 1)
serial = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
parallel = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(serial, parallel), \
"row-averaged rf image identical regardless of chunking/worker count"
+558
View File
@@ -0,0 +1,558 @@
"""Does a batch-computed FFT cache actually spare the viewer the FFT?
Storing a peak-frequency image per angle in the file is only worth doing if
displaying it is then free. The regression this module pins down is the
viewer's *dispatch* decision: it used to find the stored image only inside
ComputeWorker, so after a batch every angle change still queued a background
job behind a "Computing FFT…" popup for an image already on disk.
Both layers are covered — cached_rf_image's accept/reject rules, and the
window never reaching _start_compute for a batch-cached angle.
"""
import struct
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
from PyQt6.QtCore import QEventLoop, QTimer
from PyQt6.QtWidgets import QApplication, QDialog
import sras_compute as compute
import sras_format as fmt
from sras_compute import cache_file, cached_rf_image, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
import tools.make_test_sras as gen
_THRESHOLD_MV = 50.0 # the viewer's own default
def pump(ms: int = 200):
loop = QEventLoop()
QTimer.singleShot(ms, loop.quit)
loop.exec()
def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
waited = 0
while waited < timeout_ms:
if pred():
return True
pump(step)
waited += step
return pred()
@pytest.fixture(scope="module")
def rig(tmp_path_factory):
"""A v6 file, the FFT images a from-scratch compute gives for it, and the
same file after Batch Compute FFT has written them into its v7 cache."""
path = tmp_path_factory.mktemp("stored_cache") / "cached.sras"
gen.write(path, n_angles=4, seed=7, samples_per_frame=256)
src = SrasFile(str(path))
fresh = {a: compute_rf_image(src, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True)
for a in range(src.n_angles)}
assert src.background is not None, "the fixture file must have a background"
err = cache_file(str(path), "fft", True)
assert err == "", err
cached = SrasFile(str(path))
assert all(x is not None for x in cached.precomputed_freq_mhz)
assert all(x is None for x in cached.precomputed_dc4_mv), \
"FFT-only batch: the mask has to come from the viewer, not the file"
return SimpleNamespace(path=path, fresh=fresh, sras=cached,
n_angles=cached.n_angles)
@pytest.fixture(scope="module")
def dc_rig(tmp_path_factory):
"""A file that has been through Batch Compute DC and Store."""
path = tmp_path_factory.mktemp("stored_dc") / "dc_cached.sras"
gen.write(path, n_angles=4, seed=9, samples_per_frame=128)
err = cache_file(str(path), "dc", True)
assert err == "", err
sras = SrasFile(str(path))
assert all(x is not None for x in sras.precomputed_dc3_mv)
assert all(x is not None for x in sras.precomputed_dc4_mv)
assert len(np.unique(sras.precomputed_dc4_mv[0])) > 1, \
"a degenerate DC image would make the comparisons below vacuous"
return SimpleNamespace(path=path, sras=sras, n_angles=sras.n_angles)
@pytest.fixture
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"):
original = getattr(compute, name)
def spy(*args, _f=original, **kwargs):
calls.append(_f.__name__)
return _f(*args, **kwargs)
monkeypatch.setattr(compute, name, spy)
return calls
# ---------------------------------------------------------------------------
# cached_rf_image: when may the stored image stand in for a compute?
# ---------------------------------------------------------------------------
def test_stored_image_matches_a_fresh_compute(rig, no_fft):
for a in range(rig.n_angles):
dc4 = dc_image_mv(SrasFile(str(rig.path)), a, CH4_IDX)
img = cached_rf_image(rig.sras, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True, dc4_mv=dc4)
assert img is not None, f"angle {a} is cached in the file"
assert np.allclose(img, rig.fresh[a], atol=1e-3), \
f"angle {a} differs from a from-scratch compute"
assert not no_fft, f"the stored image was used, no FFT ran: {no_fft}"
def test_unmasked_when_no_threshold(rig):
img = cached_rf_image(rig.sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
assert img is not None and np.array_equal(img, rig.sras.precomputed_freq_mhz[0])
img[:] = -1.0
assert not np.any(rig.sras.precomputed_freq_mhz[0] == -1.0), \
"callers get a copy, never the file's own array"
def test_settings_the_stored_image_cannot_serve(rig):
"""A stored image carries one bg-sub state and one padding, so anything
else must fall through to a real compute rather than lie."""
spf = rig.sras.samples_per_frame
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
n_fft=spf * 4) is None, \
"cache was written at pad 1, a pad-4 view resolves different peaks"
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=False) is None, \
"cache was written with bg-sub on"
uncached = SrasFile(str(rig.path))
uncached.precomputed_freq_mhz[1] = None
assert cached_rf_image(uncached, 1, _THRESHOLD_MV, apply_bg_sub=True) is None
@pytest.mark.parametrize("pad", [2, 10])
def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
"""Batching at a pad factor must produce a cache that view can read back —
a pad-1-only cache is one the padded viewer can never use."""
path = tmp_path / f"pad{pad}.sras"
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) == ""
sras = SrasFile(str(path))
assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * pad) is not None, f"usable at pad {pad}"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \
"not usable unpadded"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * (pad + 1)) is None, "not usable at another pad"
# The stored numbers must be the padded ones, not pad-1 relabelled.
fresh = SrasFile(str(path))
fresh.precomputed_freq_mhz = [None] * fresh.n_angles
for a in range(sras.n_angles):
assert np.allclose(
compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad),
compute_rf_image(fresh, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad), atol=1e-3), \
f"angle {a}: stored image is the pad-{pad} answer"
def test_cach_v1_reads_as_natural_resolution(tmp_path):
"""Files cached before the pad factor existed must keep working: a v1 tail
has no pad field and is pad 1 by construction."""
path = tmp_path / "v1.sras"
gen.write(path, n_angles=2, seed=14, samples_per_frame=256)
assert cache_file(str(path), "fft", True) == ""
# Rewrite the tail as a genuine CACH v1 block (old header, no pad field).
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_pad_factor == 1
assert v1.precomputed_bg_sub is True
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True,
n_fft=v1.samples_per_frame * 10) is None
def test_mask_read_can_be_refused(rig, no_fft):
"""With no DC4 in hand, building the mask means reading a whole channel —
the GUI thread asks for None instead."""
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
allow_dc_recompute=False) is None
dc4 = dc_image_mv(SrasFile(str(rig.path)), 0, CH4_IDX)
img = cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
dc4_mv=dc4, allow_dc_recompute=False)
assert img is not None and np.allclose(img, rig.fresh[0], atol=1e-3)
assert not no_fft, f"no FFT on either branch: {no_fft}"
# ---------------------------------------------------------------------------
# The viewer: no compute job at all for a batch-cached angle
# ---------------------------------------------------------------------------
def test_viewer_shows_stored_angles_without_computing(rig, no_fft, monkeypatch):
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._load_file(str(rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
# The file carries no DC block, so the mask comes from the window's own
# background precompute — the state a user is in by the time they click.
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(rig.n_angles))), \
"DC precompute finished"
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in list(range(rig.n_angles)) + [1, 0]:
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
assert win._current_angle == a, f"angle {a} displayed"
assert np.allclose(win._current_image, rig.fresh[a], atol=1e-3), \
f"angle {a} shows the stored image"
assert dispatched == [], \
f"stored angles need no compute job, dispatched for {dispatched}"
assert not no_fft, f"no FFT ran for any stored angle: {no_fft}"
# Velocity is still a post-multiply of the same stored image.
win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
pump(120)
assert np.allclose(win._current_image,
rig.fresh[0] * win.spin_grating_um.value(), atol=1e-3)
assert dispatched == [] and not no_fft
# ...but a setting the stored image cannot serve must still recompute,
# or the fast path would be showing the wrong picture.
win.chk_bg_sub.setChecked(False)
assert wait_until(lambda: not win._job_running("compute") and bool(no_fft)), \
"bg-sub off falls through to a real FFT"
finally:
win.close()
pump(300)
def test_batch_caches_at_the_viewers_pad_factor(tmp_path, no_fft, monkeypatch):
"""The bug a pad-10 user hits: Batch Compute FFT used to store pad-1
images regardless, so the padded view recomputed every angle forever.
"""
path = tmp_path / "padded_gui.sras"
gen.write(path, n_angles=3, seed=15, samples_per_frame=256)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._fft_pad_factor = 10
win._load_file(str(path))
assert wait_until(lambda: win._sras is not None), "file loaded"
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(win._sras.n_angles))), \
"DC precompute finished"
# Convert -> Batch Compute FFT, on the open file, through the real slot.
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute("fft")
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_pad_factor == 10, \
f"cached at the viewer's pad, got {win._sras.precomputed_pad_factor}"
expected = {a: compute.cached_rf_image(
win._sras, a, dc_threshold_mv=win.spin_threshold_mv.value(),
apply_bg_sub=win.chk_bg_sub.isChecked(),
n_fft=win._current_n_fft(),
dc4_mv=win._dc_cache.get((a, CH4_IDX)))
for a in range(win._sras.n_angles)}
assert all(v is not None for v in expected.values()), "cache is readable at pad 10"
no_fft.clear()
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in range(win._sras.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert np.allclose(win._current_image, expected[a], atol=1e-3), \
f"angle {a} served from the pad-10 cache"
assert dispatched == [] and not no_fft, \
f"no recompute at pad 10 (jobs={dispatched}, fft={no_fft})"
assert "unusable" not in win.lbl_frame_warn.text()
# Change the pad and the cache legitimately stops applying — and the
# info panel has to say so rather than leave it a mystery.
win._fft_pad_factor = 4
win._update_scan_info_labels()
assert "Cached FFT unusable" in win.lbl_frame_warn.text(), \
win.lbl_frame_warn.text()
assert "pad 10x" in win.lbl_frame_warn.text()
win._refresh_display()
assert wait_until(lambda: not win._job_running("compute") and bool(no_fft)), \
"pad 4 recomputes rather than reusing the pad-10 cache"
finally:
win.close()
pump(300)
def test_viewer_shows_stored_dc_without_computing(dc_rig, monkeypatch):
"""Same for the DC half, with the background precompute silenced so the
file's stored block is the only thing that can be carrying the display."""
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
monkeypatch.setattr(type(win), "_start_dc_precompute", lambda self: None)
try:
win._load_file(str(dc_rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
pump(120)
for ch in (CH3_IDX, CH4_IDX):
win.combo_channel.setCurrentIndex(ch)
for a in range(dc_rig.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert (win._current_angle, win._current_ch) == (a, ch), \
f"angle {a} on channel {ch} displayed"
assert np.array_equal(win._current_image,
dc_rig.sras.cached_dc_mv(a, ch)), \
f"angle {a} channel {ch} shows the file's stored DC image"
assert dispatched == [], \
f"stored DC angles need no compute job, dispatched for {dispatched}"
finally:
win.close()
pump(300)
# ---------------------------------------------------------------------------
# Row-averaged FFT: cache_file("fft_rowavg", ...) and its on-disk provenance
# ---------------------------------------------------------------------------
def test_row_average_flag_and_n_round_trip(tmp_path):
path = tmp_path / "rowavg_roundtrip.sras"
gen.write(path, n_angles=2, seed=20, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert sras.version == 7
assert sras.precomputed_row_avg_n == 5
assert all(x is not None for x in sras.precomputed_freq_mhz)
def test_fft_rowavg_mode_requires_positive_n_and_threshold(tmp_path):
path = tmp_path / "rowavg_bad_args.sras"
gen.write(path, n_angles=1, seed=27, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=0.0, row_avg_n=0)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=None, row_avg_n=5)
def test_raw_and_row_averaged_caches_never_cross_served(tmp_path):
"""The central regression this feature must never allow: a raw request
served a row-averaged image (or vice versa), or a request at one window
size served a cache stored at a different one."""
path = tmp_path / "cross_serve.sras"
gen.write(path, n_angles=1, seed=21, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=0) is None, \
"a raw request must not be served a row-averaged cache"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=3) is None, \
"a request at the wrong window size must not be served either"
served = cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=5)
assert served is not None
assert np.array_equal(served, sras.precomputed_freq_mhz[0])
def test_write_v7_cache_row_avg_n_carries_forward(tmp_path):
"""A later DC-only write must leave a previously-written row-averaged
FFT block -- including its row_avg_n -- byte-for-byte unchanged."""
path = tmp_path / "carry_forward.sras"
gen.write(path, n_angles=2, seed=22, samples_per_frame=64)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=7)
assert err == "", err
before = SrasFile(str(path))
assert before.precomputed_row_avg_n == 7
freq_before = [x.copy() for x in before.precomputed_freq_mhz]
err = cache_file(str(path), "dc", True)
assert err == "", err
after = SrasFile(str(path))
assert after.precomputed_row_avg_n == 7, "row_avg_n survives a DC-only write"
assert all(np.array_equal(after.precomputed_freq_mhz[a], freq_before[a])
for a in range(after.n_angles)), \
"the row-averaged FFT block itself is untouched by a DC-only write"
def test_row_average_never_touches_dc_images(tmp_path):
path = tmp_path / "dc_untouched.sras"
gen.write(path, n_angles=2, seed=23, samples_per_frame=64)
src = SrasFile(str(path))
expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
assert cache_file(str(path), "dc", True) == ""
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=6) == ""
after = SrasFile(str(path))
assert all(np.allclose(after.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
for a in range(after.n_angles))
assert all(np.allclose(after.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
for a in range(after.n_angles))
def test_row_average_never_modifies_raw_waveform_data(tmp_path):
path = tmp_path / "waveform_untouched.sras"
gen.write(path, n_angles=2, seed=24, samples_per_frame=64)
orig = tmp_path / "waveform_untouched_orig.sras"
gen.write(orig, n_angles=2, seed=24, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=5) == ""
o, n = SrasFile(str(orig)), SrasFile(str(path))
assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
for a in range(o.n_angles)), \
"waveform data untouched by a row-averaged cache write"
def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path):
"""A v1 CACH tail predates row-averaged FFT caching entirely (no
row_avg_n byte at all) -- readers must still parse it in full, treating
it as row_avg_n=0. This is what protects an existing real-world v7
file's already-stored FFT cache from silently becoming unusable after
this change ships."""
path = tmp_path / "v1_rowavg.sras"
gen.write(path, n_angles=2, seed=25, samples_per_frame=128)
assert cache_file(str(path), "fft", True) == ""
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_row_avg_n == 0
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=0) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=5) is None, \
"a v1 tail (predating this feature) can never satisfy a row-averaged request"
def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch):
"""Driving the new 'Batch Compute Row-Averaged FFT and Store' action
end-to-end through the real menu handler: dialog values reach the
worker, the worker reaches cache_file, and the written file is
self-describing afterward. Deliberately does not assert anything about
whether viewing an angle afterward dispatches a compute job -- that is
a separate, pre-existing gap in _refresh_display shared with the plain
DC/FFT batch actions (see test_viewer_shows_stored_angles_without_computing
/ test_viewer_shows_stored_dc_without_computing above), not something
row-averaging introduces or is responsible for fixing."""
path = tmp_path / "rowavg_gui.sras"
gen.write(path, n_angles=2, seed=26, samples_per_frame=128)
class _StubDialog:
def __init__(self, *a, **k):
pass
def exec(self):
return QDialog.DialogCode.Accepted
def get_half_width(self):
return 6
def get_threshold_mv(self):
return -1e9 # mask nothing, keep the comparison simple
monkeypatch.setattr("sras_viewer.main_window.RowAverageFftOptionsDialog", _StubDialog)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
try:
win._load_file(str(path))
assert wait_until(lambda: win._sras is not None), "file loaded"
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(win._sras.n_angles))), \
"DC precompute finished"
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute_row_avg()
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_row_avg_n == 6
assert "row-averaged n=6" in win.lbl_frame_warn.text(), win.lbl_frame_warn.text()
expected = compute.cached_rf_image(win._sras, 0, dc_threshold_mv=None,
apply_bg_sub=win.chk_bg_sub.isChecked(),
row_avg_n=6)
assert expected is not None, "the batch write left a readable row-averaged cache"
finally:
win.close()
pump(300)