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

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

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

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

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

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

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales [M S E]
2026-08-12 10:36:25 -05:00
parent efecc154fe a8b962e317
commit 2937c017a7
19 changed files with 2130 additions and 617 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design.md
+66 -37
View File
@@ -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
pyproject.toml
+2 -3
View File
@@ -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",
scan_format.md
+31 -13
View File
@@ -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
sras_average.py
+154 -59
View File
@@ -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)
geo = bytearray()
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 += 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,15 +230,13 @@ 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])
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}")
print(f" Rows : {int(sras.n_rows[0])}")
print(f" Frames (actual): {n_frames_in}")
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}")
@@ -150,28 +247,26 @@ def main():
print(f"Wrote {out_path}")
return
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 available frames ({n_frames_in}). "
"The entire dataset will be averaged into a single frame.")
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.")
sras_compute.py
+163 -236
View File
@@ -1,9 +1,9 @@
#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
multiprocessing child can import it without loading Qt or matplotlib —
which matters because Python 3.14 on macOS spawns rather than forks.
Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing
child can import it without loading Qt or matplotlib — which matters
because Python 3.14 on macOS spawns rather than forks.
"""
import atexit
@@ -15,61 +15,34 @@ from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pyfftw
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
adc_to_mv)
# ---------------------------------------------------------------------------
# FFT backend
# ---------------------------------------------------------------------------
try:
import pyfftw
PYFFTW_AVAILABLE = True
except ImportError:
PYFFTW_AVAILABLE = False
try:
from threadpoolctl import threadpool_limits
except ImportError:
threadpool_limits = None
_fft_backend = "scipy" # "scipy" or "pyfftw"; set via set_fft_backend()
def set_fft_backend(name: str):
"""Select the rfft implementation. Module-level state, so it must be set
explicitly inside each multiprocessing child — it does not survive a
spawn. "numpy" is accepted as a legacy alias for "scipy"."""
global _fft_backend
if name == "numpy":
name = "scipy"
_fft_backend = name if (name == "pyfftw" and PYFFTW_AVAILABLE) else "scipy"
def get_fft_backend() -> str:
return _fft_backend
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
workers: int = 1) -> np.ndarray:
"""Dispatch rfft to the selected backend with optional multithreading."""
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
# ---------------------------------------------------------------------------
# FFT worker pool and per-thread pyFFTW plans
# ---------------------------------------------------------------------------
_FFT_BLOCK = 512 # waveforms per FFT task. The knee on a 16-core machine:
# smaller blocks serialise on GIL-held numpy dispatch,
_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.
_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_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 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,25 +598,14 @@ 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)
pool = _fft_pool() if cap > 1 else None
@@ -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()
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)
sras_format.py
+78 -24
View File
@@ -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
sras_render.py
+153
View File
@@ -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)
sras_viewer/canvases.py
+35 -17
View File
@@ -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,8 +485,21 @@ class WaveformCanvas(FigureCanvasQTAgg):
# FFT of the (possibly subtracted) waveform
power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2
power_sub[0] = 0.0
# 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
power_raw = np.abs(np.fft.rfft(waveform)) ** 2
sras_viewer/common.py
+5 -8
View File
@@ -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.
sras_viewer/dialogs.py
+1 -29
View File
@@ -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())
sras_viewer/main_window.py
+290 -55
View File
@@ -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"{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
sras_workers.py
+122 -8
View File
@@ -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.
tests/test_batch_export.py
+426
View File
@@ -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"
tests/test_compute.py
+171 -54
View File
@@ -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"
tests/test_gui.py
+48
View File
@@ -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
tests/test_row_average.py
+34 -15
View File
@@ -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, \
tests/test_stored_cache.py
+296 -4
View File
@@ -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
tools/bench_fft.py
+29 -26
View File
@@ -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,7 +24,7 @@ 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_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
@@ -38,18 +40,28 @@ 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,
@@ -57,16 +69,13 @@ def bench(sras, pads, backends):
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))))
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} {backend:>8} {label:>16} {wall:>8.2f}s "
f"{n_wf / wall:>10.0f} {util:>5.1f}x "
f"{'OK' if ok else 'MISMATCH'}")
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__":
tools/make_test_sras.py
+1 -4
View File
@@ -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