diff --git a/docs/design.md b/docs/design.md index 3ea672b..9c1ef27 100644 --- a/docs/design.md +++ b/docs/design.md @@ -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 diff --git a/scan_format.md b/scan_format.md index 3c90e16..2194f93 100644 --- a/scan_format.md +++ b/scan_format.md @@ -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) diff --git a/sras_compute.py b/sras_compute.py index 9120959..1f0f369 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -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) diff --git a/sras_format.py b/sras_format.py index 46d8003..9ce7ab8 100644 --- a/sras_format.py +++ b/sras_format.py @@ -60,7 +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) -CACH_VERSION = 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 @@ -69,9 +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 = ">4sBH" # magic, flags, n_stored +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 +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 @@ -226,6 +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 + 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.""" @@ -517,6 +529,34 @@ class SrasFile: store[angle_idx] = _read_f32_image(f, shape) 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.""" with open(self.path, "rb") as f: @@ -525,7 +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 cach_version != CACH_VERSION: + if magic != CACH_MAGIC or cach_version not in CACH_VERSIONS_READABLE: return if block_flags & CACH_FLAG_DC: @@ -535,17 +575,19 @@ class SrasFile: return if block_flags & CACH_FLAG_FFT: - flags = self._read_cache_block( - f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz]) - if flags is None: + result = self._read_sfft_block(f, cach_version) + if result is None: return + flags, row_avg_n = result self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB) + 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_bg_sub: bool | 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 @@ -553,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. @@ -565,6 +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_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] @@ -584,7 +636,9 @@ class SrasFile: if fft_entries: fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0 - payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries)) + 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_row_avg_n) for a in fft_entries: payload += struct.pack(">H", a) payload += final_freq[a].astype(">f4").tobytes() @@ -612,6 +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_row_avg_n = final_row_avg_n # ------------------------------------------------------------------ # Axes helpers diff --git a/sras_viewer/__init__.py b/sras_viewer/__init__.py index b218d72..2f87856 100644 --- a/sras_viewer/__init__.py +++ b/sras_viewer/__init__.py @@ -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 diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index 5589698..1ae09f2 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -152,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...). diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index 07c3c71..6906b88 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -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,6 +73,11 @@ 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] = [] @@ -459,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 # ------------------------------------------------------------------ @@ -598,9 +614,12 @@ 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)" + 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{bg_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") @@ -640,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)) @@ -1178,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}") @@ -1209,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 diff --git a/sras_workers.py b/sras_workers.py index 4a75dd8..7e3ebb2 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -187,9 +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). + 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 @@ -201,11 +203,14 @@ class BatchCacheWorker(QObject): file_done = pyqtSignal(str, str) finished = pyqtSignal() - def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool): + def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool, + 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._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) @@ -230,7 +235,9 @@ class BatchCacheWorker(QObject): with ProcessPoolExecutor(max_workers=n_procs) as executor: futures = { executor.submit(cache_file, p, self._mode, self._apply_bg_sub, - compute.get_fft_backend(), per_proc_workers): p + compute.get_fft_backend(), per_proc_workers, + dc_threshold_mv=self._dc_threshold, + row_avg_n=self._row_avg_n): p for p in paths } for fut in as_completed(futures): @@ -254,7 +261,9 @@ class BatchCacheWorker(QObject): try: err = cache_file(path, self._mode, self._apply_bg_sub, compute.get_fft_backend(), - compute.default_max_workers()) + compute.default_max_workers(), + dc_threshold_mv=self._dc_threshold, + row_avg_n=self._row_avg_n) except Exception as exc: err = str(exc) done += 1 diff --git a/tests/test_row_average.py b/tests/test_row_average.py new file mode 100644 index 0000000..448cee4 --- /dev/null +++ b/tests/test_row_average.py @@ -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" diff --git a/tests/test_stored_cache.py b/tests/test_stored_cache.py new file mode 100644 index 0000000..03a8793 --- /dev/null +++ b/tests/test_stored_cache.py @@ -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)