diff --git a/docs/design.md b/docs/design.md index 7b0c583..c5a6f99 100644 --- a/docs/design.md +++ b/docs/design.md @@ -32,45 +32,58 @@ share of the budget. Capping the workers alone is not enough: the chunk would still be sized against the whole budget, and N concurrent callers would each allocate all of it. -## FFT peak search: block-parallel zoom refinement (`sras_compute.py`) +## FFT peak search: block-parallel direct transform (`sras_compute.py`) The displayed RF value per pixel is the argmax of the zero-padded power -spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for -mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan -row, which is what used to collapse the old row-chunk planner to one worker -and make synthesis single-threaded. +spectrum of that pixel's CH1 waveform. This used to run through `_peak_bins_zoom`, +a coarse-rfft-plus-local-fine-DFT refinement that avoided ever materialising +a padded spectrum — at the pad factor of 40 needed for mapping resolution, a +full padded spectrum is ~9 GB per scan row, which used to collapse the old +row-chunk planner to one worker and make synthesis single-threaded. That +refinement was removed once pyFFTW became the sole, mandatory FFT backend +(SciPy dropped as a compute backend entirely): `_peak_bins` now always runs +the real full transform, and the memory problem zoom dodged is instead +solved by bounding *per-block* spectrum memory rather than avoiding full +spectra altogether. -`_peak_bins_zoom` never materialises the padded spectrum: +`_peak_bins` runs the full transform (`_block_rfft`, a cached pyFFTW +`builders.rfft` plan, FFTW_MEASURE, wisdom persisted under +`~/.cache/sras-viewer/`) and argmaxes the power spectrum, in blocks fanned +out task-parallel over a persistent thread pool (`_fft_pool()`) — each pool +thread runs one single-threaded transform at a time, so aggregate +parallelism equals the pool's worker count. This block-streaming structure +is not just about parallelism: it is also what keeps memory bounded, by +never materialising more than one block's worth of full padded spectrum at +a time, regardless of how many waveforms a chunk holds. -1. a coarse rfft at `next_fast_len(2*spf)` — 2× oversampled, so the padded - power spectrum (a trig polynomial of degree spf−1) cannot hide its global - max between coarse samples; -2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max - becomes a refinement candidate. Quarter-natural-bin scalloping at the 2× - grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide - margin. The DC-adjacent window is always refined too: the coarse DC bin - is zeroed for suppression, which would otherwise blind the scan to fine - bins closer to DC than the first coarse sample (where the leakage skirt - of an un-subtracted offset peaks); -3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every - fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated - on the exact `n_fft` grid by one small complex gemm, with np.argmax's - lowest-bin tie-break preserved across windows. +The block size is the part that has to adapt to pad factor. +`_fft_block_for(spf, n_len)` derives waveforms-per-task from a fixed +per-thread byte budget (`_FFT_PLAN_BYTES_BUDGET`, `SRAS_FFT_PLAN_BUDGET_MB`, +default 16 MB) rather than a fixed constant, because a cached pyFFTW plan's +input+output buffers are *permanent* per-thread memory (the plan cache is +never evicted) — with a fixed 512-waveform block, pad 40 at a 2500-sample +frame costs ~210 MB per pool thread (~3.3 GB total across 16 threads); +`_fft_block_for` bounds that to ~16 MB per thread (~260 MB across 16 +threads) at the same pad factor, while still reproducing the old tuned 512 +exactly at natural resolution (pad 1), where it cost nothing to begin with. +`_FFT_BLOCK_MAX` (512) and `_FFT_BLOCK_MIN` (32) cap and floor the result: +the ceiling is the measured knee on a 16-core machine at natural resolution +(smaller blocks serialise on GIL-held numpy dispatch, larger ones lose cache +residency and task granularity); the floor keeps task granularity from +collapsing at extreme pad factors, at the cost of exceeding the byte budget +there. -The selected bin is bit-identical to the full padded argmax — enforced by -`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial -spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose -baseline was captured on the old full-padded path. +The outer row-chunk sizing (`_plan_fft_rows`) needed no companion change. +It only ever budgets the raw float32 waveform *read* buffer, which this +change doesn't touch — spectrum memory is bounded independently by +`_fft_block_for`, and since the pool only ever runs as many blocks +concurrently as it has workers, peak transient spectrum memory during a +chunk's FFT phase is `_MAX_WORKERS * block * bytes_per_wf`, the same bound +whether the chunk holds 10 rows or 10,000. Queuing more rows into one chunk +to keep the read-row pool busy therefore can't blow up spectrum memory. -Work fans out over a persistent thread pool in `_FFT_BLOCK` = 512-waveform -tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose -cache residency and task granularity (measured on a 16-core machine, where -this path runs ~35× faster than the old serial padded transform at pad 40). -pyFFTW runs through per-thread `builders` plans (FFTW_MEASURE, wisdom -persisted under `~/.cache/sras-viewer/`), and `threadpoolctl` clamps BLAS to -one thread under the pool so the refinement gemm cannot oversubscribe. -`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference -full-padded path for audits. +`tools/check_equivalence.py`'s golden-hash harness remains the end-to-end +regression baseline for this path, unaffected by this change. ## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`) @@ -121,8 +134,8 @@ same edge case the masked convolution already handles. The averaging step doubles the live per-row scratch memory (a full-width `(n_frames, spf)` buffer on top of the existing compacted `waves` buffer), so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before -`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row -chunking" above. On the largest real scans `_plan_chunks` is already +`_plan_fft_rows` runs — see "Memory budget and row chunking" above. On the +largest real scans `_plan_chunks` is already clamped to its floor of one row regardless, so this costs no concurrency where it matters most; it mainly protects moderate-sized scans from an unexpected regression. @@ -144,6 +157,18 @@ by `cached_rf_image` before it hands the image back. Getting this wrong is not a slow display, it is a *wrong* display, which is why the check is a single predicate in one place rather than spread across callers. +The min peak frequency floor is the fourth recorded setting, and it sits +between the DC threshold (fully re-applicable at display time) and the +baked-in three: it is *tighten-only* re-applicable. A floor at or above the +stored one is served by masking stored pixels below it to the 0.0 "no valid +peak" sentinel — invalidated, not re-resolved; only a real recompute can +find their true above-floor peak. A floor *below* the stored one is a +genuine mismatch: the stored search never looked at those bins. Because a +served-then-masked image is a lossy stand-in for a real floored compute, +`cache_file` passes `use_stored=False` so a batch recompute always runs the +real FFT — otherwise re-batching a cached file would silently bake the +masked copy of its own old image back into the store. + Padding is the subtlest of the three, because a padded FFT looks like it should be a refinement of the unpadded one. It isn't: zero-padding interpolates between the natural bins, so it resolves a different peak @@ -170,8 +195,12 @@ recompute. The *display* path (`_stored_fft_image`) never asks bg-sub/pad/row-averaging controls — it asks whether the image matches its *own* recorded settings (`sras.precomputed_bg_sub`/`precomputed_pad_factor`/ `precomputed_row_avg_n`), which is always true whenever a stored image -exists. So presence alone decides whether it's shown; the live controls -never gate it. They still matter for two things: a genuinely never-computed +exists. Two settings are taken live instead: the DC threshold and the min +peak frequency floor, both cheaply re-appliable as masks on serve. So for +bg-sub/pad/row-avg, presence alone decides whether a stored image is shown; +the floor is the one case where a live control can gate it — a live floor +*below* the stored one can't be answered by masking, so `cached_rf_image` +declines and the caller falls through to a real compute. They still matter for two things: a genuinely never-computed angle's first live compute, and an explicit batch recompute — both of which read the live controls and produce new stored data, at which point it's the new data's *own* settings that get self-matched from then on. This is what diff --git a/pyproject.toml b/pyproject.toml index 9f7942e..7913506 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,10 +14,8 @@ dependencies = [ "scipy==1.18.0", # Angle alignment only: masked FFT phase correlation (skimage.registration). "scikit-image==0.26.0", - # Faster rfft backend; the viewer falls back to scipy.fft without it. + # Mandatory rfft backend for the peak search (no SciPy fallback). "pyFFTW==0.15.1", - # Clamps BLAS threading under the FFT worker pool. - "threadpoolctl==3.6.0", ] [project.optional-dependencies] @@ -30,6 +28,7 @@ sras-viewer = "sras_viewer.main_window:main" py-modules = [ "sras_format", "sras_compute", + "sras_render", "sras_workers", "sras_align_export", "sras_average", diff --git a/scan_format.md b/scan_format.md index 6181c02..57e4df4 100644 --- a/scan_format.md +++ b/scan_format.md @@ -255,7 +255,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 `3`; readers also accept `1` and `2` (each older tail simply lacks the fields added since — 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). | +| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `4`; readers also accept `1`–`3` (each older tail simply lacks the fields added since — 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`) @@ -292,13 +292,16 @@ never shifts an existing offset: - **`cach_version` 1**: 7 bytes, format `">4sBH"` — magic, flags, n_stored. - **`cach_version` 2**: 8 bytes, format `">4sBHB"` — + `row_avg_n`. - **`cach_version` 3**: 10 bytes, format `">4sBHBH"` — + `pad_factor`. +- **`cach_version` 4**: 14 bytes, format `">4sBHBHI"` — + `min_freq_khz`. Always written by current code. An older tail is read with its absent fields taken as the only value such a tail can describe: `row_avg_n = 0` for a `cach_version` 1 tail, which -predates row-averaged FFT caching, and `pad_factor = 1` for `cach_version` -1 or 2, which predate padded caching and are therefore natural-resolution. -Files cached before either change keep working with no recompute. +predates row-averaged FFT caching, `pad_factor = 1` for `cach_version` +1 or 2, which predate padded caching and are therefore natural-resolution, +and `min_freq_khz = 0` (no floor) for `cach_version` 1–3, which predate the +min peak frequency floor and therefore searched every bin above DC. +Files cached before any of these changes keep working with no recompute. | Offset (rel) | Size | Type | Field | Description | |--------------|------|------|-------|-------------| @@ -306,7 +309,8 @@ Files cached before either change keep working with no recompute. | 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`. | -| 8 | 2 | `u16` | `pad_factor` | *`cach_version` 3 only.* Zero-padding factor the stored `peak_freq_mhz` was resolved at: `n_fft = pad_factor × samples_per_frame`, so `1` = natural resolution. Never `0`; a `cach_version` 1 or 2 tail has no such field and is always `pad_factor = 1`. | +| 8 | 2 | `u16` | `pad_factor` | *`cach_version` ≥ 3 only.* Zero-padding factor the stored `peak_freq_mhz` was resolved at: `n_fft = pad_factor × samples_per_frame`, so `1` = natural resolution. Never `0`; a `cach_version` 1 or 2 tail has no such field and is always `pad_factor = 1`. | +| 10 | 4 | `u32` | `min_freq_khz` | *`cach_version` ≥ 4 only.* Min peak frequency floor the stored peak search excluded bins below, fixed-point in units of 0.001 MHz (kHz); `0` = no floor. Fixed-point rather than `f32` so a value that round-trips through the file compares exactly against the same value re-requested by a reader (the viewer's floor control has 0.001 MHz granularity). A `cach_version` 1–3 tail has no such field and is always `min_freq_khz = 0`. | followed by `n_stored` entries, each: @@ -341,14 +345,27 @@ Readers must fall back to real-time FFT computation (ignoring stored the reader is asking for: time-domain gating is active, the reader's requested `n_fft` doesn't equal `pad_factor × samples_per_frame`, 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; a request at one row-averaging window size must never -be served a store at another; and a request at one padding must never be -served a store at another, since a padded FFT interpolates between the -natural bins and so resolves genuinely different peak frequencies. An -`n_fft` that is not a whole multiple of `samples_per_frame` can never match -any store, because only an integer `pad_factor` is representable. +`flags.bg_sub_applied`, the reader's requested `row_avg_n` doesn't match +the stored value exactly, or the reader's requested min peak frequency +floor is *below* the stored `min_freq_khz`. A raw request must never be +served a row-averaged store, or vice versa; a request at one row-averaging +window size must never be served a store at another; and a request at one +padding must never be served a store at another, since a padded FFT +interpolates between the natural bins and so resolves genuinely different +peak frequencies. An `n_fft` that is not a whole multiple of +`samples_per_frame` can never match any store, because only an integer +`pad_factor` is representable. + +The min peak frequency floor is the one asymmetric provenance field. A +request at a floor *below* the stored one cannot be served: the stored +search never looked at bins below its floor, so the stored numbers cannot +say what a lower-floored search would have found. A request at a floor at +or *above* the stored one **is** servable — the difference is re-applied at +display time by masking every pixel whose stored `peak_freq_mhz` is below +the requested floor to `0` (the same sentinel as the DC threshold mask; +a genuine peak can never be `0`, since bin 0 is always excluded from the +search). Such masked pixels are *invalid*, not re-resolved — only a real +recompute can recover the strongest peak above the floor for them. ### In-place write ordering @@ -374,6 +391,7 @@ which has stayed `7` since the Cache Tail was introduced — this is the inner | 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. | | 3 | `SFFT` header grows a `u16` `pad_factor` — the zero-padding factor the stored `peak_freq_mhz` was resolved at, `1` = natural resolution. Before this, a padded view could never use a stored cache at all (the store was pad 1 by definition and readers rejected any `n_fft ≠ samples_per_frame`), so a user working at a pad factor got no benefit from batch-computing a file. Recording the factor lets such a view be served, while still refusing a store resolved at a *different* pad. Readers accept `cach_version` 1 and 2 tails as `pad_factor = 1`. | +| 4 | `SFFT` header grows a `u32` `min_freq_khz` — the min peak frequency floor the stored peak search excluded bins below, in 0.001 MHz units, `0` = no floor. The floor exists because a pixel that passes the DC-bias threshold but carries only weak real signal can otherwise resolve to the un-subtracted background's DC-leakage skirt — an implausibly-near-zero frequency (and so an implausibly slow velocity) for a pixel that has a genuine peak higher up. Recording the floor is what makes it enforceable against a store: without it, a stored image silently bypassed the floor entirely. Unlike the other provenance fields it is asymmetric — a *higher* requested floor is servable by masking stored pixels below it, only a *lower* one forces a recompute (see above). Readers accept `cach_version` 1–3 tails as `min_freq_khz = 0`. | A reader that does not know a `cach_version` must treat the file as uncached — not attempt a partial parse — and the file still reads as an diff --git a/sras_average.py b/sras_average.py index aed403f..861d316 100644 --- a/sras_average.py +++ b/sras_average.py @@ -3,7 +3,23 @@ sras_average.py — Waveform-averaging utility for .sras files. Reduces memory footprint by coherently averaging every N consecutive frames -along the acquisition axis, writing a new .sras file with n_frames / N frames. +along the acquisition axis, writing a new v6 .sras file with (about) +n_frames / N frames per angle. + +Handles v6/v7 only (per-angle ragged geometry). A v7 input's cache tail is +dropped — it's indexed by frame, which this changes — so the output is always +written as v6; the viewer recomputes DC/FFT on next open. Reads via memmap and +writes in row-chunks sized to a memory budget (default 1024 MB, override with +the SRAS_MEM_BUDGET_MB env var), so peak RAM stays bounded regardless of file +size — this is what makes the tool usable on multi-hundred-GB scans. + +Averaging every N frames also coarsens the physical X spacing between output +frames (each frame is a distinct stage position — see scan_format.md's +Spatial Mapping section: x_k = x_start + k * velocity_mm_s / laser_freq_hz) — +so the output header's laser_freq_hz is divided by N to keep x_axis_mm() +correct on the averaged file. This means the GUI's "Laser freq" info label +will show that adjusted value rather than the scope's real setting for an +averaged file; v6/v7 has no separate field for effective pixel pitch. Usage: python sras_average.py input.sras output.sras --n 10 @@ -13,9 +29,15 @@ Options: --n INT Number of frames to average into one (required). --discard-remainder Drop trailing frames that don't fill a complete group. Default: include a partial average for the last group. + +Environment: + SRAS_MEM_BUDGET_MB Ceiling on one row-chunk's working memory, in MB + (default 1024). Lower it on a memory-constrained + machine; the tool just takes more, smaller chunks. """ import argparse +import os import shutil import struct import sys @@ -23,14 +45,22 @@ from pathlib import Path import numpy as np -from sras_format import HDR_FMT, HDR_SIZE, SrasFile +from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile -_SUPPORTED = (2, 3, 4) +_SUPPORTED = (6, 7) +_DEFAULT_BUDGET_MB = 1024 + +# Throttle per-chunk progress printing to roughly this many lines per angle, +# so a huge angle (thousands of chunks) doesn't flood stdout while a small +# one still gets to print every chunk. +_MAX_PROGRESS_LINES = 40 def parse_args(): p = argparse.ArgumentParser( - description="Average every N waveforms in a .sras file and write a new file." + description="Average every N waveforms in a .sras file and write a new file.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, ) p.add_argument("input", help="Input .sras file") p.add_argument("output", help="Output .sras file") @@ -41,18 +71,37 @@ def parse_args(): return p.parse_args() -def read_header_sections(sras: SrasFile) -> bytes: - """The raw bytes between the header and the waveform data (angle table, - row table, preambles, background), copied through verbatim so nothing is - lost in a re-encode.""" - with open(sras.path, "rb") as f: - f.seek(HDR_SIZE) - return f.read(sras.data_offset - HDR_SIZE) +def _memory_budget_bytes() -> int: + return int(os.environ.get("SRAS_MEM_BUDGET_MB", _DEFAULT_BUDGET_MB)) * 1024 * 1024 -def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: - """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) - block. Returns int16 of shape (n_rows, n_ch, n_out, spf). +def _plan_angle(n_frames_in: int, n: int, discard_remainder: bool) -> tuple[int, int, int]: + """(n_full, remainder, n_frames_out) for averaging one angle's frames.""" + n_full = n_frames_in // n + remainder = n_frames_in % n + n_out = n_full + (1 if remainder and not discard_remainder else 0) + return n_full, remainder, n_out + + +def _chunk_rows(n_rows: int, n_channels: int, n_frames_in: int, + samples_per_frame: int, budget: int) -> int: + """How many rows to hold in RAM at once so one chunk's working buffers + (source block, mean's float64 accumulator, output block) fit the budget. + + The 4 bytes/sample below is a deliberate middle estimate, not a sum of + the three buffers: the int16 source and int16 output are 2 each, and + only mean()'s float64 result is 8, over the reduced frame axis rather + than the full block. Raise SRAS_MEM_BUDGET_MB if a machine still runs + tight at a large --n.""" + bytes_per_row = max(1, n_channels * n_frames_in * samples_per_frame * 4) + return max(1, min(n_rows, budget // bytes_per_row)) + + +def _average_block(block: np.ndarray, n: int, discard_remainder: bool, + bps: int) -> np.ndarray: + """Average every N frames of a (rows, n_ch, n_frames, spf) block along + the frame axis, returning the result already encoded in the on-disk + dtype: int8 (clipped) if bps == 1, else big-endian int16. Cast to native int16 (not float32) before calling .mean(): numpy's mean() uses a float64 accumulator by default for integer input, matching the @@ -76,38 +125,88 @@ def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarr parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) if not parts: - return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) - return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) + averaged = np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) + else: + averaged = parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) + + if bps == 1: + return np.clip(averaged, -128, 127).astype(np.int8) + return averaged.astype(">i2") -def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes, - n: int, discard_remainder: bool) -> int: - """Stream each angle through the averager, writing as we go so peak RAM - stays at one angle's block rather than the whole file.""" +def write_v6_averaged(sras: SrasFile, out_path: Path, n: int, + discard_remainder: bool, budget: int | None = None) -> list[int]: + """Write *sras* averaged every N frames to a new v6 .sras file, streaming + row-chunks per angle so peak RAM never holds more than one chunk (bounded + by *budget* bytes, default from SRAS_MEM_BUDGET_MB). + + Stages to a sibling '.part' file and os.replace()s it into position on + success, per scan_format.md's requirement that any .sras writer must + never leave a half-written file visible under its final name (a short + v6 file opens successfully with the wrong angle/frame count rather than + failing loudly). + + Returns the per-angle output frame counts. + """ + budget = _memory_budget_bytes() if budget is None else max(1, budget) + n_ch = sras.n_channels + spf = sras.samples_per_frame bps = sras.bytes_per_sample - n_frames_in = int(sras.n_frames[0]) - n_out = n_frames_in // n - if n_frames_in % n and not discard_remainder: - n_out += 1 + + plans = [_plan_angle(int(sras.n_frames[a]), n, discard_remainder) + for a in range(sras.n_angles)] + n_out_per_angle = [p[2] for p in plans] header = struct.pack( - HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]), - float(sras.x_start_mm[0]), float(sras.x_delta_mm), - sras.velocity_mm_s, sras.laser_freq_hz, - n_out, # updated frame count - sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels, + HDR_FMT_V6, b"SRAS", 6, sras.n_angles, + sras.x_start_nominal_mm, sras.y_start_nominal_mm, + sras.x_delta_nominal_mm, sras.y_delta_nominal_mm, + sras.row_spacing_mm, sras.velocity_mm_s, sras.laser_freq_hz / n, + spf, sras.sample_rate_hz, bps, n_ch, ) - with open(out_path, "wb") as f: - f.write(header) - f.write(mid_sections) - for a in range(sras.n_angles): - averaged = average_rows(sras.data[a], n, discard_remainder) - if bps == 1: - f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes()) - else: - f.write(averaged.astype(">i2").tobytes()) - return n_out + geo = bytearray() + for a in range(sras.n_angles): + geo += struct.pack(GEO_FMT_V6, float(sras.x_start_mm[a]), + float(sras.x_delta_mm_per_angle[a]), + int(n_out_per_angle[a]), int(sras.n_rows[a])) + + part_path = out_path.with_name(out_path.name + ".part") + try: + with open(part_path, "wb") as f: + f.write(header) + f.write(sras.angles_deg.astype(">f4").tobytes()) + f.write(bytes(geo)) + for a in range(sras.n_angles): + f.write(sras.y_pos_per_angle[a].astype(">f4").tobytes()) + f.write(sras.encoded_preambles()) + f.write(sras.encoded_background()) + + for a in range(sras.n_angles): + n_rows = int(sras.n_rows[a]) + n_frames_in = int(sras.n_frames[a]) + chunk_rows = _chunk_rows(n_rows, n_ch, n_frames_in, spf, budget) + total_chunks = max(1, -(-n_rows // chunk_rows)) + print_every = max(1, total_chunks // _MAX_PROGRESS_LINES) + + data = sras.data[a] + for i, r0 in enumerate(range(0, n_rows, chunk_rows)): + r1 = min(r0 + chunk_rows, n_rows) + block = np.asarray(data[r0:r1]) + out_block = _average_block(block, n, discard_remainder, bps) + f.write(out_block.tobytes()) + if total_chunks > 1 and (i % print_every == 0 or r1 == n_rows): + print(f" angle {a + 1}/{sras.n_angles}: " + f"{r1}/{n_rows} rows", flush=True) + + f.flush() + os.fsync(f.fileno()) + os.replace(part_path, out_path) + except BaseException: + part_path.unlink(missing_ok=True) + raise + + return n_out_per_angle def main(): @@ -131,18 +230,16 @@ def main(): sras = SrasFile(str(in_path)) if sras.version not in _SUPPORTED: print(f"Error: unsupported .sras version: {sras.version} " - f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})", + f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))} only)", file=sys.stderr) sys.exit(1) - n_frames_in = int(sras.n_frames[0]) - print(f" Version : v{sras.version}") - print(f" Angles : {sras.n_angles}") - print(f" Rows : {int(sras.n_rows[0])}") - print(f" Frames (actual): {n_frames_in}") - print(f" Channels : {sras.n_channels}") - print(f" Samples/frame : {sras.samples_per_frame}") - print(f" Bytes/sample : {sras.bytes_per_sample}") + aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else "" + print(f" Version : v{sras.version}") + print(f" Angles : {sras.n_angles}{aborted_note}") + print(f" Channels : {sras.n_channels}") + print(f" Samples/frame: {sras.samples_per_frame}") + print(f" Bytes/sample : {sras.bytes_per_sample}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") @@ -150,28 +247,26 @@ def main(): print(f"Wrote {out_path}") return - if args.n > n_frames_in: - print(f"Warning: --n ({args.n}) exceeds available frames ({n_frames_in}). " - "The entire dataset will be averaged into a single frame.") + print(f"\n{'idx':>4} {'angle_deg':>10} {'rows':>6} {'frames_in':>10} {'frames_out':>11}") + for a in range(sras.n_angles): + n_frames_in = int(sras.n_frames[a]) + n_full, remainder, n_out = _plan_angle(n_frames_in, args.n, args.discard_remainder) + print(f"{a:>4} {sras.angles_deg[a]:>10.4f} {int(sras.n_rows[a]):>6} " + f"{n_frames_in:>10} {n_out:>11}") + if args.n > n_frames_in: + print(f" Warning: --n ({args.n}) exceeds angle {a}'s frames " + f"({n_frames_in}); it collapses to a single frame.") - print(f"\nAveraging every {args.n} frames ...", flush=True) - print(f"\nWriting {out_path} ...", flush=True) - mid_sections = read_header_sections(sras) - n_frames_out = write_averaged(out_path, sras, mid_sections, - args.n, args.discard_remainder) + if sras.version == 7: + print("\nNote: input has a v7 cache tail; it is indexed by frame count " + "and will be dropped. The viewer will recompute DC/FFT on next open.") - n_full = n_frames_in // args.n - remainder = n_frames_in % args.n - if remainder and not args.discard_remainder: - status = f"({n_full} full groups + 1 partial group of {remainder})" - elif remainder: - status = f"({n_full} full groups, {remainder} trailing frames discarded)" - else: - status = f"({n_full} full groups)" - print(f" {n_frames_in} frames -> {n_frames_out} frames {status}") + print(f"\nAveraging every {args.n} frames and writing {out_path} ...", flush=True) + n_out_per_angle = write_v6_averaged(sras, out_path, args.n, args.discard_remainder) in_mb = in_path.stat().st_size / 1024**2 out_mb = out_path.stat().st_size / 1024**2 + print(f"\n Frames out: {n_out_per_angle}") print(f" Input size : {in_mb:.1f} MB") print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)") print("Done.") diff --git a/sras_compute.py b/sras_compute.py index ba931c6..b9c3150 100644 --- a/sras_compute.py +++ b/sras_compute.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 """Image computation and angle alignment for .sras scans. -Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a -multiprocessing child can import it without loading Qt or matplotlib — -which matters because Python 3.14 on macOS spawns rather than forks. +Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing +child can import it without loading Qt or matplotlib — which matters +because Python 3.14 on macOS spawns rather than forks. """ import atexit @@ -15,61 +15,34 @@ from dataclasses import dataclass from pathlib import Path import numpy as np +import pyfftw import scipy.fft as scipy_fft import scipy.ndimage as scipy_ndimage from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile, adc_to_mv) -# --------------------------------------------------------------------------- -# FFT backend -# --------------------------------------------------------------------------- - -try: - import pyfftw - PYFFTW_AVAILABLE = True -except ImportError: - PYFFTW_AVAILABLE = False - -try: - from threadpoolctl import threadpool_limits -except ImportError: - threadpool_limits = None - -_fft_backend = "scipy" # "scipy" or "pyfftw"; set via set_fft_backend() - - -def set_fft_backend(name: str): - """Select the rfft implementation. Module-level state, so it must be set - explicitly inside each multiprocessing child — it does not survive a - spawn. "numpy" is accepted as a legacy alias for "scipy".""" - global _fft_backend - if name == "numpy": - name = "scipy" - _fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy" - - -def get_fft_backend() -> str: - return _fft_backend - - -def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1, - workers: int = 1) -> np.ndarray: - """Dispatch rfft to the selected backend with optional multithreading.""" - if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: - return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers) - return scipy_fft.rfft(x, n=n, axis=axis, workers=workers) - - # --------------------------------------------------------------------------- # FFT worker pool and per-thread pyFFTW plans # --------------------------------------------------------------------------- -_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine: - # smaller blocks serialise on GIL-held numpy dispatch, - # larger ones lose cache residency and task granularity. -_ZOOM_MIN_PAD = 4 # zoom refinement engages at n_fft >= _ZOOM_MIN_PAD * spf -_FFT_EXACT_ENV = bool(int(os.environ.get("SRAS_FFT_EXACT", "0") or 0)) +_FFT_BLOCK_MAX = 512 # ceiling on waveforms/task: the knee measured on a + # 16-core machine at natural resolution (n_len == spf). + # Smaller blocks serialise on GIL-held numpy dispatch, + # larger ones lose cache residency and task granularity. +_FFT_BLOCK_MIN = 32 # floor, so task granularity never collapses at + # extreme pad factors — at the cost of exceeding + # _FFT_PLAN_BYTES_BUDGET there (see _fft_block_for). +_FFT_PLAN_BYTES_BUDGET = int( + os.environ.get("SRAS_FFT_PLAN_BUDGET_MB", 16)) * 1024 * 1024 + # Per-thread ceiling on one cached pyFFTW plan's resident input+output + # buffers (see _fft_block_for). This is PERMANENT memory — the plan + # cache is never evicted — so the worst case across the whole pool for + # one distinct (spf, n_len) is _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET. + # Deliberately separate from _TOTAL_BYTES_BUDGET/SRAS_MEM_BUDGET_MB, + # which bounds transient concurrently-live chunk buffers, freed at the + # end of each chunk — this one bounds a permanent per-thread tax that + # compounds with worker count, not chunk concurrency. _pool_lock = threading.Lock() _pool: ThreadPoolExecutor | None = None @@ -115,23 +88,38 @@ def _save_wisdom(): pass -def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray: +def _fft_block_for(spf: int, n_len: int) -> int: + """Waveforms per FFT task/plan at transform length n_len. Bounds a + cached pyFFTW plan's resident input+output buffers (permanent, one per + pool thread, never evicted) to _FFT_PLAN_BYTES_BUDGET regardless of pad + factor. Reproduces _FFT_BLOCK_MAX exactly at natural resolution + (n_len == spf) — see docs/design.md for the worked numbers at high pad.""" + bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1) + block = _FFT_PLAN_BYTES_BUDGET // max(1, bytes_per_wf) + return int(min(_FFT_BLOCK_MAX, max(_FFT_BLOCK_MIN, block))) + + +def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray: """rfft of a (B, spf) float32 block via a cached per-thread FFTW plan. - Plans have a fixed (_FFT_BLOCK, spf) input shape so each worker thread - plans once per transform length; a remainder block runs through the same - plan with its tail rows ignored. The returned array is the plan's output - buffer — consume it before the next call on the same thread. + Plans have a fixed (block, spf) input shape, block = _fft_block_for(spf, + n), so each worker thread plans once per (block size, transform length) + pair; a remainder block runs through the same plan with its tail rows + ignored. Single-threaded (threads=1): outer parallelism comes from the + pool, so the transform itself must not spin up threads. The returned + array is the plan's output buffer — consume it before the next call on + the same thread. """ n_wf, spf = waves.shape + block = _fft_block_for(spf, n) plans = getattr(_fftw_local, "plans", None) if plans is None: plans = _fftw_local.plans = {} - key = (_FFT_BLOCK, spf, n) + key = (block, spf, n) plan = plans.get(key) if plan is None: _load_wisdom_once() - buf = pyfftw.empty_aligned((_FFT_BLOCK, spf), dtype="float32") + buf = pyfftw.empty_aligned((block, spf), dtype="float32") # No overwrite_input: FFTW must not scribble on input_array, whose # zero-padded tail (columns spf..n) is zeroed exactly once here. plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1, @@ -144,132 +132,20 @@ def _fftw_block_rfft(waves: np.ndarray, n: int) -> np.ndarray: return plan()[:n_wf] -def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray: - """Single-threaded rfft of one block; outer parallelism comes from the - pool, so the transform itself must not spin up threads.""" - if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE: - return _fftw_block_rfft(waves, n) - return scipy_fft.rfft(waves, n=n, axis=-1, workers=1) - - # --------------------------------------------------------------------------- -# Zoom peak search: coarse rfft + local fine DFT around the winning bin +# Peak search: full transform, block-parallel # --------------------------------------------------------------------------- -_ZOOM_HALFWIDTH = 0.75 # refinement window half-width, in coarse spacings. - # Every fine bin lies within 0.5 spacings of its - # nearest coarse bin, and that bin is guaranteed to - # be a candidate (see _ZOOM_CAND_RATIO), so 0.5 - # suffices; 0.75 adds rounding margin. -_ZOOM_CAND_RATIO = 0.7 # refine every coarse bin within this power ratio of - # its row's coarse maximum. Quarter-natural-bin - # scalloping at the 2x-oversampled coarse grid can - # understate a peak by at most ~19% in power, so 0.7 - # keeps a wide margin — near-contenders are resolved - # on the fine grid, never ranked from coarse samples. - - -@dataclass -class _ZoomPlan: - """Constants for the coarse+refine peak search, built once per - compute_rf_image call and shared across worker threads. *phases* is a - lazily-filled window-start -> phase-vector cache; a benign duplicate - compute under concurrency is cheaper than locking.""" - n_fft: int - n_coarse: int - m: int # fine bins per refinement window - n_bins_fine: int - E: np.ndarray # (spf, m) complex64 fine-DFT matrix, relative bins - j: np.ndarray # (spf,) float64 sample indices - phases: dict - - -def _zoom_plan(spf: int, n_fft: int) -> _ZoomPlan: - # 2x-oversampled coarse grid: the padded power spectrum is a trig - # polynomial of degree spf-1, so at 2x sampling its global max cannot - # hide between coarse bins. - n_coarse = scipy_fft.next_fast_len(2 * spf, real=True) - n_bins_fine = n_fft // 2 + 1 - m = int(np.ceil(2 * _ZOOM_HALFWIDTH * n_fft / n_coarse)) + 1 - m = min(m, n_bins_fine - 1) - j = np.arange(spf, dtype=np.float64) - E = np.exp((-2j * np.pi / n_fft) * np.outer(j, np.arange(m))).astype(np.complex64) - return _ZoomPlan(n_fft, n_coarse, m, n_bins_fine, E, j, {}) - - -def _window_start(k_c: np.ndarray, zp: _ZoomPlan) -> np.ndarray: - """First fine bin of the refinement window around each coarse bin. - Clipped to [1, ...] so fine bin 0 stays excluded (DC suppression).""" - k0 = np.floor((k_c - _ZOOM_HALFWIDTH) * (zp.n_fft / zp.n_coarse)).astype(np.int64) - return np.clip(k0, 1, max(1, zp.n_bins_fine - zp.m)) - - -def _refine_window(waves: np.ndarray, rows: np.ndarray, k0: int, zp: _ZoomPlan, - best_pow: np.ndarray, best_bin: np.ndarray): - """Evaluate fine bins k0..k0+m-1 for the given rows and fold the result - into the per-row best (power, bin), preserving np.argmax's lowest-bin - tie-break.""" - ph = zp.phases.get(k0) - if ph is None: - ph = np.exp((-2j * np.pi * k0 / zp.n_fft) * zp.j).astype(np.complex64) - zp.phases[k0] = ph - F = (waves[rows] * ph) @ zp.E - q = F.real ** 2 - q += F.imag ** 2 - i = np.argmax(q, axis=1) - p = q[np.arange(len(rows)), i] - b = k0 + i - upd = (p > best_pow[rows]) | ((p == best_pow[rows]) & (b < best_bin[rows])) - ridx = rows[upd] - best_pow[ridx] = p[upd] - best_bin[ridx] = b[upd] - - -def _peak_bins_zoom(waves: np.ndarray, zp: _ZoomPlan) -> np.ndarray: - """Peak fine-bin per waveform without materialising the padded spectrum: - coarse rfft, then a small fine DFT (one gemm per shared window) on the - exact n_fft bin grid. Identity with the full padded argmax is enforced - by tests/test_compute.py::test_zoom_identity and the golden-hash sweep.""" - n_wf = waves.shape[0] - S = _block_rfft(waves, zp.n_coarse) - P = S.real ** 2 - P += S.imag ** 2 - P[:, 0] = 0.0 - p_c = np.max(P, axis=1) - - best_pow = np.full(n_wf, -1.0, dtype=np.float32) - best_bin = np.full(n_wf, np.iinfo(np.int64).max, dtype=np.int64) - - # For a clean signal this yields one or two windows; a noise spectrum - # (many near-equal peaks) yields a dozen or so — still a tiny fraction - # of the padded grid. - thr = np.where(p_c > 0, np.float32(_ZOOM_CAND_RATIO) * p_c, - np.float32(np.inf)) - rows_c, bins_c = np.nonzero(P >= thr[:, None]) - k0_c = _window_start(bins_c, zp) - # Zeroing the coarse DC bin (suppression) blinds the candidate scan to - # fine bins closer to DC than the first coarse sample — where the - # DC-leakage skirt of an un-subtracted offset peaks. Always refine the - # DC-adjacent window too. - rows_c = np.concatenate([rows_c, np.arange(n_wf)]) - k0_c = np.concatenate([k0_c, np.ones(n_wf, dtype=np.int64)]) - pair = np.unique(np.stack([rows_c, k0_c], axis=1), axis=0) - for k0 in np.unique(pair[:, 1]): - _refine_window(waves, pair[pair[:, 1] == k0, 0], int(k0), zp, - best_pow, best_bin) - - # An all-zero spectrum must reproduce argmax-of-zeros = bin 0. - best_bin[p_c == 0.0] = 0 - return best_bin - - -def _peak_bins_direct(waves: np.ndarray, n_len: int) -> np.ndarray: - """Peak bin per waveform via the full transform — for pad factors too - small for the zoom search to pay.""" +def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray: + """Peak bin per waveform via the full transform. *min_bin* zeroes every + bin below it (always >= 1, so true DC stays suppressed) before the + argmax, so the peak search never returns a bin below the caller's + frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks, + fanned out across _fft_pool() — see docs/design.md.""" S = _block_rfft(waves, n_len) power = S.real ** 2 power += S.imag ** 2 - power[:, 0] = 0.0 + power[:, :min_bin] = 0.0 return np.argmax(power, axis=1) @@ -300,9 +176,11 @@ def _chunk_rows_for(n_frames: int, samples_per_frame: int, def _plan_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int: """Rows per outer chunk for the block-FFT path. Only the float32 - waveform buffer scales with the chunk (in-flight block spectra total a - few MB across the whole pool), so budget it with 2x slack and let the - block fan-out saturate the pool regardless of pad factor.""" + waveform read buffer scales with the chunk; spectrum memory is bounded + independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET + regardless of chunk size or pad factor (see docs/design.md), so this + formula budgets the read buffer alone, with 2x slack, and lets the block + fan-out saturate the pool.""" bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2) return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row))) @@ -508,7 +386,8 @@ def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int: def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None, - apply_bg_sub: bool, row_avg_n: int) -> list[str]: + apply_bg_sub: bool, row_avg_n: int, + min_freq_mhz: float = 0.0) -> list[str]: """Why this file's stored FFT images can't answer a request, as human- readable phrases; empty means they can. @@ -518,8 +397,23 @@ def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None, serves a stored image iff this returns nothing, and callers that want to *explain* the miss (rather than silently recompute) format these same strings, so the two can't drift apart. + + The min peak frequency floor is the asymmetric case: a stored floor + *above* the request is a genuine mismatch (bins below the stored floor + were never searched, so the stored numbers can't answer a request that + wants them considered), but a stored floor at or *below* the request is + servable — the difference is re-applied as a display-time mask by + cached_rf_image (pixels whose stored peak falls below the requested + floor are marked invalid rather than re-resolved). Compared in whole + kHz, the header field's own fixed-point grid, so a value that + round-trips through the file can never miscompare against itself. """ reasons = [] + if round(sras.precomputed_min_freq_mhz * 1000) > round(min_freq_mhz * 1000): + reasons.append( + f"stored with a {sras.precomputed_min_freq_mhz:g} MHz " + f"min-peak-freq floor, requested {min_freq_mhz:g} MHz " + f"(bins below the stored floor were never searched)") pad = pad_factor_for(sras, n_fft) if pad != sras.precomputed_pad_factor: want = f"pad {pad}x" if pad else "a ragged n_fft" @@ -543,7 +437,8 @@ def cached_rf_image(sras: SrasFile, angle_idx: int, n_fft: int | None = None, dc4_mv: np.ndarray | None = None, row_avg_n: int = 0, - allow_dc_recompute: bool = True) -> np.ndarray | None: + allow_dc_recompute: bool = True, + min_freq_mhz: float = 0.0) -> np.ndarray | None: """The precomputed-cache fast path for compute_rf_image: a ready-to- display peak-frequency image if this file already has one matching every setting the caller cares about, else None (caller must run a @@ -571,10 +466,20 @@ def cached_rf_image(sras: SrasFile, angle_idx: int, channel on the caller's behalf; this returns None instead so a caller that wants to stay off the I/O path (e.g. a GUI thread) can choose to fall through to a real compute rather than block. + + *min_freq_mhz* above the stored floor is re-imposed here as a mask: + pixels whose stored peak is below it are set to 0.0 — the same sentinel + the DC mask uses, and unambiguous, since a genuine peak can never be + 0.0 (bin 0 is always excluded from the search). Masked-below-floor + pixels are marked invalid, not re-resolved; only a real recompute can + recover the strongest peak *above* the floor for them. A request whose + floor is *below* the stored one can't be answered at all (see + cache_mismatch_reasons) and returns None like any other mismatch. """ cached_freq = sras.precomputed_freq_mhz[angle_idx] if cached_freq is None or cache_mismatch_reasons( - sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n): + sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n, + min_freq_mhz=min_freq_mhz): return None dc4_img = None if dc_threshold_mv is not None: @@ -594,6 +499,10 @@ def cached_rf_image(sras: SrasFile, angle_idx: int, freq_img = cached_freq.copy() if dc4_img is not None: freq_img[dc4_img < dc_threshold_mv] = 0.0 + if min_freq_mhz > 0.0: + # Strict <, matching the compute path: the first bin at or above + # the floor survives searchsorted there, so it must survive here. + freq_img[freq_img < min_freq_mhz] = 0.0 return freq_img @@ -605,8 +514,9 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, max_workers: int | None = None, budget: int | None = None, should_stop=None, - exact: bool = False, - row_avg_n: int = 0) -> np.ndarray: + row_avg_n: int = 0, + min_freq_mhz: float = 0.0, + use_stored: bool = True) -> np.ndarray: """FFT of each CH1 waveform; pixel = peak frequency in MHz. Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is @@ -624,36 +534,59 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, DC-channel precompute cache), pass it as *dc4_mv* (mV, shape (n_rows, n_frames)) to reuse it instead of re-reading CH4 here. - At pad factors >= _ZOOM_MIN_PAD the padded spectrum is never - materialised: a coarse rfft finds each peak and a local fine DFT - resolves it on the exact n_fft bin grid (_peak_bins_zoom). *exact* - forces the reference full-padded transform instead — it exists for - tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and - memory-hungry at high pad. + The peak search runs the full transform (_peak_bins) in + _fft_block_for(spf, n_len)-sized blocks, fanned out across the + pyFFTW-plan-caching thread pool — see docs/design.md for how the block + size is bounded so this stays memory-safe at high pad factors. *row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n same-row neighbors (distance-weighted, valid-neighbors-only per the same dc_threshold_mv mask) before the FFT runs — see _row_average_waveforms. 0 (default) is the raw, unaveraged behavior. + *min_freq_mhz* excludes every bin below it (true DC, bin 0, is always + excluded regardless) from the peak search, on top of — not instead of — + the dc_threshold_mv mask. Without it, a pixel that passes the DC-bias + threshold but carries only weak real signal can still resolve to a + near-zero frequency: the un-subtracted background's DC-leakage skirt + then has more power than the genuine (but weak) signal peak. Raising + the floor above that skirt forces the search to report the strongest + peak that is plausibly real signal instead. 0.0 (default) disables it + (only true DC is excluded, the long-standing behavior). + Fast path: if the file has a precomputed peak-frequency image for this angle (v5 PREC or v7 CACH) matching every one of the caller's settings — including row_avg_n exactly — the stored image is used directly, no - FFT is run. See cached_rf_image. + FFT is run. See cached_rf_image. A *min_freq_mhz* at or above the + stored floor is part of that service: pixels whose stored peak falls + below it come back masked to 0.0 rather than re-resolved. A request + *below* the stored floor is a genuine mismatch and falls through to + the real FFT here. + + *use_stored=False* skips the fast path entirely and always runs the + real FFT. Batch recompute (cache_file) needs this: served-then-masked + pixels written back into the store would permanently replace peaks a + real recompute re-resolves above the floor. """ n_rows, n_frames = sras.image_shape(angle_idx) data = sras.data[angle_idx] - fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub, - n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n) - if fast is not None: - return fast + if use_stored: + fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, + apply_bg_sub=apply_bg_sub, n_fft=n_fft, + dc4_mv=dc4_mv, row_avg_n=row_avg_n, + min_freq_mhz=min_freq_mhz) + if fast is not None: + return fast # ---- Chunked FFT path -------------------------------------------------- - exact = exact or _FFT_EXACT_ENV spf = sras.samples_per_frame n_len = n_fft if n_fft is not None else spf + block = _fft_block_for(spf, n_len) freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32) + # First fine bin at or above the floor; searchsorted is exact here since + # freq32 is the very axis the floor is expressed against. + min_bin = max(1, int(np.searchsorted(freq32, min_freq_mhz))) img = np.zeros((n_rows, n_frames), dtype=np.float32) background = sras.background if (apply_bg_sub and sras.background is not None) else None cal4 = sras.cal(CH4_IDX) @@ -665,26 +598,15 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, dc4_full = dc4_mv row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None - zp = (_zoom_plan(spf, n_fft) - if not exact and n_fft is not None and n_fft >= _ZOOM_MIN_PAD * spf - else None) total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget) if row_avg_n > 0: # One extra same-sized transient buffer (the pre-averaging full-row # scratch array) is live per in-flight row; halve the budget so - # _plan_fft_rows/the exact-path sizing accounts for it rather than - # relying on _plan_fft_rows's existing 2x slack to happen to cover it. + # _plan_fft_rows accounts for it rather than relying on its existing + # 2x slack to happen to cover it. total = max(1, total // 2) cap = max_workers if max_workers is not None else _MAX_WORKERS - if exact: - # The reference path materialises the full padded spectrum, so rows - # are budgeted against it (complex64 + power + temp per bin) and the - # chunk runs as one serial transform. - bytes_per_row = max(1, n_frames * (4 * spf + 16 * (n_len // 2 + 1))) - chunk_rows = int(max(1, min(_CHUNK_ROWS_MAX, total // bytes_per_row))) - cap = 1 - else: - chunk_rows = _plan_fft_rows(n_frames, spf, total) + chunk_rows = _plan_fft_rows(n_frames, spf, total) pool = _fft_pool() if cap > 1 else None def process(r0: int, r1: int): @@ -748,23 +670,14 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, def fft_block(b0: int): if should_stop is not None and should_stop(): return - b1 = min(b0 + _FFT_BLOCK, n_wf) - w = waves[b0:b1] - bins = (_peak_bins_zoom(w, zp) if zp is not None - else _peak_bins_direct(w, n_len)) - out[b0:b1] = freq32[bins] + b1 = min(b0 + block, n_wf) + out[b0:b1] = freq32[_peak_bins(waves[b0:b1], n_len, min_bin)] - if exact: - spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1) - power = spectrum.real ** 2 - power += spectrum.imag ** 2 - power[:, 0] = 0.0 # suppress DC bin - out[:] = freq32[np.argmax(power, axis=1)] - elif pool is None: - for b0 in range(0, n_wf, _FFT_BLOCK): + if pool is None: + for b0 in range(0, n_wf, block): fft_block(b0) else: - list(pool.map(fft_block, range(0, n_wf, _FFT_BLOCK))) + list(pool.map(fft_block, range(0, n_wf, block))) # Boolean scatter is row-major, matching the read pass's # concatenation order. @@ -773,18 +686,10 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, else: img[r0:r1] = out.reshape(r1 - r0, n_frames) - # BLAS must not thread under the pool (the fine-DFT gemm would multiply - # against the pool's own workers). - limiter = (threadpool_limits(limits=1) - if pool is not None and threadpool_limits is not None else None) - try: - for r0 in range(0, n_rows, chunk_rows): - if should_stop is not None and should_stop(): - break - process(r0, min(r0 + chunk_rows, n_rows)) - finally: - if limiter is not None: - limiter.unregister() + for r0 in range(0, n_rows, chunk_rows): + if should_stop is not None and should_stop(): + break + process(r0, min(r0 + chunk_rows, n_rows)) return img @@ -793,16 +698,17 @@ def compute_rf_image(sras: SrasFile, angle_idx: int, # --------------------------------------------------------------------------- def cache_file(path: str, mode: str, apply_bg_sub: bool, - fft_backend: str = "scipy", max_workers: int = 0, + max_workers: int = 0, pad_factor: int = 1, dc_threshold_mv: float | None = None, - row_avg_n: int = 0) -> str: + row_avg_n: int = 0, + min_freq_mhz: float = 0.0) -> str: """Compute and store DC or FFT images for every angle of one file, converting v6 → v7 in place. Returns "" on success or an error message. Module-level and picklable so it can run in a ProcessPoolExecutor. The - FFT backend and worker cap are passed explicitly because module globals - do not survive a spawn. + worker cap is passed explicitly because module globals do not survive a + spawn. mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at display time), or "fft_rowavg" (same-row, distance-weighted CH1 @@ -818,6 +724,16 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, using, since a pad-1 cache is dead weight to a padded view and vice versa (cached_rf_image refuses the mismatch rather than showing peaks resolved at the wrong resolution). + + *min_freq_mhz* is the peak-search floor (see compute_rf_image), also + recorded in the SFFT block; it is quantized up front to the header + field's 0.001 MHz grid so the stored images and the recorded floor + can never disagree. Both FFT modes force use_stored=False on + compute_rf_image: a batch recompute must always run the real FFT, never + serve this file's own existing cache back to itself — with a floor (or + a changed DC threshold in fft_rowavg mode) that would silently bake a + masked copy of the old image into the store in place of a genuine + recompute. """ global _MAX_WORKERS try: @@ -828,7 +744,12 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, # between a bad argument costing nothing and costing the whole run. if not (1 <= pad_factor <= MAX_PAD_FACTOR): return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}" - set_fft_backend(fft_backend) + if not (np.isfinite(min_freq_mhz) and min_freq_mhz >= 0): + return f"min_freq_mhz must be a finite value >= 0, got {min_freq_mhz}" + # Quantize to the SFFT header's fixed-point grid (whole kHz) before + # computing, so the floor the FFT actually ran with is exactly the + # floor the header records. + min_freq_mhz = round(min_freq_mhz * 1000) / 1000.0 if max_workers: _MAX_WORKERS = max_workers @@ -860,13 +781,16 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, # internally over blocks, so angles run one at a time with the # full budget. freq = [compute_rf_image(sras, a, dc_threshold_mv=None, - apply_bg_sub=effective_bg, n_fft=n_fft) + apply_bg_sub=effective_bg, n_fft=n_fft, + min_freq_mhz=min_freq_mhz, + use_stored=False) for a in range(n)] # new_row_avg_n=0 explicitly: these images are raw per-pixel FFTs, # and carrying forward a row_avg_n left by an earlier fft_rowavg # write would label them as something they are not. sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg, - new_row_avg_n=0, new_pad_factor=pad_factor) + new_row_avg_n=0, new_pad_factor=pad_factor, + new_min_freq_mhz=min_freq_mhz) else: # "fft_rowavg" if row_avg_n <= 0: return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode" @@ -876,11 +800,14 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool, effective_bg = apply_bg_sub and sras.background is not None freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv, apply_bg_sub=effective_bg, n_fft=n_fft, - row_avg_n=row_avg_n) + row_avg_n=row_avg_n, + min_freq_mhz=min_freq_mhz, + use_stored=False) for a in range(n)] sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg, new_row_avg_n=row_avg_n, - new_pad_factor=pad_factor) + new_pad_factor=pad_factor, + new_min_freq_mhz=min_freq_mhz) return "" except Exception as exc: return str(exc) diff --git a/sras_format.py b/sras_format.py index a33a1e5..ef46264 100644 --- a/sras_format.py +++ b/sras_format.py @@ -60,15 +60,18 @@ 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 = 3 # written on every fresh write -CACH_VERSIONS_READABLE = (1, 2, 3) # accepted on read — see +CACH_VERSION = 4 # written on every fresh write +CACH_VERSIONS_READABLE = (1, 2, 3, 4) # accepted on read — see # _read_sfft_block. Each bump only # appended a field, and every older # tail has a well-defined reading: # v1 predates row-averaged FFT - # caching (row_avg_n=0) and v1/v2 + # caching (row_avg_n=0), v1/v2 # predate padded caching, so both - # are natural-resolution (pad 1). + # are natural-resolution (pad 1), + # and v1-v3 predate the min peak + # frequency floor (min_freq 0 = + # no floor). CACH_FLAG_DC = 0x01 CACH_FLAG_FFT = 0x02 @@ -79,11 +82,20 @@ SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT) SFFT_MAGIC = b"SFFT" SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored (cach_version 1) SFFT_HDR_FMT_V2 = ">4sBHB" # + row_avg_n (cach_version 2) -SFFT_HDR_FMT = ">4sBHBH" # + pad_factor (cach_version 3) +SFFT_HDR_FMT_V3 = ">4sBHBH" # + pad_factor (cach_version 3) +SFFT_HDR_FMT = ">4sBHBHI" # + min_freq_khz (cach_version 4) SFFT_HDR_SIZE_V1 = struct.calcsize(SFFT_HDR_FMT_V1) SFFT_HDR_SIZE_V2 = struct.calcsize(SFFT_HDR_FMT_V2) +SFFT_HDR_SIZE_V3 = struct.calcsize(SFFT_HDR_FMT_V3) SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT) MAX_PAD_FACTOR = 0xFFFF # the H field above +# min_freq_khz is fixed-point (u32, units of 0.001 MHz), not a float32: +# the viewer's floor spinbox has 0.001 MHz granularity, and +# round(mhz * 1000) / 1000.0 reproduces the exact float64 the user typed, +# so the accept rule in sras_compute.cache_mismatch_reasons can compare +# with plain integer ordering. A ">f" float32 of e.g. 20.1 would read back +# as 20.10000038… > 20.1 and report a spurious mismatch forever. +MAX_MIN_FREQ_KHZ = 0xFFFFFFFF # the I field above; 0 = no floor SFFT_FLAG_BG_SUB = 0x01 SFFT_FLAG_ROW_AVG = 0x02 # peak_freq_mhz came from same-row, # distance-weighted averaged CH1 @@ -134,6 +146,13 @@ def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV, return (adc - yoff_adc) * ymult_mv + yzero_mv +def _axes_extent(x_axis, y_axis, dx: float, dy: float) -> list[float]: + """Matplotlib imshow extent with half-pixel margins, Y flipped so row 0 + renders at the top.""" + return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, + y_axis[-1] + dy / 2, y_axis[0] - dy / 2] + + # --------------------------------------------------------------------------- # Binary read helpers # --------------------------------------------------------------------------- @@ -188,9 +207,10 @@ class SrasFile: 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 scalars ``precomputed_bg_sub`` / - ``precomputed_row_avg_n`` / ``precomputed_pad_factor`` record the - settings the stored FFT images were computed under, so a reader can tell - whether they answer the question it is actually asking. + ``precomputed_row_avg_n`` / ``precomputed_pad_factor`` / + ``precomputed_min_freq_mhz`` record the settings the stored FFT images + were computed under, so a reader can tell whether they answer the + question it is actually asking. """ def __init__(self, path: str): @@ -253,6 +273,13 @@ class SrasFile: # FFT resolves peaks a padded view would, and only such a view can # be served from it — see sras_compute.cached_rf_image. self.precomputed_pad_factor: int = 1 + # Min peak frequency floor (MHz) the stored peak search excluded + # bins below; 0.0 = no floor. Unlike bg-sub/pad/row-avg it is + # tighten-only re-applicable: a *higher* floor can be re-imposed on + # a stored image by masking pixels below it, but bins below the + # stored floor were never searched, so a lower floor needs a real + # recompute — see sras_compute.cache_mismatch_reasons. + self.precomputed_min_freq_mhz: float = 0.0 def encoded_preambles(self) -> bytes: """This file's Preamble Blocks section, as bytes a writer can emit. @@ -595,30 +622,35 @@ class SrasFile: store[angle_idx] = _read_f32_image(f, shape) return flags - def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int, int] | None: + def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int, int, float] | None: """Read the SFFT block header — its layout depends on cach_version, since each bump appended a trailing field (v2 row_avg_n, v3 - pad_factor) — then n_stored per-angle peak_freq_mhz entries - (unchanged across versions). + pad_factor, v4 min_freq_khz) — then n_stored per-angle + peak_freq_mhz entries (unchanged across versions). - Returns (flags, row_avg_n, pad_factor), or None if the block is - malformed. The absent fields of an older tail take the value that - describes what such a tail can only have been: row_avg_n=0 for v1, - which predates row-averaged FFT caching, and pad_factor=1 for v1/v2, - which predate padded caching and so are natural-resolution. + Returns (flags, row_avg_n, pad_factor, min_freq_mhz), or None if + the block is malformed. The absent fields of an older tail take the + value that describes what such a tail can only have been: + row_avg_n=0 for v1, which predates row-averaged FFT caching; + pad_factor=1 for v1/v2, which predate padded caching and so are + natural-resolution; and min_freq_mhz=0.0 for v1-v3, which predate + the min peak frequency floor and so searched every bin above DC. """ - hdr_fmt = {1: SFFT_HDR_FMT_V1, 2: SFFT_HDR_FMT_V2}.get( - cach_version, SFFT_HDR_FMT) + hdr_fmt = {1: SFFT_HDR_FMT_V1, 2: SFFT_HDR_FMT_V2, + 3: SFFT_HDR_FMT_V3}.get(cach_version, SFFT_HDR_FMT) raw = f.read(struct.calcsize(hdr_fmt)) if len(raw) < struct.calcsize(hdr_fmt): return None - row_avg_n, pad_factor = 0, 1 + row_avg_n, pad_factor, min_freq_khz = 0, 1, 0 if cach_version == 1: magic, flags, n_stored = struct.unpack(hdr_fmt, raw) elif cach_version == 2: magic, flags, n_stored, row_avg_n = struct.unpack(hdr_fmt, raw) - else: + elif cach_version == 3: magic, flags, n_stored, row_avg_n, pad_factor = struct.unpack(hdr_fmt, raw) + else: + (magic, flags, n_stored, row_avg_n, pad_factor, + min_freq_khz) = struct.unpack(hdr_fmt, raw) if magic != SFFT_MAGIC: return None for _ in range(n_stored): @@ -627,7 +659,7 @@ class SrasFile: break self.precomputed_freq_mhz[angle_idx] = _read_f32_image( f, self.image_shape(angle_idx)) - return flags, row_avg_n, pad_factor + return flags, row_avg_n, pad_factor, min_freq_khz / 1000.0 def _parse_cach_section(self, offset: int): """Parse the v7 CACH tail that holds precomputed DC/FFT images.""" @@ -650,10 +682,13 @@ class SrasFile: result = self._read_sfft_block(f, cach_version) if result is None: return - flags, row_avg_n, pad_factor = result + flags, row_avg_n, pad_factor, min_freq_mhz = 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 self.precomputed_pad_factor = max(1, pad_factor) + # No flag bit gates the floor: 0 (= no floor) is already the + # value every pre-v4 tail reads as. + self.precomputed_min_freq_mhz = min_freq_mhz def write_v7_cache(self, *, new_dc3_mv: list[np.ndarray | None] | None = None, @@ -661,7 +696,8 @@ class SrasFile: new_freq_mhz: list[np.ndarray | None] | None = None, new_bg_sub: bool | None = None, new_row_avg_n: int | None = None, - new_pad_factor: int | None = None): + new_pad_factor: int | None = None, + new_min_freq_mhz: float | 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 @@ -681,6 +717,13 @@ class SrasFile: different pad resolves different peaks, so recording it is what lets a reader refuse the cache instead of showing the wrong numbers. + *new_min_freq_mhz* is the min peak frequency floor the passed + *new_freq_mhz*'s peak search excluded bins below (0.0 = no floor), + carried forward the same way. It is stored fixed-point (whole kHz), + so the value is quantized to 0.001 MHz on write and + ``precomputed_min_freq_mhz`` is updated to the quantized value — + what a reload would see, never a float the header can't represent. + 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. @@ -697,11 +740,21 @@ class SrasFile: else self.precomputed_row_avg_n) final_pad_factor = (new_pad_factor if new_pad_factor is not None else self.precomputed_pad_factor) + final_min_freq_mhz = (new_min_freq_mhz if new_min_freq_mhz is not None + else self.precomputed_min_freq_mhz) 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}") if not (1 <= final_pad_factor <= MAX_PAD_FACTOR): raise ValueError( f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {final_pad_factor}") + if not (np.isfinite(final_min_freq_mhz) and final_min_freq_mhz >= 0): + raise ValueError( + f"min_freq_mhz must be a finite value >= 0, got {final_min_freq_mhz}") + final_min_freq_khz = int(round(final_min_freq_mhz * 1000)) + if final_min_freq_khz > MAX_MIN_FREQ_KHZ: + raise ValueError( + f"min_freq_mhz too large for the u32 kHz header field: " + f"{final_min_freq_mhz}") # dc3/dc4 are always populated together by every current caller, but # guard the per-angle pairing explicitly rather than assume it: an @@ -730,7 +783,7 @@ class SrasFile: 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, - final_pad_factor) + final_pad_factor, final_min_freq_khz) for a in fft_entries: payload += struct.pack(">H", a) payload += final_freq[a].astype(">f4").tobytes() @@ -770,6 +823,7 @@ class SrasFile: self.precomputed_bg_sub = final_bg_sub self.precomputed_row_avg_n = final_row_avg_n self.precomputed_pad_factor = final_pad_factor + self.precomputed_min_freq_mhz = final_min_freq_khz / 1000.0 # ------------------------------------------------------------------ # Axes helpers diff --git a/sras_render.py b/sras_render.py new file mode 100644 index 0000000..65dbcb0 --- /dev/null +++ b/sras_render.py @@ -0,0 +1,153 @@ +"""Pure-matplotlib rendering of a displayed SRAS image: imshow + colorbar + +axis/title labeling, shared by the interactive Qt canvas +(sras_viewer.canvases.ImageCanvas, which supplies its own already-Qt-backed +Figure/Axes) and the headless batch image-export worker (which builds a +throwaway Agg Figure per file and never touches Qt) -- so an exported PNG +can never quietly start looking different from what the GUI actually shows. + +Deliberately no PyQt6 import anywhere in this module: BatchExportImagesWorker +(sras_workers.py) may run export_view_image inside a spawned +ProcessPoolExecutor subprocess, exactly like sras_compute.cache_file, and +importing anything under the sras_viewer package would run its __init__.py +and pull in the whole Qt widget tree for no reason. +""" + +from pathlib import Path + +import matplotlib as mpl +import numpy as np +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure + +from sras_compute import compute_rf_image, dc_image_mv +from sras_format import CH4_IDX, CH_NAMES, SrasFile, _axes_extent + + +def draw_view_image(ax, fig, img: np.ndarray, extent: list[float], cmap, + vmin: float, vmax: float, xlabel: str, ylabel: str, + title: str, colorbar_label: str = "", cb_ticks=None, + norm=None, bad_color=None): + """imshow + colorbar + labels onto an already-created (ax, fig) pair. + + *cmap* may be a name or a Colormap instance. *norm* (which overrides + vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image + with whole-number colorbar bands instead of a continuous shade. + *bad_color*, if given, is the fill for NaN pixels -- a copy of *cmap* is + made so a shared, registered instance is never mutated. + + Shared by ImageCanvas.show_image (Qt-backed ax/fig) and + export_view_image (headless Agg ax/fig) so the two can never drift into + showing different things for the same settings. + """ + if bad_color is not None: + cmap = (cmap if hasattr(cmap, "with_extremes") + else mpl.colormaps[cmap]).with_extremes(bad=bad_color) + + kw = ({"norm": norm} if norm is not None + else {"vmin": vmin, "vmax": vmax}) + im = ax.imshow( + img, aspect="auto", origin="upper", + extent=extent, cmap=cmap, interpolation="nearest", **kw, + ) + cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, ticks=cb_ticks) + if colorbar_label: + cb.set_label(colorbar_label) + + ax.set_xlabel(xlabel) + ax.set_ylabel(ylabel) + ax.set_title(title) + return im + + +def export_view_image(path: str, *, out_dir: str, angle_idx: int, ch_idx: int, + is_fft_mode: bool, is_velocity: bool, + dc_threshold_mv: float, apply_bg_sub: bool, + pad_factor: int, min_freq_mhz: float, grating_um: float, + cmap: str, auto_scale: bool, vmin: float, vmax: float, + highlight_masked: bool, mode_str: str, + colorbar_label: str, mask_color: str = "magenta", + max_workers: int | None = None, + dpi: int = 150) -> tuple[str, str]: + """One file's contribution to Batch Export View as Images: renders + (angle_idx, ch_idx) at the given display settings to a PNG under + *out_dir*, via draw_view_image -- so a batch export is a folder of what + ImageCanvas.show_image would have put on screen for these settings, not + a raw data dump. + + Module-level and picklable, like sras_compute.cache_file, so it can run + in a ProcessPoolExecutor -- see BatchExportImagesWorker. Unlike + cache_file this never writes to *path*: export is a read of the file's + own data, not a cache conversion, so any version SrasFile can open + works, with no v6/v7 precondition. + + *pad_factor* (not n_fft) travels across files deliberately: n_fft + depends on samples_per_frame, which can differ between files in the + same batch, so n_fft is derived per file, here, from *this* file's own + value -- the same reason sras_compute.cache_file does the same thing. + + Returns (error, out_name). error is "" on success. out_name is the + filename this call targeted -- set as soon as it's known, even on most + failures -- so the caller can flag same-stem collisions across the + batch without any cross-process bookkeeping. + """ + out_name = "" + try: + sras = SrasFile(path) + if angle_idx >= sras.n_angles: + return (f"angle {angle_idx} out of range " + f"(file has {sras.n_angles} angle(s))", out_name) + + out_name = f"{Path(path).stem}_angle{angle_idx}_{CH_NAMES[ch_idx]}.png" + + if is_fft_mode: + n_fft = (sras.samples_per_frame * pad_factor + if pad_factor > 1 else None) + freq = compute_rf_image( + sras, angle_idx, dc_threshold_mv=dc_threshold_mv, + apply_bg_sub=apply_bg_sub, n_fft=n_fft, + min_freq_mhz=min_freq_mhz, max_workers=max_workers) + img = freq * grating_um if is_velocity else freq + else: + img = dc_image_mv(sras, angle_idx, ch_idx, max_workers=max_workers) + + display_img, bad_color = img, None + if highlight_masked and is_fft_mode: + # Mirrors the viewer's _redraw_image rule: DC-masked pixels and + # value-0 pixels (the "no valid peak" sentinel — DC-masked, + # below the min-freq floor, or empty spectrum; the grating + # multiply above preserves zeros, so this holds for Velocity + # too) both render in the highlight color. + dc4 = dc_image_mv(sras, angle_idx, CH4_IDX, max_workers=max_workers) + valid = dc4 >= dc_threshold_mv + if valid.shape == display_img.shape: + valid &= display_img != 0.0 + display_img = display_img.astype(np.float32, copy=True) + display_img[~valid] = np.nan + bad_color = mask_color + + if auto_scale: + v0, v1 = float(np.nanmin(display_img)), float(np.nanmax(display_img)) + if not np.isfinite(v0): + v0, v1 = 0.0, 0.0 # every pixel masked out + else: + v0, v1 = vmin, vmax + + x_axis = sras.x_axis_mm(angle_idx) + y_axis = sras.y_positions_mm(angle_idx) + dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else sras.pixel_x_mm + dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 + extent = _axes_extent(x_axis, y_axis, dx, dy) + + title = (f"{CH_NAMES[ch_idx]} | {mode_str} | " + f"{sras.angles_deg[angle_idx]:.1f}°") + + fig = Figure(figsize=(7, 5), tight_layout=True) + FigureCanvasAgg(fig) # Agg-only: never registered with pyplot + ax = fig.add_subplot(111) + draw_view_image(ax, fig, display_img, extent, cmap, v0, v1, + "X (mm)", "Y (mm)", title, colorbar_label, + bad_color=bad_color) + fig.savefig(str(Path(out_dir) / out_name), dpi=dpi) + return ("", out_name) + except Exception as exc: + return (str(exc), out_name) diff --git a/sras_viewer/canvases.py b/sras_viewer/canvases.py index 8c0a1e4..835427a 100644 --- a/sras_viewer/canvases.py +++ b/sras_viewer/canvases.py @@ -12,6 +12,7 @@ from PyQt6.QtGui import QKeyEvent from PyQt6.QtWidgets import QSizePolicy from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv +from sras_render import draw_view_image def count_colormap(n_angles: int): """(cmap, norm, ticks) for an integer "how many angles cover this pixel" @@ -166,11 +167,14 @@ class ImageCanvas(FigureCanvasQTAgg): def show_image(self, img: np.ndarray, extent: list[float], cmap, vmin: float, vmax: float, xlabel: str, ylabel: str, title: str, - colorbar_label: str = "", cb_ticks=None, norm=None): + colorbar_label: str = "", cb_ticks=None, norm=None, + bad_color=None): """*cmap* may be a name or a Colormap instance. *norm* (which overrides vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image — the wizard's overlap-count view — with whole-number colorbar bands - instead of a continuous shade.""" + instead of a continuous shade. *bad_color*, if given, is the fill for + NaN pixels — a copy of *cmap* is made so the shared, registered + instance is never mutated.""" self.figure.clf() self.ax = self.figure.add_subplot(111) # Patches and lines are destroyed by figure.clf(); drop stale refs. @@ -179,20 +183,12 @@ class ImageCanvas(FigureCanvasQTAgg): self._extent = extent self._img_shape = img.shape - kw = ({"norm": norm} if norm is not None - else {"vmin": vmin, "vmax": vmax}) - im = self.ax.imshow( - img, aspect="auto", origin="upper", - extent=extent, cmap=cmap, interpolation="nearest", **kw, - ) - cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04, - ticks=cb_ticks) - if colorbar_label: - cb.set_label(colorbar_label) - - self.ax.set_xlabel(xlabel) - self.ax.set_ylabel(ylabel) - self.ax.set_title(title) + # Shared with the headless batch image-export worker (sras_render.py) + # so an exported PNG can never quietly drift from what this canvas + # shows on screen for the same settings. + draw_view_image(self.ax, self.figure, img, extent, cmap, vmin, vmax, + xlabel, ylabel, title, colorbar_label, cb_ticks, norm, + bad_color) # Re-draw the ROI (if any) on top of the fresh image so it persists # unchanged across mode / angle / channel switches. @@ -433,13 +429,22 @@ class WaveformCanvas(FigureCanvasQTAgg): def show_rf_waveform(self, sras: SrasFile, angle_idx: int, row_idx: int, frame_idx: int, - apply_bg_sub: bool = True): + apply_bg_sub: bool = True, + min_freq_mhz: float = 0.0): """CH1 RF: time-domain + FFT spectrum. If apply_bg_sub is True and sras.background is not None, the background waveform is overlaid on the time-domain plot and the FFT is computed on the subtracted signal. The unsubtracted FFT is also shown faintly for comparison. + + *min_freq_mhz* > 0 restricts the labeled peak to bins at or above + it — the same floor the image's peak search uses, so the label + explains the map pixel instead of contradicting it — and shades the + excluded band on the spectrum. The spectrum curves themselves stay + complete (they are the evidence for choosing the floor). The peak + can still legitimately differ from a padded or row-averaged map: + this panel is always a single waveform at natural resolution. """ data = sras.data[angle_idx] waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32) @@ -480,7 +485,20 @@ class WaveformCanvas(FigureCanvasQTAgg): # FFT of the (possibly subtracted) waveform power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2 power_sub[0] = 0.0 - peak_mhz = f_mhz[int(np.argmax(power_sub))] + # First bin at or above the floor, exactly as the image peak search + # picks it (bin 0 always excluded). If the floor excludes every bin, + # fall back to the unrestricted peak rather than indexing past the + # end — the label is informational, not a mask. + lo = max(1, int(np.searchsorted(f_mhz, min_freq_mhz))) + if lo < len(power_sub): + peak_mhz = f_mhz[lo + int(np.argmax(power_sub[lo:]))] + else: + peak_mhz = f_mhz[int(np.argmax(power_sub))] + + if min_freq_mhz > 0.0: + self.ax_right.axvspan(0, min_freq_mhz, color="#888888", + alpha=0.15, zorder=0, + label=f"< {min_freq_mhz:g} MHz excluded") if bg is not None: # Also show the unsubtracted FFT for reference diff --git a/sras_viewer/common.py b/sras_viewer/common.py index f543b39..f192888 100644 --- a/sras_viewer/common.py +++ b/sras_viewer/common.py @@ -6,7 +6,7 @@ from PyQt6.QtWidgets import ( QScrollArea, QSizePolicy, QVBoxLayout, QWidget, ) -from sras_format import CH1_IDX, CH3_IDX, CH4_IDX +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, _axes_extent # noqa: F401 (re-exported) # --------------------------------------------------------------------------- # Display constants @@ -26,6 +26,10 @@ CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX) CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"] +# Fill color for masked (below-DC-threshold) pixels when "Highlight masked +# pixels" is on, chosen to stand out against every colormap in CMAPS above. +_MASKED_HIGHLIGHT_COLOR = "magenta" + # (mode_str, status-bar unit, colorbar label) per channel index _CHANNEL_DISPLAY = { CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"), @@ -109,13 +113,6 @@ def _combo(items=(), *, min_chars: int = 10) -> QComboBox: return combo -def _axes_extent(x_axis, y_axis, dx: float, dy: float) -> list[float]: - """Matplotlib imshow extent with half-pixel margins, Y flipped so row 0 - renders at the top.""" - return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2, - y_axis[-1] + dy / 2, y_axis[0] - dy / 2] - - def _wrap_label(text: str = "", css: str | None = None) -> QLabel: """A word-wrapped QLabel that reports its *wrapped* height to the layout. diff --git a/sras_viewer/dialogs.py b/sras_viewer/dialogs.py index e85bc7e..d3196cd 100644 --- a/sras_viewer/dialogs.py +++ b/sras_viewer/dialogs.py @@ -14,7 +14,6 @@ from PyQt6.QtWidgets import ( QVBoxLayout, QWidget, ) -from sras_compute import PYFFTW_AVAILABLE from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES from sras_workers import ExportChannel @@ -29,7 +28,7 @@ from .common import ( # --------------------------------------------------------------------------- class FftOptionsDialog(QDialog): - """Configure FFT backend and zero-padding. + """Configure FFT zero-padding. Changes take effect only when the user clicks Apply. Cancel discards all pending edits. The live 'frequency resolution' label updates as @@ -38,7 +37,6 @@ class FftOptionsDialog(QDialog): """ def __init__(self, parent=None, *, - current_backend: str, current_pad_factor: int, samples_per_frame: int | None, sample_rate_hz: float | None, @@ -54,29 +52,6 @@ class FftOptionsDialog(QDialog): layout = QVBoxLayout(self) - # ---- Backend --------------------------------------------------- - grp_backend = QGroupBox("FFT Backend") - bl = QVBoxLayout(grp_backend) - - self._btn_scipy = QRadioButton("SciPy FFT (pocketfft) (always available)") - self._btn_pyfftw = QRadioButton( - "pyFFTW (faster for large arrays)" if PYFFTW_AVAILABLE - else "pyFFTW (not installed — run: pip install pyfftw)") - self._btn_pyfftw.setEnabled(PYFFTW_AVAILABLE) - - self._backend_group = QButtonGroup(self) - self._backend_group.addButton(self._btn_scipy, id=0) - self._backend_group.addButton(self._btn_pyfftw, id=1) - - if current_backend == "pyfftw" and PYFFTW_AVAILABLE: - self._btn_pyfftw.setChecked(True) - else: - self._btn_scipy.setChecked(True) - - bl.addWidget(self._btn_scipy) - bl.addWidget(self._btn_pyfftw) - layout.addWidget(grp_backend) - # ---- Zero-padding ---------------------------------------------- grp_zp = QGroupBox("Zero-Padding") zl = QVBoxLayout(grp_zp) @@ -139,9 +114,6 @@ class FftOptionsDialog(QDialog): f"Velocity bin: {vel_res_ms:.3f} m/s " f"(at grating = {self._grating_um:.2f} µm)") - def get_backend(self) -> str: - return "pyfftw" if self._btn_pyfftw.isChecked() and PYFFTW_AVAILABLE else "scipy" - def get_pad_factor(self) -> int: return max(1, self._spin_pad.value()) diff --git a/sras_viewer/main_window.py b/sras_viewer/main_window.py index a9f138d..8eb4980 100644 --- a/sras_viewer/main_window.py +++ b/sras_viewer/main_window.py @@ -1,6 +1,7 @@ """The SrasViewerWindow main window and application entry point.""" import sys +from collections import Counter from pathlib import Path import numpy as np @@ -23,16 +24,16 @@ from sras_format import ( _FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC, ) from sras_workers import ( - BatchCacheWorker, BatchExportWorker, ComputeWorker, DcPrecomputeWorker, - LoadWorker, + BatchCacheWorker, BatchExportImagesWorker, BatchExportWorker, + ComputeWorker, DcPrecomputeWorker, LoadWorker, ) from .canvases import ImageCanvas, WaveformCanvas from .common import ( CH1_DERIVED_MODES, CH_LABELS, CMAPS, VELOCITY_MODE_IDX, _CHANNEL_DISPLAY, _CSS_BUSY, _axes_extent, _CSS_HINT, _CSS_INFO, _CSS_MUTED, _CSS_WARN, _LEFT_PANEL_W, - _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, _scroll_panel, - _wrap_label, + _MASKED_HIGHLIGHT_COLOR, _RIGHT_PANEL_W, Jobs, _combo, _form, _group, _make_dspin, + _scroll_panel, _wrap_label, ) from .align_wizard import AlignmentWizard from .dialogs import ( @@ -60,6 +61,7 @@ class SrasViewerWindow(QMainWindow): self._pending_ch: int = 0 self._pending_bg_sub: bool = True self._pending_threshold: float = 50.0 # mV + self._pending_min_freq_mhz: float = 0.0 self._pending_fft_pad_factor: int = 1 # Live background jobs, keyed by role — see _run_worker. @@ -72,7 +74,6 @@ class SrasViewerWindow(QMainWindow): self._settings = QSettings(QSettings.Format.IniFormat, QSettings.Scope.UserScope, "sras-viewer", "sras-viewer") - compute.set_fft_backend(str(self._settings.value("fft/backend", "scipy"))) try: pad = int(self._settings.value("fft/pad_factor", 1)) except (TypeError, ValueError): @@ -98,13 +99,13 @@ class SrasViewerWindow(QMainWindow): # 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, threshold) so revisiting the same combination - # is free. bg-sub/pad are deliberately not part of the key: once an - # angle has any FFT image (live or from the file's own stored - # cache), it stays displayed regardless of those controls — see - # _fft_cache_key. + # cached per (angle, threshold, min peak freq) so revisiting the + # same combination is free. bg-sub/pad are deliberately not part of + # the key: once an angle has any FFT image (live or from the file's + # own stored cache), it stays displayed regardless of those + # controls — see _fft_cache_key. self._dc_cache: dict[tuple[int, int], np.ndarray] = {} - self._fft_cache: dict[tuple[int, float], np.ndarray] = {} + self._fft_cache: dict[tuple[int, float, float], np.ndarray] = {} self._dc_generation: int = 0 # Angle alignment ("Fusion" menu) @@ -278,6 +279,31 @@ class SrasViewerWindow(QMainWindow): self.lbl_threshold_adc = _wrap_label( f"≈ {mv_to_adc(50.0):.1f} ADC counts", _CSS_MUTED) tl.addWidget(self.lbl_threshold_adc) + + min_freq_form = _form() + self.spin_min_freq_mhz = _make_dspin(0.0, 10000.0, 3, suffix=" MHz", + value=0.0, step=1.0) + self.spin_min_freq_mhz.setEnabled(False) + self.spin_min_freq_mhz.setToolTip( + "Excludes every FFT bin below this frequency from the peak\n" + "search (0 = off, only true DC is excluded). A pixel can pass\n" + "the DC threshold above yet still carry only weak real signal —\n" + "when that happens, the un-subtracted background's DC-leakage\n" + "skirt can have more power than the genuine signal, so the peak\n" + "search resolves to a near-zero frequency even though the pixel\n" + "is real, valid data. Raising this floor above that skirt forces\n" + "the search to report the strongest peak that is plausibly real\n" + "signal instead.\n" + "Raising the floor also re-masks images already shown or stored\n" + "in this file's cache: pixels whose stored peak falls below it\n" + "display as invalid. Lowering it below a stored cache's own\n" + "floor needs a recompute (bins below that floor were never\n" + "searched). Batch Compute FFT records the floor in the file and\n" + "re-resolves masked pixels to their true above-floor peak." + ) + self.spin_min_freq_mhz.editingFinished.connect(self._on_min_freq_changed) + min_freq_form.addRow("Min peak freq:", self.spin_min_freq_mhz) + tl.addLayout(min_freq_form) vl.addWidget(self.grp_threshold) # Background subtraction (v4+ files only) @@ -423,6 +449,22 @@ class SrasViewerWindow(QMainWindow): self.chk_auto.toggled.connect(self._on_autoscale_toggled) dl.addWidget(self.chk_auto) + self.chk_highlight_masked = QCheckBox("Highlight masked (no-data) pixels") + self.chk_highlight_masked.setChecked(True) + self.chk_highlight_masked.setEnabled(False) + self.chk_highlight_masked.setToolTip( + "RF/Velocity only. A pixel below the DC threshold has no FFT\n" + "result at all and is otherwise shown as 0 on the same grayscale\n" + "ramp as a real, valid pixel whose peak just happens to be a low\n" + "frequency — the two look identical (both near-black). When\n" + "checked, masked pixels are drawn in a distinct highlight color\n" + "instead, excluded from the colormap's data range, so real\n" + "low-frequency pixels keep their own true shade. Uncheck to\n" + "restore the old behavior where both blend together." + ) + self.chk_highlight_masked.toggled.connect(self._on_highlight_masked_toggled) + dl.addWidget(self.chk_highlight_masked) + range_form = _form() for label, attr in (("min:", "spin_vmin"), ("max:", "spin_vmax")): spin = _make_dspin(-1e9, 1e9, 4) @@ -446,7 +488,7 @@ class SrasViewerWindow(QMainWindow): fft_menu = menubar.addMenu("&FFT") fft_act = QAction("FFT &Options…", self) - fft_act.setStatusTip("Configure FFT backend and zero-padding") + fft_act.setStatusTip("Configure FFT zero-padding") fft_act.triggered.connect(self._on_fft_options) fft_menu.addAction(fft_act) @@ -486,6 +528,20 @@ class SrasViewerWindow(QMainWindow): self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg) convert_menu.addAction(self._batch_fft_rowavg_act) + convert_menu.addSeparator() + self._batch_export_images_act = QAction( + "Batch Export View as &Images…", self) + self._batch_export_images_act.setStatusTip( + "Select .sras files and export the currently selected view " + "(channel, angle, threshold, colormap, etc.) as one PNG per " + "file, rendered the same way it is shown on screen. Read-only " + "— never modifies the source files. Aligned View is ignored " + "even if enabled, since an alignment result belongs to one " + "specific file's geometry.") + self._batch_export_images_act.triggered.connect( + self._on_batch_export_images) + convert_menu.addAction(self._batch_export_images_act) + export_menu = menubar.addMenu("&Export") self._batch_export_act = QAction("&Batch Export Images…", self) self._batch_export_act.setStatusTip( @@ -638,10 +694,13 @@ class SrasViewerWindow(QMainWindow): if s.precomputed_row_avg_n else "") pad_note = (f", pad {s.precomputed_pad_factor}x" if s.precomputed_pad_factor > 1 else "") + floor_note = (f", floor ≥ {s.precomputed_min_freq_mhz:g} MHz" + if s.precomputed_min_freq_mhz > 0 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"{avg_note if n_fft else ''}{pad_note if n_fft else ''} " + f"{avg_note if n_fft else ''}{pad_note if n_fft else ''}" + f"{floor_note if n_fft else ''} " "— display is instant for cached angles") notes += self._cache_mismatch_notes() elif s.version == 7: @@ -650,17 +709,24 @@ class SrasViewerWindow(QMainWindow): def _cache_mismatch_notes(self) -> list[str]: """Informational only: whether the file's stored FFT cache was - computed under different bg-sub/pad settings than these controls - currently say. The display always shows the stored image as-is - regardless (see _stored_fft_image) — these controls only affect a - future live compute for an angle with nothing cached yet, or an - explicit batch recompute, never what's already on screen. + computed under different bg-sub/pad/min-freq settings than these + controls currently say. The display always shows the stored image + (re-masked for a raised floor) regardless — see _stored_fft_image; + these controls only affect a future live compute for an angle with + nothing cached yet, or an explicit batch recompute, never what's + already on screen. row_avg_n is compared against the file's own recorded value (a self-match), so it never contributes a reason here — there's no live control for it to diverge from, and the "Cached images" line above already reports it. + The min peak frequency floor gets two directions: a live floor + *below* the stored one is a genuine mismatch reason (the stored + search never looked below its floor), while a live floor *above* + the stored one is servable — pixels under it are shown as masked — + so that direction gets its own softer note. + Asks compute for the reasons rather than restating the accept rule, so a new provenance field can only be added in one place. """ @@ -668,15 +734,26 @@ class SrasViewerWindow(QMainWindow): if s is None or all(x is None for x in s.precomputed_freq_mhz): return [] + live_floor = self.spin_min_freq_mhz.value() + notes = [] reasons = compute.cache_mismatch_reasons( s, n_fft=self._current_n_fft(), apply_bg_sub=self.chk_bg_sub.isChecked(), - row_avg_n=s.precomputed_row_avg_n) - if not reasons: - return [] - return ["Note: current bg-sub/pad controls differ from the stored " - "cache — " + "; ".join(reasons) + ". Shown as stored; use " - "Batch Compute to recompute with these settings."] + row_avg_n=s.precomputed_row_avg_n, + min_freq_mhz=live_floor) + if reasons: + notes.append( + "Note: current bg-sub/pad/min-freq controls differ from the " + "stored cache — " + "; ".join(reasons) + ". Shown as stored; " + "use Batch Compute to recompute with these settings.") + if round(live_floor * 1000) > round(s.precomputed_min_freq_mhz * 1000): + notes.append( + f"Note: Min peak freq ({live_floor:g} MHz) is above the " + f"stored cache's floor ({s.precomputed_min_freq_mhz:g} MHz) — " + f"stored pixels whose peak falls below {live_floor:g} MHz are " + "shown as masked; Batch Compute FFT re-resolves them above " + "the floor instead.") + return notes # ------------------------------------------------------------------ # Controls @@ -699,7 +776,9 @@ class SrasViewerWindow(QMainWindow): # Threshold and bg-sub apply to all CH1 modes self.spin_threshold_mv.setEnabled(is_ch1) + self.spin_min_freq_mhz.setEnabled(is_ch1) self.chk_bg_sub.setEnabled(has_file and s.background is not None and is_ch1) + self.chk_highlight_masked.setEnabled(is_ch1) self.spin_grating_um.setEnabled(is_vel) self.grp_velocity.setVisible(is_vel) @@ -750,6 +829,16 @@ class SrasViewerWindow(QMainWindow): if self._is_fft_mode(): self._refresh_display() + def _on_min_freq_changed(self): + # Like the DC threshold, this is a genuine cache-key change that can + # be cheaply re-applied to a stored image — tighten-only: raising the + # floor masks stored pixels below it, while lowering it below the + # stored cache's own floor can only be answered by a fresh compute + # (see cached_rf_image / cache_mismatch_reasons). + if self._is_fft_mode(): + self._update_scan_info_labels() + self._refresh_display() + def _on_autoscale_toggled(self, checked: bool): manual = not checked self.spin_vmin.setEnabled(manual and self._sras is not None) @@ -757,6 +846,10 @@ class SrasViewerWindow(QMainWindow): if self._sras is not None and self._current_image is not None: self._redraw_image(self._current_image) + def _on_highlight_masked_toggled(self, checked: bool): + if self._sras is not None and self._current_image is not None: + self._redraw_image(self._current_image) + def _on_manual_range_changed(self): if not self.chk_auto.isChecked() and self._current_image is not None: self._redraw_image(self._current_image) @@ -1031,15 +1124,18 @@ class SrasViewerWindow(QMainWindow): return freq_mhz def _fft_cache_key(self, angle_idx: int) -> tuple: - """Keyed by angle and DC threshold only. Once any FFT image exists - for an angle this session — live-computed or pulled from the file's - own stored cache — it stays the displayed image for that angle - regardless of later bg-sub/pad toggles; those only affect a future - live compute for an angle with nothing cached yet, or an explicit - batch recompute (see _stored_fft_image). Threshold stays in the key - because re-masking against it is free and meant to stay interactive - (see _on_threshold_changed).""" - return (angle_idx, self.spin_threshold_mv.value()) + """Keyed by angle, DC threshold, and min peak frequency only. Once + any FFT image exists for an angle this session — live-computed or + pulled from the file's own stored cache — it stays the displayed + image for that angle regardless of later bg-sub/pad toggles; those + only affect a future live compute for an angle with nothing cached + yet, or an explicit batch recompute (see _stored_fft_image). + Threshold and min-freq are both in the key because both are cheaply + re-applied when a stored image is served (_stored_fft_image takes + them live), so the cached value genuinely reflects every component + of its key and revisiting a combination is instant.""" + return (angle_idx, self.spin_threshold_mv.value(), + self.spin_min_freq_mhz.value()) def _aligned_cache_key(self, angle_idx: int, ch_idx: int) -> tuple: """Mirrors _fft_cache's key granularity so a stale aligned image is @@ -1116,9 +1212,13 @@ class SrasViewerWindow(QMainWindow): controls never gate whether it's used, only what a *future* compute produces. See _cache_mismatch_notes for the informational (non- blocking) note when the live controls diverge from what's shown. - Only the DC threshold is taken live: re-masking a stored image - against it is free, unlike bg-sub/pad/row-averaging which are baked - irreversibly into the stored numbers. + The DC threshold and the min peak frequency floor are taken live: + re-masking a stored image against either is free, unlike + bg-sub/pad/row-averaging, which are baked irreversibly into the + stored numbers. The floor is tighten-only, though — a live floor + *below* the stored cache's own floor is the one case where a stored + image genuinely can't answer (cached_rf_image refuses it), and the + caller falls through to a real compute. allow_dc_recompute=False keeps this off the I/O path: if the mask would mean reading a whole CH4 channel, this declines and the caller @@ -1135,7 +1235,8 @@ class SrasViewerWindow(QMainWindow): n_fft=n_fft, row_avg_n=s.precomputed_row_avg_n, dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)), - allow_dc_recompute=False) + allow_dc_recompute=False, + min_freq_mhz=self.spin_min_freq_mhz.value()) def _cached_value_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None: """The already-available (no compute) image for (angle, channel): @@ -1211,6 +1312,24 @@ class SrasViewerWindow(QMainWindow): self._redraw_image(img) self._update_roi_ui() + def _dc_validity_mask(self, angle_idx: int, aligned: bool) -> np.ndarray | None: + """True where a CH1/Velocity pixel passed the DC threshold and so + actually has a real FFT result, on the same grid as display_img. + None if the CH4 DC image for this angle isn't cached yet — the + background DC precompute hasn't reached it, and this is deliberately + not worth a synchronous recompute just to redraw.""" + dc4 = self._dc_cache.get((angle_idx, CH4_IDX)) + if dc4 is None: + return None + valid = dc4 >= self.spin_threshold_mv.value() + if aligned: + # apply_alignment defaults to nearest-neighbor (order=0), so a + # 0.0/1.0 float warp stays exactly 0 or 1 - no blending at mask + # edges to second-guess with a >= 0.5 cutoff. + valid = apply_alignment(self._alignment_result, angle_idx, + valid.astype(np.float32)) >= 0.5 + return valid + def _redraw_image(self, img: np.ndarray): s = self._sras angle_idx = self._current_angle @@ -1231,8 +1350,32 @@ class SrasViewerWindow(QMainWindow): dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0 extent = _axes_extent(x_axis, y_axis, dx, dy) + # Masked-out pixels are stored as a plain 0, the same value near + # the bottom of a linear colormap as a real low-frequency pixel - + # the two are indistinguishable there. Pull masked pixels out to + # NaN (drawn in a distinct highlight color, excluded from the + # auto-scale range) so a real pixel keeps its own true shade + # instead of disappearing into the same black as "no data". In an + # FFT-derived mode a value of exactly 0 is itself the "no valid + # peak" sentinel (DC-masked, below the min-freq floor, or an empty + # spectrum — bin 0 is always excluded, so no genuine peak is ever + # 0), so those pixels are masked by value too; that covers + # floor-masked pixels even when no DC4 image is cached yet. + highlight_masked = (ch_idx in CH1_DERIVED_MODES + and self.chk_highlight_masked.isChecked()) + if highlight_masked: + mask_valid = self._dc_validity_mask(angle_idx, aligned) + if mask_valid is not None and mask_valid.shape != display_img.shape: + mask_valid = None + display_img = display_img.astype(np.float32, copy=True) + if mask_valid is not None: + display_img[~mask_valid] = np.nan + display_img[display_img == 0.0] = np.nan + if self.chk_auto.isChecked(): - vmin, vmax = float(display_img.min()), float(display_img.max()) + vmin, vmax = float(np.nanmin(display_img)), float(np.nanmax(display_img)) + if not np.isfinite(vmin): + vmin, vmax = 0.0, 0.0 # every pixel masked out for spin, val in ((self.spin_vmin, vmin), (self.spin_vmax, vmax)): with QSignalBlocker(spin): spin.setValue(val) @@ -1256,6 +1399,7 @@ class SrasViewerWindow(QMainWindow): vmin=vmin, vmax=vmax, xlabel="X (mm)", ylabel="Y (mm)", title=title, colorbar_label=colorbar_label, + bad_color=_MASKED_HIGHLIGHT_COLOR if highlight_masked else None, ) self.statusBar().showMessage( f"{s.path.name} | {ch_label} @ {angle_deg:.1f}° " @@ -1279,6 +1423,7 @@ class SrasViewerWindow(QMainWindow): self._pending_ch = ch_idx self._pending_bg_sub = self.chk_bg_sub.isChecked() self._pending_threshold = self.spin_threshold_mv.value() + self._pending_min_freq_mhz = self.spin_min_freq_mhz.value() self._pending_fft_pad_factor = self._fft_pad_factor worker = ComputeWorker( @@ -1295,6 +1440,7 @@ class SrasViewerWindow(QMainWindow): # eligible even when this fallback compute is the one that ends # up serving it (e.g. before DC precompute reaches this angle). row_avg_n=self._sras.precomputed_row_avg_n, + min_freq_mhz=self._pending_min_freq_mhz, ) if not self._run_worker( Jobs.COMPUTE, worker, @@ -1323,9 +1469,10 @@ class SrasViewerWindow(QMainWindow): already be cached.""" if (self.spin_angle.value(), self.combo_channel.currentIndex(), self.chk_bg_sub.isChecked(), self.spin_threshold_mv.value(), - self._fft_pad_factor) != ( + self.spin_min_freq_mhz.value(), self._fft_pad_factor) != ( self._pending_angle, self._pending_ch, self._pending_bg_sub, - self._pending_threshold, self._pending_fft_pad_factor): + self._pending_threshold, self._pending_min_freq_mhz, + self._pending_fft_pad_factor): self._refresh_display() def _on_compute_done(self, result): @@ -1336,7 +1483,8 @@ class SrasViewerWindow(QMainWindow): ch_idx = self._pending_ch if ch_idx in CH1_DERIVED_MODES: - self._fft_cache[(angle_idx, self._pending_threshold)] = result + key = (angle_idx, self._pending_threshold, self._pending_min_freq_mhz) + self._fft_cache[key] = result img = self._scale_for_display(result, ch_idx) else: img = result @@ -1415,7 +1563,8 @@ class SrasViewerWindow(QMainWindow): if self._current_ch in CH1_DERIVED_MODES: self.wave_canvas.show_rf_waveform( self._sras, angle_idx, row_idx, frame_idx, - apply_bg_sub=self.chk_bg_sub.isChecked()) + apply_bg_sub=self.chk_bg_sub.isChecked(), + min_freq_mhz=self.spin_min_freq_mhz.value()) else: self.wave_canvas.show_dc_waveform( self._sras, angle_idx, self._current_ch, row_idx, frame_idx) @@ -1467,7 +1616,8 @@ class SrasViewerWindow(QMainWindow): # Cache the FFT at the pad the viewer is actually displaying at, # otherwise the batch stores images this window can never use. worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked(), - pad_factor=self._fft_pad_factor) + pad_factor=self._fft_pad_factor, + min_freq_mhz=self.spin_min_freq_mhz.value()) started = self._run_worker( Jobs.BATCH, worker, connect=( @@ -1480,9 +1630,7 @@ class SrasViewerWindow(QMainWindow): if not started: return # a second trigger snuck in while the file dialog was open - self._batch_dc_act.setEnabled(False) - self._batch_fft_act.setEnabled(False) - self._batch_fft_rowavg_act.setEnabled(False) + self._set_batch_actions_enabled(False) self._show_progress( Jobs.BATCH, f"Batch computing {label} for {len(paths)} file(s)…", maximum=100) @@ -1509,7 +1657,8 @@ class SrasViewerWindow(QMainWindow): self._batch_errors = [] worker = BatchCacheWorker(paths, "fft_rowavg", self.chk_bg_sub.isChecked(), dc_threshold_mv=threshold_mv, row_avg_n=n, - pad_factor=self._fft_pad_factor) + pad_factor=self._fft_pad_factor, + min_freq_mhz=self.spin_min_freq_mhz.value()) started = self._run_worker( Jobs.BATCH, worker, connect=( @@ -1522,9 +1671,7 @@ class SrasViewerWindow(QMainWindow): 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._set_batch_actions_enabled(False) self._show_progress( Jobs.BATCH, f"Batch computing row-averaged FFT (n={n}, threshold={threshold_mv:.1f} mV) " @@ -1555,9 +1702,100 @@ class SrasViewerWindow(QMainWindow): self._load_file(str(self._sras.path)) def _after_batch(self): - self._batch_dc_act.setEnabled(True) - self._batch_fft_act.setEnabled(True) - self._batch_fft_rowavg_act.setEnabled(True) + self._set_batch_actions_enabled(True) + + def _set_batch_actions_enabled(self, enabled: bool): + """All four Convert-menu batch actions share one Jobs.BATCH slot; + grey out every sibling while one runs, rather than leaving one + clickable but silently no-op'd by the busy guard.""" + for act in (self._batch_dc_act, self._batch_fft_act, + self._batch_fft_rowavg_act, self._batch_export_images_act): + act.setEnabled(enabled) + + # ------------------------------------------------------------------ + # Batch export view as images + # ------------------------------------------------------------------ + + def _on_batch_export_images(self): + if self._job_running(Jobs.BATCH): + return + paths, _ = QFileDialog.getOpenFileNames( + self, "Select .sras files to export the current view from", "", + "SRAS files (*.sras);;All files (*)") + if not paths: + return + out_dir = QFileDialog.getExistingDirectory( + self, "Select output folder for exported images") + if not out_dir: + return + + angle_idx = self.spin_angle.value() + ch_idx = self.combo_channel.currentIndex() + is_fft_mode = ch_idx in CH1_DERIVED_MODES + mode_str, _unit, colorbar_label = _CHANNEL_DISPLAY[ch_idx] + + self._batch_errors = [] + self._batch_export_names = [] + worker = BatchExportImagesWorker( + paths, + out_dir=out_dir, angle_idx=angle_idx, ch_idx=ch_idx, + is_fft_mode=is_fft_mode, is_velocity=(ch_idx == VELOCITY_MODE_IDX), + dc_threshold_mv=self.spin_threshold_mv.value(), + apply_bg_sub=self.chk_bg_sub.isChecked(), + pad_factor=self._fft_pad_factor, + min_freq_mhz=self.spin_min_freq_mhz.value(), + grating_um=self.spin_grating_um.value(), + cmap=self.combo_cmap.currentText(), + auto_scale=self.chk_auto.isChecked(), + vmin=self.spin_vmin.value(), vmax=self.spin_vmax.value(), + highlight_masked=self.chk_highlight_masked.isChecked(), + mode_str=mode_str, colorbar_label=colorbar_label, + mask_color=_MASKED_HIGHLIGHT_COLOR, + ) + started = self._run_worker( + Jobs.BATCH, worker, + connect=( + ("progress", lambda pct: self._set_progress(Jobs.BATCH, pct)), + ("file_done", self._on_batch_export_file_done), + ("finished", + lambda p=paths, d=out_dir: self._on_batch_export_finished(p, d)), + ), + on_done=self._after_batch, + ) + if not started: + return # a second trigger snuck in while a dialog was open + + self._set_batch_actions_enabled(False) + self._show_progress( + Jobs.BATCH, f"Exporting images for {len(paths)} file(s)…", + maximum=100) + + def _on_batch_export_file_done(self, path: str, err: str, out_name: str): + if err: + self._batch_errors.append(f"{Path(path).name} — {err}") + else: + self._batch_export_names.append(out_name) + self._show_progress(Jobs.BATCH, f"Exported {Path(path).name}…") + + def _on_batch_export_finished(self, paths: list[str], out_dir: str): + self._close_progress(Jobs.BATCH) + + n_total = len(paths) + n_failed = len(self._batch_errors) + n_ok = n_total - n_failed + summary = (f"Batch export: {n_ok}/{n_total} image(s) written to " + f"{Path(out_dir).name}") + if n_failed: + summary += f", {n_failed} failed: {'; '.join(self._batch_errors)}" + dupes = sum(1 for _name, count in Counter(self._batch_export_names).items() + if count > 1) + if dupes: + summary += (f" | {dupes} filename collision(s) — later file(s) " + "overwrote earlier ones with the same output name") + self.statusBar().showMessage(summary) + self._batch_errors = [] + self._batch_export_names = [] + # No reload: unlike Batch Compute, export never touches the source files. # ------------------------------------------------------------------ # Export menu: batch image export @@ -1741,7 +1979,6 @@ class SrasViewerWindow(QMainWindow): def _on_fft_options(self): dlg = FftOptionsDialog( self, - current_backend=compute.get_fft_backend(), current_pad_factor=self._fft_pad_factor, samples_per_frame=self._sras.samples_per_frame if self._sras else None, sample_rate_hz=self._sras.sample_rate_hz if self._sras else None, @@ -1749,9 +1986,7 @@ class SrasViewerWindow(QMainWindow): ) if dlg.exec() != QDialog.DialogCode.Accepted: return - compute.set_fft_backend(dlg.get_backend()) self._fft_pad_factor = dlg.get_pad_factor() - self._settings.setValue("fft/backend", compute.get_fft_backend()) self._settings.setValue("fft/pad_factor", self._fft_pad_factor) # Pad factor no longer gates the display: it only affects a future # live compute for an angle with nothing cached yet, or an explicit diff --git a/sras_workers.py b/sras_workers.py index 1454788..ff8b3c7 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -22,6 +22,7 @@ import sras_compute as compute from sras_align_export import write_aligned_sras from sras_compute import cache_file, compute_rf_image, dc_image_mv from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile +from sras_render import export_view_image # Concurrency caps. Batch conversion runs one process per file, and each of # those processes threads internally, so the two must be divided rather than @@ -120,7 +121,8 @@ class ComputeWorker(CancellableWorker): dc_threshold_mv: float = 0.0, dc4_mv: np.ndarray | None = None, is_fft_mode: bool = False, - row_avg_n: int = 0): + row_avg_n: int = 0, + min_freq_mhz: float = 0.0): super().__init__() self._sras = sras self._angle = angle_idx @@ -131,6 +133,7 @@ class ComputeWorker(CancellableWorker): self._dc4_mv = dc4_mv self._is_fft_mode = is_fft_mode self._row_avg_n = row_avg_n + self._min_freq_mhz = min_freq_mhz def run(self): try: @@ -139,7 +142,8 @@ class ComputeWorker(CancellableWorker): self._sras, self._angle, dc_threshold_mv=self._dc_threshold, apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft, dc4_mv=self._dc4_mv, should_stop=self._stopped, - row_avg_n=self._row_avg_n) + row_avg_n=self._row_avg_n, + min_freq_mhz=self._min_freq_mhz) else: img = dc_image_mv(self._sras, self._angle, self._ch, should_stop=self._stopped) @@ -201,7 +205,9 @@ class BatchCacheWorker(QObject): Both FFT modes cache at *pad_factor*, which the caller sets from the viewer's own padding — a cache stored at a pad the user is not viewing - at is one the display can never use. + at is one the display can never use. *min_freq_mhz* travels the same + way: the viewer's live min-peak-freq floor, recorded in the store as + provenance so a reader knows which bins the peak search considered. 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 @@ -215,7 +221,7 @@ class BatchCacheWorker(QObject): def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool, dc_threshold_mv: float | None = None, row_avg_n: int = 0, - pad_factor: int = 1): + pad_factor: int = 1, min_freq_mhz: float = 0.0): super().__init__() self._paths = paths self._mode = mode @@ -223,6 +229,7 @@ class BatchCacheWorker(QObject): self._dc_threshold = dc_threshold_mv self._row_avg_n = row_avg_n self._pad_factor = pad_factor + self._min_freq_mhz = min_freq_mhz def _report(self, path: str, err: str, done: int, total: int): self.file_done.emit(path, err) @@ -247,10 +254,11 @@ class BatchCacheWorker(QObject): with ProcessPoolExecutor(max_workers=n_procs) as executor: futures = { executor.submit(cache_file, p, self._mode, self._apply_bg_sub, - compute.get_fft_backend(), per_proc_workers, + per_proc_workers, pad_factor=self._pad_factor, dc_threshold_mv=self._dc_threshold, - row_avg_n=self._row_avg_n): p + row_avg_n=self._row_avg_n, + min_freq_mhz=self._min_freq_mhz): p for p in paths } for fut in as_completed(futures): @@ -273,11 +281,11 @@ class BatchCacheWorker(QObject): for path in paths: try: err = cache_file(path, self._mode, self._apply_bg_sub, - compute.get_fft_backend(), compute.default_max_workers(), pad_factor=self._pad_factor, dc_threshold_mv=self._dc_threshold, - row_avg_n=self._row_avg_n) + row_avg_n=self._row_avg_n, + min_freq_mhz=self._min_freq_mhz) except Exception as exc: err = str(exc) done += 1 @@ -449,6 +457,112 @@ class BatchExportWorker(CancellableWorker): self.finished.emit() +class BatchExportImagesWorker(QObject): + """Renders the view settings captured by the caller at trigger time + (angle/channel/threshold/etc.) to one PNG per file in *paths*, via + sras_render.export_view_image. + + Same process-pool-with-inline-fallback strategy as BatchCacheWorker + above (same _BATCH_MAX_PROCS / _BATCH_POOL_MIN_BYTES thresholds): an + FFT-derived (CH1/Velocity) view is exactly the same expensive per-file + compute Batch Compute FFT already parallelizes this way. Never mutates + *paths* — each file is opened read-only — so unlike BatchCacheWorker + there is no version gate. + + Emits progress(int) (0-100 by files completed), file_done(str, str, str) + (path, error message or "", output filename this file targeted — set + even on most failures so the caller can flag same-stem collisions across + the batch without any cross-process bookkeeping), and finished(). + """ + progress = pyqtSignal(int) + file_done = pyqtSignal(str, str, str) + finished = pyqtSignal() + + def __init__(self, paths: list[str], **render_kwargs): + """*render_kwargs* is exactly export_view_image's keyword-only + settings (out_dir, angle_idx, ch_idx, is_fft_mode, is_velocity, + dc_threshold_mv, apply_bg_sub, pad_factor, min_freq_mhz, grating_um, + cmap, auto_scale, vmin, vmax, highlight_masked, mode_str, + colorbar_label, mask_color) — bundled rather than repeated as + positional params across __init__/_run_pooled/_run_inline.""" + super().__init__() + self._paths = paths + self._kw = render_kwargs + + def _report(self, path: str, err: str, out_name: str, done: int, total: int): + self.file_done.emit(path, err, out_name) + self.progress.emit(int(done / max(1, total) * 100)) + + def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]: + """Same contract as BatchCacheWorker._run_pooled: returns the paths + that never got a real answer because the pool itself died, so the + caller can retry them in-process.""" + per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs) + unresolved: list[str] = [] + done = 0 + + with ProcessPoolExecutor(max_workers=n_procs) as executor: + futures = { + executor.submit(export_view_image, p, + max_workers=per_proc_workers, **self._kw): p + for p in paths + } + for fut in as_completed(futures): + path = futures[fut] + try: + err, out_name = fut.result() + except BrokenProcessPool: + unresolved.append(path) + continue + except Exception as exc: + err, out_name = str(exc), "" + done += 1 + self._report(path, err, out_name, done, len(paths)) + + return unresolved + + def _run_inline(self, paths: list[str], done: int, total: int): + for path in paths: + try: + err, out_name = export_view_image( + path, max_workers=compute.default_max_workers(), **self._kw) + except Exception as exc: + err, out_name = str(exc), "" + done += 1 + self._report(path, err, out_name, done, total) + + def _worth_pooling(self, paths: list[str]) -> bool: + if len(paths) < 2: + return False + total = 0 + for p in paths: + try: + total += os.path.getsize(p) + except OSError: + pass # unreadable files are reported by export_view_image + return total >= _BATCH_POOL_MIN_BYTES + + def run(self): + paths = self._paths + n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths))) + + if not self._worth_pooling(paths): + self._run_inline(paths, 0, len(paths)) + self.finished.emit() + return + + try: + unresolved = self._run_pooled(paths, n_procs) + except Exception: + # The pool could not be created or collapsed wholesale. + unresolved = list(paths) + + if unresolved: + self._run_inline(unresolved, len(paths) - len(unresolved), len(paths)) + + self.finished.emit() + + class Ch4MaskWorker(_PooledWorker): """Fetches each requested angle's CH4 (Bias B) DC image in mV, for the alignment wizard's initial threshold-mask stack. diff --git a/tests/test_batch_export.py b/tests/test_batch_export.py new file mode 100644 index 0000000..22f8d7d --- /dev/null +++ b/tests/test_batch_export.py @@ -0,0 +1,426 @@ +"""Batch Export View as Images: does the exported PNG actually match what +the live view would show, and does the batch dispatch (menu action -> +worker -> per-file render) behave like Batch Compute's proven pattern? + +sras_render.export_view_image is tested directly (no Qt) for the plumbing +that decides *what* gets rendered -- pad_factor derived per file, masked- +pixel NaN fill, per-file auto-scale, velocity scaling -- via a +draw_view_image spy rather than pixel-diffing PNGs, the same "spy on the +seam, don't inspect the rendered artifact" approach test_highlight_masked_ +pixels (tests/test_gui.py) uses for the live canvas. + +The GUI-dispatch half drives SrasViewerWindow._on_batch_export_images() +end-to-end with patched file dialogs, the same shape as +test_stored_cache.py's test_viewer_batch_row_average_dispatch. +""" + +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pytest +from PyQt6.QtCore import QEventLoop, QTimer +from PyQt6.QtWidgets import QApplication + +import sras_render +from sras_compute import compute_rf_image, dc_image_mv +from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile +from sras_render import export_view_image +from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX +from sras_workers import BatchExportImagesWorker +import tools.make_test_sras as gen + +_THRESHOLD_MV = 50.0 + + +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() + + +# --------------------------------------------------------------------------- +# sras_render.export_view_image -- pure function, no Qt +# --------------------------------------------------------------------------- + +_DEFAULT_KW = dict( + is_fft_mode=False, is_velocity=False, dc_threshold_mv=_THRESHOLD_MV, + apply_bg_sub=True, pad_factor=1, min_freq_mhz=0.0, grating_um=1.0, + cmap="viridis", auto_scale=True, vmin=0.0, vmax=1.0, + highlight_masked=False, mode_str="DC", colorbar_label="mV", +) + + +def _kw(**overrides): + kw = dict(_DEFAULT_KW) + kw.update(overrides) + return kw + + +def _spy_draw(monkeypatch): + """Patches sras_render.draw_view_image to record the image array and + vmin/vmax/bad_color it was called with, then delegates to the real + implementation so the PNG is still written -- lets a test check *what* + export_view_image computed without depending on rendered PNG pixels.""" + orig = sras_render.draw_view_image + captured = {} + + def spy(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, title, + colorbar_label="", cb_ticks=None, norm=None, bad_color=None): + captured["img"] = np.array(img, copy=True) + captured["vmin"] = vmin + captured["vmax"] = vmax + captured["bad_color"] = bad_color + return orig(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, + title, colorbar_label, cb_ticks, norm, bad_color) + + monkeypatch.setattr(sras_render, "draw_view_image", spy) + return captured + + +def test_export_dc_channel_writes_png(tmp_path): + path = tmp_path / "dc.sras" + gen.write(path, n_angles=2, seed=1, samples_per_frame=64) + out_dir = tmp_path / "out" + out_dir.mkdir() + + err, out_name = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw()) + assert err == "" + assert out_name == f"dc_angle0_{CH_NAMES[CH4_IDX]}.png" + out_path = out_dir / out_name + assert out_path.exists() and out_path.stat().st_size > 0 + assert out_path.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n", "not a valid PNG" + + +def test_export_out_of_range_angle(tmp_path): + path = tmp_path / "short.sras" + gen.write(path, n_angles=2, seed=2, samples_per_frame=64) + out_dir = tmp_path / "out" + out_dir.mkdir() + + err, out_name = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=5, ch_idx=CH4_IDX, **_kw()) + assert err != "" and "2" in err, "error should mention the file's actual angle count" + assert out_name == "" + assert list(out_dir.iterdir()) == [], "no file written for a failed export" + + +def test_export_fft_mode_matches_compute_rf_image(tmp_path, monkeypatch): + path = tmp_path / "fft.sras" + gen.write(path, n_angles=1, seed=3, samples_per_frame=128) + out_dir = tmp_path / "out" + out_dir.mkdir() + + captured = _spy_draw(monkeypatch) + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, + **_kw(is_fft_mode=True)) + assert err == "" + + sras = SrasFile(str(path)) + expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, + apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0) + assert np.array_equal(captured["img"], expected) + + +def test_export_velocity_scales_frequency(tmp_path, monkeypatch): + path = tmp_path / "vel.sras" + gen.write(path, n_angles=1, seed=4, samples_per_frame=128) + out_dir = tmp_path / "out" + out_dir.mkdir() + + captured = _spy_draw(monkeypatch) + grating_um = 3.5 + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=VELOCITY_MODE_IDX, + **_kw(is_fft_mode=True, is_velocity=True, grating_um=grating_um)) + assert err == "" + + sras = SrasFile(str(path)) + freq = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, + apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0) + assert np.array_equal(captured["img"], freq * grating_um) + + +def test_pad_factor_uses_each_files_own_samples_per_frame(tmp_path, monkeypatch): + """n_fft must be derived per file from that file's own samples_per_frame, + never a value carried over from whichever file the caller had open -- + otherwise every file but one in a batch gets silently mis-padded.""" + orig_draw = sras_render.draw_view_image + out_dir = tmp_path / "out" + out_dir.mkdir() + + for i, spf in enumerate((64, 256)): + path = tmp_path / f"pad_{spf}.sras" + gen.write(path, n_angles=1, seed=5 + i, samples_per_frame=spf) + + captured = {} + + def spy(ax, fig, img, *a, __c=captured, **kw): + __c["img"] = np.array(img, copy=True) + return orig_draw(ax, fig, img, *a, **kw) + + monkeypatch.setattr(sras_render, "draw_view_image", spy) + + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, + **_kw(is_fft_mode=True, pad_factor=4)) + assert err == "" + + sras = SrasFile(str(path)) + expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, + apply_bg_sub=True, n_fft=spf * 4, + min_freq_mhz=0.0) + assert np.array_equal(captured["img"], expected), \ + f"samples_per_frame={spf}: n_fft must use this file's own value" + + +def test_highlight_masked_sets_nan_and_bad_color(tmp_path, monkeypatch): + path = tmp_path / "mask.sras" + gen.write(path, n_angles=1, seed=7, samples_per_frame=128) + out_dir = tmp_path / "out" + out_dir.mkdir() + + sras = SrasFile(str(path)) + dc4 = dc_image_mv(sras, 0, CH4_IDX) + threshold = float(np.median(dc4)) + expect_masked = dc4 < threshold + assert expect_masked.any() and not expect_masked.all(), \ + "fixture threshold should mask some but not all pixels" + + captured = _spy_draw(monkeypatch) + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, + **_kw(is_fft_mode=True, dc_threshold_mv=threshold, + highlight_masked=True, mask_color="magenta")) + assert err == "" + assert captured["bad_color"] == "magenta" + # The export masks by value too (0 == the "no valid peak" sentinel, same + # rule as the viewer's _redraw_image); on this fixture every + # above-threshold pixel has a nonzero peak, so the value mask coincides + # with the DC mask and the NaN set is exactly expect_masked. + assert np.array_equal(np.isnan(captured["img"]), expect_masked) + valid_vals = captured["img"][~expect_masked] + assert not np.isnan(valid_vals).any() and (valid_vals != 0).all(), \ + "fixture precondition: every valid pixel has a nonzero peak" + + captured2 = _spy_draw(monkeypatch) + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, + **_kw(is_fft_mode=True, dc_threshold_mv=threshold, highlight_masked=False)) + assert err == "" + assert captured2["bad_color"] is None + assert not np.isnan(captured2["img"]).any() + + +def test_auto_scale_uses_per_file_min_max(tmp_path, monkeypatch): + path = tmp_path / "scale.sras" + gen.write(path, n_angles=1, seed=8, samples_per_frame=64) + out_dir = tmp_path / "out" + out_dir.mkdir() + + captured = _spy_draw(monkeypatch) + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, + **_kw(auto_scale=True)) + assert err == "" + img = captured["img"] + assert captured["vmin"] == pytest.approx(float(np.nanmin(img))) + assert captured["vmax"] == pytest.approx(float(np.nanmax(img))) + + captured2 = _spy_draw(monkeypatch) + err, _ = export_view_image( + str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, + **_kw(auto_scale=False, vmin=-5.0, vmax=5.0)) + assert err == "" + assert captured2["vmin"] == -5.0 + assert captured2["vmax"] == 5.0 + + +# --------------------------------------------------------------------------- +# GUI dispatch: SrasViewerWindow._on_batch_export_images end-to-end +# --------------------------------------------------------------------------- + +def _make_window(path) -> SrasViewerWindow: + app = QApplication.instance() or QApplication([]) # noqa: F841 + win = SrasViewerWindow() + win.show() + 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" + return win + + +def test_batch_export_images_writes_one_png_per_file(tmp_path): + path = tmp_path / "src.sras" + gen.write(path, n_angles=2, seed=10, samples_per_frame=64) + paths = [str(path)] + for i in range(2): + p2 = tmp_path / f"other{i}.sras" + gen.write(p2, n_angles=2, seed=20 + i, samples_per_frame=64) + paths.append(str(p2)) + out_dir = tmp_path / "images" + out_dir.mkdir() + + win = _make_window(path) + try: + with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", + return_value=(paths, "")), \ + patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", + return_value=str(out_dir)): + win._on_batch_export_images() + assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" + + angle = win.spin_angle.value() + ch_name = CH_NAMES[win.combo_channel.currentIndex()] + expected_names = {f"{Path(p).stem}_angle{angle}_{ch_name}.png" for p in paths} + actual_names = {p.name for p in out_dir.iterdir()} + assert actual_names == expected_names + assert "Batch export: 3/3 image(s)" in win.statusBar().currentMessage() + finally: + win.close() + pump(200) + + +def test_batch_export_out_of_range_angle_reports_error_continues(tmp_path): + good_path = tmp_path / "good.sras" + short_path = tmp_path / "short.sras" + gen.write(good_path, n_angles=3, seed=30, samples_per_frame=64) + gen.write(short_path, n_angles=1, seed=31, samples_per_frame=64) + out_dir = tmp_path / "images" + out_dir.mkdir() + + win = _make_window(good_path) + try: + win.spin_angle.setValue(2) # valid for good_path, out of range for short_path + with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", + return_value=([str(good_path), str(short_path)], "")), \ + patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", + return_value=str(out_dir)): + win._on_batch_export_images() + assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" + + msg = win.statusBar().currentMessage() + assert "Batch export: 1/2 image(s)" in msg, msg + assert "1 failed" in msg, msg + assert len(list(out_dir.iterdir())) == 1, \ + "the batch must not abort — the good file still exports" + finally: + win.close() + pump(200) + + +def test_batch_export_busy_guard_skips_dialogs(tmp_path, monkeypatch): + path = tmp_path / "busy.sras" + gen.write(path, n_angles=1, seed=40, samples_per_frame=64) + win = _make_window(path) + try: + monkeypatch.setattr(win, "_job_running", lambda key: True) + with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames") as mock_dlg: + win._on_batch_export_images() + assert mock_dlg.call_count == 0, \ + "the Jobs.BATCH busy guard must return before opening any dialog" + finally: + win.close() + pump(200) + + +def test_batch_export_ignores_aligned_view_toggle(tmp_path): + """Aligned View is geometry specific to whichever single file the + Alignment Wizard last ran against and cannot be meaningfully applied + across a batch of different files -- _on_batch_export_images must not + read chk_aligned_view / self._alignment_result at all, regardless of + what's checked in the live view.""" + path = tmp_path / "aligned.sras" + gen.write(path, n_angles=1, seed=41, samples_per_frame=64) + out_dir = tmp_path / "images" + out_dir.mkdir() + + win = _make_window(path) + try: + captured_kwargs = [] + orig_init = BatchExportImagesWorker.__init__ + + def spy_init(self, paths, **kw): + captured_kwargs.append(kw) + return orig_init(self, paths, **kw) + + win.chk_aligned_view.setChecked(True) + with patch.object(BatchExportImagesWorker, "__init__", spy_init), \ + patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", + return_value=([str(path)], "")), \ + patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", + return_value=str(out_dir)): + win._on_batch_export_images() + assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" + + assert len(captured_kwargs) == 1 + assert not any("align" in k.lower() for k in captured_kwargs[0]), \ + captured_kwargs[0].keys() + finally: + win.close() + pump(200) + + +def test_batch_export_filename_collision_note(tmp_path): + dir_a, dir_b = tmp_path / "dir_a", tmp_path / "dir_b" + dir_a.mkdir() + dir_b.mkdir() + path_a, path_b = dir_a / "dup.sras", dir_b / "dup.sras" + gen.write(path_a, n_angles=1, seed=50, samples_per_frame=64) + gen.write(path_b, n_angles=1, seed=51, samples_per_frame=64) + out_dir = tmp_path / "images" + out_dir.mkdir() + + win = _make_window(path_a) + try: + with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", + return_value=([str(path_a), str(path_b)], "")), \ + patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", + return_value=str(out_dir)): + win._on_batch_export_images() + assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" + + msg = win.statusBar().currentMessage() + assert "Batch export: 2/2 image(s)" in msg, msg + assert "collision" in msg, msg + assert len(list(out_dir.iterdir())) == 1, \ + "same-stem inputs silently overwrite to one output file" + finally: + win.close() + pump(200) + + +def test_batch_export_does_not_modify_source_files(tmp_path): + path = tmp_path / "untouched.sras" + gen.write(path, n_angles=1, seed=60, samples_per_frame=64) + before = path.read_bytes() + out_dir = tmp_path / "images" + out_dir.mkdir() + + win = _make_window(path) + try: + with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", + return_value=([str(path)], "")), \ + patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", + return_value=str(out_dir)): + win._on_batch_export_images() + assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" + finally: + win.close() + pump(200) + + assert path.read_bytes() == before, "export must never write to the source file" diff --git a/tests/test_compute.py b/tests/test_compute.py index f702299..a79c2ac 100644 --- a/tests/test_compute.py +++ b/tests/test_compute.py @@ -12,6 +12,7 @@ from pathlib import Path import numpy as np import pytest +import sras_average import sras_compute as compute from sras_compute import ( cache_file, compute_dc_image, compute_rf_image, dc_image_mv, @@ -122,7 +123,7 @@ def test_parallel_identity(tmp_path, monkeypatch): # the block size so every chunk splits into many FFT tasks — the worst # case for boundary bugs. monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4) - monkeypatch.setattr(compute, "_FFT_BLOCK", 4) + monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4) fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET) assert fft_rows < n_rows, \ f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})" @@ -155,43 +156,10 @@ def test_parallel_identity(tmp_path, monkeypatch): "rf image identical (masked, padded/zoom)" -@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)]) -def test_zoom_identity(tmp_path, monkeypatch, spf, bps): - """The zoom peak search must reproduce the full padded-rfft argmax - bit-for-bit, across pad factors, masking, bg-sub, dtype, and backend.""" - path = tmp_path / f"zoom_{spf}.sras" - gen.write(path, n_angles=2, seed=6, samples_per_frame=spf, bps=bps) - sras = SrasFile(str(path)) - dc4 = dc_image_mv(sras, 0, CH4_IDX) - thr = float(np.median(dc4)) - - backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else []) - for backend in backends: - monkeypatch.setattr(compute, "_fft_backend", backend) - for pad in (4, 8, 40): - n_fft = spf * pad - for thr_v in (None, thr): - for bg in (False, True): - ref = compute_rf_image(sras, 0, dc_threshold_mv=thr_v, - apply_bg_sub=bg, n_fft=n_fft, - exact=True) - zoom = compute_rf_image(sras, 0, dc_threshold_mv=thr_v, - apply_bg_sub=bg, n_fft=n_fft) - diff = int((ref != zoom).sum()) - assert diff == 0, \ - (f"{diff} px differ: backend={backend} pad={pad} " - f"thr={thr_v} bg={bg} spf={spf}") - - # A threshold above every pixel masks everything: both paths must agree - # on an all-zero image. - all_masked = compute_rf_image(sras, 0, dc_threshold_mv=1e9, n_fft=spf * 8) - assert not all_masked.any() - - -def test_zoom_identity_fuzz(): - """Hammer _peak_bins_zoom directly with adversarial spectra: noise, +def test_peak_bins_fuzz(): + """Hammer _peak_bins directly with adversarial spectra: noise, un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs, - and all-zero rows.""" + and all-zero rows — against an independent scipy.fft reference.""" import scipy.fft as scipy_fft rng = np.random.default_rng(42) @@ -222,14 +190,102 @@ def test_zoom_identity_fuzz(): P[:, 0] = 0.0 ref = np.argmax(P, axis=1) - zp = compute._zoom_plan(spf, n_fft) - got = compute._peak_bins_zoom(w, zp) + got = compute._peak_bins(w, n_fft) bad = np.nonzero(ref != got)[0] assert not len(bad), \ (f"spf={spf} pad={pad}: rows {bad.tolist()} picked " f"{got[bad].tolist()} instead of {ref[bad].tolist()}") +def test_peak_bins_fuzz_min_freq(): + """Same adversarial-spectra fuzz as test_peak_bins_fuzz, but with a swept + min_bin floor: bins below the floor must be excluded from the argmax + exactly as the independent scipy.fft reference is, when zeroed the same + way before argmax.""" + import scipy.fft as scipy_fft + + rng = np.random.default_rng(43) + for _ in range(25): + spf = int(rng.integers(16, 220)) + pad = int(rng.choice([4, 5, 8, 16, 40])) + n_fft = spf * pad + n_wf = 24 + w = rng.normal(scale=20.0, size=(n_wf, spf)) + t = np.arange(spf) + for r in range(6): + f = rng.uniform(1.0, spf / 2 - 1) + w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2) + f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2) + w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \ + + 49.9 * np.sin(2 * np.pi * f2 * t / spf) + w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \ + + 50 * np.cos(2 * np.pi * f2 * t / spf) + w[8] = 90 + rng.normal(scale=5.0, size=spf) + w[9] = 0.0 + w = w.astype(np.float32) + + S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1) + P = S.real ** 2 + P += S.imag ** 2 + n_bins_fine = n_fft // 2 + 1 + min_bin = int(rng.integers(1, max(2, n_bins_fine // 3))) + P[:, :min_bin] = 0.0 + ref = np.argmax(P, axis=1) + + got = compute._peak_bins(w, n_fft, min_bin) + bad = np.nonzero(ref != got)[0] + assert not len(bad), \ + (f"spf={spf} pad={pad} min_bin={min_bin}: rows {bad.tolist()} picked " + f"{got[bad].tolist()} instead of {ref[bad].tolist()}") + + +@pytest.mark.parametrize("spf", [64, 500, 2500]) +def test_fft_block_for(spf): + """Block size == _FFT_BLOCK_MAX at natural resolution, shrinks and stays + >= _FFT_BLOCK_MIN as n_len grows, and the implied per-thread byte + estimate respects _FFT_PLAN_BYTES_BUDGET except when the floor is + engaged.""" + block_natural = compute._fft_block_for(spf, spf) + assert block_natural == compute._FFT_BLOCK_MAX + + prev = compute._FFT_BLOCK_MAX + for pad in (2, 4, 8, 40, 500): + n_len = spf * pad + block = compute._fft_block_for(spf, n_len) + assert compute._FFT_BLOCK_MIN <= block <= prev + bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1) + if block > compute._FFT_BLOCK_MIN: + assert block * bytes_per_wf <= compute._FFT_PLAN_BYTES_BUDGET + prev = block + + +def test_compute_rf_image_min_freq_mhz(tmp_path): + """min_freq_mhz threads through compute_rf_image end-to-end, for both + the natural-resolution and padded paths: 0.0 (default) must reproduce + the pre-existing image exactly, and a floor above every real peak must + collapse the image to bin 0 (0 MHz) — the same fallback the low-level + search uses when nothing survives the floor.""" + path = tmp_path / "floor_e2e.sras" + gen.write(path, n_angles=1, seed=12, samples_per_frame=64) + sras = SrasFile(str(path)) + huge_floor = float(sras.freq_axis_mhz(None)[-1]) + 1.0 # above Nyquist + + for n_fft in (None, 64 * 8): + unfiltered = compute_rf_image(sras, 0, dc_threshold_mv=None, + apply_bg_sub=False, n_fft=n_fft) + same = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False, + n_fft=n_fft, min_freq_mhz=0.0) + assert np.array_equal(unfiltered, same), \ + f"n_fft={n_fft}: min_freq_mhz=0.0 changed the output" + assert unfiltered.any(), \ + f"n_fft={n_fft}: fixture should have real signal" + + collapsed = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False, + n_fft=n_fft, min_freq_mhz=huge_floor) + assert not collapsed.any(), \ + f"n_fft={n_fft}: floor above Nyquist should collapse to 0 MHz" + + def test_nomask_equals_low_threshold(tmp_path): """dc_threshold_mv=None must equal a threshold below every pixel, while skipping the CH4 read.""" @@ -307,42 +363,50 @@ def test_legacy_parse(tmp_path): def test_sras_average(tmp_path): - """The sras_average.py CLI: frame averaging with remainder handling.""" - src = tmp_path / "legacy_v4.sras" - meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12, - samples_per_frame=32, seed=4) - dst = tmp_path / "legacy_v4_avg.sras" + """The sras_average.py CLI: v6 frame averaging with remainder handling, + per-angle ragged geometry, and the laser_freq_hz X-axis correction.""" + src = tmp_path / "v6.sras" + meta = gen.write(src, n_angles=2, seed=4, samples_per_frame=32, + geometry=[(4, 12)]) + src_sras = SrasFile(str(src)) + + dst = tmp_path / "v6_avg.sras" proc = subprocess.run( [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "4"], capture_output=True, text=True, cwd=REPO) assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:] avg = SrasFile(str(dst)) - assert avg.version == 4 + assert avg.version == 6 assert list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}" assert (avg.n_angles == 2 and list(avg.n_rows) == [4, 4] - and avg.n_channels == meta["n_channels"]) - assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \ + and avg.n_channels == src_sras.n_channels == 3) + assert np.allclose(avg.ch_ymult_mv, src_sras.ch_ymult_mv), \ "calibration preserved" - assert np.array_equal(avg.background, SrasFile(str(src)).background), \ + assert np.array_equal(avg.background, src_sras.background), \ "background preserved" - src_data = meta["data"] - # int16 (not float32) before .mean(): matches average_rows' own + assert avg.laser_freq_hz == pytest.approx(src_sras.laser_freq_hz / 4), \ + "laser_freq_hz divided by N keeps pixel_x_mm correct after binning" + assert np.array_equal(avg.x_start_mm, src_sras.x_start_mm), \ + "per-angle x_start unchanged" + + src_waves = meta["waveforms"] + # int16 (not float32) before .mean(): matches _average_block's own # float64-accumulator behavior for integer input, so this doesn't - # drift from what average_rows actually guarantees. - expect0 = src_data[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16) + # drift from what _average_block actually guarantees. + expect0 = src_waves[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16) assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \ "first averaged group equals the mean of its 4 source frames" # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial. - dst2 = tmp_path / "legacy_v4_avg5.sras" + dst2 = tmp_path / "v6_avg5.sras" proc2 = subprocess.run([sys.executable, str(REPO / "sras_average.py"), str(src), str(dst2), "--n", "5"], capture_output=True, text=True, cwd=REPO) assert proc2.returncode == 0, (proc2.stderr or proc2.stdout).strip()[-200:] assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \ "partial trailing group kept by default" - dst3 = tmp_path / "legacy_v4_avg5d.sras" + dst3 = tmp_path / "v6_avg5d.sras" proc3 = subprocess.run([sys.executable, str(REPO / "sras_average.py"), str(src), str(dst3), "--n", "5", "--discard-remainder"], capture_output=True, text=True, cwd=REPO) @@ -351,6 +415,59 @@ def test_sras_average(tmp_path): "--discard-remainder drops the partial group" +def test_sras_average_v7_cache_dropped(tmp_path): + """A v7 input's cache tail is indexed by frame count, so it's invalid + after averaging changes that count -- the output must always be plain + v6, never a v7 carrying a stale cache.""" + src = tmp_path / "v7.sras" + gen.write(src, n_angles=2, seed=1, samples_per_frame=32, geometry=[(3, 8)]) + src_sras = SrasFile(str(src)) + src_sras.write_v7_cache( + new_dc3_mv=[compute_dc_image(src_sras, a, CH3_IDX) for a in range(src_sras.n_angles)], + new_dc4_mv=[compute_dc_image(src_sras, a, CH4_IDX) for a in range(src_sras.n_angles)]) + assert SrasFile(str(src)).version == 7 + + dst = tmp_path / "v7_avg.sras" + proc = subprocess.run( + [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"], + capture_output=True, text=True, cwd=REPO) + assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:] + assert SrasFile(str(dst)).version == 6, "cache-bearing input still writes plain v6" + + +def test_sras_average_rejects_legacy(tmp_path): + """This tool only speaks v6/v7 now; a legacy file must fail clearly + rather than being silently misparsed.""" + src = tmp_path / "legacy_v4.sras" + gen.write_legacy(src, version=4, n_angles=1, n_rows=2, n_frames=8, + samples_per_frame=16, seed=0) + dst = tmp_path / "legacy_v4_avg.sras" + proc = subprocess.run( + [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"], + capture_output=True, text=True, cwd=REPO) + assert proc.returncode != 0 + assert "v6" in proc.stderr and "v7" in proc.stderr + + +def test_sras_average_chunking_matches_unchunked(tmp_path): + """A tiny memory budget (forcing one row per chunk) must produce + byte-identical output to a huge budget (everything in one chunk) -- the + load-bearing correctness claim of the memory-bounded rewrite: chunk + boundaries must never affect the averaged result.""" + src = tmp_path / "v6.sras" + gen.write(src, n_angles=2, seed=7, samples_per_frame=48, + geometry=[(6, 10), (5, 13)]) + sras = SrasFile(str(src)) + + dst_tiny = tmp_path / "avg_tiny.sras" + dst_big = tmp_path / "avg_big.sras" + sras_average.write_v6_averaged(sras, dst_tiny, 3, False, budget=1) + sras_average.write_v6_averaged(sras, dst_big, 3, False, budget=1 << 30) + + assert dst_tiny.read_bytes() == dst_big.read_bytes(), \ + "chunk size must not affect the averaged output" + + def test_unsupported_version_reported(tmp_path): """cache_file must report, not raise, for a file it can't handle.""" bogus = tmp_path / "bogus.sras" diff --git a/tests/test_gui.py b/tests/test_gui.py index 4b184ff..1ae89fb 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -223,6 +223,54 @@ def test_threshold_change_recomputes(ctx): f"masking zeroed some pixels ({n_zero} of {win._current_image.size})" +def test_highlight_masked_pixels(ctx): + """Masked (below-threshold) pixels are drawn as NaN - filled with a + highlight color, separate from the normal colormap - so they can't be + mistaken for a real, possibly-low-frequency pixel; unchecking restores + the old behavior where both blend into the same plain 0.""" + win = ctx.win + assert win.chk_highlight_masked.isChecked(), "on by default" + dc4 = win._dc_cache[(0, CH4_IDX)] + expect_masked = dc4 < win.spin_threshold_mv.value() + assert expect_masked.any() and not expect_masked.all(), \ + "fixture threshold should mask some but not all pixels" + + calls = [] + orig = win.image_canvas.show_image + + def spy(img, *a, **kw): + calls.append((np.array(img, copy=True), kw.get("bad_color"))) + return orig(img, *a, **kw) + + with patch.object(win.image_canvas, "show_image", side_effect=spy): + win._redraw_image(win._current_image) + shown, bad_color = calls[-1] + assert bad_color is not None, "highlight color set while checkbox is on" + # The highlight masks by value too: in an FFT mode, exactly 0 is the + # "no valid peak" sentinel (DC-masked, below the min-freq floor, or an + # empty spectrum), so the NaN set is the union of the DC mask and the + # zero-valued pixels. On this fixture every above-threshold pixel has a + # nonzero peak, so the union equals the DC mask alone. + expect_nan = expect_masked | (win._current_image == 0) + assert np.array_equal(np.isnan(shown), expect_nan), \ + "NaN where DC4 is below threshold or the value-0 sentinel, nowhere else" + assert np.array_equal(expect_nan, expect_masked), \ + "fixture precondition: every valid pixel has a nonzero peak" + + win.chk_highlight_masked.setChecked(False) + calls.clear() + with patch.object(win.image_canvas, "show_image", side_effect=spy): + win._redraw_image(win._current_image) + shown2, bad_color2 = calls[-1] + assert bad_color2 is None, "no highlight color once unchecked" + assert not np.isnan(shown2).any(), "unchecked: no pixel pulled out to NaN" + assert np.array_equal(shown2, win._current_image), \ + "unchecked: displayed array is the raw, unmodified image" + + win.chk_highlight_masked.setChecked(True) + pump(60) + + def test_bg_sub_toggle(ctx): """bg-sub no longer gates the display: it only affects a future live compute for an angle with nothing cached yet, or an explicit batch diff --git a/tests/test_row_average.py b/tests/test_row_average.py index 448cee4..dbb00d2 100644 --- a/tests/test_row_average.py +++ b/tests/test_row_average.py @@ -14,7 +14,7 @@ import pytest import sras_compute as compute from sras_compute import compute_rf_image, dc_image_mv -from sras_format import CH4_IDX, SrasFile +from sras_format import CH1_IDX, CH4_IDX, SrasFile import tools.make_test_sras as gen @@ -193,26 +193,45 @@ def test_row_average_respects_own_center_mask(tmp_path): def test_row_average_composes_with_padding(tmp_path): """row_avg_n and n_fft (zero-padding) are independent knobs: using them - together must not raise, and must still agree with the exact (non-zoom) - reference path at that pad factor -- i.e. row-averaging composes with - the zoom peak search correctly, not just with the direct one.""" + together must not raise, and must agree bit-for-bit with an independent + reference that row-averages the raw waveforms and background-subtracts + them, then runs a plain scipy rfft + argmax at the same pad factor -- + i.e. row-averaging composes correctly with the padded peak search.""" + import scipy.fft as scipy_fft + path = tmp_path / "padded_rowavg.sras" gen.write(path, n_angles=1, seed=12, samples_per_frame=64) sras = SrasFile(str(path)) spf = sras.samples_per_frame + n_rows, n_frames = sras.image_shape(0) + n_fft = spf * 40 raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True) avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, row_avg_n=3) - avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, - row_avg_n=3, n_fft=spf * 40) - avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, - row_avg_n=3, n_fft=spf * 40, exact=True) + avg_padded = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True, + row_avg_n=3, n_fft=n_fft) - assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape - assert np.all(np.isfinite(avg_padded_zoom)) - assert np.array_equal(avg_padded_zoom, avg_padded_exact), \ - "row-averaged waveforms feed the zoom and exact FFT paths identically" + assert avg_natural.shape == raw_natural.shape == avg_padded.shape + assert np.all(np.isfinite(avg_padded)) + + weights = compute._row_average_weights(3) + freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32) + data = sras.data[0] + expected = np.zeros((n_rows, n_frames), dtype=np.float32) + for r in range(n_rows): + v = np.ones(n_frames, dtype=bool) + avg = compute._row_average_waveforms( + data[r, CH1_IDX].astype(np.float32), v, weights) + avg = avg - sras.background + S = scipy_fft.rfft(avg, n=n_fft, axis=-1, workers=1) + power = S.real ** 2 + S.imag ** 2 + power[:, 0] = 0.0 + expected[r] = freq32[np.argmax(power, axis=1)] + + assert np.array_equal(avg_padded, expected), \ + "row-averaged waveforms feed the padded peak search identically " \ + "to an independent reference" def test_row_average_improves_snr_recovery(): @@ -233,8 +252,8 @@ def test_row_average_improves_snr_recovery(): weights = compute._row_average_weights(8) # wide window: lots of averaging averaged = compute._row_average_waveforms(raw, valid, weights) - raw_bins = compute._peak_bins_direct(raw, spf) - avg_bins = compute._peak_bins_direct(averaged, spf) + raw_bins = compute._peak_bins(raw, spf) + avg_bins = compute._peak_bins(averaged, spf) raw_hits = int(np.sum(raw_bins == true_bin)) avg_hits = int(np.sum(avg_bins == true_bin)) @@ -257,7 +276,7 @@ def test_row_average_parallel_identity(tmp_path, monkeypatch): sras = SrasFile(str(path)) monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4) - monkeypatch.setattr(compute, "_FFT_BLOCK", 4) + monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4) # row_avg_n > 0 halves the effective budget before chunk planning. fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2) assert fft_rows < n_rows, \ diff --git a/tests/test_stored_cache.py b/tests/test_stored_cache.py index 558baf9..1574691 100644 --- a/tests/test_stored_cache.py +++ b/tests/test_stored_cache.py @@ -90,7 +90,7 @@ def no_fft(monkeypatch): """Make any real FFT work loud: returns a list that stays empty unless a peak search actually runs.""" calls = [] - for name in ("_peak_bins_direct", "_peak_bins_zoom"): + for name in ("_peak_bins",): original = getattr(compute, name) def spy(*args, _f=original, **kwargs): @@ -147,7 +147,7 @@ def test_cache_is_stored_at_the_configured_pad(tmp_path, pad): gen.write(path, n_angles=2, seed=13, samples_per_frame=256) spf = SrasFile(str(path)).samples_per_frame - assert cache_file(str(path), "fft", True, "scipy", 0, pad) == "" + assert cache_file(str(path), "fft", True, 0, pad) == "" sras = SrasFile(str(path)) assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip" @@ -180,12 +180,18 @@ def test_cach_v1_reads_as_natural_resolution(tmp_path): # 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)) + tail_offset = v2._cache_tail_offset() 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()] + # Windows refuses to truncate a file with a live mapping (write_bytes + # opens 'wb'), and every SrasFile holds its waveform memmaps for life — + # drop the instance first. The parsed freq arrays are plain copies and + # stay usable. + del v2 + head = path.read_bytes()[:tail_offset] path.write_bytes(head + payload) v1 = SrasFile(str(path)) @@ -493,12 +499,16 @@ def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path): v2 = SrasFile(str(path)) freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles)) + tail_offset = v2._cache_tail_offset() 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()] + # See test_cach_v1_reads_as_natural_resolution: release the memmaps + # before write_bytes truncates, or Windows raises EINVAL. + del v2 + head = path.read_bytes()[:tail_offset] path.write_bytes(head + payload) v1 = SrasFile(str(path)) @@ -510,6 +520,288 @@ def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path): "a v1 tail (predating this feature) can never satisfy a row-averaged request" +# --------------------------------------------------------------------------- +# Min peak frequency floor: serve-time masking, on-disk provenance, batch +# --------------------------------------------------------------------------- + +def _biting_floor(stored: np.ndarray) -> float: + """A floor that zeroes some-but-not-all of *stored*'s positive peaks: + the median distinct positive value, so pixels below it get masked and + pixels at/above it survive (the mask is a strict <).""" + positive = np.unique(stored[stored > 0]) + assert len(positive) >= 2, "fixture must have varied peak frequencies" + return float(positive[len(positive) // 2]) + + +def test_stored_image_served_with_raised_floor_masked(rig, no_fft): + """The core of the 5 m/s bug fix: a floor above the stored one (here 0) + is re-applied when the stored image is served — pixels whose stored + peak falls below it come back as the 0.0 invalid sentinel, everything + else passes through, and no FFT runs. Both directly through + cached_rf_image and through compute_rf_image's fast path.""" + assert rig.sras.precomputed_min_freq_mhz == 0.0, "batched without a floor" + stored = rig.sras.precomputed_freq_mhz[0] + floor = _biting_floor(stored) + expect = np.where(stored < floor, np.float32(0.0), stored) + + img = cached_rf_image(rig.sras, 0, None, apply_bg_sub=True, + min_freq_mhz=floor) + assert img is not None, "an equal-or-higher floor is servable" + assert np.array_equal(img, expect) + assert (img == 0).any() and (img > 0).any(), \ + "the floor bites some pixels but not all" + + via_compute = compute_rf_image(rig.sras, 0, dc_threshold_mv=None, + apply_bg_sub=True, min_freq_mhz=floor) + assert np.array_equal(via_compute, expect), \ + "compute_rf_image's fast path applies the same serve-time mask" + assert not no_fft, f"serving + masking must not run an FFT: {no_fft}" + + +def test_min_freq_floor_round_trips_and_gates_serving(tmp_path): + """cache_file records the floor in the SFFT header and the accept rule + is asymmetric: an equal-or-higher request is servable, a lower one is + refused (the stored search never looked below its floor). 20.1 pins the + fixed-point kHz encoding — a float32 header field would read back as + 20.10000038…, above the requested 20.1, and mismatch forever.""" + path = tmp_path / "floor_roundtrip.sras" + gen.write(path, n_angles=2, seed=31, samples_per_frame=128) + floor = 20.1 + assert cache_file(str(path), "fft", True, min_freq_mhz=floor) == "" + + sras = SrasFile(str(path)) + assert sras.precomputed_min_freq_mhz == floor, "exact fixed-point round-trip" + assert all(((img == 0) | (img >= floor)).all() + for img in sras.precomputed_freq_mhz), \ + "no stored peak below the floor" + + def reasons(f): + return compute.cache_mismatch_reasons( + sras, n_fft=None, apply_bg_sub=True, row_avg_n=0, min_freq_mhz=f) + + assert reasons(floor) == [] + assert reasons(floor + 5.0) == [], "a higher request is servable (masked)" + low = reasons(0.0) + assert low and "min-peak-freq floor" in low[0], \ + "a lower request cannot be answered by the stored search" + assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \ + "default floor-0 request refused against a floored store" + assert cached_rf_image(sras, 0, None, apply_bg_sub=True, + min_freq_mhz=floor) is not None + + +def test_batch_recompute_resolves_not_masks(tmp_path, no_fft): + """Re-batching an already-cached file at a raised floor must run the + real FFT and store re-resolved peaks — never let compute_rf_image's + fast path serve the file's own stale cache back to it and bake the + masked copy in as if it were a recompute (silent, permanent data + loss: a masked pixel's true above-floor peak is unrecoverable).""" + path = tmp_path / "rebatch.sras" + gen.write(path, n_angles=2, seed=32, samples_per_frame=256) + assert cache_file(str(path), "fft", True) == "" + + first = SrasFile(str(path)) + n_angles = first.n_angles + # On the header's kHz grid up front (as the spinbox value would be), so + # the recorded floor reads back equal — cache_file quantizes whatever it + # is given, and this test wants that to be the identity. + floor = round(_biting_floor(first.precomputed_freq_mhz[0]) * 1000) / 1000.0 + bites = [img < floor for img in first.precomputed_freq_mhz] + assert bites[0].any(), "the floor must actually bite this fixture" + # What a real floored compute gives, from a view blinded to the cache. + first.precomputed_freq_mhz = [None] * n_angles + expected = [compute_rf_image(first, a, dc_threshold_mv=None, + apply_bg_sub=True, min_freq_mhz=floor) + for a in range(n_angles)] + del first # release memmaps before cache_file rewrites the tail + + no_fft.clear() + assert cache_file(str(path), "fft", True, min_freq_mhz=floor) == "" + assert no_fft, "the re-batch ran a real FFT" + + after = SrasFile(str(path)) + assert after.precomputed_min_freq_mhz == floor + for a in range(n_angles): + assert np.array_equal(after.precomputed_freq_mhz[a], expected[a]), \ + f"angle {a}: stored image is a real floored recompute" + assert (after.precomputed_freq_mhz[a][bites[a]] >= floor).all(), \ + f"angle {a}: bitten pixels re-resolved above the floor, not zeroed" + + +def test_min_freq_carries_forward_through_dc_write(tmp_path): + """A later DC-only write must leave the FFT block's recorded floor + untouched, like row_avg_n and pad_factor.""" + path = tmp_path / "floor_carry.sras" + gen.write(path, n_angles=2, seed=33, samples_per_frame=128) + assert cache_file(str(path), "fft", True, min_freq_mhz=75.0) == "" + assert cache_file(str(path), "dc", True) == "" + after = SrasFile(str(path)) + assert after.precomputed_min_freq_mhz == 75.0, \ + "floor survives a DC-only write" + + +def test_cach_v3_backward_compat_defaults_floor_zero(tmp_path): + """A v3 CACH tail predates the min peak frequency floor entirely (no + min_freq_khz field) — readers must still parse it in full, treating it + as floor 0: servable as-is at floor 0, and serve-maskable at any higher + one. This is what protects existing real-world v7 caches from silently + becoming unusable after the v4 bump ships.""" + path = tmp_path / "v3_floor.sras" + gen.write(path, n_angles=2, seed=34, samples_per_frame=128) + assert cache_file(str(path), "fft", True) == "" + + # Rewrite the tail as a genuine CACH v3 block (no min_freq field). + v4 = SrasFile(str(path)) + freq, entries = v4.precomputed_freq_mhz, list(range(v4.n_angles)) + tail_offset = v4._cache_tail_offset() + payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 3, fmt.CACH_FLAG_FFT) + payload += struct.pack(fmt.SFFT_HDR_FMT_V3, fmt.SFFT_MAGIC, + fmt.SFFT_FLAG_BG_SUB, len(entries), 0, 1) + for a in entries: + payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes() + # See test_cach_v1_reads_as_natural_resolution: release the memmaps + # before write_bytes truncates, or Windows raises EINVAL. + del v4 + head = path.read_bytes()[:tail_offset] + path.write_bytes(head + payload) + + v3 = SrasFile(str(path)) + assert v3.precomputed_min_freq_mhz == 0.0 + assert v3.precomputed_pad_factor == 1 and v3.precomputed_bg_sub is True + assert all(np.array_equal(v3.precomputed_freq_mhz[a], freq[a]) + for a in entries), "v3 images read back unchanged" + assert cached_rf_image(v3, 0, None, apply_bg_sub=True) is not None + floor = _biting_floor(freq[0]) + served = cached_rf_image(v3, 0, None, apply_bg_sub=True, + min_freq_mhz=floor) + assert served is not None + assert np.array_equal(served, + np.where(freq[0] < floor, np.float32(0.0), freq[0])) + + +def test_min_freq_validation(tmp_path): + path = tmp_path / "floor_bad.sras" + gen.write(path, n_angles=1, seed=35, samples_per_frame=64) + assert cache_file(str(path), "fft", True, min_freq_mhz=-1.0), \ + "negative floor must be an error, not a write" + assert cache_file(str(path), "fft", True, min_freq_mhz=float("nan")), \ + "NaN floor must be an error, not a write" + sras = SrasFile(str(path)) + with pytest.raises(ValueError): + sras.write_v7_cache(new_min_freq_mhz=-0.5) + + +def test_viewer_reapplies_floor_to_stored_images_without_computing( + rig, no_fft, monkeypatch): + """The session-cache poisoning bug behind the '95 MHz peak but 5 m/s' + report: changing 'Min peak freq' against a stored cache used to re-file + the identical un-floored image under a key claiming the new floor — the + UI looked updated, the pixels weren't. Now the floor really is + re-applied on serve (masked, still no compute), and clearing it + restores the unmasked image, still without computing.""" + 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" + assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache + for a in range(rig.n_angles))), \ + "DC precompute finished" + + win.combo_channel.setCurrentIndex(CH1_IDX) + assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed" + assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3) + + # Round to the spinbox's 3-decimal granularity; the chosen bin value + # still survives its own (strict-<) floor after rounding down. + floor = round(_biting_floor(rig.fresh[0]), 3) + expect = np.where(rig.fresh[0] < floor, np.float32(0.0), rig.fresh[0]) + + dispatched.clear() + no_fft.clear() + win.spin_min_freq_mhz.setValue(floor) + win._on_min_freq_changed() + pump(120) + assert np.allclose(win._current_image, expect, atol=1e-3), \ + "raised floor re-masks the stored image on serve" + assert (win._current_image == 0).any() and (win._current_image > 0).any() + assert dispatched == [] and not no_fft, \ + f"re-masked serve needs no compute (jobs={dispatched}, fft={no_fft})" + key = (0, win.spin_threshold_mv.value(), win.spin_min_freq_mhz.value()) + assert key in win._fft_cache + assert np.allclose(win._fft_cache[key], expect, atol=1e-3), \ + "the session cache holds the value its key claims" + + win.spin_min_freq_mhz.setValue(0.0) + win._on_min_freq_changed() + pump(120) + assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3), \ + "clearing the floor restores the unmasked stored image" + assert dispatched == [] and not no_fft + finally: + win.close() + pump(300) + + +def test_viewer_batch_fft_records_the_floor(tmp_path, no_fft, monkeypatch): + """Convert → Batch Compute FFT with a floor set: the live spinbox value + reaches cache_file, lands in the reloaded file's provenance and the + info panel, and the viewer then serves the floored cache without + recomputing — the tooltip's promised remedy, end to end.""" + path = tmp_path / "floor_gui.sras" + gen.write(path, n_angles=2, seed=36, 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._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" + + floor = 100.0 + win.spin_min_freq_mhz.setValue(floor) + 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_min_freq_mhz == floor, \ + "the viewer's floor reached the stored provenance" + assert "floor ≥ 100 MHz" in win.lbl_frame_warn.text() + + no_fft.clear() + dispatched.clear() + win.combo_channel.setCurrentIndex(CH1_IDX) + assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed" + pump(120) + assert dispatched == [] and not no_fft, \ + f"floored cache serves the view directly (jobs={dispatched}, fft={no_fft})" + img = win._current_image + assert ((img == 0) | (img >= floor)).all(), \ + "no displayed peak below the floor" + finally: + win.close() + pump(300) + + def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch, no_fft): """Driving the new 'Batch Compute Row-Averaged FFT and Store' action end-to-end through the real menu handler: dialog values reach the diff --git a/tools/bench_fft.py b/tools/bench_fft.py index 68a1cc3..a509b19 100644 --- a/tools/bench_fft.py +++ b/tools/bench_fft.py @@ -1,8 +1,10 @@ #!/usr/bin/env python3 -"""Benchmark the FFT peak-search path: exact vs zoom, serial vs pooled. +"""Benchmark the FFT peak-search path: serial vs pooled. Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in -cores), and verifies every variant against the exact reference image. +cores), the estimated per-thread resident pyFFTW plan footprint at each pad +factor (see sras_compute._fft_block_for), and verifies pooled output against +the serial reference. Usage: python tools/bench_fft.py # synthetic, pads 1/8/40 @@ -22,10 +24,10 @@ import numpy as np sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) import sras_compute as compute # noqa: E402 -from sras_compute import compute_rf_image, set_fft_backend # noqa: E402 -from sras_format import SrasFile # noqa: E402 -import tools.make_test_sras as gen # noqa: E402 -from tools.check_equivalence import row_slice # noqa: E402 +from sras_compute import compute_rf_image # noqa: E402 +from sras_format import SrasFile # noqa: E402 +import tools.make_test_sras as gen # noqa: E402 +from tools.check_equivalence import row_slice # noqa: E402 def _timed(fn): @@ -38,35 +40,42 @@ def _timed(fn): return out, wall, cpu / max(wall, 1e-9) -def bench(sras, pads, backends): +def _plan_mb(spf: int, n_len: int, n_workers: int) -> float: + """Estimated resident pyFFTW plan-buffer footprint across the whole + pool at this transform length (see sras_compute._fft_block_for).""" + block = compute._fft_block_for(spf, n_len) + bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1) + return block * bytes_per_wf * n_workers / (1024 * 1024) + + +def bench(sras, pads): n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a]) for a in range(sras.n_angles)) spf = sras.samples_per_frame - print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)") - print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} " - f"{'wf/s':>10} {'util':>6} match") + n_workers = compute._MAX_WORKERS + print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), " + f"{n_workers} workers") + print(f"{'pad':>4} {'variant':>8} {'wall':>9} {'wf/s':>10} {'util':>6} " + f"{'plan MB':>9} match") for pad in pads: n_fft = spf * pad if pad > 1 else None - for backend in backends: - set_fft_backend(backend) + n_len = n_fft if n_fft is not None else spf + plan_mb = _plan_mb(spf, n_len, n_workers) - def run(**kw): - imgs = [compute_rf_image(sras, a, dc_threshold_mv=None, - apply_bg_sub=True, n_fft=n_fft, **kw) - for a in range(sras.n_angles)] - return np.concatenate([i.ravel() for i in imgs]) + def run(**kw): + imgs = [compute_rf_image(sras, a, dc_threshold_mv=None, + apply_bg_sub=True, n_fft=n_fft, **kw) + for a in range(sras.n_angles)] + return np.concatenate([i.ravel() for i in imgs]) - ref, wall, util = _timed(lambda: run(exact=True)) - rows = [("exact(serial)", ref, wall, util, True)] - for label, kw in (("zoom(serial)", dict(max_workers=1)), - ("zoom(pool)", {})): - img, wall, util = _timed(lambda: run(**kw)) - rows.append((label, img, wall, util, bool(np.array_equal(img, ref)))) - for label, img, wall, util, ok in rows: - print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s " - f"{n_wf / wall:>10.0f} {util:>5.1f}x " - f"{'OK' if ok else 'MISMATCH'}") + ref, wall, util = _timed(lambda: run(max_workers=1)) + rows = [("serial", ref, wall, util, True)] + img, wall, util = _timed(lambda: run()) + rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref)))) + for label, img, wall, util, ok in rows: + print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} " + f"{util:>5.1f}x {plan_mb:>8.1f} {'OK' if ok else 'MISMATCH'}") def main(): @@ -76,30 +85,24 @@ def main(): p.add_argument("--spf", type=int, default=2500) p.add_argument("--rows", type=int, default=8) p.add_argument("--frames", type=int, default=1024) - p.add_argument("--backends", default=None, - help="comma-separated (default: scipy,pyfftw if available)") p.add_argument("--real", help="path to a real .sras file") p.add_argument("--real-rows", type=int, default=32, help="rows of angle 0 to use from the real file") args = p.parse_args() pads = [int(x) for x in args.pads.split(",")] - if args.backends: - backends = args.backends.split(",") - else: - backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else []) if args.real: sras = row_slice(SrasFile(args.real), 0, args.real_rows) sras.data = [sras.data[0]] sras.n_angles = 1 - bench(sras, pads, backends) + bench(sras, pads) else: with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp: path = Path(tmp) / "bench.sras" gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf, geometry=[(args.rows, args.frames)]) - bench(SrasFile(str(path)), pads, backends) + bench(SrasFile(str(path)), pads) if __name__ == "__main__": diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py index f27e76f..ef0d33d 100644 --- a/tools/make_test_sras.py +++ b/tools/make_test_sras.py @@ -264,10 +264,7 @@ def write_rotating(path: Path, n_angles: int = 5, samples_per_frame: int = 4, def write_legacy(path: Path, version: int = 4, n_angles: int = 2, n_rows: int = 4, n_frames: int = 10, samples_per_frame: int = 32, seed: int = 0) -> dict: - """Write a v2/v3/v4 file: uniform geometry, one flat waveform block. - - Used to exercise sras_average.py, which only handles the legacy formats. - """ + """Write a v2/v3/v4 file: uniform geometry, one flat waveform block.""" rng = np.random.default_rng(seed) n_ch, bps = 3, 1