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

Compare commits

base: tka/sras-viewer:191d1b8946f558b29d97626fa143af607bbf18b9
Branches Tags
tka/sras-viewer:main tka/sras-viewer:refactor/parallel-and-dedup tka/sras-viewer:v7-cache-format tka/sras-viewer:feature/angle-alignment tka/sras-viewer:remove-timegate-saw-add-v6
..
compare: tka/sras-viewer:refactor/parallel-and-dedup
Branches Tags
tka/sras-viewer:main tka/sras-viewer:refactor/parallel-and-dedup tka/sras-viewer:v7-cache-format tka/sras-viewer:feature/angle-alignment tka/sras-viewer:remove-timegate-saw-add-v6

1 Commits
191d1b8946 .. refactor/parallel-and-dedup

Author SHA1 Message Date
Thomas Ales 1a8187138f pre merge cleanup 2026-08-06 21:20:35 -05:00
20 changed files with 1298 additions and 5945 deletions
Expand all files Collapse all files
docs/design.md
-261
View File
@@ -72,133 +72,6 @@ one thread under the pool so the refinement gemm cannot oversubscribe.
`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
full-padded path for audits.
## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
`compute_rf_image`'s `row_avg_n` parameter averages each pixel's CH1
waveform with its up-to-n same-row neighbors before the FFT peak search, to
improve SNR on noisy scans. Never crosses rows: pixel pitch is strongly
anisotropic and varies by scan (5 µm × 50 µm on a typical scan, but as
stretched as 5 µm × 1 mm on others), so a physically meaningful "neighbor"
set can't be a fixed-shape 2-D window — but the X pitch *within one row* is
a single file-wide constant (`SrasFile.pixel_x_mm`), so restricting to the
row axis sidesteps the anisotropy question entirely rather than solving it
with an elliptical or physically-scaled 2-D kernel.
`_row_average_weights` is a Gaussian in pixel-index distance, not physical
mm distance — deliberately: within one row those are the same function up
to a fixed scale factor (`pixel_x_mm` is constant along a row), so the
kernel itself needs no pitch at all. `pixel_x_mm` is used for real exactly
once, in the GUI's options dialog, to show the window's physical width —
not in the kernel math, where it would only ever cancel out.
`_row_average_waveforms` is a masked/renormalized convolution (two
`correlate1d` calls, numerator and denominator, divided) rather than a
single fixed-normalized convolution, because a masked neighbor must
contribute *zero weight*, not a zero-amplitude sample at full weight — the
latter would bias every average near a masked run or a row's own edge
toward zero. The same two-correlation trick handles row-edge truncation for
free: `mode="constant", cval=0.0` zero-pads both the numerator and the
denominator beyond a row's own ends, so the output renormalizes by whatever
weight sum actually landed inside the row, no separate edge case.
Background subtraction stays exactly where it already was (subtracted once
from the fully-assembled `waves` buffer) rather than being threaded into the
per-neighbor gather. This is exact, not an approximation: because
`_row_average_waveforms`'s denominator is always the *actual* sum of
included, valid weights (never a fixed total), `Σwᵢ·(rawᵢ−bg) / Σwᵢ`
distributes to `avg − bg·(Σwᵢ/Σwᵢ) = avg − bg` regardless of which or how
many neighbors were included — subtracting background from the averaged
waveform is identical to subtracting it from every neighbor first, for any
window, at any row edge, with any number of masked-out neighbors.
No cross-row halo is needed: `compute_rf_image`'s chunk loop already splits
on rows only, and `read_row` already reads one row's complete
`(n_frames, spf)` slice at a time — averaging happens entirely inside that
one row's own frame axis, so a chunk boundary (which falls between rows)
can never truncate a window. Only a row's own start/end can, and that's the
same edge case the masked convolution already handles.
The averaging step doubles the live per-row scratch memory (a full-width
`(n_frames, spf)` buffer on top of the existing compacted `waves` buffer),
so `compute_rf_image` halves its byte budget when `row_avg_n > 0` before
`_plan_fft_rows`/the exact-path sizing runs — see "Memory budget and row
chunking" above. On the largest real scans `_plan_chunks` is already
clamped to its floor of one row regardless, so this costs no concurrency
where it matters most; it mainly protects moderate-sized scans from an
unexpected regression.
Persistence: `cached_rf_image` (the extracted fast-path check) requires
`sras.precomputed_row_avg_n == row_avg_n` exactly, so a raw request can
never be silently served a row-averaged cache or vice versa, and a request
at one window size can never be served a cache at another — see
`scan_format.md`'s Cache Tail / CACH tail version history sections for the
on-disk `row_avg_n` field this depends on.
## Serving a stored cache: provenance, not just presence
A stored `peak_freq_mhz` image is only interchangeable with a live compute
for the *exact* settings it was computed under. Three of them are baked
irreversibly into the numbers — background subtraction, row-averaging window,
and zero-padding — so all three are recorded in the `SFFT` block and checked
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.
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
frequency for the same waveform. `precomputed_pad_factor` exists so a padded
view can be served from a cache computed at *its* pad while still refusing
one computed at any other, including pad 1. Before it existed the store was
pad 1 by definition and any `n_fft` was rejected outright — correct, but it
meant a user working at a pad factor got nothing at all from batch-computing
a file, which is most of the point of the feature. `pad_factor_for` maps an
`n_fft` request onto the integer factor a store could have recorded, and
returns 0 for a request that is not a whole multiple of `samples_per_frame`
— unmatchable by construction, since only integer factors are representable.
The batch actions therefore have to cache at the *viewer's* current pad
factor, not a fixed one: a cache stored at a pad nobody is viewing at is
dead weight. When the two do diverge (the user changes the pad after
batching), `_cache_mismatch_notes` says so in the scan info panel, because
the symptom otherwise is just "the file I pre-computed got slow again" with
no visible cause.
That divergence note is informational only, not a warning of an impending
recompute. The *display* path (`_stored_fft_image`) never asks
`cached_rf_image` whether a stored image matches the window's live
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
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
keeps a view switch (angle, channel, or flipping bg-sub/pad) from ever
discarding precomputed data — only an explicit batch recompute does, and it
already reloads the file afterward so the new data displays immediately.
Row-averaging has no live control to diverge from in the first place (it's
only ever set inside the batch dialog), so it never appears in the
divergence note — the "Cached images" line's own `row-averaged n=…` phrase
already covers it.
### Two caches, in cost order
`_refresh_display` consults this window's in-session `_fft_cache`/`_dc_cache`
dicts first, then the open file's stored v5/v7 blocks, and only then
dispatches a `ComputeWorker`. The second tier is what makes a batch-computed
file worth having; without it every angle change queued a worker and a
progress popup for an image already sitting on disk — precisely the cost the
batch was run to avoid. (The stored image *was* reachable before, but only
from inside `ComputeWorker`, i.e. after paying for the thread and the popup.)
The file tier asks with `allow_dc_recompute=False`. If applying the DC4 mask
would mean reading a whole CH4 channel, it declines rather than blocking the
GUI thread, and the fall-through worker reaches the same stored image via
`compute_rf_image` and pays for the mask off-thread. So the GUI thread never
does I/O, and the slow path is still a fast path.
## Angle alignment coordinate frames (`sras_compute.py`)
Alignment puts every angle's images onto one shared, zero-padded pixel grid
@@ -236,140 +109,6 @@ for the same reason, and every affine maps shared-grid index → mm → undo
rotation/shift → that angle's own local mm → that angle's own raw index,
matching the output→input convention `scipy.ndimage.affine_transform` wants.
### Cropping the canvas is index translation, not a second transform
`crop_alignment_result` restricts an `AlignmentResult` to a rectangular window
of its canvas by folding the crop into each angle's existing affine rather than
composing a new one. From `_affine_out_to_src`, `matrix = D @ Rinv @ A_out`
depends only on the pitches and the rotation, and `A_out @ [row0, col0]` is
exactly the mm displacement of the new origin, so
```
matrix @ [r', c'] + (offset + matrix @ [row0, col0])
== matrix @ [r' + row0, c' + col0] + offset
```
identically. `matrix` is untouched and `offset` — which already absorbs the
origin — absorbs the crop too.
Two things follow, and both are relied on. `apply_alignment`, `reproject_mask`
and the aligned exporter all work on a cropped result with no special-casing:
resampling a cropped result is *exactly* a slice of resampling the full one
(`tests/test_align_export.py::test_crop_is_a_window_of_the_full_canvas` asserts
bit equality). And because the crop offset is a whole number of canvas pixels,
`canvas_for_params`' snap invariant — the reference angle lands on integer
canvas pixels — survives the crop, which is what keeps the reference exportable
as a verbatim block.
## Aligned export (`sras_align_export.py`)
`write_aligned_sras` bakes an alignment into a new v6 file: every angle
resampled onto the cropped shared canvas, so all of them end up with identical
geometry and the file opens already aligned. It is the only place in the
codebase that *resamples* waveform data — `sras_edit_scans` and `sras_average`
copy waveform bytes verbatim — which is why it is its own top-level module
rather than part of `sras_format` (scoped to the versioned binary spec, per the
sidecar section's own rule) or `sras_compute` (imported by every
multiprocessing child).
**Nearest neighbour, never interpolation.** Each output pixel gets exactly one
source pixel's three waveforms, verbatim. Averaging two neighbouring CH1
packets would synthesise a waveform the instrument never measured, whose FFT
peak is the peak of neither — meaningless for a technique whose entire output is
that peak frequency. The cost is that some source pixels are duplicated and
others dropped, which is the same trade `apply_alignment`'s `order=0` already
makes for the display.
**The rounding rule is `floor(x + 0.5)`, not `np.rint`.** `scipy.ndimage`'s
`order=0` rounds halves away from zero while `np.rint` rounds them to even. The
canvas is snapped to the reference's own pixel grid, so an angle whose row pitch
differs from the reference's lands on exact half-integers across whole rows —
this is the common case, not a corner case. Getting it wrong shifts those rows
by one source pixel relative to what the Aligned View drew.
**Out-of-bounds is tested on the fractional coordinate, not the rounded index.**
`scipy`'s `mode="constant"` writes `cval` wherever the coordinate leaves the
range of sample *centres*, `[0, n-1]` — a coordinate of −0.4 rounds to a
perfectly valid index 0 and is still padding. Testing the rounded index instead
puts a one-pixel rim of real data everywhere the preview shows padding.
**...but with a tolerance (`_EDGE_TOL`).** The affine is built from a chain of
mm-space multiplications, so an exactly-integer transform comes out a few times
1e-13 off: the reference angle's offset is `-20 - 7e-15`, not `-20`. A bare
`>= 0.0` therefore rejects that angle's entire first row, and `<= n-1` its last
column — for the *reference* angle, whose whole job is to pass through as an
exact integer crop. The tolerance is ~7 orders of magnitude above that noise and
~7 below the half-pixel scale at which a rounding decision means anything, so it
can only ever change pixels whose scipy answer was itself decided by noise.
**Padding is the per-channel ADC code nearest 0 mV, not 0.** Zero ADC decodes to
`(0 - yoff) * ymult + yzero`, which on real calibration is around +100 mV —
above any sensible CH4 mask threshold, so a zero fill would paint a solid
rectangle of "valid" pixels around the sample and corrupt every DC image and ROI
statistic downstream.
**Source rows are served from sliding in-RAM bands** (`_SourceReader`). A
rotated angle maps one output row to a *diagonal* across the source array, so
the pixels of a single output row come from hundreds of different source rows —
~1.4 MB each on a full-size scan. Indexing a memmap pixel-by-pixel in output
order re-faults nearly the whole angle per output row: terabytes of paging for a
gigabyte of data. Reading a contiguous band per output chunk, with the band
advancing monotonically, costs roughly 2× the source size in total reads.
**Writes go to `.part` and are `os.replace`d into position.** Not politeness: a
truncated .sras is not detectably broken, because `_parse_v6` drops incomplete
trailing angle blocks and opens what is left as an aborted scan. A half-written
export left in place would silently look like a real file with fewer angles.
**The Angle Table is carried over unchanged.** Alignment removes the *spatial*
rotation of the sample; it does not change which acoustic propagation direction
each angle measured, and that direction is the scientific content of a
multi-angle scan. Zeroing the table would make the export self-consistent for
re-registration and useless for anisotropy work. The consequence is that
re-registering an export needs `seed_deg=0.0` to put 0° inside the coarse sweep,
since `nominal_delta_deg` is still non-zero — which is exactly what the seed
parameter exists for.
## Alignment wizard (`sras_viewer/align_wizard.py`)
A `QWizard` rather than another dialog because the three steps are genuinely
sequential and the last one is destructive: correlate, choose a crop, write a
file. It replaces both former Fusion actions, so it also absorbs the old
`ManualAlignmentDialog`'s by-eye nudge editor — otherwise a scan the search
cannot fit would have no fallback at all.
Shared state lives on the wizard object, not in `registerField`: the pages pass
numpy arrays, `ManualAngleParams` and an `AlignmentResult` between them, none of
which are scalar widget properties.
`IndependentPages` is deliberately left **off**. With it set Qt never calls
`cleanupPage`, and `cleanupPage` is how the ROI page discards a crop when the
user goes back to re-correlate — a crop is indexed in canvas pixels, and a new
rotation means a different canvas, so stale indices would silently be
reinterpreted against the wrong grid. `geometry_generation` is the belt-and-
braces check for the same hazard.
The mask-stack preview shares the **final** canvas's origin and uses a pitch
that is an integer multiple of it, unlike the old manual dialog's padded,
unsnapped preview canvas. That is what lets the crop page convert a rectangle
drawn in millimetres into an exact integer window of the real canvas, with no
second coordinate frame to reconcile.
"Fit to full overlap" uses `largest_rect_at_least`, a largest-rectangle sweep,
not a bounding box of the fully-covered pixels. The full-overlap region of
several rotated scans is roughly a disc, and its bounding box has corners no
angle covers — offering that as the crop would hand the user the padding they
were trying to avoid.
Every background launch follows the two rules `_run_worker`'s docstring
establishes: disable the trigger *before* the call (so a re-entrant click cannot
start a second thread over the first), and never ignore the returned bool.
Progress is an inline `QProgressBar` on the page rather than a `QProgressDialog`
— a window-modal popup over a wizard both looks wrong and reintroduces the
event-loop pumping hazard that ordering exists to avoid. `reject()` refuses to
close while a job is in flight, since the running worker's signals are connected
to bound methods of the pages Qt would be deleting.
## Manual-alignment sidecar (`sras_compute.py`)
`<name>.sras.align.json` lives next to the scan file. The code lives in
pyproject.toml
-1
View File
@@ -31,7 +31,6 @@ py-modules = [
"sras_format",
"sras_compute",
"sras_workers",
"sras_align_export",
"sras_average",
"sras_edit_scans",
]
scan_format.md
+9 -95
View File
@@ -173,41 +173,6 @@ sum over angles a of: n_rows[a] × 3 × n_frames[a] × samples_per_frame × byte
> Table and check `file_size` against the running total before reshaping —
> a fixed `(n_angles, n_rows, ...)` reshape (as in pre-v6 readers) will not
> work since row/frame counts are no longer uniform across angles.
>
> A consequence worth stating explicitly: because a short file opens
> *successfully* as a scan with fewer angles, a truncated file is not
> detectably broken. Anything that writes a .sras must therefore stage to a
> temporary name and rename on success — the viewer's aligned export writes
> `<name>.part` and `os.replace`s it — or an interrupted write leaves behind
> something that loads without complaint and silently has the wrong angle count.
---
## Files written by the viewer's Alignment Wizard
The acquisition app is not the only producer of this format. The viewer's
`Fusion → Alignment Wizard…` writes a **v6** file holding the aligned, cropped
stack, with these properties:
* Every angle shares one grid — the cropped alignment canvas — so the
Per-Angle Geometry Table is `n_angles` identical records and the ragged Row
Table is `n_angles` identical spans. The raggedness v6 exists for is still
*expressible*, just unused, so any v6 reader works unchanged.
* `x_delta` is the reference angle's own pitch, which is exactly
`velocity_mm_s / laser_freq_hz`, so the derived X axis stays consistent with
the header.
* The `*_nominal` header fields describe the crop. Uniquely for these files they
coincide with the actual per-angle geometry, since after alignment every angle
really does scan the same box.
* The **Angle Table is unchanged**. Alignment removes the sample's spatial
rotation, not the acoustic propagation direction each angle measured — that
direction is the point of a multi-angle scan, so it is preserved.
* Output pixels with no corresponding source pixel (the canvas corners a rotated
scan cannot reach) hold the per-channel ADC code nearest **0 mV**, not zero.
Zero ADC decodes to roughly +100 mV on real calibration and would read as
signal.
* No Cache Tail is written: any cached DC/FFT is indexed by the source's grid
and would be meaningless on the new one.
---
@@ -255,7 +220,7 @@ actions.
| Offset | Size | Type | Field | Description |
|--------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `cach_magic` | `CACH` (ASCII). Missing/wrong magic → treat file as having no cache. |
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `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 `1`. Readers must treat the file as uncached if this is not a version they understand (unlike v5's `PREC` section, which read but never validated its version byte). |
| 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`)
@@ -284,29 +249,13 @@ value, same as v5's `PREC` section.
### FFT block `SFFT` (present iff `block_flags & 0x02`)
Block header layout depends on `cach_version`:
Each `cach_version` appended one trailing field, so the header grows but
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`.
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.
7-byte block header, format `">4sBH"`:
| Offset (rel) | Size | Type | Field | Description |
|--------------|------|------|-------|-------------|
| 0 | 4 | `char[4]` | `magic` | `SFFT` |
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. 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. |
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bits 1–7 reserved. |
| 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`. |
followed by `n_stored` entries, each:
@@ -315,10 +264,9 @@ u16 angle_idx — index into the angle table (0-ba
f32[n_rows[angle_idx] × n_frames[angle_idx]] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
```
**`peak_freq_mhz`** for a raw store (`row_avg_n == 0`) is computed without
any DC-threshold masking (i.e. the FFT is run on every pixel
unconditionally, same as v5's `PREC` convention). Readers apply the DC4
threshold at display time:
**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the
FFT is run on every pixel unconditionally, same as v5's `PREC` convention).
Readers apply the DC4 threshold at display time:
```
pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
@@ -328,27 +276,10 @@ display_value = peak_freq_mhz[r][f] if valid, else 0
using the DC4 image from the DC block if that angle is also cached there,
else computed on demand.
For a row-averaged store (`row_avg_n > 0`), the DC4 threshold is applied
*during* the store — a pixel below threshold is left at `0` and never
contributes to any neighbor's average — since neighbor validity can't be
deferred to display time the way plain masking can. The threshold value
itself is not recorded, only that averaging happened and at what window
size. Readers still apply their own live DC4 threshold at display time
exactly as for a raw store, using whatever mask they currently have.
Readers must fall back to real-time FFT computation (ignoring stored
`peak_freq_mhz`) whenever the store's recorded provenance doesn't match what
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.
`peak_freq_mhz`) under the same conditions as v5's PREC fast path: time-domain
gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested, or
the reader's background-subtraction setting doesn't match `flags.bg_sub_applied`.
### In-place write ordering
@@ -363,23 +294,6 @@ interrupted write leaves harmless trailing bytes rather than a corrupt file,
and the next successful write overwrites them via the same deterministic
`cache_offset`.
### CACH tail version history
Distinct from the outer `.sras` file `version` byte (top of this document),
which has stayed `7` since the Cache Tail was introduced — this is the inner
`cach_version` byte inside the `CACH` header itself.
| cach_version | Change |
|--------------|--------|
| 1 | Initial Cache Tail: `SDCB` (DC) and `SFFT` (FFT, 7-byte header) blocks. |
| 2 | `SFFT` header grows one byte, `row_avg_n` — the same-row neighbor half-width the stored `peak_freq_mhz` was averaged over before its FFT, `0` = raw. Readers still accept a `cach_version` 1 tail, treated as `row_avg_n = 0` for every angle it stores, so files cached before this change keep working without a recompute. |
| 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`. |
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
ordinary v7 (byte-identical to v6) scan, so a forward-dated tail costs a
recompute and never correctness.
---
## Acquisition Settings (fixed by sc3_aui_app.py)
sras_align_export.py
-456
View File
@@ -1,456 +0,0 @@
#!/usr/bin/env python3
"""Write an aligned, cropped .sras file from an AlignmentResult.
The alignment machinery in sras_compute never modifies a scan: it produces an
AlignmentResult, and every consumer resamples on the fly (apply_alignment for
the display, reproject_mask for the overlay). That is right for a viewer, but it
means the aligned stack cannot leave the process — no other tool can read it,
and re-opening the scan re-does the registration.
This module bakes an alignment into a new file. Each angle is resampled onto the
shared canvas that AlignmentResult already defines, cropped to the caller's
window, so every output angle ends up with *identical* geometry: same rows, same
frames, same X/Y coordinates. Rotation and translation are gone, absorbed into
where each waveform sits. The result is an ordinary v6 file that opens already
aligned, and registering it against itself returns identity.
Two deliberate choices, both about not inventing data:
* The resample is a nearest-neighbour **gather** of whole waveforms, never an
interpolation. Averaging two neighbouring pixels' CH1 packets would produce
a waveform the instrument never measured, whose FFT peak is not the peak of
either — meaningless for a technique whose entire output is that peak
frequency. So each output pixel gets exactly one source pixel's three
waveforms, verbatim, and the cost is that some source pixels are duplicated
and others dropped. This matches apply_alignment's order=0 for the same
reason.
* Output pixels with no source pixel (the canvas corners a rotated scan cannot
reach, and anything outside the crop's coverage) are filled with the ADC
code for 0 mV, not with zero. See _fill_row.
Depends only on numpy/sras_format/sras_compute — no Qt — so it is directly
unit-testable and importable from a worker thread.
"""
import os
import struct
from dataclasses import dataclass, field
from pathlib import Path
import numpy as np
import sras_compute as compute
from sras_compute import AlignmentResult
from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile, mv_to_adc
# GEO_FMT_V6 stores n_rows as ">H" and n_frames as ">I". A canvas that overflows
# either is not representable, and silently truncating would write a file whose
# geometry table disagrees with its waveform block.
_MAX_ROWS = 0xFFFF
_MAX_FRAMES = 0xFFFFFFFF
# Slack, in source pixels, on the in-bounds test at the very edge of a source
# array. Absorbs the ~1e-13 of float noise an exactly-integer affine picks up
# from being built in mm space; see _in_bounds.
_EDGE_TOL = 1e-6
# Output rows per write() call. One row is n_channels * n_cols * spf bytes —
# ~1.5 MB on a full-size scan — so a handful of rows keeps the peak buffer in
# the low tens of MB no matter how large the scan is.
_ROW_CHUNK = 8
@dataclass
class ExportPlan:
"""What write_aligned_sras would produce, without producing it.
Derived from the affine transforms alone — no waveform bytes are read — so
the wizard can call it on every ROI edit to keep a live size estimate and
per-angle coverage readout in front of the user *before* they commit to a
multi-gigabyte write.
"""
n_rows: int
n_frames: int
n_angles: int
bytes_per_angle: int
total_bytes: int
valid_px: dict[int, int] # output pixels with a source pixel
warnings: list[str] = field(default_factory=list)
def coverage_frac(self, angle_idx: int) -> float:
px = self.n_rows * self.n_frames
return (self.valid_px.get(angle_idx, 0) / px) if px else 0.0
def empty_angles(self) -> list[int]:
"""Angles that would be written as pure padding — no output pixel of
theirs has a source pixel."""
return [a for a in range(self.n_angles) if not self.valid_px.get(a, 0)]
def _src_coords(t, rows, n_cols: int) -> tuple[np.ndarray, np.ndarray]:
"""Fractional source (row, col) coordinates for whole output rows.
*rows* is an array of output row indices; both results are shaped
(len(rows), n_cols).
"""
cols = np.arange(n_cols, dtype=np.float64)
r = np.asarray(rows, dtype=np.float64)[:, None]
sr = t.matrix[0, 0] * r + t.matrix[0, 1] * cols + t.offset[0]
sc = t.matrix[1, 0] * r + t.matrix[1, 1] * cols + t.offset[1]
return sr, sc
def _round_idx(coord: np.ndarray) -> np.ndarray:
"""Nearest source index, rounding halves away from zero.
floor(x + 0.5), not np.rint: scipy.ndimage's order=0 rounds halves away
from zero while np.rint rounds them to even, and these have to be the same
source pixels the apply_alignment(order=0) preview drew. Exact halves are
not a corner case here — the canvas is snapped to the reference angle's own
pixel grid (see canvas_for_params), so an unrotated angle lands on
half-integers wherever its row pitch differs from the reference's.
"""
return np.floor(coord + 0.5).astype(np.int64)
def _in_bounds(sr: np.ndarray, sc: np.ndarray,
n_rows: int, n_frames: int) -> np.ndarray:
"""Which output pixels have a source pixel, by scipy's mode="constant" rule
plus a tolerance at the edge.
Tested on the *fractional* coordinate against the range of sample centres,
[0, n-1] inclusive — deliberately not on the rounded index. The two differ
around the whole rim: a coordinate of -0.4 rounds to a perfectly valid index
0, but scipy calls it out of bounds and writes cval there, so testing the
rounded index would put a one-pixel rim of real data everywhere the Aligned
View shows padding.
_EDGE_TOL is why this is not literally scipy's test. The affine is built
from a chain of mm-space multiplications, so an exactly-integer transform
comes out a few times 1e-13 off (the reference angle's offset lands on
-20 - 7e-15 rather than -20). Bare >= 0.0 then rejects that angle's entire
first row, and <= n-1 its last column — for the *reference* angle, whose
whole role is to pass through as an exact integer crop. The tolerance is
seven orders of magnitude above that noise and seven below the half-pixel
scale at which a rounding decision is ever meaningful, so it can only ever
change pixels whose scipy answer was itself decided by rounding noise.
"""
return ((sr >= -_EDGE_TOL) & (sr <= n_rows - 1 + _EDGE_TOL)
& (sc >= -_EDGE_TOL) & (sc <= n_frames - 1 + _EDGE_TOL))
# Output rows evaluated per numpy call when counting coverage. Counting row by
# row costs one small matmul per row (hundreds of milliseconds per angle on a
# full-size scan, on every ROI edit); counting the whole canvas at once needs
# hundreds of MB of index arrays. Blocking gets both: ~10 numpy calls per angle
# against ~30 MB of live index arrays.
_COUNT_BLOCK = 128
def _count_in_bounds(t, n_rows: int, n_cols: int,
src_rows: int, src_frames: int) -> int:
"""How many of the n_rows x n_cols output pixels have a source pixel."""
total = 0
for start in range(0, n_rows, _COUNT_BLOCK):
rows = np.arange(start, min(start + _COUNT_BLOCK, n_rows))
sr, sc = _src_coords(t, rows, n_cols)
total += int(np.count_nonzero(_in_bounds(sr, sc, src_rows, src_frames)))
return total
def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
"""Geometry, size and per-angle coverage of the file *result* would export.
Coverage is counted from the actual per-pixel index arrays rather than
approximated by the footprint parallelogram's area, because the two differ
exactly where it matters — a crop that clips one angle's scan window — and
this number is what tells the user an angle will come out mostly empty. No
waveform bytes are read, so it stays fast enough to call on every ROI edit.
"""
n_rows, n_cols = result.canvas_shape
n_angles = sras.n_angles
warnings: list[str] = []
valid_px: dict[int, int] = {}
for a in range(n_angles):
t = result.per_angle.get(a)
if t is None:
valid_px[a] = 0
warnings.append(f"Angle {a} has no transform and will be all padding.")
continue
src_rows, src_frames = sras.image_shape(a)
valid_px[a] = _count_in_bounds(t, n_rows, n_cols, src_rows, src_frames)
bytes_per_angle = (n_rows * sras.n_channels * n_cols
* sras.samples_per_frame * sras.bytes_per_sample)
# Built before the remaining warnings so they can be phrased with the
# plan's own coverage_frac rather than a second copy of the same division.
plan = ExportPlan(n_rows=n_rows, n_frames=n_cols, n_angles=n_angles,
bytes_per_angle=bytes_per_angle,
total_bytes=bytes_per_angle * n_angles,
valid_px=valid_px, warnings=warnings)
if n_rows > _MAX_ROWS:
warnings.append(
f"Crop is {n_rows} rows; the .sras geometry table caps rows at "
f"{_MAX_ROWS}. Narrow the ROI in Y.")
if n_cols > _MAX_FRAMES:
warnings.append(f"Crop is {n_cols} frames; the cap is {_MAX_FRAMES}.")
for a in range(n_angles):
frac = plan.coverage_frac(a)
if frac == 0.0:
warnings.append(
f"Angle {a} has no data inside this crop — it will be written "
f"as all padding.")
elif frac < 0.10:
warnings.append(
f"Angle {a} covers only {frac * 100:.1f}% of the crop.")
if sras.background is None:
warnings.append(
"Input has no background waveform (pre-v4 scan); a zero background "
"is written, which makes background subtraction a no-op.")
if sras.version != 6:
warnings.append(f"Input is v{sras.version}; the export is written as v6.")
if sras.scan_aborted:
warnings.append(
f"Input scan was aborted: only its {n_angles} complete angle(s) "
f"are exported.")
return plan
def _fill_row(sras: SrasFile, n_cols: int, dtype) -> np.ndarray:
"""One output row of pure padding, shape (n_channels, n_cols, spf).
Filled per channel with the ADC code for 0 mV, not with 0. Zero ADC decodes
to (0 - yoff) * ymult + yzero, which for a real scope preamble is a long way
from 0 mV — often far enough to sit above the CH4 mask threshold, which
would paint a solid rectangle of "valid" pixels around the sample and make
every DC image and every ROI statistic wrong. Rounding to the integer code
lands within half an ADC step of 0 mV, which is as close as the format can
represent.
"""
info = np.iinfo(dtype)
codes = [int(np.clip(round(mv_to_adc(0.0, *sras.cal(ch))), info.min, info.max))
for ch in range(sras.n_channels)]
row = np.empty((sras.n_channels, n_cols, sras.samples_per_frame), dtype=dtype)
for ch, code in enumerate(codes):
row[ch] = code
return row
class _SourceReader:
"""Gives the gather source rows without ever reading one twice.
This is the difference between a usable export and an unusable one, and it
is entirely about read amplification. A rotated angle maps one output row to
a *diagonal* line across the source array, so the pixels of a single output
row come from hundreds of different source rows — on a full-size scan, a
~1.4 MB source row each. Indexing a memmap pixel by pixel in output order
therefore re-faults nearly the whole angle for every output row: terabytes
of paging for a gigabyte of data.
So rows are served from a contiguous *band* held in RAM. The band for a
chunk of output rows is read in one sequential slice, and because output
rows advance monotonically through the source, consecutive chunks' bands
barely overlap: each source row is read about once, and the whole job costs
roughly 2x the source size in reads rather than a thousand times it.
Small angles skip the machinery — if the whole block fits the budget it is
materialized once and every band is a view of it.
One honest caveat: a chunk whose diagonal spans more source rows than the
budget allows still gets the band it asked for, so the budget can be
overshot. The overshoot is bounded by the span of _ROW_CHUNK output rows,
and in the worst case (an extreme rotation on a huge scan) that is the whole
angle — i.e. no worse than the in-RAM path above. Accepted deliberately:
correctness of the gather is not negotiable, and the alternative is the
memmap thrashing this class exists to avoid.
"""
def __init__(self, sras: SrasFile, angle_idx: int, budget: int):
self._src = sras.data[angle_idx]
self._n_rows = self._src.shape[0]
self._row_bytes = max(1, self._src[0].nbytes)
self._whole = np.asarray(self._src) if self._src.nbytes <= budget else None
# Leave room for the output buffer and the index arrays alongside.
self._max_band = max(1, int(budget * 0.5) // self._row_bytes)
self._band = None
self._lo = self._hi = 0
def band(self, lo: int, hi: int) -> tuple[np.ndarray, int]:
"""Rows [lo, hi) as an in-RAM array, plus the index its row 0 holds."""
lo = max(0, min(lo, self._n_rows))
hi = max(lo + 1, min(hi, self._n_rows))
if self._whole is not None:
return self._whole, 0
if self._band is None or lo < self._lo or hi > self._hi:
# Read a little more than asked so a chunk whose band creeps
# forward by a few rows does not re-read the whole span.
span = min(self._max_band, max(hi - lo, self._max_band // 2))
self._lo = lo
self._hi = min(self._n_rows, lo + span)
if self._hi < hi: # band cannot cover the ask
self._hi = hi
self._band = np.asarray(self._src[self._lo:self._hi])
return self._band, self._lo
def close(self):
self._whole = None
self._band = None
def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
*, progress_cb=None, should_stop=None,
budget: int | None = None) -> Path:
"""Write *sras*, aligned per *result* and cropped to its canvas, to a new
v6 .sras file. Returns the path written.
*result* is used exactly as given: crop the canvas first with
compute.crop_alignment_result, whose offset shift makes the cropped result
resample precisely the window the user selected.
No cache tail is written. Any DC/FFT the input had cached is indexed by the
input's grid and is meaningless on the new one, so — like sras_edit_scans —
the export drops it and lets the viewer recompute.
Writes to a sibling ".part" file and os.replace()s it into position on
success, unlinking it on error or cancellation: a half-written .sras is not
detectably broken (the v6 parser treats a short file as an aborted scan and
opens it happily), so it must never be left where the user might load it.
*should_stop* is polled once per output row chunk; returning True aborts and
raises nothing — the partial file is removed and the returned path will not
exist, so callers must check.
"""
out_path = Path(out_path)
n_rows, n_cols = result.canvas_shape
n_ch, spf = sras.n_channels, sras.samples_per_frame
n_angles = sras.n_angles
if n_rows <= 0 or n_cols <= 0:
raise ValueError(f"empty canvas: {n_rows} x {n_cols}")
if n_rows > _MAX_ROWS:
raise ValueError(
f"{n_rows} rows exceeds the .sras per-angle geometry limit of "
f"{_MAX_ROWS}; crop further in Y")
if n_cols > _MAX_FRAMES:
raise ValueError(f"{n_cols} frames exceeds the limit of {_MAX_FRAMES}")
missing = [a for a in range(n_angles) if a not in result.per_angle]
if missing:
raise ValueError(f"alignment result has no transform for angle(s) {missing}")
# The source's waveform blocks are live read-only memmaps into sras.path,
# so writing over it would corrupt the very reads the gather is making.
if out_path.exists() and out_path.samefile(sras.path):
raise ValueError(
"refusing to export onto the source scan; choose another filename")
dtype = np.dtype(np.int8 if sras.bytes_per_sample == 1 else ">i2")
x0_mm, y0_mm = result.canvas_origin_mm
y_rows = (y0_mm + np.arange(n_rows) * result.canvas_dy_mm).astype(">f4")
# Reference-only header fields. v6/v7 inputs have real ones to carry over;
# for a legacy input describe the canvas we are actually writing.
if sras.x_start_nominal_mm is not None:
nominal = (sras.x_start_nominal_mm, sras.y_start_nominal_mm,
sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
sras.row_spacing_mm)
else:
nominal = (x0_mm, y0_mm,
n_cols * sras.pixel_x_mm, n_rows * result.canvas_dy_mm,
result.canvas_dy_mm)
header = struct.pack(
HDR_FMT_V6, b"SRAS", 6, n_angles,
float(nominal[0]), float(nominal[1]), float(nominal[2]),
float(nominal[3]), float(nominal[4]),
sras.velocity_mm_s, sras.laser_freq_hz, spf, sras.sample_rate_hz,
sras.bytes_per_sample, n_ch)
# Every angle now shares one grid, so the ragged v6 tables collapse to
# n_angles copies of the same record. x_delta is the reference angle's own
# pitch (the canvas is its grid extended), which is velocity/laser_freq
# exactly, so x_axis_mm() stays self-consistent on re-read.
geo = struct.pack(GEO_FMT_V6, float(x0_mm), float(sras.pixel_x_mm),
int(n_cols), int(n_rows)) * n_angles
budget = compute.memory_budget_bytes() if budget is None else max(1, budget)
total_chunks = max(1, n_angles * ((n_rows + _ROW_CHUNK - 1) // _ROW_CHUNK))
done_chunks = 0
cancelled = False
part_path = out_path.with_name(out_path.name + ".part")
try:
with open(part_path, "wb") as fout:
fout.write(header)
fout.write(sras.angles_deg.astype(">f4").tobytes())
fout.write(geo)
fout.write(y_rows.tobytes() * n_angles)
fout.write(sras.encoded_preambles())
fout.write(sras.encoded_background())
pad = _fill_row(sras, n_cols, dtype)
for a in range(n_angles):
t = result.per_angle[a]
reader = _SourceReader(sras, a, budget)
src_rows, src_frames = sras.image_shape(a)
try:
for chunk_start in range(0, n_rows, _ROW_CHUNK):
if should_stop is not None and should_stop():
cancelled = True
break
chunk = np.arange(chunk_start,
min(chunk_start + _ROW_CHUNK, n_rows))
sr, sc = _src_coords(t, chunk, n_cols)
ok = _in_bounds(sr, sc, src_rows, src_frames)
# Clip rather than trust: _EDGE_TOL admits coordinates a
# hair outside the array, and an index off the end here
# would silently read the wrong row of the band.
idx_r = np.clip(_round_idx(sr), 0, src_rows - 1)
idx_c = np.clip(_round_idx(sc), 0, src_frames - 1)
# One band read covers the whole chunk: every source row
# any of these output rows touches, in one sequential
# slice. See _SourceReader.
if ok.any():
band, base = reader.band(int(idx_r[ok].min()),
int(idx_r[ok].max()) + 1)
else:
band, base = None, 0
for i in range(len(chunk)):
out = pad.copy()
keep = ok[i]
if keep.any():
# The two advanced indices are separated by a
# slice, so numpy puts the gathered axis first:
# (n_sel, n_ch, spf). Move it behind channels.
out[:, keep, :] = band[
idx_r[i][keep] - base, :, idx_c[i][keep], :
].transpose(1, 0, 2)
# out is C-contiguous, so the buffer protocol
# writes it straight out — .tobytes() would
# copy a full row per row written.
fout.write(out)
done_chunks += 1
if progress_cb is not None:
progress_cb(int(done_chunks / total_chunks * 100))
finally:
reader.close()
if cancelled:
break
if not cancelled:
fout.flush()
os.fsync(fout.fileno())
if cancelled:
part_path.unlink(missing_ok=True)
return out_path
os.replace(part_path, out_path)
except BaseException:
part_path.unlink(missing_ok=True)
raise
if progress_cb is not None:
progress_cb(100)
return out_path
sras_compute.py
+126 -464
View File
@@ -18,8 +18,9 @@ import numpy as np
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage
from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
adc_to_mv)
from sras_format import (
CH1_IDX, CH3_IDX, CH4_IDX, PAD_FACTOR_MAX, SrasFile, adc_to_mv,
)
# ---------------------------------------------------------------------------
# FFT backend
@@ -285,13 +286,6 @@ _TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
def memory_budget_bytes() -> int:
"""The module-wide ceiling on concurrently-live working buffers, for
callers outside this module that read the same data (the aligned exporter
sizes its source-band reader against it)."""
return _TOTAL_BYTES_BUDGET
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
budget: int = _TOTAL_BYTES_BUDGET) -> int:
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
@@ -428,172 +422,53 @@ def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
*sras.cal(ch_idx))
# ---------------------------------------------------------------------------
# Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing
#
# Averages a pixel's CH1 waveform with its same-row neighbors *before* the
# FFT peak search, to improve SNR on noisy scans. Never crosses rows: pixel
# pitch is strongly anisotropic and varies by scan, but the X pitch within
# one row is a single file-wide constant (SrasFile.pixel_x_mm), so a
# Gaussian in pixel-index distance along a row and one in true physical mm
# distance are the same function up to that constant scale factor — the
# kernel itself needs no pitch, only the GUI's physical-width hint label
# does. Rationale and the masked-average/background-subtraction proofs:
# docs/design.md ("Row-averaged FFT").
# ---------------------------------------------------------------------------
_ROW_AVG_SIGMA_FRAC = 0.5 # sigma = n * this; edge weight (distance n) is
# exp(-1/(2*frac**2)) ~= 0.135 of the center tap
def _row_average_weights(row_avg_n: int) -> np.ndarray:
"""(2n+1,) float32 Gaussian weights for distance-weighted row averaging,
symmetric around the center tap. Not pre-normalized to sum to 1 —
_row_average_waveforms renormalizes per output pixel by the actual sum
of included, valid, in-window neighbor weights, not a fixed total."""
n = max(0, int(row_avg_n))
if n == 0:
return np.ones(1, dtype=np.float32)
d = np.arange(-n, n + 1, dtype=np.float64)
sigma = n * _ROW_AVG_SIGMA_FRAC
return np.exp(-0.5 * (d / sigma) ** 2).astype(np.float32)
def _row_average_waveforms(masked_waves: np.ndarray, valid: np.ndarray,
weights: np.ndarray) -> np.ndarray:
"""Distance-weighted mean of each row position's CH1 waveform with its
same-row neighbors, counting only neighbors where *valid* is True.
masked_waves: (n_frames, spf) float32 — CH1 samples at valid[f]
positions; must already be 0.0 elsewhere (the caller must never
have read the raw memmap at an invalid position).
valid: (n_frames,) bool.
weights: (2n+1,) float32 from _row_average_weights.
Returns (n_frames, spf) float32, meaningful only where valid is True
(the caller compacts by that same mask right after, matching
compute_rf_image's existing masked-compaction contract).
Two 1-D correlations along the frame axis — the numerator against the
raw masked-zero waveforms, the denominator against the validity mask
itself — so a masked neighbor contributes *zero weight* rather than a
zero-amplitude sample at full weight, and a window truncated at a row's
edge renormalizes correctly with no separate edge case: mode="constant",
cval=0.0 pads both convolutions with zero beyond the row's own ends.
"""
num = scipy_ndimage.correlate1d(masked_waves, weights, axis=0,
mode="constant", cval=0.0)
den = scipy_ndimage.correlate1d(valid.astype(np.float32), weights, axis=0,
mode="constant", cval=0.0)
den_safe = np.where(valid, den, np.float32(1.0))
return num / den_safe[:, None]
def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int:
"""The integer zero-padding factor an *n_fft* request represents, or 0
if it represents none — i.e. it is not a whole multiple of this file's
samples_per_frame, so no stored image (which only ever records an
integer factor) can answer it.
0 rather than None so callers can compare it straight against
``sras.precomputed_pad_factor``, which is never 0.
"""
if n_fft is None:
return 1
spf = sras.samples_per_frame
if spf <= 0 or n_fft % spf:
return 0
return max(1, n_fft // spf)
def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
apply_bg_sub: bool, row_avg_n: int) -> list[str]:
"""Why this file's stored FFT images can't answer a request, as human-
readable phrases; empty means they can.
Padding, background subtraction and row-averaging are all baked into
the stored numbers, so a request that differs in any of them has to go
back through a real FFT. This is the single accept rule: cached_rf_image
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.
"""
reasons = []
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"
reasons.append(f"stored at pad {sras.precomputed_pad_factor}x, "
f"requested at {want}")
if sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None):
reasons.append("stored with background subtraction "
f"{'on' if sras.precomputed_bg_sub else 'off'}")
if sras.precomputed_row_avg_n != row_avg_n:
have = (f"row-averaged (n={sras.precomputed_row_avg_n})"
if sras.precomputed_row_avg_n else "raw per-pixel")
want = (f"row-averaged (n={row_avg_n})"
if row_avg_n else "raw per-pixel")
reasons.append(f"stored {have}, requested {want}")
return reasons
def cached_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
n_fft: int | None = None,
dc4_mv: np.ndarray | None = None,
row_avg_n: int = 0,
allow_dc_recompute: bool = True) -> np.ndarray | None:
"""The precomputed-cache fast path for compute_rf_image: a ready-to-
display peak-frequency image if this file already has one matching
every setting the caller cares about, else None (caller must run a
real FFT).
"""The stored peak-frequency image for one angle (v5 PREC or v7 CACH),
masked and ready to display — or None if no stored image can serve these
settings and a real FFT is needed.
A stored image stands in only when *all* hold: this angle actually has
a stored image (v5 PREC or v7 CACH); the requested zero-padding matches
what the store was computed at (see below); the stored bg-sub flag
matches what the caller wants; and sras.precomputed_row_avg_n ==
row_avg_n exactly (0 == "raw") — this last check is what stops a raw
request from ever being silently served a row-averaged image, or vice
versa, or a request at one window size being served a cache stored at a
different one.
Stored images are unmasked, so they serve any threshold — but they carry
one bg-sub state and one zero-padding, and a view asking for either of the
others is asking for different numbers, not a different rendering of the
same ones. Padding is the sharp edge: a pad-10 view resolves peaks on a
10x finer bin grid, so handing back a pad-1 image would quietly undo the
setting. v5 PREC and CACH v1 tails record no pad and are natural
resolution by construction.
Padding is checked the same way, and for the same reason: a padded FFT
interpolates between the natural bins, so it resolves genuinely
different peak frequencies. *n_fft* of None means natural resolution,
i.e. pad 1; anything else must be an exact whole multiple of
samples_per_frame equal to sras.precomputed_pad_factor. A ragged n_fft
that is not such a multiple can never match a stored image, since the
store only ever records an integer pad factor.
If *allow_dc_recompute* is False and no DC4 image is already cached or
supplied via *dc4_mv*, applying the mask would mean reading a whole
channel on the caller's behalf; this returns None instead so a caller
that wants to stay off the I/O path (e.g. a GUI thread) can choose to
fall through to a real compute rather than block.
Cheap enough to call on the GUI thread — a copy plus a comparison —
*except* when the CH4 mask has to be read from disk to build it. Pass
*allow_dc_recompute* False to return None in that case instead, and leave
the whole-channel read to a worker.
"""
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):
# n_fft None means "no padding", i.e. exactly samples_per_frame points.
want_n_fft = n_fft if n_fft is not None else sras.samples_per_frame
if (cached_freq is None
or want_n_fft != sras.precomputed_pad_factor * sras.samples_per_frame
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)):
return None
dc4_img = None
if dc_threshold_mv is not None:
# DC4 mask, in priority order: already-cached DC block, caller-
# supplied image, or a fresh (cheap — no FFT) recompute. Resolved
# before the copy below so the allow_dc_recompute bail-out doesn't
# allocate a full image it is about to throw away.
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
if dc4_img is None:
if dc4_mv is not None:
dc4_img = dc4_mv
elif allow_dc_recompute:
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
*sras.cal(CH4_IDX))
else:
return None
if dc_threshold_mv is None:
return cached_freq.copy()
# DC4 mask, in priority order: stored DC block, caller-supplied image,
# or a fresh (cheap — no FFT) recompute.
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
if dc4_img is None:
dc4_img = dc4_mv
if dc4_img is None:
if not allow_dc_recompute:
return None
dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
*sras.cal(CH4_IDX))
freq_img = cached_freq.copy()
if dc4_img is not None:
freq_img[dc4_img < dc_threshold_mv] = 0.0
freq_img[dc4_img < dc_threshold_mv] = 0.0
return freq_img
@@ -605,8 +480,7 @@ 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:
exact: bool = False) -> 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
@@ -631,25 +505,22 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
tests, audits, and the SRAS_FFT_EXACT=1 escape hatch, and is slow and
memory-hungry at high pad.
*row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n
same-row neighbors (distance-weighted, valid-neighbors-only per the
same dc_threshold_mv mask) before the FFT runs — see
_row_average_waveforms. 0 (default) is the raw, unaveraged behavior.
Fast path: if the file has a precomputed peak-frequency image for this
angle (v5 PREC or v7 CACH) 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.
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
matches, the stored image is used directly — no FFT is run. That test is
cached_rf_image, which the viewer also applies before it ever dispatches
a compute at all.
"""
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub,
n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n)
if fast is not None:
return fast
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
cached = cached_rf_image(sras, angle_idx, dc_threshold_mv,
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
dc4_mv=dc4_mv)
if cached is not None:
return cached
# ---- Chunked FFT path --------------------------------------------------
n_rows, n_frames = sras.image_shape(angle_idx)
data = sras.data[angle_idx]
exact = exact or _FFT_EXACT_ENV
spf = sras.samples_per_frame
n_len = n_fft if n_fft is not None else spf
@@ -657,18 +528,11 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
img = np.zeros((n_rows, n_frames), dtype=np.float32)
background = sras.background if (apply_bg_sub and sras.background is not None) else None
cal4 = sras.cal(CH4_IDX)
row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None
zp = (_zoom_plan(spf, n_fft)
if not exact and n_fft is not None and n_fft >= _ZOOM_MIN_PAD * spf
else None)
total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
if row_avg_n > 0:
# One extra same-sized transient buffer (the pre-averaging full-row
# scratch array) is live per in-flight row; halve the budget so
# _plan_fft_rows/the exact-path sizing accounts for it rather than
# relying on _plan_fft_rows's existing 2x slack to happen to cover it.
total = max(1, total // 2)
cap = max_workers if max_workers is not None else _MAX_WORKERS
if exact:
# The reference path materialises the full padded spectrum, so rows
@@ -703,31 +567,13 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
def read_row(i: int):
if should_stop is not None and should_stop():
return
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so only
# the selected elements are actually read from disk; masked-out
# pixels' pages are never paged in at all.
raw = data[r0 + i, CH1_IDX]
dst = waves[offs[i]:offs[i + 1]]
row_valid = valid[i] if valid is not None else None
if row_avg_weights is not None:
v = row_valid if row_valid is not None else np.ones(n_frames, dtype=bool)
# A plain cast + broadcast multiply, not a boolean-indexed
# scatter into a zeroed buffer: numpy's fancy indexing holds
# the GIL for its whole duration (measured zero speedup
# across threads, and net negative once threads outnumber
# physical cores), while a cast and a multiply are ordinary
# ufuncs that release it — this is what lets read_row actually
# parallelize across the pool instead of serializing on the
# scatter/gather. Trade-off: every sample in the row is read,
# valid or not (row averaging needs broad neighbor context
# regardless, unlike the plain path below, which still skips
# masked-out pixels entirely).
full = raw.astype(np.float32) * v[:, None]
avg = _row_average_waveforms(full, v, row_avg_weights)
dst[:] = avg[v] if row_valid is not None else avg
else:
# Index the raw memmap slice with the boolean mask *before*
# converting dtype — this is a lazy view until touched, so
# only the selected elements are actually read from disk;
# masked-out pixels' pages are never paged in at all.
dst[:] = raw[row_valid] if row_valid is not None else raw
dst[:] = raw[valid[i]] if valid is not None else raw
if background is not None:
dst -= background # background is 1-D (spf,)
@@ -788,40 +634,23 @@ def compute_rf_image(sras: SrasFile, angle_idx: int,
def cache_file(path: str, mode: str, apply_bg_sub: bool,
fft_backend: str = "scipy", max_workers: int = 0,
pad_factor: int = 1,
dc_threshold_mv: float | None = None,
row_avg_n: int = 0) -> str:
pad_factor: int = 1) -> 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.
FFT backend, worker cap and pad factor are 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
averaging before the FFT — see compute_rf_image's row_avg_n).
fft_rowavg needs *dc_threshold_mv* up front, unlike plain "fft":
neighbor validity is baked into the stored numbers, so it can't be
deferred to display time the way plain masking can.
*pad_factor* is the zero-padding factor to resolve the peaks at: 1 (the
default) is natural resolution, n_fft == samples_per_frame. It is
recorded in the SFFT block so a reader knows which views the stored
numbers answer for — and it has to be the pad the viewer is *actually*
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).
*pad_factor* is stored alongside the images: a cache is only usable by a
view asking for the same padding, so caching at the viewer's own setting
is the difference between a batch that pays off and one that can never be
read back (see cached_rf_image).
"""
global _MAX_WORKERS
try:
if mode not in ("dc", "fft", "fft_rowavg"):
return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')"
# write_v7_cache enforces this bound too, but only once every angle's
# FFT has already been computed. Checking up front is the difference
# 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}"
if mode not in ("dc", "fft"):
return f"unknown cache mode {mode!r} (expected 'dc' or 'fft')"
set_fft_backend(fft_backend)
if max_workers:
_MAX_WORKERS = max_workers
@@ -832,7 +661,6 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
"can be batch-cached")
n = sras.n_angles
n_fft = sras.samples_per_frame * pad_factor if pad_factor > 1 else None
n_workers, angle_budget = plan_angle_level(sras)
if mode == "dc":
dc3 = _parallel_map(
@@ -846,8 +674,10 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
budget=angle_budget), *sras.cal(CH4_IDX)),
range(n), n_workers)
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
elif mode == "fft":
else:
effective_bg = apply_bg_sub and sras.background is not None
pad = max(1, min(PAD_FACTOR_MAX, int(pad_factor)))
n_fft = None if pad <= 1 else sras.samples_per_frame * pad
# dc_threshold_mv=None: store unmasked images and mask at display
# time (same convention as v5's PREC block). Skipping the mask
# also skips reading CH4 entirely. The FFT path parallelises
@@ -856,25 +686,8 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
apply_bg_sub=effective_bg, n_fft=n_fft)
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)
else: # "fft_rowavg"
if row_avg_n <= 0:
return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode"
if dc_threshold_mv is None:
return ("fft_rowavg mode requires a DC threshold "
"(neighbor validity depends on it)")
effective_bg = apply_bg_sub and sras.background is not None
freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv,
apply_bg_sub=effective_bg, n_fft=n_fft,
row_avg_n=row_avg_n)
for a in range(n)]
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
new_row_avg_n=row_avg_n,
new_pad_factor=pad_factor)
new_pad_factor=pad)
return ""
except Exception as exc:
return str(exc)
@@ -1317,22 +1130,13 @@ def default_max_workers() -> int:
def _rotation_candidates(nominal_deg: float, search_deg: float,
step_deg: float,
signs: tuple[int, ...] = (-1, 1)) -> list[float]:
"""Coarse rotation candidates: a window around each requested sign of the
stage's reported angle change. Scoring both signs (the default) is what
makes the stage's sign convention a non-issue — the images decide which way
the stage turns, and a file whose stage reports the opposite sense
registers just as well.
*signs* exists only so a user who has established which way their own stage
turns can halve the coarse sweep from the alignment wizard. It is an
override, not an inference: nothing in the file says which sign is right.
"""
step_deg: float) -> list[float]:
"""Coarse rotation candidates: a window around *both* signs of the stage's
reported angle change. Scoring both is what makes the stage's sign
convention a non-issue — the images decide which way the stage turns, and
a file whose stage reports the opposite sense registers just as well."""
out: list[float] = []
step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below
for sign in signs:
center = sign * nominal_deg
for center in (-nominal_deg, nominal_deg):
k = int(np.floor(search_deg / step_deg))
for i in range(-k, k + 1):
out.append(center + i * step_deg)
@@ -1442,18 +1246,15 @@ def register_angle_to_reference(
dc_threshold_mv: float = 0.0,
sources: tuple[str, ...] = ("signal", "mask"),
coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
search_deg: float = 6.0, coarse_step_deg: float = 2.0,
seed_deg: float | None = None,
seed_signs: tuple[int, ...] = (-1, 1),
refine: bool = True) -> RigidFit:
search_deg: float = 6.0, coarse_step_deg: float = 2.0) -> RigidFit:
"""Rigid (rotation + translation, never scale) fit of *angle_idx* onto
*ref_angle_idx*, found entirely by cross-correlating image content.
Two stages:
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
requested signs of the stage's reported angle change (see
_rotation_candidates), for each requested source image, scored by
_overlap_ncc. Only has to pick the right basin.
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around both
signs of the stage's reported angle change (see _rotation_candidates),
for each requested source image, scored by _overlap_ncc. Only has to
pick the right basin.
2. Hill-climbing refinement of the winning (source, rotation) at
*fine_dim* with sub-pixel translation folded into every score
(_refine_rotation, _score_rotation), down to 0.05°.
@@ -1462,22 +1263,6 @@ def register_angle_to_reference(
local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
caller can surface so a bad scan is visible rather than silently fused in.
Returns identity for the reference angle itself.
The last three arguments only exist to let the alignment wizard expose the
rotation search; every default reproduces the search this function has
always done.
*seed_deg* replaces the stage's reported angle change as the center of the
coarse sweep — the pre-rotation the search starts from. None (the default)
means the stage angle from the file's Angle Table, which is nearly always
what you want: it puts the sweep within a couple of degrees of the answer.
Pass 0.0 to search around no rotation at all, which is the honest choice for
a file whose stage angles are known to be wrong.
*refine=False* stops after the coarse sweep, leaving rotation on the
*coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
it pins rotation to exactly the seed and searches translation only — for a
scan whose stage angles are trusted more than its image content.
"""
if angle_idx == ref_angle_idx:
return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
@@ -1487,10 +1272,8 @@ def register_angle_to_reference(
if not candidates:
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
if seed_deg is None else float(seed_deg))
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
seed_signs)
nominal = nominal_delta_deg(sras, angle_idx, ref_angle_idx)
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg)
# ---- Stage 1: coarse sweep, every source ------------------------------
pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
@@ -1519,10 +1302,9 @@ def register_angle_to_reference(
theta = best[1]
score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
theta, subpixel=True)
if refine:
theta, score, shift = _refine_rotation(
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
step_deg=coarse_step_deg)
theta, score, shift = _refine_rotation(
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
step_deg=coarse_step_deg)
dr, dc = shift
return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
@@ -1534,17 +1316,24 @@ def register_angle_to_reference(
def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
pitch_mm: tuple[float, float],
per_angle_params: dict[int, ManualAngleParams],
*, margin_frac: float = 0.0, snap: bool = True
) -> tuple[tuple[float, float], tuple[int, int]]:
"""Shared-canvas origin (stage mm) and (n_rows, n_cols) at *pitch_mm* that
contains every angle's footprint after its own rigid transform. Angles
missing from per_angle_params default to identity (e.g. a sidecar saved
before a rescan added more angles).
The canvas grid is aligned with the reference angle's own pixel grid, so
the reference lands on integer canvas pixels and is resampled by an exact
integer translation — the concrete meaning of "the canvas carries angle
0's X/Y coordinates". This requires pitch_mm to be the reference's own
pitch.
snap=True aligns the canvas grid with the reference angle's own pixel grid,
so the reference lands on integer canvas pixels and is resampled by an
exact integer translation — the concrete meaning of "the canvas carries
angle 0's X/Y coordinates". It requires pitch_mm to be the reference's own
pitch; the manual-alignment preview passes a coarser pitch and snap=False.
margin_frac pads the box on every side: 0 for a final canvas, nonzero for
ManualAlignmentDialog's preview canvas, which needs headroom so an ordinary
translation nudge never has to trigger a full canvas resize (an extreme
nudge can still push content past this padding; accepted, and cheap to
recover from by re-opening the dialog).
"""
dx, dy = pitch_mm
corners = np.vstack([
@@ -1555,17 +1344,27 @@ def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
for a in range(sras.n_angles)])
x_min, y_min = corners.min(axis=0)
x_max, y_max = corners.max(axis=0)
if margin_frac:
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
x_min, x_max = x_min - pad_x, x_max + pad_x
y_min, y_max = y_min - pad_y, y_max + pad_y
center = ref_center_mm(sras, ref_angle_idx)
# Express the box in the reference's own pixel indices and grow it
# outward to whole pixels, so canvas index k lands exactly where the
# reference's own pixel (k + const) does.
cy, cx = _center_idx(sras, ref_angle_idx)
cols = sorted((x_min / dx + cx, x_max / dx + cx))
rows = sorted((y_min / dy + cy, y_max / dy + cy))
col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
shape = (row1 - row0 + 1, col1 - col0 + 1)
if snap:
# Express the box in the reference's own pixel indices and grow it
# outward to whole pixels, so canvas index k lands exactly where the
# reference's own pixel (k + const) does.
cy, cx = _center_idx(sras, ref_angle_idx)
cols = sorted((x_min / dx + cx, x_max / dx + cx))
rows = sorted((y_min / dy + cy, y_max / dy + cy))
col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
shape = (row1 - row0 + 1, col1 - col0 + 1)
else:
origin_ref = np.array([x_min, y_min if dy > 0 else y_max])
shape = (int(np.ceil((y_max - y_min) / abs(dy))) + 1,
int(np.ceil((x_max - x_min) / dx)) + 1)
origin_stage = origin_ref + center
return (float(origin_stage[0]), float(origin_stage[1])), shape
@@ -1619,7 +1418,7 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
thread on every manual edit."""
pitch = pixel_pitch_mm(sras, ref_angle_idx)
canvas_origin_mm, canvas_shape = canvas_for_params(
sras, ref_angle_idx, pitch, params)
sras, ref_angle_idx, pitch, params, snap=True)
extra = extra or {}
per_angle: dict[int, AngleTransform] = {}
@@ -1636,144 +1435,6 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
pitch[0], pitch[1], canvas_origin_mm, per_angle)
def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
n_rows: int, n_cols: int) -> AlignmentResult:
"""The same alignment restricted to a rectangular window of its canvas —
canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).
Cropping folds into each angle's existing affine instead of becoming a
second transform, because a canvas crop is *pure index translation*. From
_affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
origin, so:
matrix @ [r', c'] + (offset + matrix @ [row0, col0])
== matrix @ [r' + row0, c' + col0] + offset
i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
mapping at the shifted indices, exactly. That equality is what lets
apply_alignment, reproject_mask and the aligned .sras exporter all keep
working on a cropped result with no special-casing — and it is why the crop
preview a user approves is guaranteed to be the same pixels the exporter
writes.
Bounds are the caller's responsibility (the wizard's ROI page clamps to the
canvas): an out-of-range window is geometrically well-defined here and
simply resamples padding.
"""
if n_rows <= 0 or n_cols <= 0:
raise ValueError(f"empty crop: {n_rows} x {n_cols}")
delta = np.array([float(row0), float(col0)])
per_angle = {
a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
t.offset + t.matrix @ delta, t.score, t.source)
for a, t in result.per_angle.items()
}
origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
(int(n_rows), int(n_cols)),
result.canvas_dx_mm, result.canvas_dy_mm,
origin, per_angle)
def coarse_rect_to_canvas(rect: tuple[int, int, int, int],
downsample: tuple[int, int],
canvas_shape: tuple[int, int],
*, inset_blocks: int = 0
) -> tuple[int, int, int, int]:
"""A rectangle in coarse (block-mean preview) indices as canvas pixels,
both as (row0, col0, n_rows, n_cols) and clamped to the canvas.
The single place the preview grid's relation to the real canvas is
written down: the coarse grid samples canvas pixels 0, f, 2f, …, so a
coarse pixel stands for a whole block and every caller has to agree on
which canvas pixels that block means, or a crop the user approved on the
preview lands a few pixels off in the export.
*inset_blocks* shrinks the rectangle by that many coarse blocks on each
side. A coarse pixel reported as fully covered stands for a block whose
far edge may not be, so "fit to full overlap" insets by 1 to stay honestly
inside the overlap region; a union rectangle insets by 0 because it wants
to contain the region rather than fit inside it.
"""
row0, col0, nr, nc = rect
fy, fx = downsample
n_rows, n_cols = canvas_shape
r0 = min(n_rows - 1, (row0 + inset_blocks) * fy)
c0 = min(n_cols - 1, (col0 + inset_blocks) * fx)
r1 = min(n_rows, (row0 + nr - inset_blocks) * fy)
c1 = min(n_cols, (col0 + nc - inset_blocks) * fx)
return r0, c0, max(1, r1 - r0), max(1, c1 - c0)
def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
"""Summarize a per-pixel "how many angles cover this pixel" image — the
number the alignment wizard's mask-stack view is colored by.
Separated from the drawing code because it is the actual judgement the user
makes on that screen ("do the angles land on top of each other?"), and a
plain array-in/dict-out function can be tested without Qt.
"""
counts = np.asarray(counts)
union = int(np.count_nonzero(counts))
full = int(np.count_nonzero(counts >= n_angles))
return {
"union_px": union,
"full_px": full,
"full_frac": (full / union) if union else 0.0,
"mean_count": float(counts[counts > 0].mean()) if union else 0.0,
"max_count": int(counts.max()) if counts.size else 0,
"empty": union == 0,
}
def largest_rect_at_least(counts: np.ndarray, min_count: int
) -> tuple[int, int, int, int] | None:
"""Largest axis-aligned rectangle whose every pixel has counts >= min_count,
as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.
Backs the wizard's "fit to full overlap" button. A *bounding box* of the
qualifying pixels would be the obvious thing and is wrong: the full-overlap
region of several rotated scans is roughly a disc, whose bounding box has
corners no angle covers at all. Offering that as the crop would hand the
user padding they explicitly asked to avoid, so this finds a rectangle that
is entirely inside the region.
Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
preview-sized array, so it is instant at interactive rates.
"""
good = np.asarray(counts) >= min_count
if not good.any():
return None
n_rows, n_cols = good.shape
best = (0, 0, 0, 0) # area, row0, col0, ...
best_area = 0
heights = np.zeros(n_cols, dtype=np.int64)
for r in range(n_rows):
heights = np.where(good[r], heights + 1, 0)
# Sentinel column of height 0 flushes the stack at the end of the row.
stack: list[tuple[int, int]] = [] # (start col, height)
for c in range(n_cols + 1):
h = int(heights[c]) if c < n_cols else 0
start = c
while stack and stack[-1][1] >= h:
s, sh = stack.pop()
area = sh * (c - s)
if area > best_area:
best_area = area
best = (s, sh, c - s, r)
start = s
if h:
stack.append((start, h))
col0, height, width, row_end = best
return (row_end - height + 1, col0, height, width)
def _parallel_map(fn, items, n_workers: int) -> list:
"""fn over items, in order, threaded when it pays."""
items = list(items)
@@ -1789,7 +1450,7 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
"""Top-level alignment driver: register every angle onto *ref_angle_idx* by
content, then lay them all out on that angle's own coordinate grid.
Meant for a background thread — deliberately
Runs on a background thread (see AngleAlignmentWorker) — deliberately
recomputes CH4 DC images from scratch rather than reading the GUI-thread
_dc_cache dict, since background-thread workers must not touch
GUI-thread-owned caches.
@@ -1852,9 +1513,9 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
# with no per-pixel image work at all, so build_manual_alignment is cheap enough
# to call synchronously on the GUI thread on every edit. The only genuinely
# expensive per-pixel operation anywhere in this flow is reproject_mask, and
# only the wizard's own downsampled mask stack calls that per keystroke
# — see AlignmentWizard.rebuild_stack for how it limits a nudge to reprojecting
# only the actively-edited angle.
# only ManualAlignmentDialog's own downsampled preview calls that per keystroke
# — see that class's docstring for how it limits each nudge to reprojecting only
# the actively-edited angle.
# ---------------------------------------------------------------------------
def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
@@ -1865,8 +1526,9 @@ def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
canvas_shape: tuple[int, int],
*, src_downsample: tuple[int, int] = (1, 1)) -> np.ndarray:
"""Resample one angle's binary/float mask onto an arbitrary canvas via an
explicit rotation+shift — the single building block the alignment wizard's
live mask stack repeatedly calls (once per angle per edit). *src_downsample* must match the (rows, cols)
explicit rotation+shift — the single building block ManualAlignmentDialog's
live preview repeatedly calls (once per keystroke, for only the
actively-nudged angle). *src_downsample* must match the (rows, cols)
block-mean factor already applied to *mask*, or the reprojection lands at
the wrong scale. order=0 (nearest) matches apply_alignment's own reasoning:
a binary mask must never be blended with zero-padding."""
@@ -1979,7 +1641,7 @@ def delete_manual_alignment(sras: SrasFile) -> bool:
"""Delete the sidecar if present. Returns whether a file actually existed
to delete, so Clear Alignment's status message can say so. Genuine I/O
errors (permission denied, read-only share) propagate — the caller
(the wizard's Clear path) surfaces them rather than silently
(ManualAlignmentDialog._on_clear) surfaces them rather than silently
pretending the destructive action succeeded."""
path = sidecar_path(sras.path)
try:
sras_edit_scans.py
-5
View File
@@ -14,11 +14,6 @@ Handles v2-v7. Any precomputed FFT/DC cache (v5 PREC tail, v7 CACH tail) is
dropped on write, since it's indexed by angle and would be stale/misaligned
after renumbering; the viewer just recomputes it next time the file opens.
This tool only ever *drops* angles — every kept angle's waveform bytes and
geometry are carried across verbatim. To write a file whose angles have been
resampled onto one shared aligned grid and cropped, use the viewer's
Fusion -> Alignment Wizard (sras_align_export.py) instead.
Usage:
python sras_edit_scans.py input.sras --list
python sras_edit_scans.py input.sras output.sras --drop 2,5
sras_format.py
+37 -168
View File
@@ -60,15 +60,12 @@ 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
# _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
# predate padded caching, so both
# are natural-resolution (pad 1).
# v2 added the SFFT pad factor. Writers always emit v2; readers still accept
# v1, whose FFT block is natural-resolution by construction. An older reader
# meeting a v2 tail rejects the whole section on the version check and
# recomputes — no cache, but no misreading either.
CACH_VERSION = 2
CACH_VERSION_MIN = 1
CACH_FLAG_DC = 0x01
CACH_FLAG_FFT = 0x02
@@ -77,19 +74,12 @@ SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT)
SFFT_MAGIC = b"SFFT"
SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored (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_SIZE_V1 = struct.calcsize(SFFT_HDR_FMT_V1)
SFFT_HDR_SIZE_V2 = struct.calcsize(SFFT_HDR_FMT_V2)
SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored
SFFT_HDR_FMT = ">4sBHH" # magic, flags, n_stored, pad_factor
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
MAX_PAD_FACTOR = 0xFFFF # the H field above
SFFT_FLAG_BG_SUB = 0x01
SFFT_FLAG_ROW_AVG = 0x02 # peak_freq_mhz came from same-row,
# distance-weighted averaged CH1
# waveforms, not raw per-pixel ones;
# row_avg_n is the neighbor half-width
# (pixels) used. Bits 2-7 reserved.
# The viewer clamps its pad factor to this, and the field is a uint16.
PAD_FACTOR_MAX = 256
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
@@ -187,10 +177,9 @@ class SrasFile:
``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``,
always as ragged per-angle lists (``list[np.ndarray | None]``, one entry
per angle, ``None`` where that angle was never stored) regardless of
source version. The 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.
source version. The settings the FFT images were computed under travel
with them as ``precomputed_bg_sub`` / ``precomputed_pad_factor``, since a
stored image is only usable by a view asking for the same two.
"""
def __init__(self, path: str):
@@ -247,55 +236,10 @@ class SrasFile:
self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles
self.precomputed_bg_sub: bool = False
self.precomputed_row_avg_n: int = 0
# Zero-padding factor the stored peak_freq_mhz images were resolved
# at: 1 = natural resolution (n_fft == samples_per_frame). A padded
# FFT resolves peaks a padded view would, and only such a view can
# be served from it — see sras_compute.cached_rf_image.
# Zero-padding the stored FFT images were computed at; 1 = natural
# resolution, which is all a v5 PREC or CACH v1 tail can hold.
self.precomputed_pad_factor: int = 1
def encoded_preambles(self) -> bytes:
"""This file's Preamble Blocks section, as bytes a writer can emit.
v6/v7 files kept the on-disk span verbatim, which is both cheaper and
lossless; legacy files did not keep it, and a v2 file has no preambles
at all, so those are re-encoded from the parsed strings (empty ones for
v2). Empty is not a silent downgrade: _parse_preamble("") returns {} and
_set_calibration falls back to the hardcoded scope constants, which is
exactly the calibration a v2 file already gets, so mV values round-trip
unchanged.
Lives here rather than at each writer so the version fan-out sits next
to the parser that creates it, and no writer has to probe the object
to find out which shape it got.
"""
raw = getattr(self, "preambles_raw", None)
if raw is not None:
return raw
out = bytearray()
for s in self.preambles or [""] * self.n_channels:
encoded = s.encode("utf-8")
out += struct.pack(">H", len(encoded)) + encoded
return bytes(out)
def encoded_background(self) -> bytes:
"""This file's Background Block, as bytes a writer can emit.
When there is none (v2/v3), this is samples_per_frame zeros rather than
a zero-length block. Every consumer guards on `background is not None`
and then subtracts it from a (spf,)-shaped row, so a length-0 array
would broadcast-fail at the first background-subtracted FFT; zeros make
the subtraction a correct no-op instead.
"""
raw = getattr(self, "background_raw", None)
if raw is not None:
return raw
if self.background is None:
samples = np.zeros(self.samples_per_frame, dtype=np.int8)
else:
samples = np.rint(self.background).astype(np.int8)
return struct.pack(">I", samples.size) + samples.tobytes()
def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""A stored DC image (already in mV) for (angle, channel), or None."""
store = self.precomputed_dc3_mv if ch_idx == CH3_IDX else self.precomputed_dc4_mv
@@ -326,15 +270,6 @@ class SrasFile:
self.scan_aborted = False
self.n_angles_declared = n_angles
# Pre-v6 files carry no nominal ROI. Defined as None rather than
# left absent so the object's shape does not depend on its version
# and writers can ask instead of probing with hasattr.
self.x_start_nominal_mm = None
self.y_start_nominal_mm = None
self.x_delta_nominal_mm = None
self.y_delta_nominal_mm = None
self.row_spacing_mm = None
angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32)
@@ -574,16 +509,18 @@ class SrasFile:
return end
def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
stores: list[list]) -> int | None:
stores: list[list]) -> tuple | None:
"""Read one CACH sub-block header, then its per-angle image entries
into *stores* (one list per image the block stores per angle).
Returns the header's flags byte, or None if the block is malformed.
Every sub-block header is (magic, flags, n_stored, *extras). Returns
everything after the magic, so a caller that has extras knows how to
read them, or None if the block is malformed.
"""
raw = f.read(struct.calcsize(hdr_fmt))
if len(raw) < struct.calcsize(hdr_fmt):
return None
block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
block_magic, flags, n_stored, *extras = struct.unpack(hdr_fmt, raw)
if block_magic != magic:
return None
for _ in range(n_stored):
@@ -593,41 +530,7 @@ class SrasFile:
shape = self.image_shape(angle_idx)
for store in stores:
store[angle_idx] = _read_f32_image(f, shape)
return flags
def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int, int] | 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).
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.
"""
hdr_fmt = {1: SFFT_HDR_FMT_V1, 2: SFFT_HDR_FMT_V2}.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
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:
magic, flags, n_stored, row_avg_n, pad_factor = struct.unpack(hdr_fmt, raw)
if magic != SFFT_MAGIC:
return None
for _ in range(n_stored):
(angle_idx,) = _read_struct(f, ">H")
if angle_idx >= self.n_angles:
break
self.precomputed_freq_mhz[angle_idx] = _read_f32_image(
f, self.image_shape(angle_idx))
return flags, row_avg_n, pad_factor
return (flags, *extras)
def _parse_cach_section(self, offset: int):
"""Parse the v7 CACH tail that holds precomputed DC/FFT images."""
@@ -637,7 +540,8 @@ class SrasFile:
if len(header_raw) < CACH_HDR_SIZE:
return
magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
if magic != CACH_MAGIC or cach_version not in CACH_VERSIONS_READABLE:
if (magic != CACH_MAGIC
or not CACH_VERSION_MIN <= cach_version <= CACH_VERSION):
return
if block_flags & CACH_FLAG_DC:
@@ -647,20 +551,22 @@ class SrasFile:
return
if block_flags & CACH_FLAG_FFT:
result = self._read_sfft_block(f, cach_version)
if result is None:
# v1 has no pad field: those images are natural-resolution.
hdr_fmt = SFFT_HDR_FMT if cach_version >= 2 else SFFT_HDR_FMT_V1
header = self._read_cache_block(
f, hdr_fmt, SFFT_MAGIC, [self.precomputed_freq_mhz])
if header is None:
return
flags, row_avg_n, pad_factor = result
flags, *extras = header
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)
self.precomputed_pad_factor = (
min(max(1, extras[0]), PAD_FACTOR_MAX) if extras else 1)
def write_v7_cache(self, *,
new_dc3_mv: list[np.ndarray | None] | None = None,
new_dc4_mv: list[np.ndarray | None] | None = None,
new_freq_mhz: list[np.ndarray | None] | None = None,
new_bg_sub: bool | None = None,
new_row_avg_n: int | None = None,
new_pad_factor: int | None = None):
"""Store computed DC and/or FFT images into this file's CACH tail,
in place, converting a v6 source to v7 (or updating an existing v7
@@ -669,18 +575,6 @@ class SrasFile:
``SrasFile`` already has in memory (from parsing, or a prior write
in this same session) — its bytes are never re-read from disk.
*new_row_avg_n* is the same-row neighbor half-width (pixels) the
passed *new_freq_mhz* was averaged over before its FFT, 0 for a raw
(unaveraged) compute — carried forward like *new_bg_sub* when None.
It describes the whole stored FFT block, not per-angle, mirroring
how bg-sub has never been tracked per-angle either.
*new_pad_factor* is the zero-padding factor the passed *new_freq_mhz*
was resolved at (1 = natural resolution), carried forward the same
way. Like row_avg_n it is provenance, not a hint: a view at a
different pad resolves different peaks, so recording it is what lets
a reader refuse the cache instead of showing the wrong numbers.
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.
@@ -693,15 +587,11 @@ class SrasFile:
final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv
final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz
final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub
final_row_avg_n = (new_row_avg_n if new_row_avg_n is not None
else self.precomputed_row_avg_n)
final_pad_factor = (new_pad_factor if new_pad_factor is not None
else self.precomputed_pad_factor)
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):
final_pad = (new_pad_factor if new_pad_factor is not None
else self.precomputed_pad_factor)
if not 1 <= final_pad <= PAD_FACTOR_MAX:
raise ValueError(
f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {final_pad_factor}")
f"pad factor {final_pad} outside 1..{PAD_FACTOR_MAX}")
dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]
@@ -721,10 +611,8 @@ class SrasFile:
if fft_entries:
fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
fft_flags |= SFFT_FLAG_ROW_AVG if final_row_avg_n else 0
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags,
len(fft_entries), final_row_avg_n,
final_pad_factor)
len(fft_entries), final_pad)
for a in fft_entries:
payload += struct.pack(">H", a)
payload += final_freq[a].astype(">f4").tobytes()
@@ -752,8 +640,7 @@ class SrasFile:
self.precomputed_dc4_mv = final_dc4
self.precomputed_freq_mhz = final_freq
self.precomputed_bg_sub = final_bg_sub
self.precomputed_row_avg_n = final_row_avg_n
self.precomputed_pad_factor = final_pad_factor
self.precomputed_pad_factor = final_pad
# ------------------------------------------------------------------
# Axes helpers
@@ -770,24 +657,6 @@ class SrasFile:
def y_positions_mm(self, angle_idx: int) -> np.ndarray:
return self._y_pos_per_angle[angle_idx]
def angles_share_raw_grid(self) -> bool:
"""True iff every angle's raw (x, y) pixel grid is literally the same
array as angle 0's -- the case for a .sras file the viewer's own
Alignment Wizard exported (see scan_format.md, "Files written by the
viewer's Alignment Wizard"): the writer packs one Per-Angle Geometry
record and one Row Table span and repeats those same bytes for every
angle, so re-parsed arrays are bit-identical copies rather than
independently re-derived numbers -- a bare np.array_equal is the
correct test here, no tolerance needed.
"""
if self.n_angles <= 1:
return True
x0 = self.x_axis_mm(0)
y0 = self.y_positions_mm(0)
return all(np.array_equal(self.x_axis_mm(a), x0)
and np.array_equal(self.y_positions_mm(a), y0)
for a in range(1, self.n_angles))
def time_axis_ns(self) -> np.ndarray:
return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9
sras_viewer/__init__.py
+2 -7
View File
@@ -18,12 +18,7 @@ import faulthandler
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
from .align_wizard import AlignmentWizard # noqa: E402,F401
from .canvases import ( # noqa: E402,F401
AlignOverlayCanvas, ImageCanvas, RoiQuad, WaveformCanvas,
)
from .canvases import ImageCanvas, RoiQuad, WaveformCanvas # noqa: E402,F401
from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401
from .dialogs import ( # noqa: E402,F401
FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
)
from .dialogs import FftOptionsDialog, ManualAlignmentDialog # noqa: E402,F401
from .main_window import SrasViewerWindow, main # noqa: E402,F401
sras_viewer/align_wizard.py
-1446
View File
File diff suppressed because it is too large Load Diff
sras_viewer/canvases.py
+12 -99
View File
@@ -1,9 +1,7 @@
"""Matplotlib canvases and the ROI primitive."""
import matplotlib as mpl
import numpy as np
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.figure import Figure
from matplotlib.patches import Polygon
from matplotlib.path import Path as MplPath
@@ -13,27 +11,6 @@ from PyQt6.QtWidgets import QSizePolicy
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv
def count_colormap(n_angles: int):
"""(cmap, norm, ticks) for an integer "how many angles cover this pixel"
image, 0..n_angles.
Discrete, not continuous: the judgement the wizard's stack view exists for
is "is this a plateau at N, or a fan of partial overlaps", so a region
covered by one angle too few has to read as its own band rather than a
slightly darker shade. Count 0 is fully transparent so uncovered canvas
cannot be mistaken for a low count.
Shared by both wizard pages that draw this image — they use different canvas
classes, and the same number must not change colour between them.
"""
n = max(1, int(n_angles))
base = mpl.colormaps["viridis"].resampled(n)
colors = [(0.0, 0.0, 0.0, 0.0)] + [base(i) for i in range(n)]
return (ListedColormap(colors),
BoundaryNorm(np.arange(-0.5, n + 1), len(colors)),
np.arange(0, n + 1))
# ---------------------------------------------------------------------------
# ROI (free quadrilateral in data coordinates)
# ---------------------------------------------------------------------------
@@ -123,15 +100,7 @@ class ImageCanvas(FigureCanvasQTAgg):
_HANDLE_PX = 12
_CLICK_THRESH_PX = 4 # releases within this of press count as a click
def __init__(self, parent=None, *, rect_only: bool = False):
"""*rect_only* constrains the ROI to an axis-aligned rectangle.
Used by the alignment wizard's crop page, where a free quadrilateral
would be actively misleading: v6 geometry can only express an
axis-aligned rectangle, so anything else the user drew would have to be
squared off behind their back. Default off, so the main window's
free-quad ROI is unaffected.
"""
def __init__(self, parent=None):
fig = Figure(figsize=(7, 5), tight_layout=True)
self.ax = fig.add_subplot(111)
super().__init__(fig)
@@ -139,7 +108,6 @@ class ImageCanvas(FigureCanvasQTAgg):
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
self._extent = None
self._img_shape = None
self._rect_only = rect_only
# ROI state
self._roi: RoiQuad | None = None
@@ -164,13 +132,9 @@ class ImageCanvas(FigureCanvasQTAgg):
# Public API
# ------------------------------------------------------------------
def show_image(self, img: np.ndarray, extent: list[float], cmap,
def show_image(self, img: np.ndarray, extent: list[float], cmap: str,
vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
colorbar_label: str = "", cb_ticks=None, norm=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."""
colorbar_label: str = ""):
self.figure.clf()
self.ax = self.figure.add_subplot(111)
# Patches and lines are destroyed by figure.clf(); drop stale refs.
@@ -179,14 +143,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,
extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
interpolation="nearest",
)
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
ticks=cb_ticks)
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
if colorbar_label:
cb.set_label(colorbar_label)
@@ -340,28 +302,10 @@ class ImageCanvas(FigureCanvasQTAgg):
self._roi._pts = self._snapshot.corners() + delta
elif self._state == self._DRAG_CORNER:
self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]
if self._rect_only:
self._rectify_corner(self._drag_corner_idx)
self._draw_roi()
self.draw_idle()
def _rectify_corner(self, idx: int):
"""Re-square the quad after a corner drag, anchored on the *opposite*
corner.
Anchoring on the diagonal opposite (idx ^ 2, since corners run
BL, BR, TR, TL) rather than taking the bbox of all four points is what
lets the rectangle shrink: a bbox over the three stale corners plus the
new one is the union of the old rectangle and the new point, so dragging
inward would never make it smaller.
"""
pts = self._roi.corners()
ax_, ay = pts[idx ^ 2]
bx, by = pts[idx]
self._roi._pts = RoiQuad.from_bbox(min(ax_, bx), min(ay, by),
max(ax_, bx), max(ay, by)).corners()
def _on_release(self, event):
if event.button != 1 and self._press_button != 1:
return
@@ -531,20 +475,14 @@ class WaveformCanvas(FigureCanvasQTAgg):
self.draw()
class AlignOverlayCanvas(FigureCanvasQTAgg):
"""Renders the alignment wizard's multi-angle mask views and turns keyboard
input into translate/rotate nudge requests for whichever angle is active.
Two views of the same reprojected masks, because they answer different
questions. show_counts colours each pixel by *how many* angles cover it,
which is the at-a-glance verdict on a correlation run: a good alignment is
one saturated plateau, a bad one is a fringe of low-count halos.
show_overlay gives each angle its own colour, which is what you need while
nudging a specific angle by hand.
class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
"""Renders ManualAlignmentDialog's multi-angle mask overlay and turns
keyboard input into translate/rotate nudge requests for whichever angle
the dialog currently has active.
A pure input+render widget — it holds no alignment state and never
touches SrasFile itself; the wizard page owns all of that and decides,
from these signals, whether a cheap single-layer refresh or a
touches SrasFile itself; ManualAlignmentDialog owns all of that and
decides, from these signals, whether a cheap single-layer refresh or a
full preview-canvas rebuild is needed.
FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
@@ -584,31 +522,6 @@ class AlignOverlayCanvas(FigureCanvasQTAgg):
self.figure.clf()
self.ax = self.figure.add_subplot(111)
self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
self._finish(title)
def show_counts(self, counts: np.ndarray, n_angles: int,
extent: list[float], title: str):
"""The mask stack coloured by how many angles cover each pixel.
A discrete colormap with integer-ticked colorbar rather than a
continuous one: the judgement being made is "is this a plateau at N, or
a fan of partial overlaps", and a region covered by one angle too few
has to read as its own band rather than a slightly darker shade.
Uncovered pixels are transparent so they cannot be mistaken for a low
count.
"""
self.figure.clf()
self.ax = self.figure.add_subplot(111)
cmap, norm, ticks = count_colormap(n_angles)
im = self.ax.imshow(
np.asarray(counts), extent=extent, origin="upper", aspect="auto",
interpolation="nearest", cmap=cmap, norm=norm)
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
ticks=ticks)
cb.set_label("angles overlapping")
self._finish(title)
def _finish(self, title: str):
self.ax.set_xlabel("X (mm)")
self.ax.set_ylabel("Y (mm)")
self.ax.set_title(title)
sras_viewer/common.py
+5 -32
View File
@@ -2,8 +2,8 @@
from PyQt6.QtCore import Qt
from PyQt6.QtWidgets import (
QComboBox, QDoubleSpinBox, QFormLayout, QFrame, QGroupBox, QLabel,
QScrollArea, QSizePolicy, QVBoxLayout, QWidget,
QDoubleSpinBox, QFormLayout, QFrame, QGroupBox, QLabel, QScrollArea,
QSizePolicy, QVBoxLayout, QWidget,
)
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX
@@ -60,13 +60,9 @@ class Jobs:
COMPUTE = "compute"
DC_PRECOMPUTE = "dc_precompute"
BATCH = "batch"
# The alignment wizard's three background steps: fetching each angle's CH4
# image for the mask stack, registering the angles, and writing the aligned
# export. Separate keys because a retry of one must not be blocked by
# another having run, and _run_worker's busy check is per key.
ALIGN_MASKS = "align_masks"
ALIGN_CORRELATE = "align_correlate"
ALIGN_EXPORT = "align_export"
ALIGN = "align"
MANUAL_ALIGN_MASKS = "manual_align_masks"
MANUAL_ALIGN_CORRELATE = "manual_align_correlate"
def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
@@ -85,29 +81,6 @@ def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
return spin
def _combo(items=(), *, min_chars: int = 10) -> QComboBox:
"""A combo box whose size hint does not depend on its longest entry.
By default a QComboBox asks for enough width to show its widest item. These
hold descriptive phrases, and the side panels are fixed-width — in a scroll
area with the horizontal scrollbar off (`_scroll_panel`) an unconstrained
hint pushes the inner widget past the panel and everything on the right,
including the hint text, is silently clipped instead of scrolling.
*items* is a sequence of (label, data) pairs, or of plain labels.
"""
combo = QComboBox()
combo.setSizeAdjustPolicy(
QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon)
combo.setMinimumContentsLength(min_chars)
for item in items:
if isinstance(item, tuple):
combo.addItem(item[0], item[1])
else:
combo.addItem(item)
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."""
sras_viewer/dialogs.py
+605 -265
View File
@@ -1,22 +1,33 @@
"""FFT option dialogs.
"""FFT Options and Manual Alignment dialogs."""
Angle alignment used to live here too, as ManualAlignmentDialog; it is now the
alignment wizard's first page (see align_wizard.py), which needed the same
mask-overlay editor plus the crop and export steps.
"""
from pathlib import Path
from typing import TYPE_CHECKING
import matplotlib as mpl
import numpy as np
from matplotlib.backends.backend_qtagg import NavigationToolbar2QT
from PyQt6.QtCore import QSignalBlocker, pyqtSignal
from PyQt6.QtWidgets import (
QButtonGroup, QCheckBox, QDialog, QDialogButtonBox, QFileDialog,
QGroupBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QRadioButton,
QScrollArea, QSpinBox, QVBoxLayout, QWidget,
QButtonGroup, QComboBox, QDialog, QDialogButtonBox, QGroupBox,
QHBoxLayout, QLabel, QMessageBox, QPushButton, QRadioButton, QSpinBox,
QVBoxLayout, QWidget,
)
from sras_compute import PYFFTW_AVAILABLE
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES
import sras_compute as compute
from sras_compute import (
PYFFTW_AVAILABLE, ManualAngleParams, build_manual_alignment,
delete_manual_alignment, save_manual_alignment,
)
from sras_format import SrasFile
from sras_workers import Ch4MaskWorker, CrossCorrelateWorker
from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label
from .canvases import ManualAlignOverlayCanvas
from .common import (
_CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin,
_scroll_panel, _wrap_label,
)
if TYPE_CHECKING:
from .main_window import SrasViewerWindow
# ---------------------------------------------------------------------------
@@ -37,7 +48,8 @@ class FftOptionsDialog(QDialog):
current_pad_factor: int,
samples_per_frame: int | None,
sample_rate_hz: float | None,
grating_um: float):
grating_um: float,
cached_pad_factor: int | None = None):
super().__init__(parent)
self.setWindowTitle("FFT Options")
self.setModal(True)
@@ -46,6 +58,10 @@ class FftOptionsDialog(QDialog):
self._samples_per_frame = samples_per_frame
self._sample_rate_hz = sample_rate_hz
self._grating_um = grating_um
# The pad the open file's stored FFT cache was computed at, if it has
# one — picking anything else here makes that cache unreadable, which
# is worth saying before Apply rather than after.
self._cached_pad_factor = cached_pad_factor
layout = QVBoxLayout(self)
@@ -98,6 +114,11 @@ class FftOptionsDialog(QDialog):
lbl.setStyleSheet(_CSS_HINT)
zl.addWidget(lbl)
self._lbl_cache = QLabel()
self._lbl_cache.setWordWrap(True)
self._lbl_cache.setStyleSheet(_CSS_WARN)
zl.addWidget(self._lbl_cache)
layout.addWidget(grp_zp)
# ---- Buttons ---------------------------------------------------
@@ -115,6 +136,13 @@ class FftOptionsDialog(QDialog):
sr = self._sample_rate_hz
pad = self._spin_pad.value()
self._lbl_cache.setText(
"" if self._cached_pad_factor in (None, pad) else
f"! This file's stored FFT cache was computed at pad "
f"{self._cached_pad_factor}x, so at {pad}x it can't be used and "
f"CH1/Velocity will recompute. Re-run Convert -> Batch Compute FFT "
f"to cache at {pad}x.")
if spf is None or sr is None:
self._lbl_nfft.setText("Load a file to preview FFT parameters.")
self._lbl_freq_res.setText("")
@@ -141,275 +169,587 @@ class FftOptionsDialog(QDialog):
return max(1, self._spin_pad.value())
# ---------------------------------------------------------------------------
# Row-Averaged FFT Options dialog
# ---------------------------------------------------------------------------
class RowAverageFftOptionsDialog(QDialog):
"""Configure the same-row, distance-weighted neighbor averaging applied
to each pixel's CH1 waveform before 'Batch Compute Row-Averaged FFT and
Store' re-runs the FFT peak search — a same-row SNR cleanup pass, never
mixing across rows/Y (see sras_compute._row_average_waveforms).
class ManualAlignmentDialog(QDialog):
"""Non-modal manual angle-alignment editor (Fusion -> Manual Alignment...).
Unlike the plain FFT batch action (which stores unmasked and defers
masking to display time), the DC threshold here is required up front:
it decides which same-row neighbors are eligible to contribute to a
pixel's average, so it can't be deferred.
Shows every angle's binarized CH4 (Bias B) mask overlaid in a distinct
color at partial opacity on one shared canvas, so translation/rotation
misalignment is visible by eye. Reference angle (always index 0) is
ground truth and never moves; every other angle is aligned to it. The
user picks an "active" angle and nudges its rotation+translation with
the keyboard; Auto Cross-Correlate finds every non-reference angle's
rotation *and* translation by registering its image against the
reference's (see compute.register_angle_to_reference) — meant to get every
angle stacked on top of each other so keyboard nudging only has to make
small corrections, not find an alignment from scratch; Auto De-rotate is
the weaker fallback that just seeds rotation from the stage's reported
angle, leaving translation alone. Save writes a JSON sidecar next to the
.sras file and hands a freshly-built, full-resolution AlignmentResult back
to the main window — the exact same object shape compute_angle_alignment
produces, so every existing Aligned-View code path (apply_alignment,
_aligned_canvas_axes, the pixel-inspector inverse-transform) works
completely unmodified.
Changes take effect only when the user clicks Apply. Cancel discards
all pending edits.
Non-modal by design (shown via .show(), never .exec() or setModal(True))
so the user can still interact with the main window. Talks back to
SrasViewerWindow two ways: it reuses parent._run_worker/_jobs directly
for its background mask-fetch and cross-correlate steps, so the main
window's existing shutdown/lifecycle plumbing covers both for free, and
it emits alignment_saved / alignment_cleared signals for the two moments
that should actually mutate the main window's persistent state —
everything else (nudging, Auto De-rotate, Auto Cross-Correlate, threshold
edits) stays purely local to this dialog until Save.
"""
def __init__(self, parent=None, *,
current_n: int,
current_threshold_mv: float,
pixel_x_mm: float | None):
alignment_saved = pyqtSignal(object, str) # AlignmentResult, sidecar path (str)
alignment_cleared = pyqtSignal()
_PREVIEW_MARGIN_FRAC = 0.15
_BASE_ALPHA = 0.42
_ACTIVE_ALPHA = 0.75
_MAX_PREVIEW_DIM = 1024
# (label, sources passed to compute.register_angle_to_reference). "Both"
# registers on each and keeps whichever scores higher per angle, which
# costs roughly double but removes the failure mode where the single
# chosen source is the one that happens to be uninformative for one angle.
_CORRELATE_SOURCES = (
("Both, keep best (recommended)", ("signal", "mask")),
("Raw signal", ("signal",)),
("Thresholded mask", ("mask",)),
)
def __init__(self, parent: "SrasViewerWindow", sras: SrasFile, *,
ref_angle_idx: int, dc_threshold_mv: float,
seed_per_angle: dict[int, ManualAngleParams] | None,
cached_dc4_mv: dict[int, np.ndarray]):
super().__init__(parent)
self.setWindowTitle("Row-Averaged FFT Options")
self.setModal(True)
self.setMinimumWidth(380)
self._parent = parent
self._sras = sras
self._ref_angle_idx = ref_angle_idx
self._downsample = (1, 1) # (rows, cols) block-mean factors
self._dc4_mv: dict[int, np.ndarray] = {}
self._masks_small: dict[int, np.ndarray] = {}
self._preview_layers: dict[int, np.ndarray] = {}
self._preview_origin_mm = (0.0, 0.0)
self._preview_shape = (1, 1)
self._preview_pitch_mm = (1.0, 1.0)
self._masks_ready = False
self._fit_notes: dict[int, tuple[float, str]] = {}
self._derotate_sign_flipped = False
self._pixel_x_mm = pixel_x_mm
self.setWindowTitle(f"Manual Alignment — {sras.path.name}")
self.resize(1150, 760)
layout = QVBoxLayout(self)
self._seed_initial_params(seed_per_angle)
n = sras.n_angles
cmap = mpl.colormaps["tab10"] if n <= 10 else mpl.colormaps["tab20"]
self._angle_colors = {a: cmap(a % cmap.N)[:3] for a in range(n)}
self._active_angle = 1 if ref_angle_idx == 0 and n > 1 else 0
# ---- Neighbor window ---------------------------------------------
grp_window = QGroupBox("Same-Row Neighbor Window")
wl = QVBoxLayout(grp_window)
self._build_ui(dc_threshold_mv)
self._set_controls_enabled(False) # re-enabled once masks are ready
self._start_mask_prep(cached_dc4_mv)
n_row = QHBoxLayout()
n_row.addWidget(QLabel("Neighbor half-width (n):"))
self._spin_n = QSpinBox()
self._spin_n.setRange(1, 50)
self._spin_n.setValue(max(1, current_n))
self._spin_n.setToolTip(
"Each pixel's CH1 waveform is averaged with up to n same-row\n"
"neighbors on each side, distance-weighted (Gaussian) and\n"
"counting only neighbors that already pass the DC threshold\n"
"below. Never mixes across rows/Y.")
self._spin_n.valueChanged.connect(self._update_info)
n_row.addWidget(self._spin_n)
wl.addLayout(n_row)
def showEvent(self, event):
super().showEvent(event)
self.canvas.setFocus()
self._lbl_width = QLabel()
self._lbl_width.setStyleSheet(_CSS_HINT)
wl.addWidget(self._lbl_width)
# ------------------------------------------------------------------
# Construction
# ------------------------------------------------------------------
layout.addWidget(grp_window)
def _seed_initial_params(self, seed_per_angle: dict[int, ManualAngleParams] | None):
seed = seed_per_angle or {}
self._angle_params: dict[int, ManualAngleParams] = {
a: (ManualAngleParams(seed[a].rotation_deg, seed[a].shift_mm)
if a in seed else ManualAngleParams())
for a in range(self._sras.n_angles)
}
self._angle_params[self._ref_angle_idx] = ManualAngleParams()
# ---- DC threshold ------------------------------------------------
grp_thr = QGroupBox("Neighbor Validity")
tl = QVBoxLayout(grp_thr)
thr_row = QHBoxLayout()
thr_row.addWidget(QLabel("DC threshold:"))
self._spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
value=current_threshold_mv, step=0.025)
self._spin_threshold.setToolTip(
"A same-row neighbor only contributes to a pixel's average if\n"
"its own CH4 signal is at or above this threshold -- the same\n"
"test used for RF mask display. A pixel below threshold stays\n"
"masked, exactly as today; it is never rescued by its neighbors.")
thr_row.addWidget(self._spin_threshold)
tl.addLayout(thr_row)
layout.addWidget(grp_thr)
def _build_ui(self, dc_threshold_mv: float):
root = QHBoxLayout(self)
# ---- Buttons -----------------------------------------------------
buttons = QDialogButtonBox()
buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
).clicked.connect(self.accept)
buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
).clicked.connect(self.reject)
layout.addWidget(buttons)
self.canvas = ManualAlignOverlayCanvas()
left = QWidget()
left_l = QVBoxLayout(left)
left_l.setContentsMargins(0, 0, 0, 0)
left_l.setSpacing(4)
left_l.addWidget(NavigationToolbar2QT(self.canvas, left))
left_l.addWidget(self.canvas)
root.addWidget(left, stretch=1)
self._update_info()
panel = QWidget()
panel_l = QVBoxLayout(panel)
panel_l.setContentsMargins(0, 0, 0, 0)
panel_l.setSpacing(8)
panel_l.addWidget(self._build_angle_group())
panel_l.addWidget(self._build_adjust_group())
panel_l.addWidget(self._build_step_group())
panel_l.addWidget(self._build_threshold_group(dc_threshold_mv))
panel_l.addWidget(self._build_correlate_group())
panel_l.addWidget(self._build_actions_group())
self.lbl_status = _wrap_label("", _CSS_MUTED)
panel_l.addWidget(self.lbl_status)
panel_l.addStretch()
def _update_info(self):
n = self._spin_n.value()
if self._pixel_x_mm is None:
self._lbl_width.setText("Load a file to preview the window's physical width.")
root.addWidget(_scroll_panel(panel, 320))
self._connect_controls()
def _build_angle_group(self) -> QWidget:
grp_angle, al = _group("Active Angle")
self.combo_active_angle = QComboBox()
for a in range(self._sras.n_angles):
label = f"Angle {a} ({self._sras.angles_deg[a]:.1f}°)"
if a == self._ref_angle_idx:
label += " [reference]"
self.combo_active_angle.addItem(label)
al.addWidget(self.combo_active_angle)
self.lbl_active_note = _wrap_label("", _CSS_WARN)
al.addWidget(self.lbl_active_note)
return grp_angle
def _build_adjust_group(self) -> QWidget:
self.grp_manual_adjust, mform_box = _group("Manual Adjustment")
mform = _form()
self.spin_active_rotation_deg = _make_dspin(-3600.0, 3600.0, 3, suffix=" °")
mform.addRow("Rotation:", self.spin_active_rotation_deg)
self.spin_active_shift_x_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
mform.addRow("Shift X:", self.spin_active_shift_x_mm)
self.spin_active_shift_y_mm = _make_dspin(-1e5, 1e5, 4, suffix=" mm")
mform.addRow("Shift Y:", self.spin_active_shift_y_mm)
mform_box.addLayout(mform)
return self.grp_manual_adjust
def _build_step_group(self) -> QWidget:
self.grp_step_sizes, sl = _group("Nudge Step Sizes")
sform = _form()
self.spin_step_translate_mm = _make_dspin(0.0001, 1000.0, 4,
suffix=" mm", value=0.01)
sform.addRow("Translate step:", self.spin_step_translate_mm)
self.spin_step_rotate_deg = _make_dspin(0.001, 90.0, 3,
suffix=" °", value=0.1)
sform.addRow("Rotate step:", self.spin_step_rotate_deg)
self.spin_step_multiplier = _make_dspin(1.0, 1000.0, 1, value=10.0)
sform.addRow("Coarse × (Shift):", self.spin_step_multiplier)
sl.addLayout(sform)
sl.addWidget(_wrap_label(
"Arrow keys nudge X/Y translation; Q/E nudge rotation (CCW/CW). "
"Hold Shift for the coarse step. Click the image once so it has "
"keyboard focus.", _CSS_HINT))
return self.grp_step_sizes
def _build_threshold_group(self, dc_threshold_mv: float) -> QWidget:
self.grp_mask_threshold, tl = _group("Mask Threshold")
tform = _form()
self.spin_mask_threshold_mv = _make_dspin(-500.0, 500.0, 3,
suffix=" mV", value=dc_threshold_mv)
tform.addRow("DC threshold:", self.spin_mask_threshold_mv)
tl.addLayout(tform)
return self.grp_mask_threshold
def _build_correlate_group(self) -> QWidget:
self.grp_correlate, cl = _group("Cross-Correlate (FFT)")
cform = _form()
self.combo_correlate_source = QComboBox()
for label, sources in self._CORRELATE_SOURCES:
self.combo_correlate_source.addItem(label, sources)
cform.addRow("Correlate on:", self.combo_correlate_source)
self.spin_correlate_search_deg = _make_dspin(0.0, 180.0, 1, suffix=" °",
value=6.0, step=1.0)
cform.addRow("Rotation search (±):", self.spin_correlate_search_deg)
cl.addLayout(cform)
self.btn_auto_correlate = QPushButton("Auto Cross-Correlate (vs Reference)")
cl.addWidget(self.btn_auto_correlate)
cl.addWidget(_wrap_label(
"Finds each non-reference angle's rotation *and* translation by "
"cross-correlating its image against the reference's — the stage's "
"reported angle is only the starting point of the search, and both "
"of its signs are tried. Run this first, then nudge only for small "
"corrections.", _CSS_HINT))
return self.grp_correlate
def _build_actions_group(self) -> QWidget:
grp_actions, acl = _group("Actions")
self.btn_auto_derotate = QPushButton("Auto De-rotate (use known angles)")
self.btn_save = QPushButton("Save Alignment")
self.btn_clear = QPushButton("Clear Alignment…")
self.btn_close = QPushButton("Close")
for btn in (self.btn_auto_derotate, self.btn_save, self.btn_clear, self.btn_close):
acl.addWidget(btn)
return grp_actions
def _connect_controls(self):
self.combo_active_angle.currentIndexChanged.connect(self._on_active_angle_changed)
self.spin_active_rotation_deg.editingFinished.connect(self._on_rotation_spin_edited)
self.spin_active_shift_x_mm.editingFinished.connect(self._on_shift_spin_edited)
self.spin_active_shift_y_mm.editingFinished.connect(self._on_shift_spin_edited)
self.spin_mask_threshold_mv.editingFinished.connect(self._on_mask_threshold_edited)
self.btn_auto_derotate.clicked.connect(self._on_auto_derotate)
self.btn_auto_correlate.clicked.connect(self._on_auto_correlate)
self.btn_save.clicked.connect(self._on_save)
self.btn_clear.clicked.connect(self._on_clear)
self.btn_close.clicked.connect(self.close)
self.canvas.nudge_translate.connect(self._on_nudge_translate)
self.canvas.nudge_rotate.connect(self._on_nudge_rotate)
with QSignalBlocker(self.combo_active_angle):
self.combo_active_angle.setCurrentIndex(self._active_angle)
self._on_active_angle_changed(self._active_angle)
# ------------------------------------------------------------------
# Mask preparation (initial CH4 fetch + threshold + downsample)
# ------------------------------------------------------------------
def _start_mask_prep(self, cached_dc4_mv: dict[int, np.ndarray]):
self._dc4_mv = dict(cached_dc4_mv)
missing = [a for a in range(self._sras.n_angles) if a not in self._dc4_mv]
if not missing:
self._finish_mask_prep()
return
width_um = 2 * n * self._pixel_x_mm * 1e3
self._lbl_width.setText(
f"Window: ±{n} px = {width_um:.2f} µm full width "
f"(pixel pitch {self._pixel_x_mm * 1e3:.3g} µm)")
self.lbl_status.setText(f"Preparing masks: 0/{len(missing)} angle(s) needed…")
started = self._parent._run_worker(
Jobs.MANUAL_ALIGN_MASKS, Ch4MaskWorker(self._sras, missing),
connect=(
("angle_done", self._on_mask_angle_done),
("error", lambda msg: self.lbl_status.setText(f"Mask prep error: {msg}")),
),
on_done=self._finish_mask_prep)
if not started:
self.lbl_status.setText(
"Could not start mask preparation (busy) — close and reopen.")
def get_half_width(self) -> int:
return self._spin_n.value()
def _on_mask_angle_done(self, angle_idx: int, dc4_mv: np.ndarray):
self._dc4_mv[angle_idx] = dc4_mv
self.lbl_status.setText(
f"Preparing masks: {len(self._dc4_mv)}/{self._sras.n_angles} ready…")
def get_threshold_mv(self) -> float:
return self._spin_threshold.value()
def _finish_mask_prep(self):
if len(self._dc4_mv) < self._sras.n_angles:
return # a mask-worker error left some angles unfetched
# Rows and columns get their own factor. A real scan is ~7500 frames
# wide but only ~750 rows tall, so one shared factor sized for the
# frames would throw away 8x more row detail than the preview needs and
# leave the overlay too coarse in y to judge alignment by eye.
max_rows = max(img.shape[0] for img in self._dc4_mv.values())
max_cols = max(img.shape[1] for img in self._dc4_mv.values())
self._downsample = (
max(1, int(np.ceil(max_rows / self._MAX_PREVIEW_DIM))),
max(1, int(np.ceil(max_cols / self._MAX_PREVIEW_DIM))))
self._recompute_masks_small()
self._rebuild_preview_canvas()
self._set_controls_enabled(True)
self.lbl_status.setText("Ready.")
def _recompute_masks_small(self):
"""Threshold + downsample every angle's already-in-memory full-res
CH4 mV image. Cheap (a compare + block-mean), so this re-runs in
full whenever the mask-threshold spin box changes — no re-fetch.
Purely for the overlay's visuals: no alignment geometry depends on this
threshold, only which pixels the overlay paints."""
threshold = self.spin_mask_threshold_mv.value()
fy, fx = self._downsample
self._masks_small = {
a: compute.block_mean_2d((img >= threshold).astype(np.float32), fy, fx)
for a, img in self._dc4_mv.items()
}
# ------------------------------------------------------------------
# Preview canvas: full rebuild vs. incremental single-layer refresh
# ------------------------------------------------------------------
def _rebuild_preview_canvas(self):
"""Full geometry rebuild: recomputes the shared preview canvas's
origin/shape (rotation can grow the union bbox — translation alone
cannot, per the padding baked in via _PREVIEW_MARGIN_FRAC) and every
angle's reprojected mask layer. Triggered by: dialog open,
mask-threshold change, Auto De-rotate, a rotation nudge/edit of the
active angle. NOT triggered by a translation-only nudge — see
_refresh_active_preview_layer."""
dx_ref, dy_ref = compute.pixel_pitch_mm(self._sras, self._ref_angle_idx)
fy, fx = self._downsample
pitch = (dx_ref * fx, dy_ref * fy)
origin, shape = compute.canvas_for_params(
self._sras, self._ref_angle_idx, pitch, self._angle_params,
margin_frac=self._PREVIEW_MARGIN_FRAC, snap=False)
self._preview_origin_mm, self._preview_shape = origin, shape
self._preview_pitch_mm = pitch
self._preview_layers = {
a: self._reproject(a) for a in range(self._sras.n_angles)
}
self._redraw_overlay()
def _reproject(self, angle_idx: int) -> np.ndarray:
"""One angle's downsampled mask on the current preview canvas.
src_downsample must match _masks_small's block-mean factors, or the
layer lands magnified and offset instead of where the alignment
actually puts it."""
p = self._angle_params[angle_idx]
return compute.reproject_mask(
self._sras, angle_idx, self._ref_angle_idx,
self._masks_small[angle_idx], p.rotation_deg, p.shift_mm,
self._preview_pitch_mm, self._preview_origin_mm, self._preview_shape,
src_downsample=self._downsample)
def _refresh_active_preview_layer(self):
"""Cheap path for a translation-only nudge/edit of the active angle:
reproject just that one angle's downsampled mask onto the *existing*
preview canvas — every other angle's cached layer is untouched."""
self._preview_layers[self._active_angle] = self._reproject(self._active_angle)
self._redraw_overlay()
def _redraw_overlay(self):
"""Alpha-composite every angle's colored mask layer into one RGBA
image ("all thresholds overlaid with varying opacity"). Each angle
keeps a fixed, distinct color regardless of which is active; the
active angle is drawn last (on top) at a visibly higher alpha so
it's easy to track while nudging."""
if not self._preview_layers:
return # mask prep hasn't finished yet — nothing to draw
n_rows, n_cols = self._preview_shape
rgba = np.zeros((n_rows, n_cols, 4), dtype=np.float32)
order = sorted(range(self._sras.n_angles), key=lambda a: a == self._active_angle)
for a in order:
layer = self._preview_layers.get(a)
if layer is None:
continue
alpha = self._ACTIVE_ALPHA if a == self._active_angle else self._BASE_ALPHA
color = self._angle_colors[a]
fg_a = layer * alpha
for c in range(3):
rgba[..., c] = color[c] * fg_a + rgba[..., c] * rgba[..., 3] * (1 - fg_a)
rgba[..., 3] = fg_a + rgba[..., 3] * (1 - fg_a)
x0, y0 = self._preview_origin_mm
dx, dy = self._preview_pitch_mm
x_axis = x0 + np.arange(n_cols) * dx
y_axis = y0 + np.arange(n_rows) * dy
extent = _axes_extent(x_axis, y_axis, dx, dy)
title = (f"Angle {self._active_angle} active "
f"({self._sras.angles_deg[self._active_angle]:.1f}°)")
self.canvas.show_overlay(rgba, extent, title)
# ------------------------------------------------------------------
# Angle selection / nudge / edit handlers
# ------------------------------------------------------------------
def _on_active_angle_changed(self, angle_idx: int):
self._active_angle = angle_idx
is_ref = angle_idx == self._ref_angle_idx
self.grp_manual_adjust.setEnabled(self._masks_ready and not is_ref)
self.lbl_active_note.setText(
"Reference angle — defines the shared origin, not adjustable." if is_ref else "")
self._sync_active_spinboxes()
self._redraw_overlay()
def _sync_active_spinboxes(self):
p = self._angle_params[self._active_angle]
for spin, val in ((self.spin_active_rotation_deg, p.rotation_deg),
(self.spin_active_shift_x_mm, p.shift_mm[0]),
(self.spin_active_shift_y_mm, p.shift_mm[1])):
with QSignalBlocker(spin):
spin.setValue(val)
def _on_nudge_translate(self, dir_x: int, dir_y: int, coarse: bool):
if not self._masks_ready or self._active_angle == self._ref_angle_idx:
return
step = self.spin_step_translate_mm.value()
if coarse:
step *= self.spin_step_multiplier.value()
p = self._angle_params[self._active_angle]
p.shift_mm = (p.shift_mm[0] + dir_x * step, p.shift_mm[1] + dir_y * step)
self._sync_active_spinboxes()
self._refresh_active_preview_layer()
def _on_nudge_rotate(self, direction: int, coarse: bool):
if not self._masks_ready or self._active_angle == self._ref_angle_idx:
return
step = self.spin_step_rotate_deg.value()
if coarse:
step *= self.spin_step_multiplier.value()
self._angle_params[self._active_angle].rotation_deg += direction * step
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
def _on_rotation_spin_edited(self):
if self._active_angle == self._ref_angle_idx:
return
self._angle_params[self._active_angle].rotation_deg = self.spin_active_rotation_deg.value()
self._rebuild_preview_canvas()
def _on_shift_spin_edited(self):
if self._active_angle == self._ref_angle_idx:
return
p = self._angle_params[self._active_angle]
p.shift_mm = (self.spin_active_shift_x_mm.value(), self.spin_active_shift_y_mm.value())
self._refresh_active_preview_layer()
def _on_mask_threshold_edited(self):
if not self._masks_ready:
return
self._recompute_masks_small()
self._rebuild_preview_canvas()
# ------------------------------------------------------------------
# Actions
# ------------------------------------------------------------------
def _on_auto_derotate(self):
"""Seed every angle's rotation from the stage's reported angle.
A starting point for nudging by eye, not an alignment: the stage's
sign convention relative to this module's is not knowable from the
file, so the sign that lines the scans up is whichever of the two looks
right in the overlay. Auto Cross-Correlate decides that from the images
instead, and is the button to reach for first.
"""
sign = -1.0 if self._derotate_sign_flipped else 1.0
self._derotate_sign_flipped = not self._derotate_sign_flipped
n_changed = 0
for a in range(self._sras.n_angles):
if a == self._ref_angle_idx:
continue
self._angle_params[a].rotation_deg = sign * compute.nominal_delta_deg(
self._sras, a, self._ref_angle_idx)
n_changed += 1
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self.lbl_status.setText(
f"Rotation set to the stage angle ({'−' if sign < 0 else '+'}delta) "
f"for {n_changed} angle(s); translation untouched. Click again to "
"try the opposite sign.")
def _on_auto_correlate(self):
if not self._masks_ready:
return
angles = [a for a in range(self._sras.n_angles) if a != self._ref_angle_idx]
if not angles:
return
worker = CrossCorrelateWorker(
self._sras, self._ref_angle_idx, angles, self._dc4_mv,
sources=self.combo_correlate_source.currentData(),
dc_threshold_mv=self.spin_mask_threshold_mv.value(),
search_deg=self.spin_correlate_search_deg.value())
self._correlate_done_count = 0
self._correlate_total = len(angles)
self._fit_notes = {}
self._set_controls_enabled(False)
self.lbl_status.setText(f"Cross-correlating: 0/{self._correlate_total} angle(s)…")
started = self._parent._run_worker(
Jobs.MANUAL_ALIGN_CORRELATE, worker,
connect=(
("angle_done", self._on_correlate_angle_done),
("error", self._on_correlate_error),
),
on_done=self._finish_auto_correlate)
if not started:
self._set_controls_enabled(True)
self.lbl_status.setText("Could not start cross-correlation (busy) — try again.")
def _on_correlate_angle_done(self, angle_idx: int, rotation_deg: float,
shift_x_mm: float, shift_y_mm: float,
score: float, source: str):
self._angle_params[angle_idx] = ManualAngleParams(rotation_deg, (shift_x_mm, shift_y_mm))
self._fit_notes[angle_idx] = (score, source)
self._correlate_done_count += 1
self.lbl_status.setText(
f"Cross-correlating: {self._correlate_done_count}/{self._correlate_total} angle(s)…")
def _on_correlate_error(self, msg: str):
self.lbl_status.setText(f"Cross-correlation error: {msg}")
def _finish_auto_correlate(self):
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self._set_controls_enabled(True)
self.lbl_status.setText(
f"Cross-correlated {self._correlate_done_count} angle(s) against "
f"Angle {self._ref_angle_idx}.\n" + self._fit_report())
def _fit_report(self) -> str:
"""Per-angle registration quality, worst first.
Surfaced rather than buried because a single bad acquisition (stage
glitch, laser dropout) registers poorly and would otherwise be fused in
silently — seeing which angle it is, is what makes dropping it with
sras_edit_scans.py actionable. The deviation from the stage's own
reported angle is shown alongside: a large one means the search and the
stage disagree, which is either a genuine mechanical error or a sign
that this angle's fit is not to be trusted.
"""
if not self._fit_notes:
return ""
rows = sorted(self._fit_notes.items(), key=lambda kv: kv[1][0])
worst = rows[0]
lines = [f"Worst fit: angle {worst[0]} (score {worst[1][0]:.3f}, "
f"{worst[1][1]})."]
drifted = []
for a, _note in rows:
nominal = compute.nominal_delta_deg(self._sras, a, self._ref_angle_idx)
got = self._angle_params[a].rotation_deg
dev = min(abs(got - nominal), abs(got + nominal))
if dev > 1.0:
drifted.append(f"{a} ({dev:.2f}°)")
if drifted:
lines.append("Rotation differs from the stage angle by >1° for "
"angle(s) " + ", ".join(drifted) + ".")
lines.append("Nudge from here for any remaining fine correction.")
return " ".join(lines)
def _on_save(self):
threshold = self.spin_mask_threshold_mv.value()
resolved = dict(self._angle_params) # already concrete floats
try:
path = save_manual_alignment(self._sras, self._ref_angle_idx, threshold, resolved)
result = build_manual_alignment(self._sras, self._ref_angle_idx,
threshold, resolved)
except OSError as exc:
QMessageBox.warning(self, "Save Alignment Failed", str(exc))
return
self.lbl_status.setText(f"Saved to {path.name}.")
self.alignment_saved.emit(result, str(path))
def _on_clear(self):
reply = QMessageBox.question(
self, "Clear Alignment",
"This resets every angle back to raw/unaligned (0° rotation, no "
"shift) and deletes the saved alignment file for this scan, if "
"any. This cannot be undone. Continue?",
QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
QMessageBox.StandardButton.No)
if reply != QMessageBox.StandardButton.Yes:
return
try:
existed = delete_manual_alignment(self._sras)
except OSError as exc:
QMessageBox.warning(self, "Clear Alignment Failed",
f"Could not delete the saved alignment file: {exc}")
return
self._angle_params = {a: ManualAngleParams() for a in range(self._sras.n_angles)}
self._fit_notes = {}
self._sync_active_spinboxes()
self._rebuild_preview_canvas()
self.lbl_status.setText(
"Alignment cleared; saved file removed." if existed
else "Alignment cleared (there was no saved file).")
self.alignment_cleared.emit()
def _set_controls_enabled(self, enabled: bool):
self._masks_ready = enabled
self.combo_active_angle.setEnabled(enabled)
self.grp_manual_adjust.setEnabled(enabled and self._active_angle != self._ref_angle_idx)
self.grp_step_sizes.setEnabled(enabled)
self.grp_mask_threshold.setEnabled(enabled)
self.grp_correlate.setEnabled(enabled)
self.btn_auto_derotate.setEnabled(enabled)
self.btn_save.setEnabled(enabled)
self.btn_clear.setEnabled(enabled)
# ---------------------------------------------------------------------------
# Export Fused ROI dialog
# ---------------------------------------------------------------------------
class FusedRoiExportDialog(QDialog):
"""Choose a value type and which angles to fuse for Export Fused ROI.
Every angle in the file is listed (a live AlignmentResult's per_angle
always covers every angle, and the raw-shared-grid path needs no per-
angle transform at all). Each checkbox is disabled — and auto-unchecked
— whenever *availability_fn(angle_idx, ch_idx)* is False for the
currently selected value type; switching the value-type radio
re-evaluates every checkbox live, since availability is per (angle,
value type) rather than just per angle — e.g. DC may be ready
everywhere while FFT is ready nowhere.
"""
_VALUE_MODES = (CH1_IDX, CH3_IDX, CH4_IDX, VELOCITY_MODE_IDX)
def __init__(self, parent=None, *,
angles: list[tuple[int, float]],
availability_fn,
default_ch_idx: int,
out_dir: str,
stem: str,
grid_note: str):
super().__init__(parent)
self.setWindowTitle("Export Fused ROI")
self.setModal(True)
self.setMinimumWidth(420)
self._availability_fn = availability_fn
self._out_dir = out_dir
self._stem = stem
self._path_user_chosen = False
layout = QVBoxLayout(self)
layout.addWidget(_wrap_label(grid_note, _CSS_HINT))
# ---- Value type --------------------------------------------------
grp_val, vl = _group("Value to Export")
self._val_group = QButtonGroup(self)
self._val_buttons: dict[int, QRadioButton] = {}
for ch_idx, label in zip(self._VALUE_MODES, CH_LABELS):
rb = QRadioButton(label)
self._val_group.addButton(rb, id=ch_idx)
self._val_buttons[ch_idx] = rb
vl.addWidget(rb)
self._val_buttons[default_ch_idx].setChecked(True)
self._val_group.idClicked.connect(self._on_value_type_changed)
layout.addWidget(grp_val)
# ---- Angles --------------------------------------------------
grp_ang, al = _group("Angles to Include")
sel_row = QHBoxLayout()
btn_all = QPushButton("Select All Available")
btn_none = QPushButton("Select None")
btn_all.clicked.connect(self._on_select_all_available)
btn_none.clicked.connect(self._on_select_none)
sel_row.addWidget(btn_all)
sel_row.addWidget(btn_none)
al.addLayout(sel_row)
scroll_inner = QWidget()
scroll_layout = QVBoxLayout(scroll_inner)
self._angle_checks: dict[int, QCheckBox] = {}
for angle_idx, angle_deg in angles:
cb = QCheckBox(f"{angle_deg:.1f}° (angle {angle_idx})")
self._angle_checks[angle_idx] = cb
cb.toggled.connect(self._update_accept_enabled)
scroll_layout.addWidget(cb)
scroll = QScrollArea()
scroll.setWidget(scroll_inner)
scroll.setWidgetResizable(True)
scroll.setMaximumHeight(220)
al.addWidget(scroll)
self._lbl_none_available = _wrap_label("", _CSS_WARN)
al.addWidget(self._lbl_none_available)
layout.addWidget(grp_ang)
# ---- Output path --------------------------------------------------
grp_out, ol = _group("Output File")
path_row = QHBoxLayout()
self._edit_path = QLineEdit()
self._edit_path.setReadOnly(True)
path_row.addWidget(self._edit_path, 1)
btn_browse = QPushButton("Browse…")
btn_browse.clicked.connect(self._on_browse)
path_row.addWidget(btn_browse)
ol.addLayout(path_row)
layout.addWidget(grp_out)
# ---- Buttons -----------------------------------------------------
buttons = QDialogButtonBox()
self._btn_export = buttons.addButton(
"Export", QDialogButtonBox.ButtonRole.AcceptRole)
self._btn_export.clicked.connect(self.accept)
buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
).clicked.connect(self.reject)
layout.addWidget(buttons)
self._refresh_default_path()
self._apply_availability()
# ---- internals -----------------------------------------------------
def _current_ch_idx(self) -> int:
return self._val_group.checkedId()
def _apply_availability(self):
ch_idx = self._current_ch_idx()
n_ok = 0
for angle_idx, cb in self._angle_checks.items():
ok = self._availability_fn(angle_idx, ch_idx)
cb.setEnabled(ok)
if ok:
n_ok += 1
cb.setToolTip("")
else:
cb.setChecked(False)
cb.setToolTip(
f"No cached/stored {CH_LABELS[ch_idx]} data for this "
"angle yet — view it in the main window (or run "
"Batch Compute) first.")
self._lbl_none_available.setText(
"" if n_ok else "No angle has this value type ready yet.")
self._update_accept_enabled()
def _on_value_type_changed(self, _id: int):
self._apply_availability()
if not self._path_user_chosen:
self._refresh_default_path()
def _on_select_all_available(self):
for cb in self._angle_checks.values():
if cb.isEnabled():
cb.setChecked(True)
def _on_select_none(self):
for cb in self._angle_checks.values():
cb.setChecked(False)
def _refresh_default_path(self):
ch_idx = self._current_ch_idx()
name = f"{self._stem}_fused_roi_{CH_NAMES[ch_idx]}.csv"
self._edit_path.setText(str(Path(self._out_dir) / name))
self._update_accept_enabled()
def _on_browse(self):
path, _ = QFileDialog.getSaveFileName(
self, "Export Fused ROI as CSV", self._edit_path.text(),
"CSV files (*.csv);;All files (*)")
if path:
self._edit_path.setText(path)
self._path_user_chosen = True
self._update_accept_enabled()
def _update_accept_enabled(self):
any_checked = any(cb.isChecked() for cb in self._angle_checks.values())
self._btn_export.setEnabled(any_checked and bool(self._edit_path.text()))
# ---- getters ---------------------------------------------------------
def get_ch_idx(self) -> int:
return self._current_ch_idx()
def get_selected_angles(self) -> list[int]:
return sorted(a for a, cb in self._angle_checks.items() if cb.isChecked())
def get_output_path(self) -> str:
return self._edit_path.text()
sras_viewer/main_window.py
+241 -437
View File
@@ -16,27 +16,24 @@ from PyQt6.QtWidgets import (
import sras_compute as compute
from sras_compute import (
ManualAngleParams, apply_alignment, build_manual_alignment,
load_manual_alignment, save_manual_alignment, sidecar_path,
load_manual_alignment, sidecar_path,
)
from sras_format import (
CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, mv_to_adc,
_FALLBACK_YMULT_MV, _FALLBACK_YOFF_ADC,
)
from sras_workers import (
BatchCacheWorker, ComputeWorker, DcPrecomputeWorker, LoadWorker,
AngleAlignmentWorker, BatchCacheWorker, 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,
)
from .align_wizard import AlignmentWizard
from .dialogs import (
FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
_RIGHT_PANEL_W, Jobs, _form, _group, _make_dspin, _scroll_panel, _wrap_label,
)
from .dialogs import FftOptionsDialog, ManualAlignmentDialog
# ---------------------------------------------------------------------------
# Main window
@@ -76,35 +73,32 @@ class SrasViewerWindow(QMainWindow):
except (TypeError, ValueError):
pad = 1
self._fft_pad_factor: int = max(1, min(256, pad)) # 1 = no padding
try:
row_avg_n = int(self._settings.value("fft/row_avg_n", 3))
except (TypeError, ValueError):
row_avg_n = 3
self._pending_row_avg_n: int = max(1, min(50, row_avg_n))
# Convert menu: batch DC/FFT compute-and-store (v6 -> v7)
self._batch_errors: list[str] = []
# Display-only settings (colormap, grating) never trigger a
# recompute — they're applied to cached data on redraw. DC images
# (CH3/CH4) are cheap and precomputed for every angle in the
# background right after load. CH1/Velocity FFT images are
# computed lazily (with a progress popup) the first time an
# angle/threshold combination is viewed — using the cached DC4
# image to skip the FFT entirely for masked-out pixels — and
# cached per (angle, 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.
# (CH3/CH4) are cheap: they come straight from the file's stored
# cache when Batch Compute DC has been run for it (see
# _stored_dc_image), and are otherwise precomputed for every angle
# in the background right after load. CH1/Velocity FFT images come
# from the file's own stored cache when Batch Compute FFT has been
# run for it (see _stored_fft_image), and are otherwise computed
# lazily (with a progress popup) the first time an angle/threshold
# combination is viewed — using the cached DC4 image to skip the
# FFT entirely for masked-out pixels. Either way the result is
# cached per (angle, bg_sub, n_fft, threshold) so revisiting the
# same combination is free.
self._dc_cache: dict[tuple[int, int], np.ndarray] = {}
self._fft_cache: dict[tuple[int, float], np.ndarray] = {}
self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {}
self._dc_generation: int = 0
# Angle alignment ("Fusion" menu)
self._alignment_result = None
self._alignment_generation: int = 0
self._aligned_cache: dict[tuple, np.ndarray] = {}
self._align_wizard: AlignmentWizard | None = None
self._manual_align_dialog: ManualAlignmentDialog | None = None
self._build_ui()
@@ -229,12 +223,12 @@ class SrasViewerWindow(QMainWindow):
self.spin_angle.setRange(0, 0)
self.spin_angle.setEnabled(False)
self.spin_angle.setMinimumWidth(64)
# valueChanged with keyboard tracking off, not editingFinished: the
# latter fires only on Return or focus-out, so stepping the angle (an
# arrow click or Up/Down, the ordinary way to walk a scan) changed the
# number and left the image behind. Tracking off is what keeps
# valueChanged from also firing per keystroke mid-typing, which on a
# large scan would launch a compute for every intermediate angle.
# Keyboard tracking off is what makes valueChanged safe to act on: it
# stops the signal firing per keystroke, so typing "12" is one change
# to angle 12 (on Return or focus-out) rather than a trip to angle 1
# first. Arrow clicks and the wheel still emit immediately, which
# editingFinished alone did not — that is why stepping the angle used
# to leave the plot on the previous one until the box lost focus.
self.spin_angle.setKeyboardTracking(False)
self.spin_angle.valueChanged.connect(self._on_view_changed)
self.lbl_angle_deg = QLabel("—")
@@ -247,10 +241,14 @@ class SrasViewerWindow(QMainWindow):
ar.addStretch()
view_form.addRow("Angle:", angle_field)
self.combo_channel = _combo(CH_LABELS, min_chars=12)
self.combo_channel = QComboBox()
self.combo_channel.addItems(CH_LABELS)
self.combo_channel.setEnabled(False)
self.combo_channel.setSizePolicy(QSizePolicy.Policy.Expanding,
QSizePolicy.Policy.Fixed)
self.combo_channel.setSizeAdjustPolicy(
QComboBox.SizeAdjustPolicy.AdjustToMinimumContentsLengthWithIcon)
self.combo_channel.setMinimumContentsLength(12)
self.combo_channel.currentIndexChanged.connect(self._on_channel_changed)
view_form.addRow("Channel:", self.combo_channel)
vl.addLayout(view_form)
@@ -280,9 +278,7 @@ class SrasViewerWindow(QMainWindow):
self.chk_bg_sub.setEnabled(False)
self.chk_bg_sub.setToolTip(
"Subtract the stored background waveform from each CH1 frame\n"
"before computing the FFT (v4+ files only). Applies to angles\n"
"not yet computed and to future batch recomputes — it does not\n"
"change an image already shown or already stored in the file."
"before computing the FFT (v4+ files only)."
)
self.chk_bg_sub.toggled.connect(self._on_bg_sub_toggled)
vl.addWidget(self.chk_bg_sub)
@@ -337,11 +333,6 @@ class SrasViewerWindow(QMainWindow):
self.btn_export_roi.clicked.connect(self._on_export_roi_csv)
rl.addWidget(self.btn_export_roi)
self.btn_export_fused_roi = QPushButton("Export Fused ROI…")
self.btn_export_fused_roi.setEnabled(False)
self.btn_export_fused_roi.clicked.connect(self._on_export_fused_roi_csv)
rl.addWidget(self.btn_export_fused_roi)
self.lbl_roi_center = _wrap_label("centroid: —", _CSS_HINT)
self.lbl_roi_size = _wrap_label("bbox: —", _CSS_HINT)
self.lbl_roi_npix = _wrap_label("pixels inside: —", _CSS_HINT)
@@ -445,13 +436,23 @@ class SrasViewerWindow(QMainWindow):
fft_menu.addAction(fft_act)
fusion_menu = menubar.addMenu("&Fusion")
self._wizard_act = QAction("Alignment &Wizard…", self)
self._wizard_act.setStatusTip(
"Align the angles, crop to a region of interest, and save the "
"aligned data as a new .sras file. Requires >1 angle.")
self._wizard_act.setEnabled(False)
self._wizard_act.triggered.connect(self._on_alignment_wizard)
fusion_menu.addAction(self._wizard_act)
self._alignment_act = QAction("Angle &Alignment", self)
self._alignment_act.setStatusTip(
"Compute a rotation+translation alignment across all angles "
"(from CH4 masks) and enable Aligned View. Requires >1 angle.")
self._alignment_act.setEnabled(False)
self._alignment_act.triggered.connect(self._on_angle_alignment)
fusion_menu.addAction(self._alignment_act)
self._manual_align_act = QAction("&Manual Alignment…", self)
self._manual_align_act.setStatusTip(
"Open an interactive dialog to align angles by eye: overlaid CH4 "
"threshold masks, keyboard nudge (translate + rotate), auto "
"de-rotate to the known scan angles, and save/clear a persistent "
"alignment.")
self._manual_align_act.setEnabled(False)
self._manual_align_act.triggered.connect(self._on_manual_alignment)
fusion_menu.addAction(self._manual_align_act)
convert_menu = menubar.addMenu("&Convert")
self._batch_dc_act = QAction("Batch Compute DC and &Store…", self)
@@ -469,17 +470,6 @@ class SrasViewerWindow(QMainWindow):
self._batch_fft_act.triggered.connect(lambda: self._on_batch_compute("fft"))
convert_menu.addAction(self._batch_fft_act)
self._batch_fft_rowavg_act = QAction(
"Batch Compute Row-A&veraged FFT and Store…", self)
self._batch_fft_rowavg_act.setStatusTip(
"Select .sras files and compute+store a same-row, distance-weighted "
"smoothed FFT peak-frequency image for every angle — improves SNR "
"on noisy regions. DC images and raw waveform data are never "
"touched; the raw (unsmoothed) FFT is always a recompute away. "
"Converts v6 files to v7 in place.")
self._batch_fft_rowavg_act.triggered.connect(self._on_batch_compute_row_avg)
convert_menu.addAction(self._batch_fft_rowavg_act)
# ------------------------------------------------------------------
# Drag-and-drop
# ------------------------------------------------------------------
@@ -527,9 +517,9 @@ class SrasViewerWindow(QMainWindow):
# A manual-alignment dialog bound to the previous file must not
# survive a reload — its per-angle state (and the sras it was
# constructed against) no longer matches the new file's geometry.
if self._align_wizard is not None:
self._align_wizard.close()
self._align_wizard = None
if self._manual_align_dialog is not None:
self._manual_align_dialog.close()
self._manual_align_dialog = None
# Caches (and any in-flight DC precompute) belong to the previous
# file's geometry — discard and start fresh. Bumping the generation
@@ -619,50 +609,34 @@ class SrasViewerWindow(QMainWindow):
n_fft = sum(1 for x in s.precomputed_freq_mhz if x is not None)
if n_dc or n_fft:
bg_note = " (bg-sub)" if s.precomputed_bg_sub else " (no bg-sub)"
avg_note = (f", row-averaged n={s.precomputed_row_avg_n}"
if s.precomputed_row_avg_n else "")
pad_note = (f", pad {s.precomputed_pad_factor}x"
if s.precomputed_pad_factor > 1 else "")
if n_fft and s.precomputed_pad_factor > 1 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"FFT {n_fft}/{s.n_angles} angles"
f"{bg_note + pad_note if n_fft else ''} "
"— display is instant for cached angles")
notes += self._cache_mismatch_notes()
elif s.version == 7:
notes.append("v7 format: no cache blocks stored yet")
# A stored FFT cache the current settings can't read back is the one
# failure here with no other symptom: the images are right there in
# the file and every angle still recomputes. Say why.
if n_fft:
reasons = []
if s.precomputed_pad_factor != self._fft_pad_factor:
reasons.append(f"stored at pad {s.precomputed_pad_factor}x, "
f"FFT Options is set to {self._fft_pad_factor}x")
if s.precomputed_bg_sub != (self.chk_bg_sub.isChecked()
and s.background is not None):
reasons.append("stored with bg-sub "
f"{'on' if s.precomputed_bg_sub else 'off'}")
if reasons:
notes.append(f"! Cached FFT unusable as configured ({'; '.join(reasons)})"
" — CH1/Velocity will recompute. Re-run Convert -> "
"Batch Compute FFT to cache at the current settings.")
self.lbl_frame_warn.setText("\n".join(notes))
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.
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.
Asks compute for the reasons rather than restating the accept rule,
so a new provenance field can only be added in one place.
"""
s = self._sras
if s is None or all(x is None for x in s.precomputed_freq_mhz):
return []
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."]
# ------------------------------------------------------------------
# Controls
# ------------------------------------------------------------------
@@ -697,10 +671,11 @@ class SrasViewerWindow(QMainWindow):
can_batch = not self._job_running(Jobs.BATCH)
self._batch_dc_act.setEnabled(can_batch)
self._batch_fft_act.setEnabled(can_batch)
self._batch_fft_rowavg_act.setEnabled(can_batch)
self._wizard_act.setEnabled(
has_file and s.n_angles > 1 and self._align_wizard is None)
self._alignment_act.setEnabled(
has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN))
self._manual_align_act.setEnabled(
has_file and s.n_angles > 1 and not self._job_running(Jobs.ALIGN))
self.chk_aligned_view.setEnabled(enabled and self._alignment_result is not None)
self._update_roi_ui()
@@ -709,12 +684,14 @@ class SrasViewerWindow(QMainWindow):
self._on_view_changed()
def _on_bg_sub_toggled(self):
# 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
# recompute — never what's already shown. Just keep the info panel's
# divergence note current.
# Background subtraction changes the FFT input, so it genuinely
# invalidates the cached raw FFT (the cache key includes it) —
# _refresh_display() recomputes only on a miss for the new state.
# It can also decide whether the file's stored cache is readable, so
# the info panel's verdict on that has to be redrawn with it.
self._update_scan_info_labels()
if self._is_fft_mode():
self._update_scan_info_labels()
self._refresh_display()
def _on_grating_changed(self):
# Grating is a pure post-multiply on the cached frequency image —
@@ -833,91 +810,6 @@ class SrasViewerWindow(QMainWindow):
self.statusBar().showMessage(
f"Exported ROI ({n_pix} pixels) to {Path(path).name}")
def _on_export_fused_roi_csv(self):
s = self._sras
if s is None:
return
roi = self.image_canvas.get_roi()
if roi is None or not self._fused_grid_ready():
return # button is disabled in these states; defensive no-op
angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
grid_note = (
"Using the live Fusion alignment canvas."
if self._alignment_result is not None else
"Using the file's shared raw grid (no live alignment result — "
"this file's angles already share one grid).")
dlg = FusedRoiExportDialog(
self, angles=angles,
availability_fn=lambda a, c: self._cached_value_image(a, c) is not None,
default_ch_idx=self.combo_channel.currentIndex(),
out_dir=str(s.path.parent), stem=s.path.stem, grid_note=grid_note)
if dlg.exec() != QDialog.DialogCode.Accepted:
return
self._write_fused_roi_csv(
roi, dlg.get_ch_idx(), dlg.get_selected_angles(), dlg.get_output_path())
def _write_fused_roi_csv(self, roi, ch_idx: int, angle_idxs: list[int],
out_path: str):
"""Mask once on the fused grid, then one value column per selected
angle. No row/frame index columns: once angles are fused onto one
shared canvas there is no single meaningful raw (row, frame) per
output pixel, unlike the single-angle _on_export_roi_csv above."""
s = self._sras
if s is None or not angle_idxs or not out_path:
return
x_axis, y_axis = self._fused_export_axes()
mask = roi.mask_for_grid(x_axis, y_axis)
if not mask.any():
self.statusBar().showMessage(
"ROI does not overlap any pixel on the fused grid")
return
columns: list[np.ndarray] = []
used: list[int] = []
skipped: list[int] = []
for a in angle_idxs:
img = self._fused_value_image(a, ch_idx)
if img is None or img.shape != mask.shape:
# Defensive only: nothing else can mutate the caches between
# dialog-accept and this synchronous call in a single-
# threaded GUI callback, so this should never trigger — skip
# the angle rather than abort the whole export.
skipped.append(a)
continue
columns.append(img[mask].astype(np.float64))
used.append(a)
if not columns:
self.statusBar().showMessage(
"Nothing to export — no selected angle had data")
return
X, Y = np.meshgrid(np.asarray(x_axis, dtype=np.float64),
np.asarray(y_axis, dtype=np.float64))
data = np.column_stack([X[mask], Y[mask], *columns])
ch_name = CH_NAMES[ch_idx]
corners_str = " ".join(f"({p[0]:.6g},{p[1]:.6g})" for p in roi.corners())
angle_list_str = ", ".join(f"{s.angles_deg[a]:.4g}°" for a in used)
header_cols = ["x_mm", "y_mm"] + [f"v_{s.angles_deg[a]:.4g}deg" for a in used]
header = (
f"# ROI quad corners (BL BR TR TL) mm: {corners_str}\n"
f"# source: {s.path.name}, value={ch_name}, "
f"grid={'aligned canvas' if self._alignment_result is not None else 'shared raw grid'}\n"
f"# angles (deg): {angle_list_str}\n"
f"# n_pixels={int(mask.sum())}\n"
+ ",".join(header_cols)
)
np.savetxt(out_path, data, delimiter=",",
fmt=["%.6g"] * data.shape[1], header=header, comments="")
note = (f" ({len(skipped)} angle(s) skipped — no longer available)"
if skipped else "")
self.statusBar().showMessage(
f"Exported fused ROI ({int(mask.sum())} pixels, {len(used)} "
f"angle(s)) to {Path(out_path).name}{note}")
# ------------------------------------------------------------------
# ROI
# ------------------------------------------------------------------
@@ -947,7 +839,6 @@ class SrasViewerWindow(QMainWindow):
self.lbl_roi_npix.setText("pixels inside: —")
self.btn_clear_roi.setEnabled(False)
self.btn_export_roi.setEnabled(False)
self._update_fused_export_enabled()
return
cen = roi.centroid()
@@ -971,26 +862,6 @@ class SrasViewerWindow(QMainWindow):
self.lbl_roi_npix.setText(f"pixels inside: {npix}")
self.btn_clear_roi.setEnabled(True)
self.btn_export_roi.setEnabled(self._current_image is not None and npix > 0)
self._update_fused_export_enabled()
def _update_fused_export_enabled(self):
s = self._sras
roi = self.image_canvas.get_roi() if s is not None else None
if s is None:
ready, reason = False, "Open a file first."
elif roi is None:
ready, reason = False, "Draw an ROI first."
elif not self._fused_grid_ready():
ready, reason = False, (
"Angles need a common grid to fuse: run Fusion → Alignment "
"Wizard, or open a file the wizard already exported.")
else:
ready, reason = True, (
"Export the ROI as one CSV with a column per selected "
"angle's value (choose which value and which angles in "
"the dialog).")
self.btn_export_fused_roi.setEnabled(ready)
self.btn_export_fused_roi.setToolTip(reason)
# ------------------------------------------------------------------
# Display
@@ -1001,6 +872,14 @@ class SrasViewerWindow(QMainWindow):
return None
return self._sras.samples_per_frame * self._fft_pad_factor
def _cached_fft_pad_factor(self) -> int | None:
"""The pad factor the open file's stored FFT images were computed at,
or None if it has none stored."""
if self._sras is None or not any(x is not None
for x in self._sras.precomputed_freq_mhz):
return None
return self._sras.precomputed_pad_factor
def _is_fft_mode(self) -> bool:
"""Is the selected channel an FFT-derived (CH1/Velocity) mode?"""
return (self._sras is not None
@@ -1014,15 +893,8 @@ 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())
return (angle_idx, self.chk_bg_sub.isChecked(), self._current_n_fft(),
self.spin_threshold_mv.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
@@ -1047,137 +919,69 @@ class SrasViewerWindow(QMainWindow):
self._aligned_cache[key] = cached
return cached
# ------------------------------------------------------------------
# Fused ROI export (Export Fused ROI…)
# ------------------------------------------------------------------
def _fused_grid_ready(self) -> bool:
"""Is there a common (x, y) grid to fuse angles onto right now?
Either a live alignment result exists (angles are reconciled onto its
canvas via apply_alignment), or — with no alignment run this session
— the open file's raw per-angle grids already coincide, which is true
for a .sras file the Alignment Wizard itself previously exported (see
scan_format.md, "Files written by the viewer's Alignment Wizard")."""
return (self._sras is not None
and (self._alignment_result is not None
or self._sras.angles_share_raw_grid()))
def _fused_export_axes(self) -> tuple[np.ndarray, np.ndarray]:
"""(x_axis, y_axis) a fused ROI export masks and labels against: the
alignment canvas when a live result exists — the more current,
deliberate source of truth even if the raw grids happen to already
match too — else the grid every angle already shares."""
if self._alignment_result is not None:
return self._aligned_canvas_axes()
s = self._sras
return s.x_axis_mm(0), s.y_positions_mm(0)
def _fused_value_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""One angle's image on the fused grid for (angle, channel), scaled
exactly like the on-screen display — so a Velocity export and the
on-screen Velocity image share the same scaling code. None if
nothing is cached/stored for this angle/channel without a compute."""
raw = self._cached_value_image(angle_idx, ch_idx)
if raw is None:
return None
img = (self._scale_for_display(raw, ch_idx)
if ch_idx in CH1_DERIVED_MODES else raw)
if self._alignment_result is not None:
return self._get_aligned_display_image(img, angle_idx, ch_idx)
return img # angles already share the raw grid — no resample needed
def _stored_fft_image(self, angle_idx: int) -> np.ndarray | None:
"""The open file's own stored peak-frequency image for this angle,
masked and ready to display, or None if the file has nothing stored
for it.
"""The file's own batch-computed FFT image for this angle, masked for
the current threshold — or None if it can't serve the current settings.
Asks for the image at the settings it was actually computed under
(sras.precomputed_bg_sub / precomputed_pad_factor /
precomputed_row_avg_n) rather than the window's live bg-sub/pad
controls, so a stored image is always shown once present — those
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.
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
falls through to the background worker, which reaches the same stored
image via compute_rf_image and pays for the mask off the GUI thread.
Batch Compute FFT stores a peak-frequency image per angle in the file
itself, and the whole point of paying for that once is that display is
then instant. Without this check the viewer only ever found the stored
image deep inside ComputeWorker, so every angle change after a batch
still dispatched a background job behind a "Computing FFT…" popup for
an image that was already sitting on disk.
"""
s = self._sras
n_fft = (s.samples_per_frame * s.precomputed_pad_factor
if s.precomputed_pad_factor > 1 else None)
return compute.cached_rf_image(
s, angle_idx,
img = compute.cached_rf_image(
self._sras, angle_idx,
dc_threshold_mv=self.spin_threshold_mv.value(),
apply_bg_sub=s.precomputed_bg_sub,
n_fft=n_fft,
row_avg_n=s.precomputed_row_avg_n,
apply_bg_sub=self.chk_bg_sub.isChecked(),
n_fft=self._current_n_fft(),
dc4_mv=self._dc_cache.get((angle_idx, CH4_IDX)),
# Masking is a copy plus a comparison, but *reading* CH4 to build
# the mask is a whole-channel read — that case belongs on a worker,
# so it returns None here and falls through to _start_compute.
allow_dc_recompute=False)
if img is not None:
self._fft_cache[self._fft_cache_key(angle_idx)] = img
return img
def _cached_value_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""The already-available (no compute) image for (angle, channel):
this window's own in-session dicts first, then the file's stored
v5/v7 cache blocks — the same two tiers, in the same cost order, that
_refresh_display and Export Fused ROI both need, so both share this
one lookup rather than drifting apart. None means genuinely nothing
is cached/stored, i.e. only a real compute could produce it — and
neither caller here is allowed to trigger one: _refresh_display falls
back to _start_compute() itself, and a fused export simply treats the
angle as unavailable.
def _stored_dc_image(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
"""The file's own batch-computed DC image for (angle, channel), if it
has one — the CH3/CH4 half of _stored_fft_image.
For CH1/Velocity this is the unscaled, already-masked peak-frequency
(MHz) image; callers displaying/exporting Velocity must still run it
through _scale_for_display.
The worker this saves would have handed back the very same stored
array (dc_image_mv prefers it over recomputing), so all it ever cost
was a thread and a progress popup — but that is exactly what made
"Batch Compute DC and Store" look like it had done nothing.
"""
s = self._sras
if s is None:
return None
if ch_idx in CH1_DERIVED_MODES:
key = self._fft_cache_key(angle_idx)
raw = self._fft_cache.get(key)
if raw is None:
raw = self._stored_fft_image(angle_idx)
if raw is not None:
# Masking the stored image is cheap but not free; keep the
# result so revisiting this angle costs nothing at all.
self._fft_cache[key] = raw
return raw
# A stored DC image needs no post-processing, so the file's own
# parsed array is served directly — as the compute path already
# does, _current_image is treated as read-only by every consumer.
cached = self._dc_cache.get((angle_idx, ch_idx))
if cached is None:
cached = s.cached_dc_mv(angle_idx, ch_idx)
return cached
img = self._sras.cached_dc_mv(angle_idx, ch_idx)
if img is not None:
self._dc_cache[(angle_idx, ch_idx)] = img
return img
def _refresh_display(self):
"""Show the image for the current angle/channel/threshold, using
cached data whenever possible and only falling back to a background
compute (with progress popup) when genuinely nothing is cached yet.
Two caches are consulted, in cost order: this window's own in-session
dicts, then the file's stored v5/v7 cache blocks. The second is what
makes a batch-computed file worth having — without it every angle
change queued a worker and a progress popup for an image already on
disk, which is exactly the cost the batch was run to avoid.
"""
compute (with progress popup) when genuinely nothing is cached yet."""
if self._sras is None:
return
angle_idx = self.spin_angle.value()
ch_idx = self.combo_channel.currentIndex()
raw = self._cached_value_image(angle_idx, ch_idx)
if raw is not None:
img = (self._scale_for_display(raw, ch_idx)
if ch_idx in CH1_DERIVED_MODES else raw)
self._show_image_now(img, angle_idx, ch_idx)
return
if ch_idx in CH1_DERIVED_MODES:
raw = self._fft_cache.get(self._fft_cache_key(angle_idx))
if raw is None:
raw = self._stored_fft_image(angle_idx)
if raw is not None:
self._show_image_now(self._scale_for_display(raw, ch_idx),
angle_idx, ch_idx)
return
else:
cached = self._dc_cache.get((angle_idx, ch_idx))
if cached is None:
cached = self._stored_dc_image(angle_idx, ch_idx)
if cached is not None:
self._show_image_now(cached, angle_idx, ch_idx)
return
# Nothing cached for these settings — need a real compute. Changing
# the DC threshold changes *which* pixels get an FFT at all, so it
@@ -1315,7 +1119,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
self._fft_cache[(angle_idx, self._pending_bg_sub,
self._current_n_fft(), self._pending_threshold)] = result
img = self._scale_for_display(result, ch_idx)
else:
img = result
@@ -1443,10 +1248,11 @@ class SrasViewerWindow(QMainWindow):
return
self._batch_errors = []
# Cache the FFT at the pad the viewer is actually displaying at,
# otherwise the batch stores images this window can never use.
# Cache at the pad factor currently configured, not a fixed pad 1 —
# a cache stored at any other padding is one this viewer can never
# read back (see _stored_fft_image).
worker = BatchCacheWorker(paths, mode, self.chk_bg_sub.isChecked(),
pad_factor=self._fft_pad_factor)
self._fft_pad_factor)
started = self._run_worker(
Jobs.BATCH, worker,
connect=(
@@ -1461,54 +1267,10 @@ class SrasViewerWindow(QMainWindow):
self._batch_dc_act.setEnabled(False)
self._batch_fft_act.setEnabled(False)
self._batch_fft_rowavg_act.setEnabled(False)
self._show_progress(
Jobs.BATCH, f"Batch computing {label} for {len(paths)} file(s)…",
maximum=100)
def _on_batch_compute_row_avg(self):
if self._job_running(Jobs.BATCH):
return
dlg = RowAverageFftOptionsDialog(
self, current_n=self._pending_row_avg_n,
current_threshold_mv=self.spin_threshold_mv.value(),
pixel_x_mm=self._sras.pixel_x_mm if self._sras is not None else None)
if dlg.exec() != QDialog.DialogCode.Accepted:
return
n, threshold_mv = dlg.get_half_width(), dlg.get_threshold_mv()
self._pending_row_avg_n = n
self._settings.setValue("fft/row_avg_n", n)
paths, _ = QFileDialog.getOpenFileNames(
self, "Select .sras files to batch-compute Row-Averaged FFT", "",
"SRAS files (*.sras);;All files (*)")
if not paths:
return
self._batch_errors = []
worker = BatchCacheWorker(paths, "fft_rowavg", self.chk_bg_sub.isChecked(),
dc_threshold_mv=threshold_mv, row_avg_n=n,
pad_factor=self._fft_pad_factor)
started = self._run_worker(
Jobs.BATCH, worker,
connect=(
("progress", lambda pct: self._set_progress(Jobs.BATCH, pct)),
("file_done", self._on_batch_file_done),
("finished", lambda p=paths: self._on_batch_finished(p)),
),
on_done=self._after_batch,
)
if not started:
return # a second trigger snuck in while a dialog was open
self._batch_dc_act.setEnabled(False)
self._batch_fft_act.setEnabled(False)
self._batch_fft_rowavg_act.setEnabled(False)
self._show_progress(
Jobs.BATCH,
f"Batch computing row-averaged FFT (n={n}, threshold={threshold_mv:.1f} mV) "
f"for {len(paths)} file(s)…", maximum=100)
def _on_batch_file_done(self, path: str, err: str):
if err:
self._batch_errors.append(f"{Path(path).name} — {err}")
@@ -1536,83 +1298,123 @@ class SrasViewerWindow(QMainWindow):
def _after_batch(self):
self._batch_dc_act.setEnabled(True)
self._batch_fft_act.setEnabled(True)
self._batch_fft_rowavg_act.setEnabled(True)
# ------------------------------------------------------------------
# Fusion: angle alignment
# ------------------------------------------------------------------
def _on_alignment_wizard(self):
def _on_angle_alignment(self):
if self._sras is None or self._sras.n_angles <= 1:
return
if self._align_wizard is not None:
self._align_wizard.raise_()
self._align_wizard.activateWindow()
ref_idx = 0
threshold_mv = self.spin_threshold_mv.value()
generation = self._alignment_generation
started = self._run_worker(
Jobs.ALIGN, AngleAlignmentWorker(self._sras, ref_idx, threshold_mv),
connect=(
("progress", lambda pct: self._set_progress("main", pct)),
("finished", lambda result, err, g=generation:
self._on_alignment_done(g, result, err)),
),
on_done=lambda: self._update_controls_enabled(self._sras is not None),
)
if not started:
return
self._alignment_act.setEnabled(False)
self._show_progress(
"main",
f"Computing angle alignment ({self._sras.n_angles} angles, "
f"ref=angle 0, CH4 mask ≥ {threshold_mv:.3f} mV)…",
maximum=100)
def _on_alignment_done(self, generation: int, result, error_msg: str):
self._close_progress("main")
if generation != self._alignment_generation:
return # a new file was loaded while this was computing — discard
if error_msg:
self.statusBar().showMessage(f"Angle alignment failed: {error_msg}")
return
# No generation bump: this result *is* the current generation's.
self._apply_alignment_result(result, view_checked=True,
bump_generation=False)
nr, nc = result.canvas_shape
self.statusBar().showMessage(
f"Angle alignment computed ({self._sras.n_angles} angles, "
f"canvas {nc}×{nr} px).")
self._refresh_display()
# ------------------------------------------------------------------
# Fusion: manual alignment
# ------------------------------------------------------------------
def _on_manual_alignment(self):
if self._sras is None or self._sras.n_angles <= 1:
return
if self._manual_align_dialog is not None:
self._manual_align_dialog.raise_()
self._manual_align_dialog.activateWindow()
return
ref_idx = 0
threshold_mv = self.spin_threshold_mv.value()
# Only the threshold carries over from a saved alignment. The wizard's
# first page starts every angle pre-rotated from the stage angles and
# owns the parameters from there, so inheriting saved per-angle values
# would make "Reset to pre-rotation only" mean something different
# each time.
seed: dict[int, ManualAngleParams] = {}
# Seed only from a previously *saved manual* alignment (this dialog's
# own Save also writes this sidecar) -- never from self._alignment_result
# when it holds the automatic Fusion -> Angle Alignment's output. That
# path's translation comes from FFT phase correlation, which is the
# very thing manual mode exists to work around; inheriting it here
# would silently reintroduce the same bad translations under a
# "manual" label, on top of the (correct) analytic rotation, which is
# exactly what makes manual mode look like it "still does the same
# thing" the automatic one does.
sidecar = load_manual_alignment(self._sras)
if sidecar is not None and sidecar.ref_angle_idx == ref_idx:
seed = dict(sidecar.per_angle)
threshold_mv = sidecar.dc_threshold_mv
cached_dc4 = {a: img for (a, ch), img in self._dc_cache.items() if ch == CH4_IDX}
wiz = AlignmentWizard(
dlg = ManualAlignmentDialog(
self, self._sras, ref_angle_idx=ref_idx, dc_threshold_mv=threshold_mv,
cached_dc4_mv=cached_dc4)
wiz.alignment_ready.connect(self._on_wizard_finished)
wiz.finished.connect(self._on_wizard_closed)
wiz.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
self._align_wizard = wiz
self._update_controls_enabled(True)
wiz.show()
seed_per_angle=seed, cached_dc4_mv=cached_dc4)
dlg.alignment_saved.connect(self._on_manual_alignment_saved)
dlg.alignment_cleared.connect(self._on_manual_alignment_cleared)
dlg.finished.connect(self._on_manual_align_dialog_closed)
dlg.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
self._manual_align_dialog = dlg
dlg.show()
def _on_wizard_closed(self, _result_code: int):
self._align_wizard = None
self._update_controls_enabled(self._sras is not None)
def _on_manual_align_dialog_closed(self, _result_code: int):
self._manual_align_dialog = None
def _on_wizard_finished(self, result, out_path: str):
"""The wizard exported a file; make the session match what was written.
Applies the *cropped* result, so Aligned View shows exactly the extent
that went into the file rather than the wider uncropped canvas, and
saves the sidecar for the input scan so reopening it lands in the same
place.
"""
if result is None:
return
self._apply_alignment_result(result, view_checked=True)
note = ""
try:
save_manual_alignment(
self._sras, result.ref_angle_idx, result.dc_threshold_mv,
{a: ManualAngleParams(t.rotation_deg, t.shift_mm)
for a, t in result.per_angle.items()})
except OSError as exc:
note = f" (could not save the sidecar: {exc})"
self._update_controls_enabled(self._sras is not None)
nr, nc = result.canvas_shape
self.statusBar().showMessage(
f"Aligned scan written to {Path(out_path).name}; Aligned View now "
f"shows the exported {nc}×{nr} px region.{note}")
if self._current_image is not None:
self._refresh_display()
def _apply_alignment_result(self, result, *, view_checked: bool):
def _apply_alignment_result(self, result, *, view_checked: bool,
bump_generation: bool = True):
"""Install (or clear, with result=None) the active alignment: reset
the aligned-image cache and set the Aligned View checkbox without
firing its change signal."""
self._alignment_result = result
self._aligned_cache = {}
if bump_generation:
self._alignment_generation += 1
with QSignalBlocker(self.chk_aligned_view):
self.chk_aligned_view.setChecked(view_checked)
self.chk_aligned_view.setEnabled(result is not None)
def _manual_alignment_changed(self, result, message: str):
self._apply_alignment_result(result, view_checked=result is not None)
self._update_controls_enabled(self._sras is not None)
self.statusBar().showMessage(message)
if self._current_image is not None:
self._refresh_display()
def _on_manual_alignment_saved(self, result, sidecar_path_str: str):
self._manual_alignment_changed(
result, f"Manual alignment saved to {Path(sidecar_path_str).name}")
def _on_manual_alignment_cleared(self):
self._manual_alignment_changed(None, "Manual alignment cleared.")
# ------------------------------------------------------------------
# FFT Options
# ------------------------------------------------------------------
@@ -1625,6 +1427,7 @@ class SrasViewerWindow(QMainWindow):
samples_per_frame=self._sras.samples_per_frame if self._sras else None,
sample_rate_hz=self._sras.sample_rate_hz if self._sras else None,
grating_um=self.spin_grating_um.value(),
cached_pad_factor=self._cached_fft_pad_factor(),
)
if dlg.exec() != QDialog.DialogCode.Accepted:
return
@@ -1632,18 +1435,19 @@ class SrasViewerWindow(QMainWindow):
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
# batch recompute — never what's already shown. Just keep the info
# panel's divergence note current.
# Pad factor changes the FFT bin count, so it genuinely invalidates
# the cached raw FFT (part of the cache key) — _refresh_display()
# recomputes only on a cache miss. It also decides whether the file's
# stored cache can be read back at all, which the info panel reports.
self._update_scan_info_labels()
if self._is_fft_mode():
self._update_scan_info_labels()
self._refresh_display()
# ------------------------------------------------------------------
def closeEvent(self, event):
if self._align_wizard is not None:
self._align_wizard.close()
if self._manual_align_dialog is not None:
self._manual_align_dialog.close()
# Signal every cancellable worker first, then wait. Waiting without
# signalling means sitting out whatever is in flight — on a large
sras_workers.py
+50 -69
View File
@@ -15,8 +15,9 @@ import numpy as np
from PyQt6.QtCore import QObject, pyqtSignal
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_compute import (
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
)
from sras_format import CH3_IDX, CH4_IDX, SrasFile
# Concurrency caps. Batch conversion runs one process per file, and each of
@@ -186,15 +187,12 @@ class BatchCacheWorker(QObject):
an existing v7 file's cache blocks without disturbing whatever the other
block already holds.
*mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
images, unmasked — masking is applied at display time, same as v5's PREC
convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
averaging before the FFT — needs *dc_threshold_mv* and a positive
*row_avg_n*; see ``sras_compute.cache_file``).
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.
convention). FFT images are computed and stored at *pad_factor*, the
viewer's own zero-padding setting: a stored image only serves a view
asking for the same padding, so caching at any other one produces a file
the viewer will never read back.
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
@@ -207,14 +205,11 @@ class BatchCacheWorker(QObject):
finished = pyqtSignal()
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):
super().__init__()
self._paths = paths
self._mode = mode
self._apply_bg_sub = apply_bg_sub
self._dc_threshold = dc_threshold_mv
self._row_avg_n = row_avg_n
self._pad_factor = pad_factor
def _report(self, path: str, err: str, done: int, total: int):
@@ -241,9 +236,7 @@ class BatchCacheWorker(QObject):
futures = {
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), per_proc_workers,
pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n): p
self._pad_factor): p
for p in paths
}
for fut in as_completed(futures):
@@ -268,9 +261,7 @@ class BatchCacheWorker(QObject):
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)
self._pad_factor)
except Exception as exc:
err = str(exc)
done += 1
@@ -308,9 +299,34 @@ class BatchCacheWorker(QObject):
self.finished.emit()
class AngleAlignmentWorker(QObject):
"""Computes the rigid (rotation + translation, never scale) alignment for
every angle in *sras* against *ref_angle_idx*, by cross-correlating each
angle's CH4 image against the reference's. Both the rotation and the
translation are found from image content — see compute_angle_alignment.
"""
progress = pyqtSignal(int) # 0–100
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._threshold = dc_threshold_mv
def run(self):
try:
result = compute_angle_alignment(
self._sras, self._ref, self._threshold,
progress_cb=self.progress.emit)
self.finished.emit(result, "")
except Exception as exc:
self.finished.emit(None, str(exc))
class Ch4MaskWorker(_PooledWorker):
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
alignment wizard's initial threshold-mask stack.
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
ManualAlignmentDialog's initial threshold-mask overlay.
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
from raw waveforms, so this only does real work for a file that hasn't
@@ -319,7 +335,7 @@ class Ch4MaskWorker(_PooledWorker):
every file load) hasn't reached yet. In the common case — the user opens
Fusion -> Manual Alignment after DC precompute has already finished —
*angle_indices* is empty and this worker is never even constructed (see
CorrelatePage._start_mask_prep).
ManualAlignmentDialog._start_mask_prep).
"""
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
@@ -346,8 +362,8 @@ class Ch4MaskWorker(_PooledWorker):
class CrossCorrelateWorker(_PooledWorker):
"""Rigid registration (rotation + translation, never scale) of each of
*angle_indices* against *ref_angle_idx*, for the alignment wizard's
Run/Re-run Correlation button.
*angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto
Cross-Correlate button.
Runs on a background thread — registering a real many-angle,
high-resolution scan takes long enough that doing it on the GUI thread
@@ -361,18 +377,17 @@ class CrossCorrelateWorker(_PooledWorker):
angle_done = pyqtSignal(int, float, float, float, float, str)
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
dc4_mv: dict[int, np.ndarray], *, reg_kwargs: dict | None = None):
"""*reg_kwargs* is splatted into register_angle_to_reference — every
registration setting the wizard exposes (sources, threshold, search
width, seed, signs, refine, grid sizes) travels in it, so this class
holds no opinion about which knobs exist and exposing another needs no
change here."""
dc4_mv: dict[int, np.ndarray], *,
sources: tuple[str, ...], dc_threshold_mv: float,
search_deg: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._angles = angle_indices
self._dc4_mv = dc4_mv
self._reg_kwargs = dict(reg_kwargs or {})
self._sources = sources
self._threshold = dc_threshold_mv
self._search_deg = search_deg
def _plan(self) -> int:
return compute.registration_workers(self._sras)
@@ -382,45 +397,11 @@ class CrossCorrelateWorker(_PooledWorker):
def _one(self, a: int) -> tuple[int, compute.RigidFit]:
return a, compute.register_angle_to_reference(
self._sras, a, self._ref, self._dc4_mv, **self._reg_kwargs)
self._sras, a, self._ref, self._dc4_mv,
dc_threshold_mv=self._threshold, sources=self._sources,
search_deg=self._search_deg)
def _emit(self, result):
a, fit = result
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
fit.shift_mm[1], fit.score, fit.source)
class AlignedExportWorker(CancellableWorker):
"""Writes the aligned, cropped .sras on a background thread.
Unlike every other worker here this one produces a *file*, which changes
what cancellation has to mean: write_aligned_sras stages into a ".part"
sibling and removes it when should_stop() fires, so a cancelled or crashed
export leaves nothing behind. That matters more than it sounds — a
truncated .sras is not detectably broken, since the v6 parser reads a short
file as an aborted scan and opens it happily.
Cancellation is polled per output row chunk, the same granularity
CancellableWorker's docstring justifies, so closing the window never waits
on a multi-gigabyte write.
"""
progress = pyqtSignal(int) # 0-100
finished = pyqtSignal(str, str) # written path ("" = none), error
def __init__(self, sras: SrasFile, result, out_path: str):
super().__init__()
self._sras = sras
self._result = result
self._out_path = out_path
def run(self):
try:
written = write_aligned_sras(
self._sras, self._result, self._out_path,
progress_cb=self.progress.emit, should_stop=self._stopped)
if self._stopped():
self.finished.emit("", "") # cancelled: no file, no error
else:
self.finished.emit(str(written), "")
except Exception as exc:
self.finished.emit("", str(exc))
tests/test_align_export.py
-646
View File
@@ -1,646 +0,0 @@
"""Aligned/cropped .sras export: does the written file actually hold the
alignment the viewer showed?
The export is the one place an alignment stops being a transform applied on the
fly and becomes bytes on disk, so these tests care about two things above all:
the file's geometry describes what was written, and the pixels in it are the
same pixels apply_alignment would have drawn. The strongest check is the
round-trip — register the exported file against itself and demand identity,
which no amount of self-consistent-but-wrong index math can fake.
No Qt: this exercises sras_align_export and sras_compute directly.
"""
import struct
import numpy as np
import pytest
import sras_align_export as export
import sras_compute as compute
from sras_format import CH3_IDX, CH4_IDX, HDR_SIZE_V6, SrasFile, adc_to_mv, mv_to_adc
import tools.make_test_sras as gen
_THRESHOLD_MV = 80.0
# Same reasoning as tests/test_alignment.py: a quarter degree is already
# sub-pixel for this sample at the registration pitch.
_ROT_TOL_DEG = 0.5
_SHIFT_TOL_MM = 0.02
def dc_mv(sras: SrasFile, angle_idx: int, ch: int = CH4_IDX) -> np.ndarray:
return adc_to_mv(compute.compute_dc_image(sras, angle_idx, ch), *sras.cal(ch))
@pytest.fixture(scope="module")
def rig(tmp_path_factory):
"""The rotating-sample scan, its truth alignment, and its export."""
tmpdir = tmp_path_factory.mktemp("sras_export")
src_path = tmpdir / "rotating.sras"
meta = gen.write_rotating(src_path, n_angles=4)
sras = SrasFile(str(src_path))
params = {a: compute.ManualAngleParams(rot, shift)
for a, (rot, shift) in meta["truth"].items()}
result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
out_path = tmpdir / "rotating_aligned.sras"
export.write_aligned_sras(sras, result, out_path)
return type("Rig", (), dict(
tmpdir=tmpdir, src_path=src_path, sras=sras, meta=meta,
result=result, out_path=out_path, out=SrasFile(str(out_path))))
# ---------------------------------------------------------------------------
# Geometry and file structure
# ---------------------------------------------------------------------------
def test_output_is_v6_with_uniform_geometry(rig):
out, result = rig.out, rig.result
n_rows, n_cols = result.canvas_shape
assert out.version == 6
assert out.n_angles == rig.sras.n_angles
assert set(out.n_rows) == {n_rows}, "every angle must share the canvas rows"
assert set(out.n_frames) == {n_cols}, "every angle must share the canvas frames"
assert np.allclose(out.x_start_mm, result.canvas_origin_mm[0])
# x_delta must stay velocity/laser_freq or x_axis_mm() contradicts the
# geometry table; the canvas pitch is the reference angle's own pitch, so
# this is exact rather than approximate.
assert np.allclose(out.x_delta_mm_per_angle, rig.sras.pixel_x_mm)
assert out.pixel_x_mm == pytest.approx(rig.sras.pixel_x_mm)
def test_row_table_matches_the_canvas(rig):
expected = (rig.result.canvas_origin_mm[1]
+ np.arange(rig.result.canvas_shape[0]) * rig.result.canvas_dy_mm)
for a in range(rig.out.n_angles):
assert rig.out.y_positions_mm(a) == pytest.approx(expected, abs=1e-4)
def test_angle_table_and_calibration_round_trip(rig):
assert rig.out.angles_deg == pytest.approx(rig.sras.angles_deg)
for ch in range(rig.sras.n_channels):
assert rig.out.cal(ch) == pytest.approx(rig.sras.cal(ch))
assert rig.out.samples_per_frame == rig.sras.samples_per_frame
assert rig.out.bytes_per_sample == rig.sras.bytes_per_sample
assert rig.out.n_channels == rig.sras.n_channels
assert rig.out.background == pytest.approx(rig.sras.background)
def test_no_cache_tail(rig):
"""File ends exactly at the waveform data — nothing trailing.
A stale cache tail would be indexed by the *input's* grid, so the export
must not carry one; asserting on the exact file size is what proves it,
since a v7 tail would simply be ignored by a v6 parser.
"""
end = max(off + n for _, off, n in rig.out.iter_angle_blocks())
assert rig.out_path.stat().st_size == end
assert all(img is None for img in rig.out.precomputed_dc4_mv)
def test_declared_header_size_is_v6(rig):
raw = rig.out_path.read_bytes()[:HDR_SIZE_V6]
magic, version, n_angles = struct.unpack(">4sBH", raw[:7])
assert (magic, version, n_angles) == (b"SRAS", 6, rig.sras.n_angles)
# ---------------------------------------------------------------------------
# The pixels themselves
# ---------------------------------------------------------------------------
def test_export_matches_apply_alignment(rig):
"""The exported waveforms decode to the same DC image the viewer drew —
over the *whole* canvas, padding included.
Two rules have to be exactly right for this, and each fails differently:
* rounding must be floor(x + 0.5), not np.rint, or pixels on exact
half-integer boundaries pick the neighbouring source pixel;
* out-of-bounds must be tested on the fractional coordinate against
[0, n-1], not on the rounded index, or a one-pixel rim gets real data
where the preview shows padding.
Comparing every pixel rather than only the interior is what catches the
second one, since a rim discrepancy hides inside a `preview != 0` mask.
"""
for a in range(rig.sras.n_angles):
preview = compute.apply_alignment(rig.result, a, dc_mv(rig.sras, a))
actual = dc_mv(rig.out, a)
assert actual.shape == preview.shape
# Padding matches to within half an ADC step: apply_alignment pads with
# literal 0.0 mV, the export with the nearest integer ADC code to 0 mV.
tol = abs(rig.sras.cal(CH4_IDX)[0]) / 2.0 + 1e-4
# Exclude the epsilon rim the export deliberately keeps and scipy drops
# (see test_edge_tolerance_only_affects_the_epsilon_rim).
sr, sc = export._src_coords(rig.result.per_angle[a],
np.arange(preview.shape[0]), preview.shape[1])
rim = export._in_bounds(sr, sc, *rig.sras.image_shape(a)) & (preview == 0.0)
cmp = ~rim
assert actual[cmp] == pytest.approx(preview[cmp], abs=tol), \
f"angle {a}: exported pixels differ from the aligned preview"
# And exactly, wherever there is real data.
inside = (preview != 0.0)
assert inside.any(), f"angle {a}: preview is entirely padding"
assert actual[inside] == pytest.approx(preview[inside], abs=1e-6), \
f"angle {a}: exported data pixels are not bit-equal to the preview"
def test_reference_angle_is_exported_whole(rig):
"""The reference angle must survive as a complete, exact integer crop.
It is the coordinate authority — its transform is the identity with an
integer offset by construction — so every one of its source pixels has to
appear in the export. This is what _EDGE_TOL exists for: that offset comes
out of the mm-space affine chain as -20 - 7e-15, and a bare `>= 0` bounds
test silently drops the angle's entire first row and last column.
"""
ref = rig.result.ref_angle_idx
src_rows, src_frames = rig.sras.image_shape(ref)
plan = export.plan_export(rig.sras, rig.result)
assert plan.valid_px[ref] == src_rows * src_frames, \
"reference angle lost pixels to the in-bounds test"
# And the values themselves land as an exact, unrotated block.
src_img = dc_mv(rig.sras, ref)
out_img = dc_mv(rig.out, ref)
t = rig.result.per_angle[ref]
row0, col0 = (int(round(-t.offset[0])), int(round(-t.offset[1])))
assert np.array_equal(out_img[row0:row0 + src_rows, col0:col0 + src_frames],
src_img), \
"reference angle is not a verbatim block in the export"
def test_edge_tolerance_only_affects_the_epsilon_rim(rig):
"""Where the export's bounds test and scipy's disagree, the coordinate must
be within _EDGE_TOL of the boundary — i.e. only pixels whose scipy answer
was itself decided by float noise, never a real half-pixel decision."""
for a in range(rig.sras.n_angles):
t = rig.result.per_angle[a]
src_rows, src_frames = rig.sras.image_shape(a)
n_rows, n_cols = rig.result.canvas_shape
ones = np.ones((src_rows, src_frames), dtype=np.float32)
scipy_valid = compute.apply_alignment(rig.result, a, ones) > 0.5
sr, sc = export._src_coords(t, np.arange(n_rows), n_cols)
ours = export._in_bounds(sr, sc, src_rows, src_frames)
differ = ours != scipy_valid
assert not (scipy_valid & ~ours).any(), \
f"angle {a}: export drops pixels scipy keeps"
if differ.any():
# Every disagreement sits within the tolerance of an edge.
near = (np.abs(sr) <= export._EDGE_TOL)
near |= (np.abs(sr - (src_rows - 1)) <= export._EDGE_TOL)
near |= (np.abs(sc) <= export._EDGE_TOL)
near |= (np.abs(sc - (src_frames - 1)) <= export._EDGE_TOL)
assert near[differ].all(), \
f"angle {a}: bounds differ away from the epsilon rim"
def test_export_matches_apply_alignment_on_ch3(rig):
"""Channel-agnostic: the gather moves whole pixels, not per-channel images."""
for a in range(rig.sras.n_angles):
preview = compute.apply_alignment(rig.result, a,
dc_mv(rig.sras, a, CH3_IDX))
actual = dc_mv(rig.out, a, CH3_IDX)
inside = preview != 0.0
assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
def test_padding_is_zero_mv_not_zero_adc(rig):
"""Unreachable canvas pixels must read as ~0 mV on every channel.
Filling with literal zero ADC would decode to (0 - yoff) * ymult + yzero —
for this fixture's CH4 calibration that is +100 mV, well above any sensible
mask threshold, so the padding would masquerade as valid sample everywhere.
"""
a = rig.sras.n_angles - 1
t = rig.result.per_angle[a]
n_rows, n_cols = rig.result.canvas_shape
src_rows, src_frames = rig.sras.image_shape(a)
sr, sc = export._src_coords(t, np.arange(n_rows), n_cols)
outside = ~export._in_bounds(sr, sc, src_rows, src_frames)
assert outside.any(), "rotated angle should leave unreachable canvas corners"
for ch in (CH3_IDX, CH4_IDX):
img = dc_mv(rig.out, a, ch)
half_step = abs(rig.sras.cal(ch)[0]) / 2.0
assert np.abs(img[outside]).max() <= half_step + 1e-6, \
f"CH{ch} padding is not within half an ADC step of 0 mV"
# And the sanity check that makes the above meaningful: zero ADC would not
# have passed it.
assert abs(adc_to_mv(0, *rig.sras.cal(CH4_IDX))) > 10.0
def test_reregistering_the_export_is_identity(rig):
"""The export really is aligned: registering it against its own angle 0
recovers no rotation and no shift.
The end-to-end check — it fails for any index error, sign flip, wrong pivot
or origin mistake anywhere in crop/affine/gather, in a way the
self-consistency tests above cannot.
"""
dc4 = {a: dc_mv(rig.out, a) for a in range(rig.out.n_angles)}
for a in range(1, rig.out.n_angles):
fit = compute.register_angle_to_reference(
rig.out, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
seed_deg=0.0, seed_signs=(1,))
assert abs(fit.rotation_deg) <= _ROT_TOL_DEG, \
f"angle {a} still rotated by {fit.rotation_deg:.3f}° after export"
assert float(np.hypot(*fit.shift_mm)) <= _SHIFT_TOL_MM, \
f"angle {a} still shifted by {fit.shift_mm} mm after export"
def test_export_of_int16_input(rig, tmp_path):
"""bps=2 inputs keep their big-endian int16 dtype through the gather."""
src_path = tmp_path / "i16.sras"
gen.write(src_path, n_angles=2, samples_per_frame=16, bps=2)
sras = SrasFile(str(src_path))
result = compute.build_manual_alignment(sras, 0, 0.0, {})
out_path = tmp_path / "i16_aligned.sras"
export.write_aligned_sras(sras, result, out_path)
out = SrasFile(str(out_path))
assert out.bytes_per_sample == 2
assert out.data[0].dtype == np.dtype(">i2")
for a in range(sras.n_angles):
preview = compute.apply_alignment(result, a, dc_mv(sras, a))
actual = dc_mv(out, a)
inside = preview != 0.0
assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
# ---------------------------------------------------------------------------
# Cropping
# ---------------------------------------------------------------------------
def test_crop_is_a_window_of_the_full_canvas(rig):
"""crop_alignment_result must resample exactly the sub-rectangle it names.
Asserted as bit-exact equality, not approximately: the crop composes into
the affine's offset by an integer number of canvas pixels, so anything but
an exact match means the composition is wrong.
"""
n_rows, n_cols = rig.result.canvas_shape
row0, col0 = n_rows // 5, n_cols // 4
nr, nc = n_rows // 2, n_cols // 3
cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc)
assert cropped.canvas_shape == (nr, nc)
assert cropped.canvas_origin_mm[0] == pytest.approx(
rig.result.canvas_origin_mm[0] + col0 * rig.result.canvas_dx_mm)
assert cropped.canvas_origin_mm[1] == pytest.approx(
rig.result.canvas_origin_mm[1] + row0 * rig.result.canvas_dy_mm)
for a in range(rig.sras.n_angles):
img = dc_mv(rig.sras, a)
full = compute.apply_alignment(rig.result, a, img)
assert np.array_equal(
compute.apply_alignment(cropped, a, img),
full[row0:row0 + nr, col0:col0 + nc]), \
f"angle {a}: cropped resample is not the same window"
# Rotation/shift are properties of the angle, not of the canvas.
assert cropped.per_angle[a].rotation_deg == rig.result.per_angle[a].rotation_deg
assert cropped.per_angle[a].shift_mm == rig.result.per_angle[a].shift_mm
def test_cropped_export_round_trips(rig, tmp_path):
n_rows, n_cols = rig.result.canvas_shape
row0, col0, nr, nc = n_rows // 4, n_cols // 4, n_rows // 2, n_cols // 2
cropped = compute.crop_alignment_result(rig.result, row0, col0, nr, nc)
out_path = tmp_path / "cropped.sras"
export.write_aligned_sras(rig.sras, cropped, out_path)
out = SrasFile(str(out_path))
assert set(out.n_rows) == {nr} and set(out.n_frames) == {nc}
assert out.x_start_mm[0] == pytest.approx(cropped.canvas_origin_mm[0], abs=1e-4)
for a in range(rig.sras.n_angles):
preview = compute.apply_alignment(cropped, a, dc_mv(rig.sras, a))
actual = dc_mv(out, a)
inside = preview != 0.0
if inside.any():
assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
def test_crop_rejects_empty_window(rig):
with pytest.raises(ValueError, match="empty crop"):
compute.crop_alignment_result(rig.result, 0, 0, 0, 10)
with pytest.raises(ValueError, match="empty crop"):
compute.crop_alignment_result(rig.result, 0, 0, 10, -1)
# ---------------------------------------------------------------------------
# plan_export and overlap_stats
# ---------------------------------------------------------------------------
def test_plan_export_matches_what_was_written(rig):
plan = export.plan_export(rig.sras, rig.result)
n_rows, n_cols = rig.result.canvas_shape
assert (plan.n_rows, plan.n_frames) == (n_rows, n_cols)
assert plan.n_angles == rig.sras.n_angles
data_bytes = sum(n for _, _, n in rig.out.iter_angle_blocks())
assert plan.total_bytes == data_bytes
assert plan.bytes_per_angle * plan.n_angles == plan.total_bytes
# Coverage must agree with the pixels that actually carry data. The
# reference angle is unrotated, so its whole footprint lands inside.
ref_px = np.prod(rig.sras.image_shape(0))
assert plan.valid_px[0] == ref_px
for a in range(1, rig.sras.n_angles):
assert 0 < plan.valid_px[a] <= n_rows * n_cols
assert 0.0 < plan.coverage_frac(a) < 1.0
def test_plan_export_flags_a_crop_that_misses_an_angle(rig):
"""A crop over a corner the rotated angles cannot reach must warn, and the
export must still succeed by writing that angle as padding."""
n_rows, n_cols = rig.result.canvas_shape
corner = compute.crop_alignment_result(rig.result, 0, 0,
max(1, n_rows // 12),
max(1, n_cols // 12))
plan = export.plan_export(rig.sras, corner)
empty = [a for a in range(rig.sras.n_angles) if plan.valid_px[a] == 0]
assert empty, "top-left canvas corner should be unreachable for some angle"
assert any("all padding" in w for w in plan.warnings)
def test_overlap_stats():
counts = np.array([[0, 1, 2], [3, 3, 0], [0, 2, 3]])
stats = compute.overlap_stats(counts, 3)
assert stats["union_px"] == 6
assert stats["full_px"] == 3
assert stats["full_frac"] == pytest.approx(0.5)
assert stats["max_count"] == 3
assert stats["mean_count"] == pytest.approx((1 + 2 + 3 + 3 + 2 + 3) / 6)
assert stats["empty"] is False
empty = compute.overlap_stats(np.zeros((4, 4), dtype=int), 3)
assert empty["empty"] is True
assert empty["full_frac"] == 0.0 and empty["mean_count"] == 0.0
def test_largest_rect_at_least():
# A 2x3 block of 3s with a notch that a bounding box would swallow.
counts = np.array([
[0, 0, 0, 0, 0],
[0, 3, 3, 3, 0],
[0, 3, 3, 3, 0],
[0, 3, 0, 3, 0],
])
row0, col0, nr, nc = compute.largest_rect_at_least(counts, 3)
assert (nr * nc) == 6 and (row0, col0, nr, nc) == (1, 1, 2, 3)
assert (counts[row0:row0 + nr, col0:col0 + nc] >= 3).all()
# A column taller than the wide block is the better rectangle.
tall = np.array([[3, 3], [3, 0], [3, 0], [3, 0]])
r0, c0, nr2, nc2 = compute.largest_rect_at_least(tall, 3)
assert (r0, c0, nr2, nc2) == (0, 0, 4, 1)
assert compute.largest_rect_at_least(np.zeros((3, 3), dtype=int), 1) is None
# Whole-array case: no notch, so the answer is the array itself.
assert compute.largest_rect_at_least(np.full((3, 4), 2), 2) == (0, 0, 3, 4)
def test_largest_rect_matches_brute_force():
"""Randomized check against an O(n^4) reference.
The histogram sweep is short and easy to get subtly wrong — an off-by-one in
the stack unwind yields rectangles that are merely large, and "large but not
maximal" is invisible by eye on real data.
"""
def brute(good):
n_rows, n_cols = good.shape
best = 0
for r0 in range(n_rows):
for r1 in range(r0 + 1, n_rows + 1):
run = 0
for g in good[r0:r1].all(axis=0):
run = run + 1 if g else 0
best = max(best, run * (r1 - r0))
return best
rng = np.random.default_rng(0)
for _ in range(200):
counts = rng.integers(0, 3, size=(int(rng.integers(1, 9)),
int(rng.integers(1, 9))))
got = compute.largest_rect_at_least(counts, 2)
expected = brute(counts >= 2)
if got is None:
assert expected == 0
continue
row0, col0, nr, nc = got
assert (counts[row0:row0 + nr, col0:col0 + nc] >= 2).all(), \
f"rectangle is not pure:\n{counts}\n{got}"
assert nr * nc == expected, \
f"not maximal ({nr * nc} < {expected}):\n{counts}\n{got}"
def test_largest_rect_is_pure_on_the_real_fixture(rig):
"""On real overlap counts the returned rectangle must contain only
full-overlap pixels — the property a bounding box would violate."""
n = rig.sras.n_angles
masks = {a: (dc_mv(rig.sras, a) >= _THRESHOLD_MV).astype(np.float32)
for a in range(n)}
counts = sum(compute.apply_alignment(rig.result, a, masks[a]) > 0.5
for a in range(n)).astype(int)
assert counts.max() == n, "fixture alignment should have a full-overlap region"
rect = compute.largest_rect_at_least(counts, n)
assert rect is not None
row0, col0, nr, nc = rect
assert (counts[row0:row0 + nr, col0:col0 + nc] == n).all(), \
"convenience crop must not include pixels some angle misses"
# And it must beat the naive bounding box, which here is impure.
rr, cc = np.nonzero(counts == n)
bbox_pure = (counts[rr.min():rr.max() + 1, cc.min():cc.max() + 1] == n).all()
assert not bbox_pure, "fixture no longer exercises the bounding-box hazard"
# ---------------------------------------------------------------------------
# Legacy inputs, validation and durability
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("version", [2, 4])
def test_legacy_input_exports_as_v6(version, tmp_path):
"""v2-v5 inputs keep no verbatim preamble/background spans, so those
sections have to be re-encoded. v2 additionally has neither."""
src_path = tmp_path / f"legacy_v{version}.sras"
gen.write_legacy(src_path, version=version, n_angles=2)
sras = SrasFile(str(src_path))
result = compute.build_manual_alignment(sras, 0, 0.0, {})
out_path = tmp_path / f"legacy_v{version}_aligned.sras"
export.write_aligned_sras(sras, result, out_path)
out = SrasFile(str(out_path))
assert out.version == 6
assert out.n_angles == sras.n_angles
# A zero background rather than a zero-length one: consumers subtract it
# from a (spf,)-shaped row, which a length-0 array cannot broadcast against.
assert out.background is not None
assert out.background.size == sras.samples_per_frame
if sras.background is None:
assert np.all(out.background == 0)
assert any("no background" in w for w in
export.plan_export(sras, result).warnings)
# Calibration must survive: v2 has no preambles and falls back to the
# hardcoded scope constants, and the re-encoded empty preambles must land on
# exactly the same fallback.
for ch in range(sras.n_channels):
assert out.cal(ch) == pytest.approx(sras.cal(ch))
for a in range(sras.n_angles):
preview = compute.apply_alignment(result, a, dc_mv(sras, a))
actual = dc_mv(out, a)
inside = preview != 0.0
assert actual[inside] == pytest.approx(preview[inside], abs=1e-3)
def test_too_many_rows_is_rejected_before_writing(rig, tmp_path):
"""The geometry table stores n_rows as a u16; silently truncating would
write a file whose header disagrees with its own waveform block."""
huge = compute.crop_alignment_result(rig.result, 0, 0, 70000, 4)
out_path = tmp_path / "huge.sras"
with pytest.raises(ValueError, match="exceeds the .sras per-angle geometry"):
export.write_aligned_sras(rig.sras, huge, out_path)
assert not out_path.exists()
assert not out_path.with_name(out_path.name + ".part").exists()
def test_missing_transform_is_rejected(rig, tmp_path):
broken = compute.crop_alignment_result(rig.result, 0, 0,
*rig.result.canvas_shape)
del broken.per_angle[1]
with pytest.raises(ValueError, match="no transform for angle"):
export.write_aligned_sras(rig.sras, broken, tmp_path / "broken.sras")
def test_cancelled_export_leaves_nothing_behind(rig, tmp_path):
out_path = tmp_path / "cancelled.sras"
written = export.write_aligned_sras(rig.sras, rig.result, out_path,
should_stop=lambda: True)
assert written == out_path
assert not out_path.exists(), "cancelled export must not leave an output file"
assert not out_path.with_name(out_path.name + ".part").exists()
def test_failed_write_leaves_nothing_behind(rig, tmp_path):
"""An exception mid-write must remove the partial file: a short .sras is
not detectably broken — the v6 parser reads it as an aborted scan."""
out_path = tmp_path / "boom.sras"
def explode(_pct):
raise RuntimeError("boom")
with pytest.raises(RuntimeError, match="boom"):
export.write_aligned_sras(rig.sras, rig.result, out_path,
progress_cb=explode)
assert not out_path.exists()
assert not out_path.with_name(out_path.name + ".part").exists()
def test_progress_is_monotonic_and_completes(rig, tmp_path):
seen: list[int] = []
export.write_aligned_sras(rig.sras, rig.result, tmp_path / "prog.sras",
progress_cb=seen.append)
assert seen and seen[-1] == 100
assert seen == sorted(seen)
assert all(0 <= p <= 100 for p in seen)
def test_band_reader_path_is_byte_identical(rig, tmp_path, monkeypatch):
"""A source block too large to hold in RAM is served from sliding bands
instead. That path only runs on multi-gigabyte scans, so force it with a
tiny budget and demand the same bytes — otherwise the one code path that
matters on real data is the one never tested."""
whole = tmp_path / "whole.sras"
export.write_aligned_sras(rig.sras, rig.result, whole)
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 4096)
banded = tmp_path / "banded.sras"
export.write_aligned_sras(rig.sras, rig.result, banded)
assert banded.read_bytes() == whole.read_bytes()
def test_row_chunking_is_invariant(rig, tmp_path, monkeypatch):
"""Output must not depend on how many rows are buffered per write."""
base = tmp_path / "base.sras"
export.write_aligned_sras(rig.sras, rig.result, base)
monkeypatch.setattr(export, "_ROW_CHUNK", 1)
one = tmp_path / "one.sras"
export.write_aligned_sras(rig.sras, rig.result, one)
assert one.read_bytes() == base.read_bytes()
def test_refuses_to_overwrite_the_source(rig):
"""The source's waveform blocks are live read-only memmaps; writing over
the file would corrupt the reads the gather is making from it."""
with pytest.raises(ValueError, match="refusing to export onto the source"):
export.write_aligned_sras(rig.sras, rig.result, rig.src_path)
assert SrasFile(str(rig.src_path)).n_angles == rig.sras.n_angles
def test_overwrites_an_existing_file(rig, tmp_path):
out_path = tmp_path / "existing.sras"
out_path.write_bytes(b"not a scan")
export.write_aligned_sras(rig.sras, rig.result, out_path)
assert SrasFile(str(out_path)).version == 6
# ---------------------------------------------------------------------------
# The registration knobs the wizard exposes
# ---------------------------------------------------------------------------
def test_locked_rotation_returns_exactly_the_seed(rig):
"""search_deg=0 + one sign + refine=False pins rotation to the seed, which
is what "lock rotation to the stage angle" means on the wizard's first
page. Only the translation may be searched."""
dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)}
for a in range(1, rig.sras.n_angles):
nominal = compute.nominal_delta_deg(rig.sras, a, 0)
fit = compute.register_angle_to_reference(
rig.sras, a, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
search_deg=0.0, coarse_step_deg=2.0, seed_signs=(-1,), refine=False)
assert fit.rotation_deg == pytest.approx(-nominal)
def test_seed_deg_overrides_the_stage_angle(rig):
"""seed_deg=0.0 searches around no rotation at all, so a scan whose angles
are genuinely ~37° apart must fail to find them within a ±2° window —
proving the seed is what positions the search."""
dc4 = {a: dc_mv(rig.sras, a) for a in range(rig.sras.n_angles)}
fit = compute.register_angle_to_reference(
rig.sras, 1, 0, dc4, dc_threshold_mv=_THRESHOLD_MV,
search_deg=2.0, seed_deg=0.0, seed_signs=(1,), refine=False)
truth_rot = rig.meta["truth"][1][0]
assert abs(fit.rotation_deg) <= 2.0
assert abs(fit.rotation_deg - truth_rot) > 10.0
def test_rotation_candidates_signs():
both = compute._rotation_candidates(10.0, 2.0, 2.0)
assert both == compute._rotation_candidates(10.0, 2.0, 2.0, (-1, 1)), \
"default must stay the both-signs sweep"
assert compute._rotation_candidates(10.0, 0.0, 2.0, (1,)) == [10.0]
assert compute._rotation_candidates(10.0, 0.0, 2.0, (-1,)) == [-10.0]
# A zero seed collapses the two windows; the dedupe must keep one copy.
assert compute._rotation_candidates(0.0, 2.0, 2.0) == [-2.0, 0.0, 2.0]
def test_zero_mv_fill_code_is_clipped_to_dtype():
"""mv_to_adc is unclamped, so the fill code must be clipped or the int8
cast wraps around to a large-magnitude value."""
fake = type("S", (), dict(
n_channels=1, samples_per_frame=2,
cal=lambda self, ch: (1e-6, 0.0, 5000.0)))()
row = export._fill_row(fake, 3, np.dtype(np.int8))
assert row.shape == (1, 3, 2)
assert row.min() == row.max() == np.iinfo(np.int8).min
assert mv_to_adc(0.0, 1e-6, 0.0, 5000.0) < np.iinfo(np.int8).min
tests/test_alignment.py
+1 -1
View File
@@ -155,7 +155,7 @@ def test_all_angles_stack(rig):
def test_downsampled_preview_lands_with_full_res(rig):
# The wizard reprojects block-mean-downsampled masks, so the
# ManualAlignmentDialog reprojects block-mean-downsampled masks, so the
# affine has to account for the factor. When it did not, every preview
# layer came out magnified by that factor and offset — the overlay showed a
# blown-up crop of each mask, which is not something you can align by eye.
tests/test_fused_export.py
-75
View File
@@ -1,75 +0,0 @@
"""SrasFile.angles_share_raw_grid(): the no-alignment-needed gating path for
Export Fused ROI.
A plain multi-angle scan gives each angle its own bounding box and stage
x_start (scan_format.md's whole reason v6 geometry is per-angle), so it must
read as "not shareable" without a live alignment. A file the viewer's own
Alignment Wizard exported repeats one Per-Angle Geometry record and one Row
Table span for every angle (scan_format.md, "Files written by the viewer's
Alignment Wizard"), so it must read as "shareable" with no alignment needed
at all.
No Qt: this exercises sras_format/sras_compute/sras_align_export directly,
mirroring tests/test_align_export.py.
"""
import sras_align_export as export
import sras_compute as compute
from sras_format import SrasFile
import tools.make_test_sras as gen
_THRESHOLD_MV = 80.0
def test_single_angle_file_always_shares_its_grid(tmp_path):
path = tmp_path / "one_angle.sras"
gen.write(path, n_angles=1)
sras = SrasFile(str(path))
assert sras.angles_share_raw_grid()
def test_plain_multi_angle_file_does_not_share_its_grid(tmp_path):
"""tools.make_test_sras.write gives every angle its own geometry and
stage x_start (build()'s `x_start = -0.5 + 0.1 * a`), matching how real
v6 scans vary per angle — so this must read as "not shareable"."""
path = tmp_path / "plain.sras"
gen.write(path, n_angles=3)
sras = SrasFile(str(path))
assert not sras.angles_share_raw_grid()
def test_wizard_exported_file_shares_its_grid(tmp_path):
src_path = tmp_path / "rotating.sras"
meta = gen.write_rotating(src_path, n_angles=4)
sras = SrasFile(str(src_path))
params = {a: compute.ManualAngleParams(rot, shift)
for a, (rot, shift) in meta["truth"].items()}
result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
out_path = tmp_path / "rotating_aligned.sras"
export.write_aligned_sras(sras, result, out_path)
out = SrasFile(str(out_path))
assert not sras.angles_share_raw_grid(), (
"sanity check: the *source* rotating scan must NOT already share a "
"grid, or this test would not actually exercise the wizard export")
assert out.angles_share_raw_grid()
def test_mutated_angle_breaks_the_shared_grid(tmp_path):
"""A loaded SrasFile's x_start_mm is a public per-angle array (as
tests/test_gui.py::test_alignment_geometry_is_stage_independent also
relies on) -- mutating one angle's start must be visible here too."""
src_path = tmp_path / "rotating.sras"
meta = gen.write_rotating(src_path, n_angles=3)
sras = SrasFile(str(src_path))
params = {a: compute.ManualAngleParams(rot, shift)
for a, (rot, shift) in meta["truth"].items()}
result = compute.build_manual_alignment(sras, 0, _THRESHOLD_MV, params)
out_path = tmp_path / "rotating_aligned.sras"
export.write_aligned_sras(sras, result, out_path)
out = SrasFile(str(out_path))
assert out.angles_share_raw_grid()
out.x_start_mm[1] += 1.0
assert not out.angles_share_raw_grid()
tests/test_gui.py
+210 -545
View File
@@ -3,8 +3,8 @@ signals and worker threads under the offscreen platform plugin.
Covers the interactions a manual smoke test would: load, switch angles and
channels, background DC precompute, lazy FFT compute, threshold and bg-sub
changes, the alignment wizard end to end (pre-rotation, correlation, manual
nudging, crop, export), aligned view, ROI draw/move, and CSV export.
changes, angle alignment, manual angle alignment, aligned view, ROI
draw/move, and CSV export.
NOTE: this module is one ordered integration sequence over a single shared
window — the tests build on each other's state and must run in definition
@@ -20,13 +20,11 @@ import numpy as np
import pytest
from PyQt6.QtCore import QEventLoop, Qt, QTimer
from PyQt6.QtTest import QTest
from PyQt6.QtWidgets import QApplication, QDialog, QMessageBox
from PyQt6.QtWidgets import QApplication, QMessageBox
import sras_compute as compute
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
from sras_viewer import (
FusedRoiExportDialog, RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX,
)
from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX
import tools.make_test_sras as gen
@@ -96,11 +94,18 @@ def test_dc_precompute_all_angles(ctx):
def test_angle_switching_from_cache(ctx):
"""Setting the angle must redraw on its own — no nudge from the test.
The spinbox used to be wired to editingFinished alone, so stepping it with
the arrows changed the number and left the plot on the previous angle until
the box happened to lose focus.
"""
win, s = ctx.win, ctx.s
for a in range(s.n_angles):
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
assert win._current_angle == a, \
f"angle {a} redrawn from the spinbox alone, showing {win._current_angle}"
expected = win._sras.image_shape(a)
assert win._current_image.shape == expected, \
f"angle {a} shows its own geometry {expected}, got {win._current_image.shape}"
@@ -108,69 +113,42 @@ def test_angle_switching_from_cache(ctx):
"no compute job needed for cached DC angles"
def test_stepping_the_angle_spinbox_redraws(ctx):
"""Clicking the angle spinbox's arrows (or pressing Up/Down in it) must
move the display, not just the number.
This is the ordinary way to walk a scan, and it used to do nothing: the
spinbox was wired on editingFinished, which QAbstractSpinBox emits only
on Return or focus-out — never on a step. Every other test in this module
called _on_view_changed() by hand and so could not have caught it.
"""
def test_angle_steps_and_typing(ctx):
"""The arrows step and redraw one angle at a time; typing commits once,
without a stop at every intermediate value."""
win, s = ctx.win, ctx.s
assert s.n_angles >= 3, "need room to step in both directions"
assert s.n_angles > 2, "fixture needs several angles"
win.spin_angle.setValue(0)
assert wait_until(lambda: win._current_angle == 0), "settled on angle 0"
for expected in range(1, s.n_angles):
win.spin_angle.stepUp()
assert wait_until(lambda e=expected: win._current_angle == e), \
f"stepping up to angle {expected} redrew the display"
win.spin_angle.stepDown()
assert wait_until(lambda: win._current_angle == s.n_angles - 2), \
"stepping down redraws too"
# Keyboard stepping goes through the same signal, so it must work as well.
QTest.keyClick(win.spin_angle, Qt.Key.Key_Down)
assert wait_until(lambda: win._current_angle == s.n_angles - 3), \
"Key_Down redraws"
win.spin_angle.setValue(0)
assert wait_until(lambda: win._current_angle == 0), "back to angle 0"
def test_typing_an_angle_does_not_compute_intermediate_angles(ctx):
"""Keyboard tracking must stay off: with it on, valueChanged fires per
keystroke, so typing "12" would dispatch a compute for angle 1 first —
on a real scan, a whole wasted FFT for an angle the user never asked for.
"""
win, s = ctx.win, ctx.s
assert not win.spin_angle.keyboardTracking(), \
"keyboard tracking off is what makes valueChanged safe to connect"
target = s.n_angles - 1
assert target >= 2, "need a multi-digit-ish range to make the point"
win.spin_angle.setValue(0)
wait_until(lambda: win._current_angle == 0)
pump(60)
seen = []
win.spin_angle.valueChanged.connect(seen.append)
try:
win.spin_angle.lineEdit().selectAll()
QTest.keyClicks(win.spin_angle, str(target))
win.spin_angle.stepUp()
pump(60)
assert seen == [], f"no signal while typing, got {seen}"
assert win._current_angle == 1, f"arrow stepped to 1, showing {win._current_angle}"
win.spin_angle.stepDown()
pump(60)
assert win._current_angle == 0, f"arrow stepped to 0, showing {win._current_angle}"
# Typing must not redraw per keystroke — that is what editingFinished
# bought and what keyboard tracking has to keep buying.
seen.clear()
last = s.n_angles - 1
win.spin_angle.lineEdit().selectAll()
QTest.keyClicks(win.spin_angle.lineEdit(), str(last))
pump(30)
assert seen == [] and win._current_angle == 0, \
f"half-typed text must not redraw ({seen}, showing {win._current_angle})"
QTest.keyClick(win.spin_angle, Qt.Key.Key_Return)
pump(60)
assert seen == [target], f"one signal on commit, got {seen}"
assert seen == [last], f"committed exactly once on Return: {seen}"
assert win._current_angle == last
finally:
win.spin_angle.valueChanged.disconnect(seen.append)
assert wait_until(lambda: win._current_angle == target), "committed angle shown"
win.spin_angle.setValue(0)
assert wait_until(lambda: win._current_angle == 0), "back to angle 0"
win.spin_angle.clearFocus()
win.spin_angle.setValue(0)
pump(60)
def test_channel_switching(ctx):
@@ -224,25 +202,16 @@ def test_threshold_change_recomputes(ctx):
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
recompute. Toggling it on an angle that already has an FFT image must
leave that image on screen, untouched."""
win = ctx.win
n_before = len(win._fft_cache)
img_before = win._current_image
win.chk_bg_sub.setChecked(False)
pump(200)
assert not win._job_running("compute"), \
"toggling bg-sub alone must not dispatch a recompute"
assert len(win._fft_cache) == n_before, "no new cache entry from the toggle"
assert np.array_equal(win._current_image, img_before), \
"displayed image unchanged by the bg-sub toggle"
assert wait_until(
lambda: not win._job_running("compute") and len(win._fft_cache) > n_before), \
"recomputed without bg-sub"
win.chk_bg_sub.setChecked(True)
pump(200)
assert not win._job_running("compute")
assert len(win._fft_cache) == n_before
assert np.array_equal(win._current_image, img_before)
assert not win._job_running("compute"), \
"returning to bg-sub was a cache hit (no recompute)"
def test_roi_and_csv_export(ctx):
@@ -294,16 +263,44 @@ def test_roi_survives_switches(ctx):
"ROI still present after channel switch"
def test_alignment_geometry_is_stage_independent(ctx):
def test_angle_alignment(ctx):
win, s = ctx.win, ctx.s
win.spin_angle.setValue(0)
win._on_view_changed()
wait_until(lambda: not win._job_running("compute"))
assert win._alignment_act.isEnabled(), "alignment action enabled"
win._on_angle_alignment()
assert wait_until(
lambda: win._alignment_result is not None and not win._job_running("align"),
timeout_ms=60000), "alignment completed"
r = win._alignment_result
assert len(r.per_angle) == s.n_angles, "transform for every angle"
assert all(r.canvas_shape[0] >= int(s.n_rows[a])
and r.canvas_shape[1] >= int(s.n_frames[a])
for a in range(s.n_angles)), \
f"canvas is at least as large as any single angle: {r.canvas_shape}"
assert r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0), \
"reference angle has zero shift"
assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked(), \
"Aligned View auto-enabled and checked"
pump(200)
assert win.image_canvas._img_shape == r.canvas_shape, \
f"{win.image_canvas._img_shape} vs {r.canvas_shape}"
win.chk_aligned_view.setChecked(False)
pump(200)
assert win.image_canvas._img_shape == s.image_shape(0), \
f"unchecking returns to the raw per-angle grid: {win.image_canvas._img_shape}"
def test_manual_alignment_geometry(ctx):
"""Local mm is anchored on each angle's array center, not its stage
position: that is what makes a scan's placement independent of where its
window happened to sit. (Registration accuracy itself is covered by
tests/test_alignment.py, which has a synthetic sample to register.)"""
win, s = ctx.win, ctx.s
win.spin_angle.setValue(0)
win._on_view_changed()
wait_until(lambda: not win._job_running("compute"))
assert win._wizard_act.isEnabled(), "alignment wizard action enabled"
assert win._manual_align_act.isEnabled(), "manual alignment action enabled"
n_rows, n_frames = s.image_shape(0)
assert np.allclose(compute._center_idx(s, 0),
@@ -324,7 +321,7 @@ def test_alignment_geometry_is_stage_independent(ctx):
("moving every non-reference angle's scan window must leave the canvas "
f"unchanged: {origin_a} {shape_a} vs {origin_b} {shape_b}")
# Both signs of the stage's reported angle are searched by default.
# Both signs of the stage's reported angle are searched.
cands = compute._rotation_candidates(30.0, 6.0, 2.0)
assert min(cands) < -29.0 and max(cands) > 29.0, f"{min(cands)}..{max(cands)}"
@@ -337,295 +334,128 @@ def test_alignment_geometry_is_stage_independent(ctx):
"_shift_into moves content by exactly the requested offset"
def test_wizard_opens_prerotated(ctx):
"""The wizard shows a mask stack before any correlation has run, built from
the stage angles in the file — the "pre-rotate" step. Nothing may seed from
the still-live automatic result; only a saved sidecar."""
def test_manual_dialog_opens_at_identity(ctx):
"""Open must NOT seed from the still-live automatic AlignmentResult.
Manual mode exists to fix up whatever the automatic registration got
wrong, so it must start from identity (every angle centered on the
reference, no rotation) regardless of whatever the automatic run last
computed. Only a previously *saved manual* alignment (sidecar) should
ever seed this dialog."""
win, s = ctx.win, ctx.s
win._on_alignment_wizard()
assert win._align_wizard is not None, "wizard opened"
ctx.wiz = wiz = win._align_wizard
ctx.p1 = p1 = wiz.page(wiz.PAGE_CORRELATE)
assert not win._job_running("align_masks"), \
win._on_manual_alignment()
assert win._manual_align_dialog is not None, "dialog opened"
ctx.dlg = dlg = win._manual_align_dialog
assert not win._job_running("manual_align_masks"), \
"mask prep needed no background worker (already DC-cached)"
assert wait_until(lambda: p1.isComplete()), "masks ready, Next enabled"
assert not win._wizard_act.isEnabled(), \
"wizard action disabled while a wizard is open"
st = wiz.state
assert st.result is not None, "an AlignmentResult exists from pre-rotation alone"
assert st.counts is not None and st.counts.shape == st.preview_shape
assert 0 <= st.counts.max() <= s.n_angles
assert not st.fits, "no fits before a correlation has run"
for a in range(s.n_angles):
nominal = compute.nominal_delta_deg(s, a, st.ref_angle_idx)
expected = 0.0 if a == st.ref_angle_idx else nominal
assert abs(st.params[a].rotation_deg - expected) < 1e-9, \
f"angle {a} not pre-rotated to its stage angle"
assert st.params[a].shift_mm == (0.0, 0.0), \
"pre-rotation must not invent a translation"
assert all(dlg._angle_params[a] == compute.ManualAngleParams()
for a in range(s.n_angles)), \
"no manual sidecar yet -> dialog starts at identity, not the automatic result"
def test_wizard_reference_angle_is_locked(ctx):
p1, wiz = ctx.p1, ctx.wiz
p1.combo_active.setCurrentIndex(wiz.state.ref_angle_idx)
def test_reference_angle_is_locked(ctx):
dlg = ctx.dlg
dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx)
pump(30)
before = wiz.state.params[wiz.state.ref_angle_idx]
p1._on_nudge_translate(1, 0, False)
p1._on_nudge_rotate(1, False)
assert wiz.state.params[wiz.state.ref_angle_idx] == before, \
before_ref = dlg._angle_params[dlg._ref_angle_idx]
dlg._on_nudge_translate(1, 0, False)
dlg._on_nudge_rotate(1, False)
assert not dlg.grp_manual_adjust.isEnabled(), "reference angle group disabled"
assert dlg._angle_params[dlg._ref_angle_idx] == before_ref, \
"reference angle untouched by nudge attempts"
def test_wizard_nudges(ctx):
"""Manual correction, which the wizard absorbed from the old dialog."""
p1, wiz, s = ctx.p1, ctx.wiz, ctx.s
def test_nudges(ctx):
"""Nudging a real angle (fine + coarse, translate + rotate)."""
dlg, s = ctx.dlg, ctx.s
ctx.active = active = 1 if s.n_angles > 1 else 0
p1.combo_active.setCurrentIndex(active)
dlg.combo_active_angle.setCurrentIndex(active)
pump(30)
before = wiz.state.params[active].shift_mm
p1._on_nudge_translate(1, 0, False)
fine = p1.spin_step_translate.value()
assert abs(wiz.state.params[active].shift_mm[0] - (before[0] + fine)) < 1e-9, \
before = dlg._angle_params[active].shift_mm
dlg._on_nudge_translate(1, 0, False) # fine +X
fine_step = dlg.spin_step_translate_mm.value()
assert abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9, \
"fine translate nudge moved shift_x by exactly one fine step"
before = wiz.state.params[active].shift_mm
p1._on_nudge_translate(0, -1, True)
coarse = fine * p1.spin_step_mult.value()
assert abs(wiz.state.params[active].shift_mm[1] - (before[1] - coarse)) < 1e-9, \
before = dlg._angle_params[active].shift_mm
dlg._on_nudge_translate(0, -1, True) # coarse -Y
coarse_step = fine_step * dlg.spin_step_multiplier.value()
assert abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9, \
"coarse translate nudge uses the multiplier"
before_rot = wiz.state.params[active].rotation_deg
p1._on_nudge_rotate(1, False)
assert wiz.state.params[active].rotation_deg != before_rot
assert len(wiz.state.layers) == s.n_angles, \
"stack rebuilt for every angle after a rotation nudge"
before_rot = dlg._angle_params[active].rotation_deg
dlg._on_nudge_rotate(1, False)
assert dlg._angle_params[active].rotation_deg != before_rot, \
"rotate nudge changed rotation_deg"
assert len(dlg._preview_layers) == s.n_angles, \
"preview canvas rebuilt for every angle after a rotation nudge"
# A nudge only reprojects the angle that moved, patching the overlap counts
# in place. That shortcut is only sound if it lands on exactly what a full
# rebuild would have produced.
incremental = wiz.state.counts.copy()
wiz.rebuild_stack()
assert np.array_equal(wiz.state.counts, incremental), \
"incremental nudge update matches a full stack rebuild"
# Real key-event wiring (keyPressEvent -> signal -> slot).
before = wiz.state.params[active].shift_mm
QTest.keyClick(p1.canvas, Qt.Key.Key_Right)
assert wiz.state.params[active].shift_mm[0] > before[0], \
# Real key-event wiring (proves keyPressEvent -> signal -> slot).
before = dlg._angle_params[active].shift_mm
QTest.keyClick(dlg.canvas, Qt.Key.Key_Right)
assert dlg._angle_params[active].shift_mm[0] > before[0], \
"a real Right-arrow key event nudged shift_x"
# Both views render from the same reprojected layers.
p1.combo_view.setCurrentIndex(1)
pump(50)
p1.combo_view.setCurrentIndex(0)
pump(50)
def test_auto_derotate(ctx):
"""Auto De-rotate: seeds rotation from the stage angle, no translation."""
dlg, s, active = ctx.dlg, ctx.s, ctx.active
shift_before_derotate = dlg._angle_params[active].shift_mm
dlg._on_auto_derotate()
nominal = compute.nominal_delta_deg(s, active, dlg._ref_angle_idx)
assert abs(dlg._angle_params[active].rotation_deg - nominal) < 1e-6, \
"auto de-rotate seeded rotation from the stage's reported angle"
assert dlg._angle_params[active].shift_mm == shift_before_derotate, \
"auto de-rotate left translation untouched"
assert dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0, \
"reference angle stays identity after auto de-rotate"
# Clicking again offers the other sign, since which one lines the scans up
# is not knowable from the file.
dlg._on_auto_derotate()
assert abs(dlg._angle_params[active].rotation_deg + nominal) < 1e-6, \
"auto de-rotate offers the opposite sign on a second click"
def test_wizard_correlate(ctx):
"""Cross-correlation, for every source option, retryable."""
win, p1, wiz, s = ctx.win, ctx.p1, ctx.wiz, ctx.s
from sras_viewer.align_wizard import _CORRELATE_SOURCES
for idx, (label, _sources) in enumerate(_CORRELATE_SOURCES):
p1.combo_source.setCurrentIndex(idx)
p1.btn_correlate.click()
assert not p1.isComplete(), \
f"Next must be disabled while correlating ({label})"
assert wait_until(lambda: not win._job_running("align_correlate"),
timeout_ms=60000), f"correlation finished ({label})"
assert p1.isComplete(), f"Next re-enabled ({label})"
assert all(a in wiz.state.fits for a in range(s.n_angles)
if a != wiz.state.ref_angle_idx), \
def test_auto_cross_correlate(ctx):
"""Auto Cross-Correlate: searches rotation *and* translation."""
win, dlg, s = ctx.win, ctx.dlg, ctx.s
assert dlg.btn_auto_correlate.isEnabled(), \
"cross-correlate action enabled once masks are ready"
for label_idx, (label, _sources) in enumerate(dlg._CORRELATE_SOURCES):
dlg.combo_correlate_source.setCurrentIndex(label_idx)
dlg._on_auto_correlate()
assert wait_until(
lambda: not win._job_running("manual_align_correlate"),
timeout_ms=60000), f"auto cross-correlate completed ({label})"
assert all(a in dlg._fit_notes for a in range(s.n_angles)
if a != dlg._ref_angle_idx), \
f"every non-reference angle got a fit ({label})"
assert wiz.state.params[wiz.state.ref_angle_idx] == compute.ManualAngleParams(), \
"reference angle stays identity after correlation"
assert p1.btn_correlate.isEnabled(), "controls re-enabled when done"
assert p1.table.rowCount() == s.n_angles and p1.table.item(0, 0) is not None, \
"per-angle fit table populated"
assert p1.lbl_overlap.text(), "overlap summary reported"
r = wiz.state.result
assert len(r.per_angle) == s.n_angles, "transform for every angle"
assert r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0), \
"reference angle has zero shift"
assert dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams(), \
"auto cross-correlate reference angle stays identity"
assert dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled(), \
"auto cross-correlate re-enabled controls when done"
assert len(dlg._preview_layers) == s.n_angles, \
"preview canvas rebuilt after cross-correlate"
assert dlg._fit_report(), "fit quality is reported per angle"
def test_wizard_retry_changes_geometry(ctx):
"""Editing a parameter and re-running is the retry path, and it must
invalidate anything indexed against the old canvas."""
p1, wiz = ctx.p1, ctx.wiz
gen_before = wiz.state.geometry_generation
p1.spin_threshold.setValue(p1.spin_threshold.value() + 5.0)
p1.spin_threshold.editingFinished.emit()
pump(60)
assert wiz.state.geometry_generation > gen_before, \
"a threshold change rebuilt the geometry"
# Reset drops the fits and returns to pre-rotation only.
p1.btn_reset.click()
pump(60)
assert not wiz.state.fits, "reset cleared the fits"
nominal = compute.nominal_delta_deg(ctx.s, ctx.active, wiz.state.ref_angle_idx)
assert abs(wiz.state.params[ctx.active].rotation_deg - nominal) < 1e-9
assert wiz.state.params[ctx.active].shift_mm == (0.0, 0.0), \
"reset also drops nudged translation"
# Put a real correlation back for the pages that follow.
p1.btn_correlate.click()
assert wait_until(lambda: not ctx.win._job_running("align_correlate"),
timeout_ms=60000)
def test_wizard_roi_page(ctx):
"""The crop page: presets, and two-way sync between the drawn rectangle and
the numeric canvas-pixel boxes."""
wiz = ctx.wiz
wiz.next()
pump(150)
assert wiz.currentId() == wiz.PAGE_ROI, "advanced to the ROI page"
ctx.p2 = p2 = wiz.page(wiz.PAGE_ROI)
st = wiz.state
assert st.crop is not None and p2.isComplete(), \
"a default crop is offered on entry"
n_rows, n_cols = st.result.canvas_shape
assert st.crop[2] > 1 or n_rows == 1, \
f"default crop must not collapse to a single row: {st.crop}"
p2.btn_whole.click()
pump(50)
assert st.crop == (0, 0, n_rows, n_cols), "whole-canvas preset"
p2.btn_fit_union.click()
pump(50)
assert st.counts[st.crop[0]:st.crop[0] + st.crop[2],
st.crop[1]:st.crop[1] + st.crop[3]].sum() == st.counts.sum(), \
"fit-to-union must keep every covered pixel"
if p2.btn_fit_overlap.isEnabled():
p2.btn_fit_overlap.click()
pump(50)
row0, col0, nr, nc = st.crop
assert (st.counts[row0:row0 + nr, col0:col0 + nc] >= 1).all(), \
"full-overlap crop must not include uncovered pixels"
# Numeric -> drawn rectangle.
p2.btn_whole.click()
pump(50)
target = (0, 0, max(1, n_rows // 2), max(1, n_cols // 2))
p2.spin_rows.setValue(target[2])
p2.spin_cols.setValue(target[3])
pump(50)
assert st.crop == target, f"spin boxes drive the crop: {st.crop} vs {target}"
# Drawn rectangle -> numeric, round-tripping exactly.
x0, y0 = wiz.canvas_to_mm(target[1] - 0.5, target[0] - 0.5)
x1, y1 = wiz.canvas_to_mm(target[1] + target[3] - 0.5,
target[0] + target[2] - 0.5)
p2.canvas.set_roi(RoiQuad.from_bbox(min(x0, x1), min(y0, y1),
max(x0, x1), max(y0, y1)))
pump(80)
assert st.crop == target, \
f"drawn rectangle round-trips to the same crop: {st.crop} vs {target}"
# A degenerate crop blocks Next.
st.crop = None
p2.completeChanged.emit()
assert not p2.isComplete(), "an absent crop blocks Next"
p2._set_crop(*target)
assert p2.isComplete()
ctx.crop = target
def test_wizard_crop_dropped_when_going_back(ctx):
"""A crop is canvas-pixel indexed, so it cannot survive a re-correlation."""
wiz, p2 = ctx.wiz, ctx.p2
wiz.back()
pump(120)
assert wiz.currentId() == wiz.PAGE_CORRELATE
assert wiz.state.crop is None, "cleanupPage discarded the stale crop"
assert wiz.cropped_plan() == (None, None), \
"nothing derived from the dropped crop survives either"
wiz.next()
pump(150)
assert wiz.state.crop is not None, "a fresh default crop is offered again"
p2._set_crop(*ctx.crop)
def test_wizard_export(ctx):
"""Writing the file: Finish stays unavailable until a write succeeds."""
win, wiz = ctx.win, ctx.wiz
out = ctx.tmpdir / "wizard_aligned.sras"
with patch("sras_viewer.align_wizard.QMessageBox.question",
return_value=QMessageBox.StandardButton.Yes):
wiz.next()
pump(150)
assert wiz.currentId() == wiz.PAGE_SAVE, "advanced to the save page"
ctx.p3 = p3 = wiz.page(wiz.PAGE_SAVE)
cropped, _ = wiz.cropped_plan()
assert cropped is not None, "crop applied on leaving page 2"
assert cropped.canvas_shape == ctx.crop[2:], \
"cropped result carries the chosen shape"
assert not p3.isComplete(), "Finish unavailable before anything is written"
assert p3.lbl_summary.text(), "a summary of what will be written is shown"
with patch("sras_viewer.align_wizard.QFileDialog.getSaveFileName",
return_value=(str(out), "")):
p3.btn_browse.click()
assert wiz.state.out_path == str(out)
p3.btn_export.click()
assert wait_until(lambda: not win._job_running("align_export"),
timeout_ms=60000), "export finished"
assert wiz.state.exported_path == str(out), p3.lbl_status.text()
assert p3.isComplete(), "Finish available once the file exists"
assert out.exists()
written = SrasFile(str(out))
ctx.written = written
assert written.version == 6, "export is a v6 file"
assert written.n_angles == ctx.s.n_angles
assert all(written.image_shape(a) == ctx.crop[2:]
for a in range(written.n_angles)), \
"every angle shares the cropped grid"
assert not out.with_name(out.name + ".part").exists(), \
"no staging file left behind"
def test_wizard_finish_applies_and_persists(ctx):
"""Finish makes the session match the file: Aligned View shows the exported
extent, and the sidecar records it for the input scan."""
win, wiz = ctx.win, ctx.wiz
wiz.accept()
pump(250)
assert win._align_wizard is None, "wizard reference released"
assert win._wizard_act.isEnabled(), "wizard action available again"
assert win._alignment_result is not None
assert win._alignment_result.canvas_shape == ctx.crop[2:], \
"the *cropped* result is what the view now uses"
assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked()
pump(200)
assert win.image_canvas._img_shape == ctx.crop[2:], \
f"canvas shows the cropped extent: {win.image_canvas._img_shape}"
sidecar = compute.sidecar_path(ctx.s.path)
assert sidecar.exists(), "sidecar written for the input scan"
def test_save_sidecar(ctx):
win, dlg, s, active = ctx.win, ctx.dlg, ctx.s, ctx.active
dlg._on_save()
sidecar = compute.sidecar_path(s.path)
assert sidecar.exists(), "sidecar file written"
ctx.sidecar = sidecar
ctx.sidecar_raw = raw = json.loads(sidecar.read_text())
assert raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION
assert all(raw["per_angle"][str(a)]["rotation_deg"]
== win._alignment_result.per_angle[a].rotation_deg
for a in range(ctx.s.n_angles)), \
"sidecar round-trips the applied rotations"
win.chk_aligned_view.setChecked(False)
pump(200)
assert win.image_canvas._img_shape == ctx.s.image_shape(0), \
f"unchecking returns to the raw per-angle grid: {win.image_canvas._img_shape}"
assert raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION, \
"sidecar schema_version is current"
assert all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg")
== dlg._angle_params[a].rotation_deg for a in range(s.n_angles)), \
"sidecar per_angle round-trips the dialog's resolved params"
assert (win._alignment_result is not None
and win._alignment_result.per_angle[active].rotation_deg
== dlg._angle_params[active].rotation_deg), \
"main window's alignment_result replaced by the manual build"
assert win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked(), \
"Aligned View auto-enabled after Save"
def test_stale_schema_sidecar_ignored(ctx):
@@ -639,221 +469,56 @@ def test_stale_schema_sidecar_ignored(ctx):
sidecar.write_text(json.dumps(raw)) # restore for the rest of the sequence
def test_clear_with_confirmation(ctx):
win, dlg, s = ctx.win, ctx.dlg, ctx.s
with patch("sras_viewer.dialogs.QMessageBox.question",
return_value=QMessageBox.StandardButton.Yes):
dlg._on_clear()
assert not ctx.sidecar.exists(), "sidecar file deleted"
assert all(dlg._angle_params[a] == compute.ManualAngleParams()
for a in range(s.n_angles)), "dialog params reset to identity"
assert win._alignment_result is None, "main window alignment_result cleared"
assert (not win.chk_aligned_view.isEnabled()
and not win.chk_aligned_view.isChecked()), \
"Aligned View disabled after Clear"
dlg.close()
pump(150)
assert win._manual_align_dialog is None, "dialog reference released on close"
def test_sidecar_restored_on_reload(ctx):
win, active = ctx.win, ctx.active
saved = json.loads(ctx.sidecar.read_text())["per_angle"][str(active)]
win._on_manual_alignment()
dlg = win._manual_align_dialog
dlg.combo_active_angle.setCurrentIndex(active)
pump(30)
dlg._on_auto_derotate()
dlg._on_nudge_translate(1, 1, True)
saved_rotation = dlg._angle_params[active].rotation_deg
saved_shift = dlg._angle_params[active].shift_mm
dlg._on_save()
dlg.close()
pump(150)
old_sras_id = id(win._sras)
win._load_file(str(ctx.path)) # reload the same file fresh
assert wait_until(
lambda: win._sras is not None and id(win._sras) != old_sras_id), \
"file reloaded"
ctx.s = win._sras
assert win._align_wizard is None, "no wizard left open across a reload"
assert win._manual_align_dialog is None, \
"manual dialog force-closed by a reload"
assert win._alignment_result is not None, \
"reload restores the saved alignment automatically"
"reload restores the saved manual alignment automatically"
assert abs(win._alignment_result.per_angle[active].rotation_deg
- saved["rotation_deg"]) < 1e-9, \
"restored rotation matches what was saved"
- saved_rotation) < 1e-9, "restored rotation matches what was saved"
assert win._alignment_result.per_angle[active].shift_mm == saved_shift, \
"restored shift matches what was saved"
assert win.chk_aligned_view.isChecked(), \
"Aligned View auto-checked after restoring a saved alignment"
def test_wizard_closes_with_a_reload(ctx):
"""An open wizard belongs to the file it was opened on."""
win = ctx.win
win._on_alignment_wizard()
assert win._align_wizard is not None
assert wait_until(
lambda: win._align_wizard.page(win._align_wizard.PAGE_CORRELATE).isComplete())
win._load_file(str(ctx.path))
assert wait_until(lambda: not win._job_running("load"))
pump(200)
assert win._align_wizard is None, "wizard force-closed by a reload"
ctx.s = win._sras
# ---------------------------------------------------------------------------
# Export Fused ROI
# ---------------------------------------------------------------------------
#
# Two independent ways angles can end up sharing one (x, y) grid to fuse
# onto: a live alignment result (case a, exercised on ctx.win -- the reload
# above restored one from the sidecar), or a file that is itself a previous
# Alignment Wizard export, whose angles already share a grid on disk with no
# alignment result needed at all (case b, exercised on a second window
# opened on ctx.written from test_wizard_export).
def test_fused_export_gating_case_a(ctx):
"""A live alignment result bridges the raw scan's per-angle grids --
angles_share_raw_grid() alone would be False here."""
win, s = ctx.win, ctx.s
assert win._alignment_result is not None, "alignment restored from sidecar"
assert not s.angles_share_raw_grid(), \
"sanity check: the raw (un-aligned) scan must not already share a grid"
x, y = win._aligned_canvas_axes()
roi = RoiQuad.from_bbox(float(x[1]), float(y[1]), float(x[-2]), float(y[-2]))
win.image_canvas.set_roi(roi)
pump(120)
assert win._fused_grid_ready()
assert win.btn_export_fused_roi.isEnabled()
def test_write_fused_roi_csv_case_a_content(ctx):
win, s = ctx.win, ctx.s
roi = win.image_canvas.get_roi()
assert roi is not None, "ROI drawn by test_fused_export_gating_case_a"
csv_path = ctx.tmpdir / "fused_roi_case_a.csv"
win._write_fused_roi_csv(roi, CH4_IDX, [0, 1], str(csv_path))
assert csv_path.exists()
lines = csv_path.read_text().splitlines()
body = [l for l in lines if not l.startswith("#")]
header, *data_lines = body
assert header == (
f"x_mm,y_mm,v_{s.angles_deg[0]:.4g}deg,v_{s.angles_deg[1]:.4g}deg")
x, y = win._fused_export_axes()
mask = roi.mask_for_grid(x, y)
assert len(data_lines) == int(mask.sum())
data = np.array([[float(v) for v in line.split(",")] for line in data_lines])
expect0 = win._fused_value_image(0, CH4_IDX)[mask]
expect1 = win._fused_value_image(1, CH4_IDX)[mask]
assert np.allclose(data[:, 2], expect0, rtol=1e-5, atol=1e-4)
assert np.allclose(data[:, 3], expect1, rtol=1e-5, atol=1e-4)
def test_fused_roi_dialog_availability_live_updates(ctx):
"""Switching the value-type radio re-evaluates every angle checkbox,
disabling/auto-unchecking whichever ones are no longer available --
independent of what actually backs availability_fn, so a synthetic
stand-in keeps this a fast, deterministic test of the dialog itself."""
win, s = ctx.win, ctx.s
angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
only_angle0_has_ch1 = lambda a, c: (a == 0) if c == CH1_IDX else True
dlg = FusedRoiExportDialog(
win, angles=angles, availability_fn=only_angle0_has_ch1,
default_ch_idx=CH1_IDX, out_dir=str(s.path.parent), stem=s.path.stem,
grid_note="test")
try:
assert dlg._angle_checks[0].isEnabled()
assert all(not dlg._angle_checks[a].isEnabled()
for a in range(1, s.n_angles))
dlg._angle_checks[0].setChecked(True)
dlg._val_buttons[CH4_IDX].click()
assert all(dlg._angle_checks[a].isEnabled() for a in range(s.n_angles)), \
"CH4 is available for every angle"
assert dlg._angle_checks[0].isChecked(), \
"stays checked -- still available under CH4"
if s.n_angles > 1:
dlg._angle_checks[1].setChecked(True)
dlg._val_buttons[CH1_IDX].click()
assert dlg._angle_checks[0].isChecked()
if s.n_angles > 1:
assert not dlg._angle_checks[1].isEnabled()
assert not dlg._angle_checks[1].isChecked(), \
"auto-unchecked: angle 1 has no data under CH1"
finally:
dlg.close()
def test_fused_roi_dialog_select_all_none(ctx):
win, s = ctx.win, ctx.s
angles = [(a, float(s.angles_deg[a])) for a in range(s.n_angles)]
dlg = FusedRoiExportDialog(
win, angles=angles, availability_fn=lambda a, c: c == CH4_IDX,
default_ch_idx=CH4_IDX, out_dir=str(s.path.parent), stem=s.path.stem,
grid_note="test")
try:
assert not dlg._btn_export.isEnabled(), "nothing checked yet"
dlg._on_select_all_available()
assert all(cb.isChecked() for cb in dlg._angle_checks.values())
assert dlg._btn_export.isEnabled()
dlg._on_select_none()
assert not any(cb.isChecked() for cb in dlg._angle_checks.values())
assert not dlg._btn_export.isEnabled()
finally:
dlg.close()
def test_on_export_fused_roi_csv_end_to_end(ctx):
win, s = ctx.win, ctx.s
roi = win.image_canvas.get_roi()
assert roi is not None, "ROI from the earlier fused-export tests is still set"
csv_path = ctx.tmpdir / "fused_roi_e2e.csv"
with patch("sras_viewer.main_window.FusedRoiExportDialog") as MockDlg:
inst = MockDlg.return_value
inst.exec.return_value = QDialog.DialogCode.Accepted
inst.get_ch_idx.return_value = CH4_IDX
inst.get_selected_angles.return_value = [0, 1]
inst.get_output_path.return_value = str(csv_path)
win.btn_export_fused_roi.click()
assert csv_path.exists()
kwargs = MockDlg.call_args.kwargs
assert kwargs["default_ch_idx"] == win.combo_channel.currentIndex()
assert kwargs["out_dir"] == str(s.path.parent)
assert kwargs["stem"] == s.path.stem
def test_fused_export_no_alignment_shared_grid_path(ctx):
"""ctx.written (from test_wizard_export) is itself a previous Alignment
Wizard export: opened fresh with no sidecar for its own path, so no
alignment result is ever restored -- but its angles already share one
grid on disk, so the export must work through the no-resample path."""
win2 = SrasViewerWindow()
try:
win2._load_file(str(ctx.written.path))
assert wait_until(lambda: win2._sras is not None)
s2 = win2._sras
assert win2._alignment_result is None, \
"no sidecar exists for this path -- nothing auto-restored"
assert s2.angles_share_raw_grid(), \
"a wizard export already shares one grid across angles"
assert wait_until(lambda: all((a, CH4_IDX) in win2._dc_cache
for a in range(s2.n_angles))), \
"DC precomputed for every angle"
x, y = s2.x_axis_mm(0), s2.y_positions_mm(0)
roi = RoiQuad.from_bbox(float(x[1]), float(y[1]), float(x[-2]), float(y[-2]))
win2.image_canvas.set_roi(roi)
pump(120)
assert win2._fused_grid_ready()
assert win2.btn_export_fused_roi.isEnabled()
aligned_cache_before = len(win2._aligned_cache)
angle_idxs = list(range(min(2, s2.n_angles)))
csv_path = ctx.tmpdir / "fused_roi_case_b.csv"
win2._write_fused_roi_csv(roi, CH4_IDX, angle_idxs, str(csv_path))
assert csv_path.exists()
assert len(win2._aligned_cache) == aligned_cache_before, \
"no-resample path must never touch apply_alignment"
header = next(l for l in csv_path.read_text().splitlines()
if not l.startswith("#"))
expected_header = "x_mm,y_mm," + ",".join(
f"v_{s2.angles_deg[a]:.4g}deg" for a in angle_idxs)
assert header == expected_header
# Break the shared grid and confirm gating flips off.
mutate_idx = 1 if s2.n_angles > 1 else 0
s2.x_start_mm[mutate_idx] += 1.0
assert not s2.angles_share_raw_grid()
win2._update_fused_export_enabled()
assert not win2._fused_grid_ready()
assert not win2.btn_export_fused_roi.isEnabled()
assert "Alignment Wizard" in win2.btn_export_fused_roi.toolTip()
finally:
win2.close()
pump(200)
def test_pixel_inspector(ctx):
win = ctx.win
win.chk_aligned_view.setChecked(False)
tests/test_row_average.py
-278
View File
@@ -1,278 +0,0 @@
"""Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing.
Covers the properties the design depends on: the kernel is symmetric and
n=0 is a true no-op; the masked/renormalized convolution matches an
independent brute-force reference and gives masked neighbors exactly zero
weight regardless of their content; background subtraction after averaging
is algebraically identical to subtracting before; chunking/worker count
never changes the result; and a pixel that's itself masked is never
"rescued" by averaging.
"""
import numpy as np
import pytest
import sras_compute as compute
from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH4_IDX, SrasFile
import tools.make_test_sras as gen
def _reference_row_average(masked_waves: np.ndarray, valid: np.ndarray,
weights: np.ndarray) -> np.ndarray:
"""Independent, unvectorized reference for _row_average_waveforms: for
each row position, sum weighted valid neighbors within the kernel's
radius and normalize by the actual included weight sum. Same
definition, computed by brute-force nested loops instead of
correlate1d, so it can't share a bug with the implementation."""
n_frames, spf = masked_waves.shape
n = len(weights) // 2
out = np.zeros_like(masked_waves)
for i in range(n_frames):
num = np.zeros(spf, dtype=np.float64)
den = 0.0
for d in range(-n, n + 1):
j = i + d
if 0 <= j < n_frames and valid[j]:
w = float(weights[d + n])
num += w * masked_waves[j].astype(np.float64)
den += w
out[i] = num / den if den > 0 else 0.0
return out
# ---------------------------------------------------------------------------
# _row_average_weights
# ---------------------------------------------------------------------------
def test_row_average_weights_shape_and_symmetry():
w0 = compute._row_average_weights(0)
assert w0.shape == (1,) and w0[0] == 1.0
for n in (1, 2, 5):
w = compute._row_average_weights(n)
assert w.shape == (2 * n + 1,)
assert w[n] == pytest.approx(1.0), "center tap is the peak weight"
assert np.allclose(w, w[::-1]), "symmetric about the center"
half = w[n:]
assert np.all(np.diff(half) < 0), "strictly decreasing away from center"
# ---------------------------------------------------------------------------
# _row_average_waveforms
# ---------------------------------------------------------------------------
def test_row_average_matches_hand_rolled_reference():
rng = np.random.default_rng(0)
n_frames, spf = 15, 6
raw = rng.integers(-50, 51, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[2, 3, 9]] = False # a run of two invalid, plus a lone invalid
masked = raw.copy()
masked[~valid] = 0.0
weights = compute._row_average_weights(3)
got = compute._row_average_waveforms(masked, valid, weights)
ref = _reference_row_average(masked, valid, weights)
assert np.allclose(got[valid], ref[valid], atol=1e-4)
def test_row_average_edge_of_row():
"""A window wider than the row itself must still renormalize correctly
at both ends -- mode='constant', cval=0.0 zero-pads both the numerator
and denominator, so this is not a special case, but it's the one most
likely to break if that padding were ever mismatched between the two."""
n_frames, spf = 6, 3
raw = np.arange(n_frames * spf, dtype=np.float32).reshape(n_frames, spf)
valid = np.ones(n_frames, dtype=bool)
weights = compute._row_average_weights(4) # window (9 taps) > n_frames (6)
got = compute._row_average_waveforms(raw, valid, weights)
ref = _reference_row_average(raw, valid, weights)
assert np.allclose(got, ref, atol=1e-4)
def test_row_average_excludes_masked_neighbor_from_normalization():
"""A masked neighbor must contribute zero *weight* to the normalization,
not participate as a legitimate zero-valued sample at full weight --
the two give different answers, and only the former is correct. (Note:
masked_waves must already be 0 at invalid positions per
_row_average_waveforms's contract -- that's what read_row's zero-filled
scratch buffer guarantees in production -- so the only way to vary
"what a masked position looks like" while respecting that contract is
whether its weight is excluded from the denominator at all.)"""
n_frames, spf = 9, 4
weights = compute._row_average_weights(2)
# A: position 4 is masked -- excluded from the weight sum entirely.
valid_a = np.ones(n_frames, dtype=bool)
valid_a[4] = False
masked_a = np.zeros((n_frames, spf), dtype=np.float32)
masked_a[valid_a] = 1.0
got_a = compute._row_average_waveforms(masked_a, valid_a, weights)
# B: position 4 is valid but genuinely zero-valued -- included in the
# weight sum, diluting neighbors' averages.
valid_b = np.ones(n_frames, dtype=bool)
masked_b = np.ones((n_frames, spf), dtype=np.float32)
masked_b[4] = 0.0
got_b = compute._row_average_waveforms(masked_b, valid_b, weights)
# Every position whose window reaches index 4 must average *higher* in
# A (excluded from the denominator) than in B (included as a real zero).
affected = [2, 3, 5, 6]
assert np.all(got_a[affected] > got_b[affected]), \
"masking must exclude a neighbor from normalization, not just zero its value"
# Positions outside the window (radius 2) are unaffected either way.
assert np.allclose(got_a[[0, 1, 7, 8]], got_b[[0, 1, 7, 8]])
def test_background_subtracted_once_equals_subtract_then_average():
"""Algebraic identity the implementation relies on: subtracting a fixed
background from the already-averaged waveform equals subtracting it
from every valid neighbor first, because the denominator is always the
*actual* included weight sum (never a fixed total)."""
rng = np.random.default_rng(1)
n_frames, spf = 11, 8
raw = rng.integers(-40, 41, size=(n_frames, spf)).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
valid[[1, 7]] = False
masked = raw.copy()
masked[~valid] = 0.0
background = rng.integers(-5, 6, size=spf).astype(np.float32)
weights = compute._row_average_weights(3)
# Order A (what the code does): average first, subtract background once.
order_a = compute._row_average_waveforms(masked, valid, weights) - background
# Order B: subtract background from every valid neighbor first (restoring
# the "0 at invalid positions" contract afterward), then average.
bg_subbed = masked - background
bg_subbed[~valid] = 0.0
order_b = compute._row_average_waveforms(bg_subbed, valid, weights)
assert np.allclose(order_a[valid], order_b[valid], atol=1e-3)
# ---------------------------------------------------------------------------
# compute_rf_image(row_avg_n=...) integration
# ---------------------------------------------------------------------------
def test_row_average_zero_is_identity(tmp_path):
"""row_avg_n=0 must take the exact same code path as before this
feature existed (row_avg_weights stays None), not a single-tap kernel
that merely computes to the same answer."""
path = tmp_path / "zero.sras"
gen.write(path, n_angles=1, seed=10, samples_per_frame=64)
sras = SrasFile(str(path))
plain = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
explicit_zero = compute_rf_image(sras, 0, dc_threshold_mv=None,
apply_bg_sub=True, row_avg_n=0)
assert np.array_equal(plain, explicit_zero)
def test_row_average_respects_own_center_mask(tmp_path):
"""A pixel that's itself below threshold stays masked (0) after row
averaging -- averaging never rescues a masked pixel, matching the
'valid neighbors only' design (masked pixels are excluded from other
pixels' averages, and are never themselves smoothed)."""
path = tmp_path / "center_mask.sras"
gen.write(path, n_angles=1, seed=13, samples_per_frame=64)
sras = SrasFile(str(path))
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.percentile(dc4, 50))
mask = dc4 >= thr
assert mask.any() and not mask.all(), "threshold actually splits the image"
img = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(img == 0, ~mask), \
"masked pixels stay exactly 0 after row averaging; valid ones don't"
def test_row_average_composes_with_padding(tmp_path):
"""row_avg_n and n_fft (zero-padding) are independent knobs: using them
together must not raise, and must still agree with the exact (non-zoom)
reference path at that pad factor -- i.e. row-averaging composes with
the zoom peak search correctly, not just with the direct one."""
path = tmp_path / "padded_rowavg.sras"
gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
sras = SrasFile(str(path))
spf = sras.samples_per_frame
raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3)
avg_padded_zoom = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40)
avg_padded_exact = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
row_avg_n=3, n_fft=spf * 40, exact=True)
assert avg_natural.shape == raw_natural.shape == avg_padded_zoom.shape
assert np.all(np.isfinite(avg_padded_zoom))
assert np.array_equal(avg_padded_zoom, avg_padded_exact), \
"row-averaged waveforms feed the zoom and exact FFT paths identically"
def test_row_average_improves_snr_recovery():
"""The actual point of the feature: averaging same-row waveforms that
share a true underlying tone but carry independent noise recovers that
tone far more reliably than any single raw (unaveraged) waveform does."""
rng = np.random.default_rng(42)
n_frames, spf = 21, 128
true_bin = 9
t = np.arange(spf)
tone = 15.0 * np.sin(2 * np.pi * true_bin * t / spf) # same true signal
# at every position
noise_sigma = 40.0 # much larger than the tone -- deliberately poor SNR
raw = (tone[None, :] + rng.normal(scale=noise_sigma, size=(n_frames, spf))
).astype(np.float32)
valid = np.ones(n_frames, dtype=bool)
weights = compute._row_average_weights(8) # wide window: lots of averaging
averaged = compute._row_average_waveforms(raw, valid, weights)
raw_bins = compute._peak_bins_direct(raw, spf)
avg_bins = compute._peak_bins_direct(averaged, spf)
raw_hits = int(np.sum(raw_bins == true_bin))
avg_hits = int(np.sum(avg_bins == true_bin))
assert avg_hits > raw_hits, (
f"row averaging should recover the true bin ({true_bin}) more often "
f"than raw per-pixel estimates: raw {raw_hits}/{n_frames}, "
f"averaged {avg_hits}/{n_frames}")
assert avg_hits >= n_frames * 0.7, \
f"averaged recovery should be reliable, not just barely better: {avg_hits}/{n_frames}"
def test_row_average_parallel_identity(tmp_path, monkeypatch):
"""Forcing 1 worker vs many must give an identical row-averaged image --
catches chunk-boundary bugs (there should be none, since averaging never
crosses rows, but this is the empirical proof, not just inspection)."""
path = tmp_path / "parallel_rowavg.sras"
n_rows, n_frames, spf = 40, 13, 128
gen.write(path, n_angles=1, seed=11, samples_per_frame=spf,
geometry=[(n_rows, n_frames)])
sras = SrasFile(str(path))
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
# row_avg_n > 0 halves the effective budget before chunk planning.
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
assert fft_rows < n_rows, \
f"row-averaged FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
dc4 = dc_image_mv(sras, 0, CH4_IDX)
thr = float(np.median(dc4))
monkeypatch.setattr(compute, "_MAX_WORKERS", 1)
serial = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
parallel = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True,
row_avg_n=4)
assert np.array_equal(serial, parallel), \
"row-averaged rf image identical regardless of chunking/worker count"
tests/test_stored_cache.py
-595
View File
@@ -1,595 +0,0 @@
"""Does a batch-computed FFT cache actually spare the viewer the FFT?
Storing a peak-frequency image per angle in the file is only worth doing if
displaying it is then free. The regression this module pins down is the
viewer's *dispatch* decision: it used to find the stored image only inside
ComputeWorker, so after a batch every angle change still queued a background
job behind a "Computing FFT…" popup for an image already on disk.
Both layers are covered — cached_rf_image's accept/reject rules, and the
window never reaching _start_compute for a batch-cached angle.
"""
import struct
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import pytest
from PyQt6.QtCore import QEventLoop, QTimer
from PyQt6.QtWidgets import QApplication, QDialog
import sras_compute as compute
import sras_format as fmt
from sras_compute import cache_file, cached_rf_image, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
import tools.make_test_sras as gen
_THRESHOLD_MV = 50.0 # the viewer's own default
def pump(ms: int = 200):
loop = QEventLoop()
QTimer.singleShot(ms, loop.quit)
loop.exec()
def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
waited = 0
while waited < timeout_ms:
if pred():
return True
pump(step)
waited += step
return pred()
@pytest.fixture(scope="module")
def rig(tmp_path_factory):
"""A v6 file, the FFT images a from-scratch compute gives for it, and the
same file after Batch Compute FFT has written them into its v7 cache."""
path = tmp_path_factory.mktemp("stored_cache") / "cached.sras"
gen.write(path, n_angles=4, seed=7, samples_per_frame=256)
src = SrasFile(str(path))
fresh = {a: compute_rf_image(src, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True)
for a in range(src.n_angles)}
assert src.background is not None, "the fixture file must have a background"
err = cache_file(str(path), "fft", True)
assert err == "", err
cached = SrasFile(str(path))
assert all(x is not None for x in cached.precomputed_freq_mhz)
assert all(x is None for x in cached.precomputed_dc4_mv), \
"FFT-only batch: the mask has to come from the viewer, not the file"
return SimpleNamespace(path=path, fresh=fresh, sras=cached,
n_angles=cached.n_angles)
@pytest.fixture(scope="module")
def dc_rig(tmp_path_factory):
"""A file that has been through Batch Compute DC and Store."""
path = tmp_path_factory.mktemp("stored_dc") / "dc_cached.sras"
gen.write(path, n_angles=4, seed=9, samples_per_frame=128)
err = cache_file(str(path), "dc", True)
assert err == "", err
sras = SrasFile(str(path))
assert all(x is not None for x in sras.precomputed_dc3_mv)
assert all(x is not None for x in sras.precomputed_dc4_mv)
assert len(np.unique(sras.precomputed_dc4_mv[0])) > 1, \
"a degenerate DC image would make the comparisons below vacuous"
return SimpleNamespace(path=path, sras=sras, n_angles=sras.n_angles)
@pytest.fixture
def no_fft(monkeypatch):
"""Make any real FFT work loud: returns a list that stays empty unless a
peak search actually runs."""
calls = []
for name in ("_peak_bins_direct", "_peak_bins_zoom"):
original = getattr(compute, name)
def spy(*args, _f=original, **kwargs):
calls.append(_f.__name__)
return _f(*args, **kwargs)
monkeypatch.setattr(compute, name, spy)
return calls
# ---------------------------------------------------------------------------
# cached_rf_image: when may the stored image stand in for a compute?
# ---------------------------------------------------------------------------
def test_stored_image_matches_a_fresh_compute(rig, no_fft):
for a in range(rig.n_angles):
dc4 = dc_image_mv(SrasFile(str(rig.path)), a, CH4_IDX)
img = cached_rf_image(rig.sras, a, dc_threshold_mv=_THRESHOLD_MV,
apply_bg_sub=True, dc4_mv=dc4)
assert img is not None, f"angle {a} is cached in the file"
assert np.allclose(img, rig.fresh[a], atol=1e-3), \
f"angle {a} differs from a from-scratch compute"
assert not no_fft, f"the stored image was used, no FFT ran: {no_fft}"
def test_unmasked_when_no_threshold(rig):
img = cached_rf_image(rig.sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
assert img is not None and np.array_equal(img, rig.sras.precomputed_freq_mhz[0])
img[:] = -1.0
assert not np.any(rig.sras.precomputed_freq_mhz[0] == -1.0), \
"callers get a copy, never the file's own array"
def test_settings_the_stored_image_cannot_serve(rig):
"""A stored image carries one bg-sub state and one padding, so anything
else must fall through to a real compute rather than lie."""
spf = rig.sras.samples_per_frame
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
n_fft=spf * 4) is None, \
"cache was written at pad 1, a pad-4 view resolves different peaks"
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=False) is None, \
"cache was written with bg-sub on"
uncached = SrasFile(str(rig.path))
uncached.precomputed_freq_mhz[1] = None
assert cached_rf_image(uncached, 1, _THRESHOLD_MV, apply_bg_sub=True) is None
@pytest.mark.parametrize("pad", [2, 10])
def test_cache_is_stored_at_the_configured_pad(tmp_path, pad):
"""Batching at a pad factor must produce a cache that view can read back —
a pad-1-only cache is one the padded viewer can never use."""
path = tmp_path / f"pad{pad}.sras"
gen.write(path, n_angles=2, seed=13, samples_per_frame=256)
spf = SrasFile(str(path)).samples_per_frame
assert cache_file(str(path), "fft", True, "scipy", 0, pad) == ""
sras = SrasFile(str(path))
assert sras.precomputed_pad_factor == pad, "pad factor survives the round trip"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * pad) is not None, f"usable at pad {pad}"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \
"not usable unpadded"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
n_fft=spf * (pad + 1)) is None, "not usable at another pad"
# The stored numbers must be the padded ones, not pad-1 relabelled.
fresh = SrasFile(str(path))
fresh.precomputed_freq_mhz = [None] * fresh.n_angles
for a in range(sras.n_angles):
assert np.allclose(
compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad),
compute_rf_image(fresh, a, dc_threshold_mv=None, apply_bg_sub=True,
n_fft=spf * pad), atol=1e-3), \
f"angle {a}: stored image is the pad-{pad} answer"
def test_cach_v1_reads_as_natural_resolution(tmp_path):
"""Files cached before the pad factor existed must keep working: a v1 tail
has no pad field and is pad 1 by construction."""
path = tmp_path / "v1.sras"
gen.write(path, n_angles=2, seed=14, samples_per_frame=256)
assert cache_file(str(path), "fft", True) == ""
# Rewrite the tail as a genuine CACH v1 block (old header, no pad field).
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_pad_factor == 1
assert v1.precomputed_bg_sub is True
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True,
n_fft=v1.samples_per_frame * 10) is None
def test_mask_read_can_be_refused(rig, no_fft):
"""With no DC4 in hand, building the mask means reading a whole channel —
the GUI thread asks for None instead."""
assert cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
allow_dc_recompute=False) is None
dc4 = dc_image_mv(SrasFile(str(rig.path)), 0, CH4_IDX)
img = cached_rf_image(rig.sras, 0, _THRESHOLD_MV, apply_bg_sub=True,
dc4_mv=dc4, allow_dc_recompute=False)
assert img is not None and np.allclose(img, rig.fresh[0], atol=1e-3)
assert not no_fft, f"no FFT on either branch: {no_fft}"
# ---------------------------------------------------------------------------
# The viewer: no compute job at all for a batch-cached angle
# ---------------------------------------------------------------------------
def test_viewer_shows_stored_angles_without_computing(rig, no_fft, monkeypatch):
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._load_file(str(rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
# The file carries no DC block, so the mask comes from the window's own
# background precompute — the state a user is in by the time they click.
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(rig.n_angles))), \
"DC precompute finished"
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in list(range(rig.n_angles)) + [1, 0]:
win.spin_angle.setValue(a)
win._on_view_changed()
pump(60)
assert win._current_angle == a, f"angle {a} displayed"
assert np.allclose(win._current_image, rig.fresh[a], atol=1e-3), \
f"angle {a} shows the stored image"
assert dispatched == [], \
f"stored angles need no compute job, dispatched for {dispatched}"
assert not no_fft, f"no FFT ran for any stored angle: {no_fft}"
# Velocity is still a post-multiply of the same stored image.
win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
pump(120)
assert np.allclose(win._current_image,
rig.fresh[0] * win.spin_grating_um.value(), atol=1e-3)
assert dispatched == [] and not no_fft
# ...and a live control that no longer matches the stored image's own
# provenance must NOT force a recompute either — the stored image is
# shown as-is; only an explicit batch recompute changes what's shown.
win.combo_channel.setCurrentIndex(CH1_IDX)
pump(60)
win.chk_bg_sub.setChecked(False)
pump(200)
assert not win._job_running("compute") and not no_fft, \
"bg-sub off still shows the stored image, no real FFT"
assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3)
finally:
win.close()
pump(300)
def test_batch_caches_at_the_viewers_pad_factor(tmp_path, no_fft, monkeypatch):
"""The bug a pad-10 user hits: Batch Compute FFT used to store pad-1
images regardless, so the padded view recomputed every angle forever.
"""
path = tmp_path / "padded_gui.sras"
gen.write(path, n_angles=3, seed=15, samples_per_frame=256)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
try:
win._fft_pad_factor = 10
win._load_file(str(path))
assert wait_until(lambda: win._sras is not None), "file loaded"
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
for a in range(win._sras.n_angles))), \
"DC precompute finished"
# Convert -> Batch Compute FFT, on the open file, through the real slot.
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute("fft")
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_pad_factor == 10, \
f"cached at the viewer's pad, got {win._sras.precomputed_pad_factor}"
expected = {a: compute.cached_rf_image(
win._sras, a, dc_threshold_mv=win.spin_threshold_mv.value(),
apply_bg_sub=win.chk_bg_sub.isChecked(),
n_fft=win._current_n_fft(),
dc4_mv=win._dc_cache.get((a, CH4_IDX)))
for a in range(win._sras.n_angles)}
assert all(v is not None for v in expected.values()), "cache is readable at pad 10"
no_fft.clear()
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in range(win._sras.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert np.allclose(win._current_image, expected[a], atol=1e-3), \
f"angle {a} served from the pad-10 cache"
assert dispatched == [] and not no_fft, \
f"no recompute at pad 10 (jobs={dispatched}, fft={no_fft})"
assert "unusable" not in win.lbl_frame_warn.text()
# Change the live pad control so it no longer matches the stored
# image's own provenance — the info panel has to say so rather than
# leave it a mystery, but the stored pad-10 image keeps displaying;
# only an explicit batch recompute would ever produce a pad-4 one.
win._fft_pad_factor = 4
win._update_scan_info_labels()
assert "differ from the stored cache" in win.lbl_frame_warn.text(), \
win.lbl_frame_warn.text()
assert "pad 10x" in win.lbl_frame_warn.text()
last_angle = win._current_angle
win._refresh_display()
assert wait_until(lambda: not win._job_running("compute")), "settled"
assert not no_fft, "no real FFT ran — the pad-10 cache still served the view"
assert np.allclose(win._current_image, expected[last_angle], atol=1e-3), \
"pad-10 cache still shown after the live pad control diverged"
finally:
win.close()
pump(300)
def test_viewer_shows_stored_dc_without_computing(dc_rig, monkeypatch):
"""Same for the DC half, with the background precompute silenced so the
file's stored block is the only thing that can be carrying the display."""
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
dispatched = []
original_start = type(win)._start_compute
monkeypatch.setattr(type(win), "_start_compute",
lambda self: (dispatched.append(self.spin_angle.value()),
original_start(self))[1])
monkeypatch.setattr(type(win), "_start_dc_precompute", lambda self: None)
try:
win._load_file(str(dc_rig.path))
assert wait_until(lambda: win._sras is not None), "file loaded"
pump(120)
for ch in (CH3_IDX, CH4_IDX):
win.combo_channel.setCurrentIndex(ch)
for a in range(dc_rig.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert (win._current_angle, win._current_ch) == (a, ch), \
f"angle {a} on channel {ch} displayed"
assert np.array_equal(win._current_image,
dc_rig.sras.cached_dc_mv(a, ch)), \
f"angle {a} channel {ch} shows the file's stored DC image"
assert dispatched == [], \
f"stored DC angles need no compute job, dispatched for {dispatched}"
finally:
win.close()
pump(300)
# ---------------------------------------------------------------------------
# Row-averaged FFT: cache_file("fft_rowavg", ...) and its on-disk provenance
# ---------------------------------------------------------------------------
def test_row_average_flag_and_n_round_trip(tmp_path):
path = tmp_path / "rowavg_roundtrip.sras"
gen.write(path, n_angles=2, seed=20, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert sras.version == 7
assert sras.precomputed_row_avg_n == 5
assert all(x is not None for x in sras.precomputed_freq_mhz)
def test_fft_rowavg_mode_requires_positive_n_and_threshold(tmp_path):
path = tmp_path / "rowavg_bad_args.sras"
gen.write(path, n_angles=1, seed=27, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=0.0, row_avg_n=0)
assert cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=None, row_avg_n=5)
def test_raw_and_row_averaged_caches_never_cross_served(tmp_path):
"""The central regression this feature must never allow: a raw request
served a row-averaged image (or vice versa), or a request at one window
size served a cache stored at a different one."""
path = tmp_path / "cross_serve.sras"
gen.write(path, n_angles=1, seed=21, samples_per_frame=128)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
assert err == "", err
sras = SrasFile(str(path))
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=0) is None, \
"a raw request must not be served a row-averaged cache"
assert cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=3) is None, \
"a request at the wrong window size must not be served either"
served = cached_rf_image(sras, 0, None, apply_bg_sub=True, row_avg_n=5)
assert served is not None
assert np.array_equal(served, sras.precomputed_freq_mhz[0])
def test_write_v7_cache_row_avg_n_carries_forward(tmp_path):
"""A later DC-only write must leave a previously-written row-averaged
FFT block -- including its row_avg_n -- byte-for-byte unchanged."""
path = tmp_path / "carry_forward.sras"
gen.write(path, n_angles=2, seed=22, samples_per_frame=64)
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=7)
assert err == "", err
before = SrasFile(str(path))
assert before.precomputed_row_avg_n == 7
freq_before = [x.copy() for x in before.precomputed_freq_mhz]
err = cache_file(str(path), "dc", True)
assert err == "", err
after = SrasFile(str(path))
assert after.precomputed_row_avg_n == 7, "row_avg_n survives a DC-only write"
assert all(np.array_equal(after.precomputed_freq_mhz[a], freq_before[a])
for a in range(after.n_angles)), \
"the row-averaged FFT block itself is untouched by a DC-only write"
def test_row_average_never_touches_dc_images(tmp_path):
path = tmp_path / "dc_untouched.sras"
gen.write(path, n_angles=2, seed=23, samples_per_frame=64)
src = SrasFile(str(path))
expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
assert cache_file(str(path), "dc", True) == ""
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=6) == ""
after = SrasFile(str(path))
assert all(np.allclose(after.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
for a in range(after.n_angles))
assert all(np.allclose(after.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
for a in range(after.n_angles))
def test_row_average_never_modifies_raw_waveform_data(tmp_path):
path = tmp_path / "waveform_untouched.sras"
gen.write(path, n_angles=2, seed=24, samples_per_frame=64)
orig = tmp_path / "waveform_untouched_orig.sras"
gen.write(orig, n_angles=2, seed=24, samples_per_frame=64)
assert cache_file(str(path), "fft_rowavg", True,
dc_threshold_mv=-1e9, row_avg_n=5) == ""
o, n = SrasFile(str(orig)), SrasFile(str(path))
assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
for a in range(o.n_angles)), \
"waveform data untouched by a row-averaged cache write"
def test_cach_v1_backward_compat_defaults_row_avg_n_zero(tmp_path):
"""A v1 CACH tail predates row-averaged FFT caching entirely (no
row_avg_n byte at all) -- readers must still parse it in full, treating
it as row_avg_n=0. This is what protects an existing real-world v7
file's already-stored FFT cache from silently becoming unusable after
this change ships."""
path = tmp_path / "v1_rowavg.sras"
gen.write(path, n_angles=2, seed=25, samples_per_frame=128)
assert cache_file(str(path), "fft", True) == ""
v2 = SrasFile(str(path))
freq, entries = v2.precomputed_freq_mhz, list(range(v2.n_angles))
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
fmt.SFFT_FLAG_BG_SUB, len(entries))
for a in entries:
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
head = path.read_bytes()[:v2._cache_tail_offset()]
path.write_bytes(head + payload)
v1 = SrasFile(str(path))
assert v1.precomputed_row_avg_n == 0
assert all(np.array_equal(v1.precomputed_freq_mhz[a], freq[a]) for a in entries), \
"v1 images read back unchanged"
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=0) is not None
assert cached_rf_image(v1, 0, None, apply_bg_sub=True, row_avg_n=5) is None, \
"a v1 tail (predating this feature) can never satisfy a row-averaged request"
def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch, 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
worker, the worker reaches cache_file, and the written file is
self-describing afterward. Also the exact scenario the row-averaged-FFT
recompute bug reported: before the fix, the display always asked for
row_avg_n=0 regardless of what the file actually had stored, so viewing
an angle after this batch action saw a phantom mismatch and launched a
full raw recompute on every view switch. This asserts that no longer
happens -- the stored row-averaged image is shown directly, with no
dispatched compute job and no real FFT."""
path = tmp_path / "rowavg_gui.sras"
gen.write(path, n_angles=2, seed=26, samples_per_frame=128)
class _StubDialog:
def __init__(self, *a, **k):
pass
def exec(self):
return QDialog.DialogCode.Accepted
def get_half_width(self):
return 6
def get_threshold_mv(self):
return -1e9 # mask nothing, keep the comparison simple
monkeypatch.setattr("sras_viewer.main_window.RowAverageFftOptionsDialog", _StubDialog)
app = QApplication.instance() or QApplication([]) # noqa: F841
win = SrasViewerWindow()
win.show()
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"
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
return_value=([str(path)], "")):
win._on_batch_compute_row_avg()
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
assert wait_until(lambda: win._sras is not None
and win._sras.version == 7), "file reloaded as v7"
pump(200)
assert win._sras.precomputed_row_avg_n == 6
assert "row-averaged n=6" in win.lbl_frame_warn.text(), win.lbl_frame_warn.text()
expected = compute.cached_rf_image(win._sras, 0, dc_threshold_mv=None,
apply_bg_sub=win.chk_bg_sub.isChecked(),
row_avg_n=6)
assert expected is not None, "the batch write left a readable row-averaged cache"
# The regression this batch action used to leave unfixed: viewing an
# angle afterward must show the stored row-averaged image directly,
# never fall through to a real (raw) recompute.
no_fft.clear()
dispatched.clear()
win.combo_channel.setCurrentIndex(CH1_IDX)
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
for a in range(win._sras.n_angles):
win.spin_angle.setValue(a)
pump(60)
assert np.allclose(win._current_image,
compute.cached_rf_image(
win._sras, a,
dc_threshold_mv=win.spin_threshold_mv.value(),
apply_bg_sub=win.chk_bg_sub.isChecked(),
row_avg_n=6), atol=1e-3), \
f"angle {a} shows the stored row-averaged image"
assert dispatched == [] and not no_fft, \
f"no recompute for row-averaged angles (jobs={dispatched}, fft={no_fft})"
finally:
win.close()
pump(300)