Compare commits
19 Commits
eecb0e3a82
..
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 48d560fb00 | |||
| 2937c017a7 | |||
| a8b962e317 | |||
| efecc154fe | |||
| ec0d322871 | |||
| aca501ae5c | |||
| 105b9514a5 | |||
| 4212b8c313 | |||
| 6976cbf767 | |||
| c54cce453c | |||
| 1caf6373cb | |||
| 191d1b8946 | |||
| d5914b5793 | |||
| c30c8b1815 | |||
| 8348ad313c | |||
| f40c965b74 | |||
| 6d0e30b9ce | |||
| 3989c2a1b8 | |||
| 1a8187138f |
+323
-33
@@ -32,45 +32,201 @@ share of the budget. Capping the workers alone is not enough: the chunk would
|
||||
still be sized against the whole budget, and N concurrent callers would each
|
||||
allocate all of it.
|
||||
|
||||
## FFT peak search: block-parallel zoom refinement (`sras_compute.py`)
|
||||
## FFT peak search: block-parallel direct transform (`sras_compute.py`)
|
||||
|
||||
The displayed RF value per pixel is the argmax of the zero-padded power
|
||||
spectrum of that pixel's CH1 waveform. At the pad factor of 40 needed for
|
||||
mapping resolution, materialising padded spectra is hopeless: ~9 GB per scan
|
||||
row, which is what used to collapse the old row-chunk planner to one worker
|
||||
and make synthesis single-threaded.
|
||||
spectrum of that pixel's CH1 waveform. This used to run through `_peak_bins_zoom`,
|
||||
a coarse-rfft-plus-local-fine-DFT refinement that avoided ever materialising
|
||||
a padded spectrum — at the pad factor of 40 needed for mapping resolution, a
|
||||
full padded spectrum is ~9 GB per scan row, which used to collapse the old
|
||||
row-chunk planner to one worker and make synthesis single-threaded. That
|
||||
refinement was removed once pyFFTW became the sole, mandatory FFT backend
|
||||
(SciPy dropped as a compute backend entirely): `_peak_bins` now always runs
|
||||
the real full transform, and the memory problem zoom dodged is instead
|
||||
solved by bounding *per-block* spectrum memory rather than avoiding full
|
||||
spectra altogether.
|
||||
|
||||
`_peak_bins_zoom` never materialises the padded spectrum:
|
||||
`_peak_bins` runs the full transform (`_block_rfft`, a cached pyFFTW
|
||||
`builders.rfft` plan, FFTW_MEASURE, wisdom persisted under
|
||||
`~/.cache/sras-viewer/`) and argmaxes the power spectrum, in blocks fanned
|
||||
out task-parallel over a persistent thread pool (`_fft_pool()`) — each pool
|
||||
thread runs one single-threaded transform at a time, so aggregate
|
||||
parallelism equals the pool's worker count. This block-streaming structure
|
||||
is not just about parallelism: it is also what keeps memory bounded, by
|
||||
never materialising more than one block's worth of full padded spectrum at
|
||||
a time, regardless of how many waveforms a chunk holds.
|
||||
|
||||
1. a coarse rfft at `next_fast_len(2*spf)` — 2× oversampled, so the padded
|
||||
power spectrum (a trig polynomial of degree spf−1) cannot hide its global
|
||||
max between coarse samples;
|
||||
2. every coarse bin within `_ZOOM_CAND_RATIO` (0.7) of its row's coarse max
|
||||
becomes a refinement candidate. Quarter-natural-bin scalloping at the 2×
|
||||
grid can understate a peak's power by at most ~19%, so 0.7 keeps a wide
|
||||
margin. The DC-adjacent window is always refined too: the coarse DC bin
|
||||
is zeroed for suppression, which would otherwise blind the scan to fine
|
||||
bins closer to DC than the first coarse sample (where the leakage skirt
|
||||
of an un-subtracted offset peaks);
|
||||
3. each candidate window (±`_ZOOM_HALFWIDTH` = 0.75 coarse spacings; every
|
||||
fine bin lies within 0.5 spacings of its nearest coarse bin) is evaluated
|
||||
on the exact `n_fft` grid by one small complex gemm, with np.argmax's
|
||||
lowest-bin tie-break preserved across windows.
|
||||
The block size is the part that has to adapt to pad factor.
|
||||
`_fft_block_for(spf, n_len)` derives waveforms-per-task from a fixed
|
||||
per-thread byte budget (`_FFT_PLAN_BYTES_BUDGET`, `SRAS_FFT_PLAN_BUDGET_MB`,
|
||||
default 16 MB) rather than a fixed constant, because a cached pyFFTW plan's
|
||||
input+output buffers are *permanent* per-thread memory (the plan cache is
|
||||
never evicted) — with a fixed 512-waveform block, pad 40 at a 2500-sample
|
||||
frame costs ~210 MB per pool thread (~3.3 GB total across 16 threads);
|
||||
`_fft_block_for` bounds that to ~16 MB per thread (~260 MB across 16
|
||||
threads) at the same pad factor, while still reproducing the old tuned 512
|
||||
exactly at natural resolution (pad 1), where it cost nothing to begin with.
|
||||
`_FFT_BLOCK_MAX` (512) and `_FFT_BLOCK_MIN` (32) cap and floor the result:
|
||||
the ceiling is the measured knee on a 16-core machine at natural resolution
|
||||
(smaller blocks serialise on GIL-held numpy dispatch, larger ones lose cache
|
||||
residency and task granularity); the floor keeps task granularity from
|
||||
collapsing at extreme pad factors, at the cost of exceeding the byte budget
|
||||
there.
|
||||
|
||||
The selected bin is bit-identical to the full padded argmax — enforced by
|
||||
`tests/test_compute.py::test_zoom_identity`, a fuzz test over adversarial
|
||||
spectra, and the golden-hash harness (`tools/check_equivalence.py`), whose
|
||||
baseline was captured on the old full-padded path.
|
||||
The outer row-chunk sizing (`_plan_fft_rows`) needed no companion change.
|
||||
It only ever budgets the raw float32 waveform *read* buffer, which this
|
||||
change doesn't touch — spectrum memory is bounded independently by
|
||||
`_fft_block_for`, and since the pool only ever runs as many blocks
|
||||
concurrently as it has workers, peak transient spectrum memory during a
|
||||
chunk's FFT phase is `_MAX_WORKERS * block * bytes_per_wf`, the same bound
|
||||
whether the chunk holds 10 rows or 10,000. Queuing more rows into one chunk
|
||||
to keep the read-row pool busy therefore can't blow up spectrum memory.
|
||||
|
||||
Work fans out over a persistent thread pool in `_FFT_BLOCK` = 512-waveform
|
||||
tasks: smaller blocks serialise on GIL-held numpy dispatch, larger ones lose
|
||||
cache residency and task granularity (measured on a 16-core machine, where
|
||||
this path runs ~35× faster than the old serial padded transform at pad 40).
|
||||
pyFFTW runs through per-thread `builders` plans (FFTW_MEASURE, wisdom
|
||||
persisted under `~/.cache/sras-viewer/`), and `threadpoolctl` clamps BLAS to
|
||||
one thread under the pool so the refinement gemm cannot oversubscribe.
|
||||
`compute_rf_image(exact=True)` (or `SRAS_FFT_EXACT=1`) keeps the reference
|
||||
full-padded path for audits.
|
||||
`tools/check_equivalence.py`'s golden-hash harness remains the end-to-end
|
||||
regression baseline for this path, unaffected by this change.
|
||||
|
||||
## Row-averaged FFT: same-row, distance-weighted SNR cleanup (`sras_compute.py`)
|
||||
|
||||
`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` 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.
|
||||
|
||||
The min peak frequency floor is the fourth recorded setting, and it sits
|
||||
between the DC threshold (fully re-applicable at display time) and the
|
||||
baked-in three: it is *tighten-only* re-applicable. A floor at or above the
|
||||
stored one is served by masking stored pixels below it to the 0.0 "no valid
|
||||
peak" sentinel — invalidated, not re-resolved; only a real recompute can
|
||||
find their true above-floor peak. A floor *below* the stored one is a
|
||||
genuine mismatch: the stored search never looked at those bins. Because a
|
||||
served-then-masked image is a lossy stand-in for a real floored compute,
|
||||
`cache_file` passes `use_stored=False` so a batch recompute always runs the
|
||||
real FFT — otherwise re-batching a cached file would silently bake the
|
||||
masked copy of its own old image back into the store.
|
||||
|
||||
Padding is the subtlest of the three, because a padded FFT looks like it
|
||||
should be a refinement of the unpadded one. It isn't: zero-padding
|
||||
interpolates between the natural bins, so it resolves a different peak
|
||||
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. Two settings are taken live instead: the DC threshold and the min
|
||||
peak frequency floor, both cheaply re-appliable as masks on serve. So for
|
||||
bg-sub/pad/row-avg, presence alone decides whether a stored image is shown;
|
||||
the floor is the one case where a live control can gate it — a live floor
|
||||
*below* the stored one can't be answered by masking, so `cached_rf_image`
|
||||
declines and the caller falls through to a real compute. They still matter for two things: a genuinely never-computed
|
||||
angle's first live compute, and an explicit batch recompute — both of which
|
||||
read the live controls and produce new stored data, at which point it's the
|
||||
new data's *own* settings that get self-matched from then on. This is what
|
||||
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`)
|
||||
|
||||
@@ -109,6 +265,140 @@ 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
|
||||
|
||||
+3
-3
@@ -14,10 +14,8 @@ dependencies = [
|
||||
"scipy==1.18.0",
|
||||
# Angle alignment only: masked FFT phase correlation (skimage.registration).
|
||||
"scikit-image==0.26.0",
|
||||
# Faster rfft backend; the viewer falls back to scipy.fft without it.
|
||||
# Mandatory rfft backend for the peak search (no SciPy fallback).
|
||||
"pyFFTW==0.15.1",
|
||||
# Clamps BLAS threading under the FFT worker pool.
|
||||
"threadpoolctl==3.6.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
@@ -30,7 +28,9 @@ sras-viewer = "sras_viewer.main_window:main"
|
||||
py-modules = [
|
||||
"sras_format",
|
||||
"sras_compute",
|
||||
"sras_render",
|
||||
"sras_workers",
|
||||
"sras_align_export",
|
||||
"sras_average",
|
||||
"sras_edit_scans",
|
||||
]
|
||||
|
||||
+113
-9
@@ -173,6 +173,41 @@ 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.
|
||||
|
||||
---
|
||||
|
||||
@@ -220,7 +255,7 @@ actions.
|
||||
| Offset | Size | Type | Field | Description |
|
||||
|--------|------|------|-------|-------------|
|
||||
| 0 | 4 | `char[4]` | `cach_magic` | `CACH` (ASCII). Missing/wrong magic → treat file as having no cache. |
|
||||
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `1`. Readers must treat the file as uncached if this is not a version they understand (unlike v5's `PREC` section, which read but never validated its version byte). |
|
||||
| 4 | 1 | `u8` | `cach_version` | Cache format version. Currently `4`; readers also accept `1`–`3` (each older tail simply lacks the fields added since — see the `SFFT` block below and [CACH tail version history](#cach-tail-version-history)). Any other value → treat the file as uncached (unlike v5's `PREC` section, which read but never validated its version byte). |
|
||||
| 5 | 1 | `u8` | `block_flags` | Bit 0 = DC block (`SDCB`) follows. Bit 1 = FFT block (`SFFT`) follows, immediately after the DC block if both are present. Bits 2–7 reserved, must be zero on write. |
|
||||
|
||||
### DC block `SDCB` (present iff `block_flags & 0x01`)
|
||||
@@ -249,13 +284,33 @@ value, same as v5's `PREC` section.
|
||||
|
||||
### FFT block `SFFT` (present iff `block_flags & 0x02`)
|
||||
|
||||
7-byte block header, format `">4sBH"`:
|
||||
Block header layout depends on `cach_version`:
|
||||
|
||||
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`.
|
||||
- **`cach_version` 4**: 14 bytes, format `">4sBHBHI"` — + `min_freq_khz`.
|
||||
Always written by current code.
|
||||
|
||||
An older tail is read with its absent fields taken as the only value such a
|
||||
tail can describe: `row_avg_n = 0` for a `cach_version` 1 tail, which
|
||||
predates row-averaged FFT caching, `pad_factor = 1` for `cach_version`
|
||||
1 or 2, which predate padded caching and are therefore natural-resolution,
|
||||
and `min_freq_khz = 0` (no floor) for `cach_version` 1–3, which predate the
|
||||
min peak frequency floor and therefore searched every bin above DC.
|
||||
Files cached before any of these changes keep working with no recompute.
|
||||
|
||||
| Offset (rel) | Size | Type | Field | Description |
|
||||
|--------------|------|------|-------|-------------|
|
||||
| 0 | 4 | `char[4]` | `magic` | `SFFT` |
|
||||
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bits 1–7 reserved. |
|
||||
| 4 | 1 | `u8` | `flags` | Bit 0 = `bg_sub_applied` — background waveform was subtracted from CH1 before the FFT when these images were computed. Bit 1 = `row_averaged` — `peak_freq_mhz` came from same-row, distance-weighted averaged CH1 waveforms rather than raw per-pixel ones; `row_avg_n` (below) is the neighbor half-width used. Bits 2–7 reserved. |
|
||||
| 5 | 2 | `u16` | `n_stored` | Number of angle entries that follow |
|
||||
| 7 | 1 | `u8` | `row_avg_n` | *`cach_version` ≥ 2 only.* Same-row neighbor half-width, in pixels, that `peak_freq_mhz` was averaged over before its FFT; `0` = raw (unaveraged). Meaningful only when `flags` bit 1 is set — a `cach_version` 1 tail has no such byte and is always `row_avg_n = 0`. |
|
||||
| 8 | 2 | `u16` | `pad_factor` | *`cach_version` ≥ 3 only.* Zero-padding factor the stored `peak_freq_mhz` was resolved at: `n_fft = pad_factor × samples_per_frame`, so `1` = natural resolution. Never `0`; a `cach_version` 1 or 2 tail has no such field and is always `pad_factor = 1`. |
|
||||
| 10 | 4 | `u32` | `min_freq_khz` | *`cach_version` ≥ 4 only.* Min peak frequency floor the stored peak search excluded bins below, fixed-point in units of 0.001 MHz (kHz); `0` = no floor. Fixed-point rather than `f32` so a value that round-trips through the file compares exactly against the same value re-requested by a reader (the viewer's floor control has 0.001 MHz granularity). A `cach_version` 1–3 tail has no such field and is always `min_freq_khz = 0`. |
|
||||
|
||||
followed by `n_stored` entries, each:
|
||||
|
||||
@@ -264,9 +319,10 @@ u16 angle_idx — index into the angle table (0-ba
|
||||
f32[n_rows[angle_idx] × n_frames[angle_idx]] peak_freq_mhz — CH1 FFT peak frequency, MHz, row-major
|
||||
```
|
||||
|
||||
**`peak_freq_mhz`** is computed without any DC-threshold masking (i.e. the
|
||||
FFT is run on every pixel unconditionally, same as v5's `PREC` convention).
|
||||
Readers apply the DC4 threshold at display time:
|
||||
**`peak_freq_mhz`** for a raw store (`row_avg_n == 0`) is computed without
|
||||
any DC-threshold masking (i.e. the FFT is run on every pixel
|
||||
unconditionally, same as v5's `PREC` convention). Readers apply the DC4
|
||||
threshold at display time:
|
||||
|
||||
```
|
||||
pixel is valid ⟺ dc4_mv[r][f] ≥ threshold_mv
|
||||
@@ -276,10 +332,40 @@ display_value = peak_freq_mhz[r][f] if valid, else 0
|
||||
using the DC4 image from the DC block if that angle is also cached there,
|
||||
else computed on demand.
|
||||
|
||||
For a row-averaged store (`row_avg_n > 0`), the DC4 threshold is applied
|
||||
*during* the store — a pixel below threshold is left at `0` and never
|
||||
contributes to any neighbor's average — since neighbor validity can't be
|
||||
deferred to display time the way plain masking can. The threshold value
|
||||
itself is not recorded, only that averaging happened and at what window
|
||||
size. Readers still apply their own live DC4 threshold at display time
|
||||
exactly as for a raw store, using whatever mask they currently have.
|
||||
|
||||
Readers must fall back to real-time FFT computation (ignoring stored
|
||||
`peak_freq_mhz`) under the same conditions as v5's PREC fast path: time-domain
|
||||
gating is active, zero-padding (`n_fft ≠ samples_per_frame`) is requested, or
|
||||
the reader's background-subtraction setting doesn't match `flags.bg_sub_applied`.
|
||||
`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`, the reader's requested `row_avg_n` doesn't match
|
||||
the stored value exactly, or the reader's requested min peak frequency
|
||||
floor is *below* the stored `min_freq_khz`. A raw request must never be
|
||||
served a row-averaged store, or vice versa; a request at one row-averaging
|
||||
window size must never be served a store at another; and a request at one
|
||||
padding must never be served a store at another, since a padded FFT
|
||||
interpolates between the natural bins and so resolves genuinely different
|
||||
peak frequencies. An `n_fft` that is not a whole multiple of
|
||||
`samples_per_frame` can never match any store, because only an integer
|
||||
`pad_factor` is representable.
|
||||
|
||||
The min peak frequency floor is the one asymmetric provenance field. A
|
||||
request at a floor *below* the stored one cannot be served: the stored
|
||||
search never looked at bins below its floor, so the stored numbers cannot
|
||||
say what a lower-floored search would have found. A request at a floor at
|
||||
or *above* the stored one **is** servable — the difference is re-applied at
|
||||
display time by masking every pixel whose stored `peak_freq_mhz` is below
|
||||
the requested floor to `0` (the same sentinel as the DC threshold mask;
|
||||
a genuine peak can never be `0`, since bin 0 is always excluded from the
|
||||
search). Such masked pixels are *invalid*, not re-resolved — only a real
|
||||
recompute can recover the strongest peak above the floor for them.
|
||||
|
||||
### In-place write ordering
|
||||
|
||||
@@ -294,6 +380,24 @@ 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`. |
|
||||
| 4 | `SFFT` header grows a `u32` `min_freq_khz` — the min peak frequency floor the stored peak search excluded bins below, in 0.001 MHz units, `0` = no floor. The floor exists because a pixel that passes the DC-bias threshold but carries only weak real signal can otherwise resolve to the un-subtracted background's DC-leakage skirt — an implausibly-near-zero frequency (and so an implausibly slow velocity) for a pixel that has a genuine peak higher up. Recording the floor is what makes it enforceable against a store: without it, a stored image silently bypassed the floor entirely. Unlike the other provenance fields it is asymmetric — a *higher* requested floor is servable by masking stored pixels below it, only a *lower* one forces a recompute (see above). Readers accept `cach_version` 1–3 tails as `min_freq_khz = 0`. |
|
||||
|
||||
A reader that does not know a `cach_version` must treat the file as
|
||||
uncached — not attempt a partial parse — and the file still reads as an
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,456 @@
|
||||
#!/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
|
||||
+169
-69
@@ -3,7 +3,23 @@
|
||||
sras_average.py — Waveform-averaging utility for .sras files.
|
||||
|
||||
Reduces memory footprint by coherently averaging every N consecutive frames
|
||||
along the acquisition axis, writing a new .sras file with n_frames / N frames.
|
||||
along the acquisition axis, writing a new v6 .sras file with (about)
|
||||
n_frames / N frames per angle.
|
||||
|
||||
Handles v6/v7 only (per-angle ragged geometry). A v7 input's cache tail is
|
||||
dropped — it's indexed by frame, which this changes — so the output is always
|
||||
written as v6; the viewer recomputes DC/FFT on next open. Reads via memmap and
|
||||
writes in row-chunks sized to a memory budget (default 1024 MB, override with
|
||||
the SRAS_MEM_BUDGET_MB env var), so peak RAM stays bounded regardless of file
|
||||
size — this is what makes the tool usable on multi-hundred-GB scans.
|
||||
|
||||
Averaging every N frames also coarsens the physical X spacing between output
|
||||
frames (each frame is a distinct stage position — see scan_format.md's
|
||||
Spatial Mapping section: x_k = x_start + k * velocity_mm_s / laser_freq_hz) —
|
||||
so the output header's laser_freq_hz is divided by N to keep x_axis_mm()
|
||||
correct on the averaged file. This means the GUI's "Laser freq" info label
|
||||
will show that adjusted value rather than the scope's real setting for an
|
||||
averaged file; v6/v7 has no separate field for effective pixel pitch.
|
||||
|
||||
Usage:
|
||||
python sras_average.py input.sras output.sras --n 10
|
||||
@@ -13,9 +29,15 @@ Options:
|
||||
--n INT Number of frames to average into one (required).
|
||||
--discard-remainder Drop trailing frames that don't fill a complete group.
|
||||
Default: include a partial average for the last group.
|
||||
|
||||
Environment:
|
||||
SRAS_MEM_BUDGET_MB Ceiling on one row-chunk's working memory, in MB
|
||||
(default 1024). Lower it on a memory-constrained
|
||||
machine; the tool just takes more, smaller chunks.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import shutil
|
||||
import struct
|
||||
import sys
|
||||
@@ -23,14 +45,22 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sras_format import HDR_FMT, HDR_SIZE, SrasFile
|
||||
from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile
|
||||
|
||||
_SUPPORTED = (2, 3, 4)
|
||||
_SUPPORTED = (6, 7)
|
||||
_DEFAULT_BUDGET_MB = 1024
|
||||
|
||||
# Throttle per-chunk progress printing to roughly this many lines per angle,
|
||||
# so a huge angle (thousands of chunks) doesn't flood stdout while a small
|
||||
# one still gets to print every chunk.
|
||||
_MAX_PROGRESS_LINES = 40
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(
|
||||
description="Average every N waveforms in a .sras file and write a new file."
|
||||
description="Average every N waveforms in a .sras file and write a new file.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
p.add_argument("input", help="Input .sras file")
|
||||
p.add_argument("output", help="Output .sras file")
|
||||
@@ -41,21 +71,45 @@ def parse_args():
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def read_header_sections(sras: SrasFile) -> bytes:
|
||||
"""The raw bytes between the header and the waveform data (angle table,
|
||||
row table, preambles, background), copied through verbatim so nothing is
|
||||
lost in a re-encode."""
|
||||
with open(sras.path, "rb") as f:
|
||||
f.seek(HDR_SIZE)
|
||||
return f.read(sras.data_offset - HDR_SIZE)
|
||||
def _memory_budget_bytes() -> int:
|
||||
return int(os.environ.get("SRAS_MEM_BUDGET_MB", _DEFAULT_BUDGET_MB)) * 1024 * 1024
|
||||
|
||||
|
||||
def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
|
||||
"""Average every N frames of one angle's (n_rows, n_ch, n_frames, spf)
|
||||
block. Returns int16 of shape (n_rows, n_ch, n_out, spf).
|
||||
def _plan_angle(n_frames_in: int, n: int, discard_remainder: bool) -> tuple[int, int, int]:
|
||||
"""(n_full, remainder, n_frames_out) for averaging one angle's frames."""
|
||||
n_full = n_frames_in // n
|
||||
remainder = n_frames_in % n
|
||||
n_out = n_full + (1 if remainder and not discard_remainder else 0)
|
||||
return n_full, remainder, n_out
|
||||
|
||||
Averaging is done in float32 and rounded on cast, matching numpy's mean
|
||||
followed by an int16 cast in the original implementation.
|
||||
|
||||
def _chunk_rows(n_rows: int, n_channels: int, n_frames_in: int,
|
||||
samples_per_frame: int, budget: int) -> int:
|
||||
"""How many rows to hold in RAM at once so one chunk's working buffers
|
||||
(source block, mean's float64 accumulator, output block) fit the budget.
|
||||
|
||||
The 4 bytes/sample below is a deliberate middle estimate, not a sum of
|
||||
the three buffers: the int16 source and int16 output are 2 each, and
|
||||
only mean()'s float64 result is 8, over the reduced frame axis rather
|
||||
than the full block. Raise SRAS_MEM_BUDGET_MB if a machine still runs
|
||||
tight at a large --n."""
|
||||
bytes_per_row = max(1, n_channels * n_frames_in * samples_per_frame * 4)
|
||||
return max(1, min(n_rows, budget // bytes_per_row))
|
||||
|
||||
|
||||
def _average_block(block: np.ndarray, n: int, discard_remainder: bool,
|
||||
bps: int) -> np.ndarray:
|
||||
"""Average every N frames of a (rows, n_ch, n_frames, spf) block along
|
||||
the frame axis, returning the result already encoded in the on-disk
|
||||
dtype: int8 (clipped) if bps == 1, else big-endian int16.
|
||||
|
||||
Cast to native int16 (not float32) before calling .mean(): numpy's mean()
|
||||
uses a float64 accumulator by default for integer input, matching the
|
||||
original implementation exactly (which kept the whole file as int16 and
|
||||
called .mean() directly). A float32 cast here would use a float32
|
||||
accumulator instead — for large group sizes that can round the sum
|
||||
differently than float64 and, after the int16 cast below, occasionally
|
||||
land on a value 1 ADC count away from the original tool's output.
|
||||
"""
|
||||
n_frames = block.shape[2]
|
||||
n_full = n_frames // n
|
||||
@@ -63,46 +117,96 @@ def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarr
|
||||
|
||||
parts = []
|
||||
if n_full:
|
||||
full = block[:, :, :n_full * n, :].astype(np.float32)
|
||||
full = block[:, :, :n_full * n, :].astype(np.int16)
|
||||
full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3])
|
||||
parts.append(full.mean(axis=3).astype(np.int16))
|
||||
if remainder and not discard_remainder:
|
||||
tail = block[:, :, n_full * n:, :].astype(np.float32)
|
||||
tail = block[:, :, n_full * n:, :].astype(np.int16)
|
||||
parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16))
|
||||
|
||||
if not parts:
|
||||
return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
|
||||
return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
|
||||
averaged = np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
|
||||
else:
|
||||
averaged = parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
|
||||
|
||||
if bps == 1:
|
||||
return np.clip(averaged, -128, 127).astype(np.int8)
|
||||
return averaged.astype(">i2")
|
||||
|
||||
|
||||
def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes,
|
||||
n: int, discard_remainder: bool) -> int:
|
||||
"""Stream each angle through the averager, writing as we go so peak RAM
|
||||
stays at one angle's block rather than the whole file."""
|
||||
def write_v6_averaged(sras: SrasFile, out_path: Path, n: int,
|
||||
discard_remainder: bool, budget: int | None = None) -> list[int]:
|
||||
"""Write *sras* averaged every N frames to a new v6 .sras file, streaming
|
||||
row-chunks per angle so peak RAM never holds more than one chunk (bounded
|
||||
by *budget* bytes, default from SRAS_MEM_BUDGET_MB).
|
||||
|
||||
Stages to a sibling '.part' file and os.replace()s it into position on
|
||||
success, per scan_format.md's requirement that any .sras writer must
|
||||
never leave a half-written file visible under its final name (a short
|
||||
v6 file opens successfully with the wrong angle/frame count rather than
|
||||
failing loudly).
|
||||
|
||||
Returns the per-angle output frame counts.
|
||||
"""
|
||||
budget = _memory_budget_bytes() if budget is None else max(1, budget)
|
||||
n_ch = sras.n_channels
|
||||
spf = sras.samples_per_frame
|
||||
bps = sras.bytes_per_sample
|
||||
n_frames_in = int(sras.n_frames[0])
|
||||
n_out = n_frames_in // n
|
||||
if n_frames_in % n and not discard_remainder:
|
||||
n_out += 1
|
||||
|
||||
plans = [_plan_angle(int(sras.n_frames[a]), n, discard_remainder)
|
||||
for a in range(sras.n_angles)]
|
||||
n_out_per_angle = [p[2] for p in plans]
|
||||
|
||||
header = struct.pack(
|
||||
HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]),
|
||||
float(sras.x_start_mm[0]), float(sras.x_delta_mm),
|
||||
sras.velocity_mm_s, sras.laser_freq_hz,
|
||||
n_out, # updated frame count
|
||||
sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels,
|
||||
HDR_FMT_V6, b"SRAS", 6, sras.n_angles,
|
||||
sras.x_start_nominal_mm, sras.y_start_nominal_mm,
|
||||
sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
|
||||
sras.row_spacing_mm, sras.velocity_mm_s, sras.laser_freq_hz / n,
|
||||
spf, sras.sample_rate_hz, bps, n_ch,
|
||||
)
|
||||
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(header)
|
||||
f.write(mid_sections)
|
||||
for a in range(sras.n_angles):
|
||||
averaged = average_rows(sras.data[a], n, discard_remainder)
|
||||
if bps == 1:
|
||||
f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes())
|
||||
else:
|
||||
f.write(averaged.astype(">i2").tobytes())
|
||||
return n_out
|
||||
geo = bytearray()
|
||||
for a in range(sras.n_angles):
|
||||
geo += struct.pack(GEO_FMT_V6, float(sras.x_start_mm[a]),
|
||||
float(sras.x_delta_mm_per_angle[a]),
|
||||
int(n_out_per_angle[a]), int(sras.n_rows[a]))
|
||||
|
||||
part_path = out_path.with_name(out_path.name + ".part")
|
||||
try:
|
||||
with open(part_path, "wb") as f:
|
||||
f.write(header)
|
||||
f.write(sras.angles_deg.astype(">f4").tobytes())
|
||||
f.write(bytes(geo))
|
||||
for a in range(sras.n_angles):
|
||||
f.write(sras.y_pos_per_angle[a].astype(">f4").tobytes())
|
||||
f.write(sras.encoded_preambles())
|
||||
f.write(sras.encoded_background())
|
||||
|
||||
for a in range(sras.n_angles):
|
||||
n_rows = int(sras.n_rows[a])
|
||||
n_frames_in = int(sras.n_frames[a])
|
||||
chunk_rows = _chunk_rows(n_rows, n_ch, n_frames_in, spf, budget)
|
||||
total_chunks = max(1, -(-n_rows // chunk_rows))
|
||||
print_every = max(1, total_chunks // _MAX_PROGRESS_LINES)
|
||||
|
||||
data = sras.data[a]
|
||||
for i, r0 in enumerate(range(0, n_rows, chunk_rows)):
|
||||
r1 = min(r0 + chunk_rows, n_rows)
|
||||
block = np.asarray(data[r0:r1])
|
||||
out_block = _average_block(block, n, discard_remainder, bps)
|
||||
f.write(out_block.tobytes())
|
||||
if total_chunks > 1 and (i % print_every == 0 or r1 == n_rows):
|
||||
print(f" angle {a + 1}/{sras.n_angles}: "
|
||||
f"{r1}/{n_rows} rows", flush=True)
|
||||
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(part_path, out_path)
|
||||
except BaseException:
|
||||
part_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
return n_out_per_angle
|
||||
|
||||
|
||||
def main():
|
||||
@@ -126,18 +230,16 @@ def main():
|
||||
sras = SrasFile(str(in_path))
|
||||
if sras.version not in _SUPPORTED:
|
||||
print(f"Error: unsupported .sras version: {sras.version} "
|
||||
f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})",
|
||||
f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))} only)",
|
||||
file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
n_frames_in = int(sras.n_frames[0])
|
||||
print(f" Version : v{sras.version}")
|
||||
print(f" Angles : {sras.n_angles}")
|
||||
print(f" Rows : {int(sras.n_rows[0])}")
|
||||
print(f" Frames (actual): {n_frames_in}")
|
||||
print(f" Channels : {sras.n_channels}")
|
||||
print(f" Samples/frame : {sras.samples_per_frame}")
|
||||
print(f" Bytes/sample : {sras.bytes_per_sample}")
|
||||
aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else ""
|
||||
print(f" Version : v{sras.version}")
|
||||
print(f" Angles : {sras.n_angles}{aborted_note}")
|
||||
print(f" Channels : {sras.n_channels}")
|
||||
print(f" Samples/frame: {sras.samples_per_frame}")
|
||||
print(f" Bytes/sample : {sras.bytes_per_sample}")
|
||||
|
||||
if args.n == 1:
|
||||
print("--n 1: no averaging needed; copying file as-is.")
|
||||
@@ -145,28 +247,26 @@ def main():
|
||||
print(f"Wrote {out_path}")
|
||||
return
|
||||
|
||||
if args.n > n_frames_in:
|
||||
print(f"Warning: --n ({args.n}) exceeds available frames ({n_frames_in}). "
|
||||
"The entire dataset will be averaged into a single frame.")
|
||||
print(f"\n{'idx':>4} {'angle_deg':>10} {'rows':>6} {'frames_in':>10} {'frames_out':>11}")
|
||||
for a in range(sras.n_angles):
|
||||
n_frames_in = int(sras.n_frames[a])
|
||||
n_full, remainder, n_out = _plan_angle(n_frames_in, args.n, args.discard_remainder)
|
||||
print(f"{a:>4} {sras.angles_deg[a]:>10.4f} {int(sras.n_rows[a]):>6} "
|
||||
f"{n_frames_in:>10} {n_out:>11}")
|
||||
if args.n > n_frames_in:
|
||||
print(f" Warning: --n ({args.n}) exceeds angle {a}'s frames "
|
||||
f"({n_frames_in}); it collapses to a single frame.")
|
||||
|
||||
print(f"\nAveraging every {args.n} frames ...", flush=True)
|
||||
print(f"\nWriting {out_path} ...", flush=True)
|
||||
mid_sections = read_header_sections(sras)
|
||||
n_frames_out = write_averaged(out_path, sras, mid_sections,
|
||||
args.n, args.discard_remainder)
|
||||
if sras.version == 7:
|
||||
print("\nNote: input has a v7 cache tail; it is indexed by frame count "
|
||||
"and will be dropped. The viewer will recompute DC/FFT on next open.")
|
||||
|
||||
n_full = n_frames_in // args.n
|
||||
remainder = n_frames_in % args.n
|
||||
if remainder and not args.discard_remainder:
|
||||
status = f"({n_full} full groups + 1 partial group of {remainder})"
|
||||
elif remainder:
|
||||
status = f"({n_full} full groups, {remainder} trailing frames discarded)"
|
||||
else:
|
||||
status = f"({n_full} full groups)"
|
||||
print(f" {n_frames_in} frames -> {n_frames_out} frames {status}")
|
||||
print(f"\nAveraging every {args.n} frames and writing {out_path} ...", flush=True)
|
||||
n_out_per_angle = write_v6_averaged(sras, out_path, args.n, args.discard_remainder)
|
||||
|
||||
in_mb = in_path.stat().st_size / 1024**2
|
||||
out_mb = out_path.stat().st_size / 1024**2
|
||||
print(f"\n Frames out: {n_out_per_angle}")
|
||||
print(f" Input size : {in_mb:.1f} MB")
|
||||
print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)")
|
||||
print("Done.")
|
||||
|
||||
+645
-319
File diff suppressed because it is too large
Load Diff
@@ -14,6 +14,11 @@ 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
|
||||
|
||||
+247
-17
@@ -60,7 +60,18 @@ PREC_FLAG_BG_SUB = 0x01
|
||||
CACH_MAGIC = b"CACH"
|
||||
CACH_HDR_FMT = ">4sBB" # magic, cach_version, block_flags
|
||||
CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT)
|
||||
CACH_VERSION = 1
|
||||
CACH_VERSION = 4 # written on every fresh write
|
||||
CACH_VERSIONS_READABLE = (1, 2, 3, 4) # accepted on read — see
|
||||
# _read_sfft_block. Each bump only
|
||||
# appended a field, and every older
|
||||
# tail has a well-defined reading:
|
||||
# v1 predates row-averaged FFT
|
||||
# caching (row_avg_n=0), v1/v2
|
||||
# predate padded caching, so both
|
||||
# are natural-resolution (pad 1),
|
||||
# and v1-v3 predate the min peak
|
||||
# frequency floor (min_freq 0 =
|
||||
# no floor).
|
||||
CACH_FLAG_DC = 0x01
|
||||
CACH_FLAG_FFT = 0x02
|
||||
|
||||
@@ -69,9 +80,28 @@ SDCB_HDR_FMT = ">4sBH" # magic, reserved, n_stored
|
||||
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT)
|
||||
|
||||
SFFT_MAGIC = b"SFFT"
|
||||
SFFT_HDR_FMT = ">4sBH" # magic, flags, n_stored
|
||||
SFFT_HDR_FMT_V1 = ">4sBH" # magic, flags, n_stored (cach_version 1)
|
||||
SFFT_HDR_FMT_V2 = ">4sBHB" # + row_avg_n (cach_version 2)
|
||||
SFFT_HDR_FMT_V3 = ">4sBHBH" # + pad_factor (cach_version 3)
|
||||
SFFT_HDR_FMT = ">4sBHBHI" # + min_freq_khz (cach_version 4)
|
||||
SFFT_HDR_SIZE_V1 = struct.calcsize(SFFT_HDR_FMT_V1)
|
||||
SFFT_HDR_SIZE_V2 = struct.calcsize(SFFT_HDR_FMT_V2)
|
||||
SFFT_HDR_SIZE_V3 = struct.calcsize(SFFT_HDR_FMT_V3)
|
||||
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
|
||||
MAX_PAD_FACTOR = 0xFFFF # the H field above
|
||||
# min_freq_khz is fixed-point (u32, units of 0.001 MHz), not a float32:
|
||||
# the viewer's floor spinbox has 0.001 MHz granularity, and
|
||||
# round(mhz * 1000) / 1000.0 reproduces the exact float64 the user typed,
|
||||
# so the accept rule in sras_compute.cache_mismatch_reasons can compare
|
||||
# with plain integer ordering. A ">f" float32 of e.g. 20.1 would read back
|
||||
# as 20.10000038… > 20.1 and report a spurious mismatch forever.
|
||||
MAX_MIN_FREQ_KHZ = 0xFFFFFFFF # the I field above; 0 = no floor
|
||||
SFFT_FLAG_BG_SUB = 0x01
|
||||
SFFT_FLAG_ROW_AVG = 0x02 # peak_freq_mhz came from same-row,
|
||||
# distance-weighted averaged CH1
|
||||
# waveforms, not raw per-pixel ones;
|
||||
# row_avg_n is the neighbor half-width
|
||||
# (pixels) used. Bits 2-7 reserved.
|
||||
|
||||
# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
|
||||
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
|
||||
@@ -116,6 +146,13 @@ def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV,
|
||||
return (adc - yoff_adc) * ymult_mv + yzero_mv
|
||||
|
||||
|
||||
def _axes_extent(x_axis, y_axis, dx: float, dy: float) -> list[float]:
|
||||
"""Matplotlib imshow extent with half-pixel margins, Y flipped so row 0
|
||||
renders at the top."""
|
||||
return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
|
||||
y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Binary read helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -169,7 +206,11 @@ 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.
|
||||
source version. The scalars ``precomputed_bg_sub`` /
|
||||
``precomputed_row_avg_n`` / ``precomputed_pad_factor`` /
|
||||
``precomputed_min_freq_mhz`` record the settings the stored FFT images
|
||||
were computed under, so a reader can tell whether they answer the
|
||||
question it is actually asking.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str):
|
||||
@@ -226,6 +267,61 @@ 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.
|
||||
self.precomputed_pad_factor: int = 1
|
||||
# Min peak frequency floor (MHz) the stored peak search excluded
|
||||
# bins below; 0.0 = no floor. Unlike bg-sub/pad/row-avg it is
|
||||
# tighten-only re-applicable: a *higher* floor can be re-imposed on
|
||||
# a stored image by masking pixels below it, but bins below the
|
||||
# stored floor were never searched, so a lower floor needs a real
|
||||
# recompute — see sras_compute.cache_mismatch_reasons.
|
||||
self.precomputed_min_freq_mhz: float = 0.0
|
||||
|
||||
def encoded_preambles(self) -> bytes:
|
||||
"""This file's Preamble Blocks section, as bytes a writer can emit.
|
||||
|
||||
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."""
|
||||
@@ -257,6 +353,15 @@ 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)
|
||||
|
||||
@@ -517,6 +622,45 @@ class SrasFile:
|
||||
store[angle_idx] = _read_f32_image(f, shape)
|
||||
return flags
|
||||
|
||||
def _read_sfft_block(self, f, cach_version: int) -> tuple[int, int, int, float] | None:
|
||||
"""Read the SFFT block header — its layout depends on cach_version,
|
||||
since each bump appended a trailing field (v2 row_avg_n, v3
|
||||
pad_factor, v4 min_freq_khz) — then n_stored per-angle
|
||||
peak_freq_mhz entries (unchanged across versions).
|
||||
|
||||
Returns (flags, row_avg_n, pad_factor, min_freq_mhz), or None if
|
||||
the block is malformed. The absent fields of an older tail take the
|
||||
value that describes what such a tail can only have been:
|
||||
row_avg_n=0 for v1, which predates row-averaged FFT caching;
|
||||
pad_factor=1 for v1/v2, which predate padded caching and so are
|
||||
natural-resolution; and min_freq_mhz=0.0 for v1-v3, which predate
|
||||
the min peak frequency floor and so searched every bin above DC.
|
||||
"""
|
||||
hdr_fmt = {1: SFFT_HDR_FMT_V1, 2: SFFT_HDR_FMT_V2,
|
||||
3: SFFT_HDR_FMT_V3}.get(cach_version, SFFT_HDR_FMT)
|
||||
raw = f.read(struct.calcsize(hdr_fmt))
|
||||
if len(raw) < struct.calcsize(hdr_fmt):
|
||||
return None
|
||||
row_avg_n, pad_factor, min_freq_khz = 0, 1, 0
|
||||
if cach_version == 1:
|
||||
magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
|
||||
elif cach_version == 2:
|
||||
magic, flags, n_stored, row_avg_n = struct.unpack(hdr_fmt, raw)
|
||||
elif cach_version == 3:
|
||||
magic, flags, n_stored, row_avg_n, pad_factor = struct.unpack(hdr_fmt, raw)
|
||||
else:
|
||||
(magic, flags, n_stored, row_avg_n, pad_factor,
|
||||
min_freq_khz) = struct.unpack(hdr_fmt, raw)
|
||||
if magic != SFFT_MAGIC:
|
||||
return None
|
||||
for _ in range(n_stored):
|
||||
(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, min_freq_khz / 1000.0
|
||||
|
||||
def _parse_cach_section(self, offset: int):
|
||||
"""Parse the v7 CACH tail that holds precomputed DC/FFT images."""
|
||||
with open(self.path, "rb") as f:
|
||||
@@ -525,7 +669,7 @@ class SrasFile:
|
||||
if len(header_raw) < CACH_HDR_SIZE:
|
||||
return
|
||||
magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
|
||||
if magic != CACH_MAGIC or cach_version != CACH_VERSION:
|
||||
if magic != CACH_MAGIC or cach_version not in CACH_VERSIONS_READABLE:
|
||||
return
|
||||
|
||||
if block_flags & CACH_FLAG_DC:
|
||||
@@ -535,17 +679,25 @@ class SrasFile:
|
||||
return
|
||||
|
||||
if block_flags & CACH_FLAG_FFT:
|
||||
flags = self._read_cache_block(
|
||||
f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz])
|
||||
if flags is None:
|
||||
result = self._read_sfft_block(f, cach_version)
|
||||
if result is None:
|
||||
return
|
||||
flags, row_avg_n, pad_factor, min_freq_mhz = result
|
||||
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
|
||||
self.precomputed_row_avg_n = row_avg_n if (flags & SFFT_FLAG_ROW_AVG) else 0
|
||||
self.precomputed_pad_factor = max(1, pad_factor)
|
||||
# No flag bit gates the floor: 0 (= no floor) is already the
|
||||
# value every pre-v4 tail reads as.
|
||||
self.precomputed_min_freq_mhz = min_freq_mhz
|
||||
|
||||
def write_v7_cache(self, *,
|
||||
new_dc3_mv: list[np.ndarray | None] | None = None,
|
||||
new_dc4_mv: list[np.ndarray | None] | None = None,
|
||||
new_freq_mhz: list[np.ndarray | None] | None = None,
|
||||
new_bg_sub: bool | None = None):
|
||||
new_bg_sub: bool | None = None,
|
||||
new_row_avg_n: int | None = None,
|
||||
new_pad_factor: int | None = None,
|
||||
new_min_freq_mhz: float | None = None):
|
||||
"""Store computed DC and/or FFT images into this file's CACH tail,
|
||||
in place, converting a v6 source to v7 (or updating an existing v7
|
||||
file). Only the block(s) passed in are recomputed; whichever block
|
||||
@@ -553,6 +705,25 @@ 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.
|
||||
|
||||
*new_min_freq_mhz* is the min peak frequency floor the passed
|
||||
*new_freq_mhz*'s peak search excluded bins below (0.0 = no floor),
|
||||
carried forward the same way. It is stored fixed-point (whole kHz),
|
||||
so the value is quantized to 0.001 MHz on write and
|
||||
``precomputed_min_freq_mhz`` is updated to the quantized value —
|
||||
what a reload would see, never a float the header can't represent.
|
||||
|
||||
The waveform data itself is never touched: the cache tail always
|
||||
starts at ``_cache_tail_offset()``, a fixed offset derived from the
|
||||
header and geometry table alone.
|
||||
@@ -565,8 +736,33 @@ 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)
|
||||
final_min_freq_mhz = (new_min_freq_mhz if new_min_freq_mhz is not None
|
||||
else self.precomputed_min_freq_mhz)
|
||||
if not (0 <= final_row_avg_n <= 255):
|
||||
raise ValueError(f"row_avg_n must fit in a byte (0-255), got {final_row_avg_n}")
|
||||
if not (1 <= final_pad_factor <= MAX_PAD_FACTOR):
|
||||
raise ValueError(
|
||||
f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {final_pad_factor}")
|
||||
if not (np.isfinite(final_min_freq_mhz) and final_min_freq_mhz >= 0):
|
||||
raise ValueError(
|
||||
f"min_freq_mhz must be a finite value >= 0, got {final_min_freq_mhz}")
|
||||
final_min_freq_khz = int(round(final_min_freq_mhz * 1000))
|
||||
if final_min_freq_khz > MAX_MIN_FREQ_KHZ:
|
||||
raise ValueError(
|
||||
f"min_freq_mhz too large for the u32 kHz header field: "
|
||||
f"{final_min_freq_mhz}")
|
||||
|
||||
dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
|
||||
# dc3/dc4 are always populated together by every current caller, but
|
||||
# guard the per-angle pairing explicitly rather than assume it: an
|
||||
# angle present in only one of the two arrays would otherwise crash
|
||||
# below on final_dc4[a].astype(...) (or silently store the wrong
|
||||
# dc3/dc4 pairing for that angle).
|
||||
dc_entries = [a for a in range(self.n_angles)
|
||||
if final_dc3[a] is not None and final_dc4[a] is not None]
|
||||
fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]
|
||||
|
||||
block_flags = ((CACH_FLAG_DC if dc_entries else 0)
|
||||
@@ -584,7 +780,10 @@ class SrasFile:
|
||||
|
||||
if fft_entries:
|
||||
fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
|
||||
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries))
|
||||
fft_flags |= SFFT_FLAG_ROW_AVG if final_row_avg_n else 0
|
||||
payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags,
|
||||
len(fft_entries), final_row_avg_n,
|
||||
final_pad_factor, final_min_freq_khz)
|
||||
for a in fft_entries:
|
||||
payload += struct.pack(">H", a)
|
||||
payload += final_freq[a].astype(">f4").tobytes()
|
||||
@@ -595,13 +794,23 @@ class SrasFile:
|
||||
f.truncate()
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
# Version-byte flip last: if the process dies before this point,
|
||||
# the file is still readable as plain v6 (v6 parsing only
|
||||
# bounds-checks per-angle offset+nbytes <= file_size, it never
|
||||
# asserts exactly how many bytes follow the last angle) — so an
|
||||
# interrupted write can never corrupt the file, only leave
|
||||
# harmless trailing bytes that the next successful write
|
||||
# overwrites via this same deterministic cache offset.
|
||||
# Version-byte flip last: when converting a v6 source (self.version
|
||||
# was 6 on entry), if the process dies before this point the file
|
||||
# is still readable as plain v6 (v6 parsing only bounds-checks
|
||||
# per-angle offset+nbytes <= file_size, it never asserts exactly
|
||||
# how many bytes follow the last angle) — so an interrupted write
|
||||
# can never corrupt the file, only leave harmless trailing bytes
|
||||
# that the next successful write overwrites via this same
|
||||
# deterministic cache offset.
|
||||
#
|
||||
# That guarantee does NOT extend to updating an already-v7 file:
|
||||
# the version byte here is already 7 before this call, so a crash
|
||||
# during the payload write above (before flush/fsync/truncate)
|
||||
# can leave a cache tail that mixes a prefix of the new payload
|
||||
# with a stale suffix of the old one, and this method has no
|
||||
# protection against that case (no atomic rename — the tail is
|
||||
# rewritten in place to avoid copying the, potentially huge,
|
||||
# waveform data that precedes it).
|
||||
f.seek(4)
|
||||
f.write(struct.pack("B", 7))
|
||||
f.flush()
|
||||
@@ -612,6 +821,9 @@ 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_min_freq_mhz = final_min_freq_khz / 1000.0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Axes helpers
|
||||
@@ -628,6 +840,24 @@ 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
|
||||
|
||||
|
||||
+204
@@ -0,0 +1,204 @@
|
||||
"""Pure-matplotlib rendering of a displayed SRAS image: imshow + colorbar +
|
||||
axis/title labeling, shared by the interactive Qt canvas
|
||||
(sras_viewer.canvases.ImageCanvas, which supplies its own already-Qt-backed
|
||||
Figure/Axes) and the headless batch image-export worker (which builds a
|
||||
throwaway Agg Figure per file and never touches Qt) -- so an exported PNG
|
||||
can never quietly start looking different from what the GUI actually shows.
|
||||
|
||||
Deliberately no PyQt6 import anywhere in this module: BatchExportImagesWorker
|
||||
(sras_workers.py) may run export_view_image inside a spawned
|
||||
ProcessPoolExecutor subprocess, exactly like sras_compute.cache_file, and
|
||||
importing anything under the sras_viewer package would run its __init__.py
|
||||
and pull in the whole Qt widget tree for no reason.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib as mpl
|
||||
import numpy as np
|
||||
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
||||
from matplotlib.figure import Figure
|
||||
|
||||
from sras_compute import compute_rf_image, dc_image_mv
|
||||
from sras_format import CH4_IDX, CH_NAMES, SrasFile, _axes_extent
|
||||
|
||||
# Figure size (inches) to fall back on when a caller has no live view to
|
||||
# match: the size ImageCanvas is constructed at, so a headless render with
|
||||
# no canvas behind it still looks like the viewer's starting layout.
|
||||
DEFAULT_FIGSIZE = (7.0, 5.0)
|
||||
|
||||
# The canvas size travels from the GUI as-is, so it can be degenerate (a
|
||||
# collapsed splitter pane, a window minimized mid-batch) or, on a very wide
|
||||
# multi-monitor window, big enough that width * dpi approaches matplotlib's
|
||||
# 2**16 pixel limit. Clamp rather than fail: an export's aspect ratio is a
|
||||
# presentation detail and must never be the reason a batch loses an image.
|
||||
_MIN_FIG_IN = 1.0
|
||||
_MAX_FIG_IN = 100.0
|
||||
|
||||
|
||||
def sanitize_figsize(figsize) -> tuple[float, float]:
|
||||
"""(width, height) in inches, clamped to something renderable.
|
||||
|
||||
*figsize* is None (use DEFAULT_FIGSIZE) or any 2-sequence of numbers --
|
||||
including the float64 pair Figure.get_size_inches returns, which is how
|
||||
the GUI hands over the live canvas's current size.
|
||||
"""
|
||||
try:
|
||||
w, h = float(figsize[0]), float(figsize[1])
|
||||
except (TypeError, ValueError, IndexError, KeyError):
|
||||
return DEFAULT_FIGSIZE
|
||||
if not (np.isfinite(w) and np.isfinite(h)):
|
||||
return DEFAULT_FIGSIZE
|
||||
return (min(max(w, _MIN_FIG_IN), _MAX_FIG_IN),
|
||||
min(max(h, _MIN_FIG_IN), _MAX_FIG_IN))
|
||||
|
||||
|
||||
def draw_view_image(ax, fig, img: np.ndarray, extent: list[float], cmap,
|
||||
vmin: float, vmax: float, xlabel: str, ylabel: str,
|
||||
title: str, colorbar_label: str = "", cb_ticks=None,
|
||||
norm=None, bad_color=None):
|
||||
"""imshow + colorbar + labels onto an already-created (ax, fig) pair.
|
||||
|
||||
*cmap* may be a name or a Colormap instance. *norm* (which overrides
|
||||
vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image
|
||||
with whole-number colorbar bands instead of a continuous shade.
|
||||
*bad_color*, if given, is the fill for NaN pixels -- a copy of *cmap* is
|
||||
made so a shared, registered instance is never mutated.
|
||||
|
||||
Shared by ImageCanvas.show_image (Qt-backed ax/fig) and
|
||||
export_view_image (headless Agg ax/fig) so the two can never drift into
|
||||
showing different things for the same settings.
|
||||
"""
|
||||
if bad_color is not None:
|
||||
cmap = (cmap if hasattr(cmap, "with_extremes")
|
||||
else mpl.colormaps[cmap]).with_extremes(bad=bad_color)
|
||||
|
||||
kw = ({"norm": norm} if norm is not None
|
||||
else {"vmin": vmin, "vmax": vmax})
|
||||
im = ax.imshow(
|
||||
img, aspect="auto", origin="upper",
|
||||
extent=extent, cmap=cmap, interpolation="nearest", **kw,
|
||||
)
|
||||
cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, ticks=cb_ticks)
|
||||
if colorbar_label:
|
||||
cb.set_label(colorbar_label)
|
||||
|
||||
ax.set_xlabel(xlabel)
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_title(title)
|
||||
return im
|
||||
|
||||
|
||||
def export_view_image(path: str, *, out_dir: str, angle_idx: int, ch_idx: int,
|
||||
is_fft_mode: bool, is_velocity: bool,
|
||||
dc_threshold_mv: float, apply_bg_sub: bool,
|
||||
pad_factor: int, min_freq_mhz: float, grating_um: float,
|
||||
cmap: str, auto_scale: bool, vmin: float, vmax: float,
|
||||
highlight_masked: bool, mode_str: str,
|
||||
colorbar_label: str, mask_color: str = "magenta",
|
||||
max_workers: int | None = None,
|
||||
figsize: tuple[float, float] | None = None,
|
||||
dpi: int = 150) -> tuple[str, str]:
|
||||
"""One file's contribution to Batch Export View as Images: renders
|
||||
(angle_idx, ch_idx) at the given display settings to a PNG under
|
||||
*out_dir*, via draw_view_image -- so a batch export is a folder of what
|
||||
ImageCanvas.show_image would have put on screen for these settings, not
|
||||
a raw data dump.
|
||||
|
||||
Module-level and picklable, like sras_compute.cache_file, so it can run
|
||||
in a ProcessPoolExecutor -- see BatchExportImagesWorker. Unlike
|
||||
cache_file this never writes to *path*: export is a read of the file's
|
||||
own data, not a cache conversion, so any version SrasFile can open
|
||||
works, with no v6/v7 precondition.
|
||||
|
||||
*figsize* is the on-screen ImageCanvas's current size in inches, so the
|
||||
PNG carries the aspect ratio the view was actually being read at. It
|
||||
matters more here than it would for a plot with fixed data aspect:
|
||||
draw_view_image uses aspect="auto", so the image stretches to whatever
|
||||
box it is given -- rendering a view the user has sized wide into a
|
||||
hard-coded 7x5 squeezes the map into a different shape than the one
|
||||
they judged it by. Passing the full size, not just the ratio, also
|
||||
keeps titles, tick labels and the colorbar in the same proportion to
|
||||
the map as on screen; *dpi* alone then sets the output resolution.
|
||||
|
||||
*pad_factor* (not n_fft) travels across files deliberately: n_fft
|
||||
depends on samples_per_frame, which can differ between files in the
|
||||
same batch, so n_fft is derived per file, here, from *this* file's own
|
||||
value -- the same reason sras_compute.cache_file does the same thing.
|
||||
|
||||
row_avg_n is per file for the same reason and so is not a parameter at
|
||||
all: it comes from *this* file's stored cache, mirroring
|
||||
SrasViewerWindow._stored_fft_image and _start_compute, which both ask at
|
||||
the window size the store was written at. Leaving it at compute_rf_image's
|
||||
"raw per-pixel" default would make a row-averaged cache a mismatch, so a
|
||||
file the viewer displays from its store would export as a full raw
|
||||
recompute instead -- a different (and much slower) image than the one the
|
||||
batch was triggered to reproduce.
|
||||
|
||||
Returns (error, out_name). error is "" on success. out_name is the
|
||||
filename this call targeted -- set as soon as it's known, even on most
|
||||
failures -- so the caller can flag same-stem collisions across the
|
||||
batch without any cross-process bookkeeping.
|
||||
"""
|
||||
out_name = ""
|
||||
try:
|
||||
sras = SrasFile(path)
|
||||
if angle_idx >= sras.n_angles:
|
||||
return (f"angle {angle_idx} out of range "
|
||||
f"(file has {sras.n_angles} angle(s))", out_name)
|
||||
|
||||
out_name = f"{Path(path).stem}_angle{angle_idx}_{CH_NAMES[ch_idx]}.png"
|
||||
|
||||
if is_fft_mode:
|
||||
n_fft = (sras.samples_per_frame * pad_factor
|
||||
if pad_factor > 1 else None)
|
||||
freq = compute_rf_image(
|
||||
sras, angle_idx, dc_threshold_mv=dc_threshold_mv,
|
||||
apply_bg_sub=apply_bg_sub, n_fft=n_fft,
|
||||
row_avg_n=sras.precomputed_row_avg_n,
|
||||
min_freq_mhz=min_freq_mhz, max_workers=max_workers)
|
||||
img = freq * grating_um if is_velocity else freq
|
||||
else:
|
||||
img = dc_image_mv(sras, angle_idx, ch_idx, max_workers=max_workers)
|
||||
|
||||
display_img, bad_color = img, None
|
||||
if highlight_masked and is_fft_mode:
|
||||
# Mirrors the viewer's _redraw_image rule: DC-masked pixels and
|
||||
# value-0 pixels (the "no valid peak" sentinel — DC-masked,
|
||||
# below the min-freq floor, or empty spectrum; the grating
|
||||
# multiply above preserves zeros, so this holds for Velocity
|
||||
# too) both render in the highlight color.
|
||||
dc4 = dc_image_mv(sras, angle_idx, CH4_IDX, max_workers=max_workers)
|
||||
valid = dc4 >= dc_threshold_mv
|
||||
if valid.shape == display_img.shape:
|
||||
valid &= display_img != 0.0
|
||||
display_img = display_img.astype(np.float32, copy=True)
|
||||
display_img[~valid] = np.nan
|
||||
bad_color = mask_color
|
||||
|
||||
if auto_scale:
|
||||
v0, v1 = float(np.nanmin(display_img)), float(np.nanmax(display_img))
|
||||
if not np.isfinite(v0):
|
||||
v0, v1 = 0.0, 0.0 # every pixel masked out
|
||||
else:
|
||||
v0, v1 = vmin, vmax
|
||||
|
||||
x_axis = sras.x_axis_mm(angle_idx)
|
||||
y_axis = sras.y_positions_mm(angle_idx)
|
||||
dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else sras.pixel_x_mm
|
||||
dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
|
||||
extent = _axes_extent(x_axis, y_axis, dx, dy)
|
||||
|
||||
title = (f"{CH_NAMES[ch_idx]} | {mode_str} | "
|
||||
f"{sras.angles_deg[angle_idx]:.1f}°")
|
||||
|
||||
fig = Figure(figsize=sanitize_figsize(figsize), tight_layout=True)
|
||||
FigureCanvasAgg(fig) # Agg-only: never registered with pyplot
|
||||
ax = fig.add_subplot(111)
|
||||
draw_view_image(ax, fig, display_img, extent, cmap, v0, v1,
|
||||
"X (mm)", "Y (mm)", title, colorbar_label,
|
||||
bad_color=bad_color)
|
||||
fig.savefig(str(Path(out_dir) / out_name), dpi=dpi)
|
||||
return ("", out_name)
|
||||
except Exception as exc:
|
||||
return (str(exc), out_name)
|
||||
@@ -18,7 +18,12 @@ import faulthandler
|
||||
|
||||
faulthandler.enable() # print a native stack trace on SIGSEGV/SIGABRT/etc.
|
||||
|
||||
from .canvases import ImageCanvas, RoiQuad, WaveformCanvas # noqa: E402,F401
|
||||
from .align_wizard import AlignmentWizard # noqa: E402,F401
|
||||
from .canvases import ( # noqa: E402,F401
|
||||
AlignOverlayCanvas, ImageCanvas, RoiQuad, WaveformCanvas,
|
||||
)
|
||||
from .common import CH_LABELS, CMAPS, VELOCITY_MODE_IDX # noqa: E402,F401
|
||||
from .dialogs import FftOptionsDialog, ManualAlignmentDialog # noqa: E402,F401
|
||||
from .dialogs import ( # noqa: E402,F401
|
||||
FftOptionsDialog, FusedRoiExportDialog, RowAverageFftOptionsDialog,
|
||||
)
|
||||
from .main_window import SrasViewerWindow, main # noqa: E402,F401
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+128
-23
@@ -1,7 +1,9 @@
|
||||
"""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
|
||||
@@ -10,6 +12,28 @@ from PyQt6.QtGui import QKeyEvent
|
||||
from PyQt6.QtWidgets import QSizePolicy
|
||||
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv
|
||||
from sras_render import draw_view_image
|
||||
|
||||
def count_colormap(n_angles: int):
|
||||
"""(cmap, norm, ticks) for an integer "how many angles cover this pixel"
|
||||
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)
|
||||
@@ -100,7 +124,15 @@ class ImageCanvas(FigureCanvasQTAgg):
|
||||
_HANDLE_PX = 12
|
||||
_CLICK_THRESH_PX = 4 # releases within this of press count as a click
|
||||
|
||||
def __init__(self, parent=None):
|
||||
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.
|
||||
"""
|
||||
fig = Figure(figsize=(7, 5), tight_layout=True)
|
||||
self.ax = fig.add_subplot(111)
|
||||
super().__init__(fig)
|
||||
@@ -108,6 +140,7 @@ 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
|
||||
@@ -132,9 +165,16 @@ class ImageCanvas(FigureCanvasQTAgg):
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def show_image(self, img: np.ndarray, extent: list[float], cmap: str,
|
||||
def show_image(self, img: np.ndarray, extent: list[float], cmap,
|
||||
vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
|
||||
colorbar_label: str = ""):
|
||||
colorbar_label: str = "", cb_ticks=None, norm=None,
|
||||
bad_color=None):
|
||||
"""*cmap* may be a name or a Colormap instance. *norm* (which overrides
|
||||
vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image —
|
||||
the wizard's overlap-count view — with whole-number colorbar bands
|
||||
instead of a continuous shade. *bad_color*, if given, is the fill for
|
||||
NaN pixels — a copy of *cmap* is made so the shared, registered
|
||||
instance is never mutated."""
|
||||
self.figure.clf()
|
||||
self.ax = self.figure.add_subplot(111)
|
||||
# Patches and lines are destroyed by figure.clf(); drop stale refs.
|
||||
@@ -143,18 +183,12 @@ class ImageCanvas(FigureCanvasQTAgg):
|
||||
self._extent = extent
|
||||
self._img_shape = img.shape
|
||||
|
||||
im = self.ax.imshow(
|
||||
img, aspect="auto", origin="upper",
|
||||
extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
|
||||
interpolation="nearest",
|
||||
)
|
||||
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
|
||||
if colorbar_label:
|
||||
cb.set_label(colorbar_label)
|
||||
|
||||
self.ax.set_xlabel(xlabel)
|
||||
self.ax.set_ylabel(ylabel)
|
||||
self.ax.set_title(title)
|
||||
# Shared with the headless batch image-export worker (sras_render.py)
|
||||
# so an exported PNG can never quietly drift from what this canvas
|
||||
# shows on screen for the same settings.
|
||||
draw_view_image(self.ax, self.figure, img, extent, cmap, vmin, vmax,
|
||||
xlabel, ylabel, title, colorbar_label, cb_ticks, norm,
|
||||
bad_color)
|
||||
|
||||
# Re-draw the ROI (if any) on top of the fresh image so it persists
|
||||
# unchanged across mode / angle / channel switches.
|
||||
@@ -302,10 +336,28 @@ 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
|
||||
@@ -377,13 +429,22 @@ class WaveformCanvas(FigureCanvasQTAgg):
|
||||
|
||||
def show_rf_waveform(self, sras: SrasFile, angle_idx: int,
|
||||
row_idx: int, frame_idx: int,
|
||||
apply_bg_sub: bool = True):
|
||||
apply_bg_sub: bool = True,
|
||||
min_freq_mhz: float = 0.0):
|
||||
"""CH1 RF: time-domain + FFT spectrum.
|
||||
|
||||
If apply_bg_sub is True and sras.background is not None, the background
|
||||
waveform is overlaid on the time-domain plot and the FFT is computed
|
||||
on the subtracted signal. The unsubtracted FFT is also shown faintly
|
||||
for comparison.
|
||||
|
||||
*min_freq_mhz* > 0 restricts the labeled peak to bins at or above
|
||||
it — the same floor the image's peak search uses, so the label
|
||||
explains the map pixel instead of contradicting it — and shades the
|
||||
excluded band on the spectrum. The spectrum curves themselves stay
|
||||
complete (they are the evidence for choosing the floor). The peak
|
||||
can still legitimately differ from a padded or row-averaged map:
|
||||
this panel is always a single waveform at natural resolution.
|
||||
"""
|
||||
data = sras.data[angle_idx]
|
||||
waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32)
|
||||
@@ -424,7 +485,20 @@ class WaveformCanvas(FigureCanvasQTAgg):
|
||||
# FFT of the (possibly subtracted) waveform
|
||||
power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2
|
||||
power_sub[0] = 0.0
|
||||
peak_mhz = f_mhz[int(np.argmax(power_sub))]
|
||||
# First bin at or above the floor, exactly as the image peak search
|
||||
# picks it (bin 0 always excluded). If the floor excludes every bin,
|
||||
# fall back to the unrestricted peak rather than indexing past the
|
||||
# end — the label is informational, not a mask.
|
||||
lo = max(1, int(np.searchsorted(f_mhz, min_freq_mhz)))
|
||||
if lo < len(power_sub):
|
||||
peak_mhz = f_mhz[lo + int(np.argmax(power_sub[lo:]))]
|
||||
else:
|
||||
peak_mhz = f_mhz[int(np.argmax(power_sub))]
|
||||
|
||||
if min_freq_mhz > 0.0:
|
||||
self.ax_right.axvspan(0, min_freq_mhz, color="#888888",
|
||||
alpha=0.15, zorder=0,
|
||||
label=f"< {min_freq_mhz:g} MHz excluded")
|
||||
|
||||
if bg is not None:
|
||||
# Also show the unsubtracted FFT for reference
|
||||
@@ -475,14 +549,20 @@ class WaveformCanvas(FigureCanvasQTAgg):
|
||||
self.draw()
|
||||
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
A pure input+render widget — it holds no alignment state and never
|
||||
touches SrasFile itself; ManualAlignmentDialog owns all of that and
|
||||
decides, from these signals, whether a cheap single-layer refresh or a
|
||||
touches SrasFile itself; the wizard page 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
|
||||
@@ -522,6 +602,31 @@ class ManualAlignOverlayCanvas(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)
|
||||
|
||||
+36
-11
@@ -2,11 +2,11 @@
|
||||
|
||||
from PyQt6.QtCore import Qt
|
||||
from PyQt6.QtWidgets import (
|
||||
QDoubleSpinBox, QFormLayout, QFrame, QGroupBox, QLabel, QScrollArea,
|
||||
QSizePolicy, QVBoxLayout, QWidget,
|
||||
QComboBox, QDoubleSpinBox, QFormLayout, QFrame, QGroupBox, QLabel,
|
||||
QScrollArea, QSizePolicy, QVBoxLayout, QWidget,
|
||||
)
|
||||
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, _axes_extent # noqa: F401 (re-exported)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Display constants
|
||||
@@ -26,6 +26,10 @@ CH1_DERIVED_MODES = (CH1_IDX, VELOCITY_MODE_IDX)
|
||||
|
||||
CMAPS = ["gray", "viridis", "plasma", "inferno", "hot", "jet", "RdBu_r", "seismic"]
|
||||
|
||||
# Fill color for masked (below-DC-threshold) pixels when "Highlight masked
|
||||
# pixels" is on, chosen to stand out against every colormap in CMAPS above.
|
||||
_MASKED_HIGHLIGHT_COLOR = "magenta"
|
||||
|
||||
# (mode_str, status-bar unit, colorbar label) per channel index
|
||||
_CHANNEL_DISPLAY = {
|
||||
CH1_IDX: ("RF", "Peak frequency (MHz)", "MHz"),
|
||||
@@ -60,9 +64,14 @@ class Jobs:
|
||||
COMPUTE = "compute"
|
||||
DC_PRECOMPUTE = "dc_precompute"
|
||||
BATCH = "batch"
|
||||
ALIGN = "align"
|
||||
MANUAL_ALIGN_MASKS = "manual_align_masks"
|
||||
MANUAL_ALIGN_CORRELATE = "manual_align_correlate"
|
||||
EXPORT = "export"
|
||||
# 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"
|
||||
|
||||
|
||||
def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
|
||||
@@ -81,11 +90,27 @@ def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
|
||||
return spin
|
||||
|
||||
|
||||
def _axes_extent(x_axis, y_axis, dx: float, dy: float) -> list[float]:
|
||||
"""Matplotlib imshow extent with half-pixel margins, Y flipped so row 0
|
||||
renders at the top."""
|
||||
return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
|
||||
y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
|
||||
def _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 _wrap_label(text: str = "", css: str | None = None) -> QLabel:
|
||||
|
||||
+379
-594
File diff suppressed because it is too large
Load Diff
+843
-174
File diff suppressed because it is too large
Load Diff
+345
-42
@@ -10,15 +10,19 @@ everything they need through their constructor and hand results back by signal.
|
||||
import os
|
||||
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
|
||||
from concurrent.futures.process import BrokenProcessPool
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from matplotlib.backends.backend_agg import FigureCanvasAgg
|
||||
from matplotlib.figure import Figure
|
||||
from PyQt6.QtCore import QObject, pyqtSignal
|
||||
|
||||
import sras_compute as compute
|
||||
from sras_compute import (
|
||||
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
|
||||
)
|
||||
from sras_format import CH3_IDX, CH4_IDX, SrasFile
|
||||
from sras_align_export import write_aligned_sras
|
||||
from sras_compute import cache_file, compute_rf_image, dc_image_mv
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
|
||||
from sras_render import export_view_image, sanitize_figsize
|
||||
|
||||
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||||
# those processes threads internally, so the two must be divided rather than
|
||||
@@ -116,7 +120,9 @@ class ComputeWorker(CancellableWorker):
|
||||
apply_bg_sub: bool = True, n_fft: int | None = None,
|
||||
dc_threshold_mv: float = 0.0,
|
||||
dc4_mv: np.ndarray | None = None,
|
||||
is_fft_mode: bool = False):
|
||||
is_fft_mode: bool = False,
|
||||
row_avg_n: int = 0,
|
||||
min_freq_mhz: float = 0.0):
|
||||
super().__init__()
|
||||
self._sras = sras
|
||||
self._angle = angle_idx
|
||||
@@ -126,6 +132,8 @@ class ComputeWorker(CancellableWorker):
|
||||
self._dc_threshold = dc_threshold_mv
|
||||
self._dc4_mv = dc4_mv
|
||||
self._is_fft_mode = is_fft_mode
|
||||
self._row_avg_n = row_avg_n
|
||||
self._min_freq_mhz = min_freq_mhz
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
@@ -133,7 +141,9 @@ class ComputeWorker(CancellableWorker):
|
||||
img = compute_rf_image(
|
||||
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||||
dc4_mv=self._dc4_mv, should_stop=self._stopped)
|
||||
dc4_mv=self._dc4_mv, should_stop=self._stopped,
|
||||
row_avg_n=self._row_avg_n,
|
||||
min_freq_mhz=self._min_freq_mhz)
|
||||
else:
|
||||
img = dc_image_mv(self._sras, self._angle, self._ch,
|
||||
should_stop=self._stopped)
|
||||
@@ -187,9 +197,17 @@ class BatchCacheWorker(QObject):
|
||||
an existing v7 file's cache blocks without disturbing whatever the other
|
||||
block already holds.
|
||||
|
||||
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
|
||||
*mode* is ``"dc"`` (CH3/CH4 mean images), ``"fft"`` (CH1 peak-frequency
|
||||
images, unmasked — masking is applied at display time, same as v5's PREC
|
||||
convention).
|
||||
convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
|
||||
averaging before the FFT — needs *dc_threshold_mv* and a positive
|
||||
*row_avg_n*; see ``sras_compute.cache_file``).
|
||||
|
||||
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. *min_freq_mhz* travels the same
|
||||
way: the viewer's live min-peak-freq floor, recorded in the store as
|
||||
provenance so a reader knows which bins the peak search considered.
|
||||
|
||||
Files are processed one per subprocess: they are fully independent, each
|
||||
opens its own memmap and writes only its own bytes, and only path strings
|
||||
@@ -201,11 +219,17 @@ class BatchCacheWorker(QObject):
|
||||
file_done = pyqtSignal(str, str)
|
||||
finished = pyqtSignal()
|
||||
|
||||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
|
||||
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
|
||||
dc_threshold_mv: float | None = None, row_avg_n: int = 0,
|
||||
pad_factor: int = 1, min_freq_mhz: float = 0.0):
|
||||
super().__init__()
|
||||
self._paths = paths
|
||||
self._mode = mode
|
||||
self._apply_bg_sub = apply_bg_sub
|
||||
self._dc_threshold = dc_threshold_mv
|
||||
self._row_avg_n = row_avg_n
|
||||
self._pad_factor = pad_factor
|
||||
self._min_freq_mhz = min_freq_mhz
|
||||
|
||||
def _report(self, path: str, err: str, done: int, total: int):
|
||||
self.file_done.emit(path, err)
|
||||
@@ -230,7 +254,11 @@ class BatchCacheWorker(QObject):
|
||||
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||||
futures = {
|
||||
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
|
||||
compute.get_fft_backend(), per_proc_workers): p
|
||||
per_proc_workers,
|
||||
pad_factor=self._pad_factor,
|
||||
dc_threshold_mv=self._dc_threshold,
|
||||
row_avg_n=self._row_avg_n,
|
||||
min_freq_mhz=self._min_freq_mhz): p
|
||||
for p in paths
|
||||
}
|
||||
for fut in as_completed(futures):
|
||||
@@ -253,8 +281,11 @@ class BatchCacheWorker(QObject):
|
||||
for path in paths:
|
||||
try:
|
||||
err = cache_file(path, self._mode, self._apply_bg_sub,
|
||||
compute.get_fft_backend(),
|
||||
compute.default_max_workers())
|
||||
compute.default_max_workers(),
|
||||
pad_factor=self._pad_factor,
|
||||
dc_threshold_mv=self._dc_threshold,
|
||||
row_avg_n=self._row_avg_n,
|
||||
min_freq_mhz=self._min_freq_mhz)
|
||||
except Exception as exc:
|
||||
err = str(exc)
|
||||
done += 1
|
||||
@@ -292,34 +323,271 @@ 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)
|
||||
@dataclass
|
||||
class ExportChannel:
|
||||
"""One row of a batch image export: which raw channel to read, whether
|
||||
to apply the Velocity post-multiply, and the fixed display range/labels
|
||||
to render it with.
|
||||
|
||||
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||||
Kept free of sras_viewer's display constants (CH_LABELS, VELOCITY_MODE_IDX,
|
||||
etc.) so BatchExportWorker has no dependency on the GUI module — the
|
||||
caller resolves labels/units/filename tags once, up front.
|
||||
"""
|
||||
ch_idx: int # CH1_IDX, CH3_IDX, or CH4_IDX -- which raw data to read
|
||||
is_velocity: bool # True only for the derived Velocity map (post-multiply of CH1 freq)
|
||||
vmin: float
|
||||
vmax: float
|
||||
label: str # e.g. "CH3 -- Bias A (DC mean)"
|
||||
unit: str # colorbar units, e.g. "mV"
|
||||
tag: str # filename tag: "CH1", "CH3", "CH4", "VEL"
|
||||
|
||||
|
||||
def _render_map_png(img: np.ndarray, extent: list[float], *, cmap: str,
|
||||
vmin: float, vmax: float, title: str, colorbar_label: str,
|
||||
out_path: str, figsize: tuple[float, float] | None = None):
|
||||
"""Render one map image to *out_path* with a fixed vmin/vmax, using a
|
||||
headless Agg canvas so this never touches the GUI thread's interactive
|
||||
matplotlib backend. Layout mirrors ImageCanvas.show_image.
|
||||
|
||||
*figsize* is the live canvas's size in inches, so the PNG comes out at
|
||||
the shape the view was being read at instead of a fixed 7x5 -- with
|
||||
aspect="auto" below, the figure box is what sets the map's proportions.
|
||||
None falls back to sras_render.DEFAULT_FIGSIZE."""
|
||||
fig = Figure(figsize=sanitize_figsize(figsize), tight_layout=True)
|
||||
FigureCanvasAgg(fig)
|
||||
ax = fig.add_subplot(111)
|
||||
im = ax.imshow(img, aspect="auto", origin="upper", extent=extent,
|
||||
cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
|
||||
cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||||
if colorbar_label:
|
||||
cb.set_label(colorbar_label)
|
||||
ax.set_xlabel("X (mm)")
|
||||
ax.set_ylabel("Y (mm)")
|
||||
ax.set_title(title)
|
||||
fig.savefig(out_path, dpi=150)
|
||||
|
||||
|
||||
class BatchExportWorker(CancellableWorker):
|
||||
"""Renders and saves one PNG per (angle, selected channel) for an
|
||||
already-open SrasFile, with each channel's vmin/vmax held fixed across
|
||||
every angle so the colorbar is directly comparable image to image.
|
||||
|
||||
Takes the SrasFile object directly (like ComputeWorker/DcPrecomputeWorker)
|
||||
rather than a path -- this runs against the file already open in the GUI,
|
||||
not an arbitrary batch of files, so there is no need to reopen it in a
|
||||
subprocess the way BatchCacheWorker does.
|
||||
|
||||
Every setting that decides what an FFT-derived pixel *is* has to arrive
|
||||
here explicitly, because compute_rf_image defaults each one to "off" and
|
||||
an omitted argument is therefore not a no-op -- it silently exports a
|
||||
different image than the one on screen. *min_freq_mhz* is the one with
|
||||
teeth: dropping it does not just skip a mask, it also makes this file's
|
||||
stored cache (written by Batch Compute FFT at some earlier floor) look
|
||||
like a match, so the export hands back the pre-floor peaks the user
|
||||
raised the floor to get rid of. It mirrors
|
||||
SrasViewerWindow._stored_fft_image, which takes the floor live for the
|
||||
same reason -- re-masking a stored image against a *higher* floor is
|
||||
free, so the displayed map always honors the spin box.
|
||||
|
||||
Emits progress(int) (0-100 over all angle x channel pairs), file_done(str,
|
||||
str) (output path, error message or ""), and finished() -- same shape as
|
||||
BatchCacheWorker.
|
||||
"""
|
||||
progress = pyqtSignal(int)
|
||||
file_done = pyqtSignal(str, str)
|
||||
finished = pyqtSignal()
|
||||
|
||||
def __init__(self, sras: SrasFile, channels: list[ExportChannel],
|
||||
output_dir: str, prefix: str, *, cmap: str,
|
||||
apply_bg_sub: bool, dc_threshold_mv: float,
|
||||
n_fft: int | None, grating_um: float,
|
||||
min_freq_mhz: float = 0.0,
|
||||
figsize: tuple[float, float] | None = None):
|
||||
super().__init__()
|
||||
self._sras = sras
|
||||
self._ref = ref_angle_idx
|
||||
self._threshold = dc_threshold_mv
|
||||
self._channels = channels
|
||||
self._output_dir = Path(output_dir)
|
||||
self._prefix = prefix
|
||||
self._cmap = cmap
|
||||
self._apply_bg_sub = apply_bg_sub
|
||||
self._dc_threshold_mv = dc_threshold_mv
|
||||
self._n_fft = n_fft
|
||||
self._grating_um = grating_um
|
||||
self._min_freq_mhz = min_freq_mhz
|
||||
self._figsize = figsize
|
||||
|
||||
def _angle_extent(self, angle_idx: int) -> list[float]:
|
||||
s = self._sras
|
||||
x_axis = s.x_axis_mm(angle_idx)
|
||||
y_axis = s.y_positions_mm(angle_idx)
|
||||
dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else s.pixel_x_mm
|
||||
dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
|
||||
return [x_axis[0] - dx / 2, x_axis[-1] + dx / 2,
|
||||
y_axis[-1] + dy / 2, y_axis[0] - dy / 2]
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
result = compute_angle_alignment(
|
||||
self._sras, self._ref, self._threshold,
|
||||
progress_cb=self.progress.emit)
|
||||
self.finished.emit(result, "")
|
||||
s = self._sras
|
||||
n_angles = s.n_angles
|
||||
total = n_angles * len(self._channels)
|
||||
done = 0
|
||||
needs_freq = any(c.ch_idx == CH1_IDX for c in self._channels)
|
||||
|
||||
for angle_idx in range(n_angles):
|
||||
if self._stop:
|
||||
break
|
||||
extent = self._angle_extent(angle_idx)
|
||||
angle_deg = s.angles_deg[angle_idx]
|
||||
|
||||
freq_mhz = None
|
||||
if needs_freq:
|
||||
freq_mhz = compute_rf_image(
|
||||
s, angle_idx, dc_threshold_mv=self._dc_threshold_mv,
|
||||
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||||
should_stop=self._stopped, row_avg_n=s.precomputed_row_avg_n,
|
||||
min_freq_mhz=self._min_freq_mhz)
|
||||
|
||||
for channel in self._channels:
|
||||
if self._stop:
|
||||
break
|
||||
out_path = (self._output_dir
|
||||
/ f"{self._prefix}_angle{angle_idx:02d}_{channel.tag}.png")
|
||||
try:
|
||||
if channel.ch_idx == CH1_IDX:
|
||||
img = (freq_mhz * self._grating_um if channel.is_velocity
|
||||
else freq_mhz)
|
||||
else:
|
||||
img = dc_image_mv(s, angle_idx, channel.ch_idx,
|
||||
should_stop=self._stopped)
|
||||
title = f"{channel.tag} | {angle_deg:.1f}°"
|
||||
colorbar_label = (f"{channel.label} ({channel.unit})"
|
||||
if channel.unit else channel.label)
|
||||
_render_map_png(
|
||||
img, extent, cmap=self._cmap,
|
||||
vmin=channel.vmin, vmax=channel.vmax,
|
||||
title=title, colorbar_label=colorbar_label,
|
||||
out_path=str(out_path), figsize=self._figsize)
|
||||
self.file_done.emit(str(out_path), "")
|
||||
except Exception as exc:
|
||||
self.file_done.emit(str(out_path), str(exc))
|
||||
done += 1
|
||||
self.progress.emit(int(done / max(1, total) * 100))
|
||||
|
||||
self.finished.emit()
|
||||
except Exception as exc:
|
||||
self.finished.emit(None, str(exc))
|
||||
self.file_done.emit("", str(exc))
|
||||
self.finished.emit()
|
||||
|
||||
|
||||
class BatchExportImagesWorker(QObject):
|
||||
"""Renders the view settings captured by the caller at trigger time
|
||||
(angle/channel/threshold/etc.) to one PNG per file in *paths*, via
|
||||
sras_render.export_view_image.
|
||||
|
||||
Same process-pool-with-inline-fallback strategy as BatchCacheWorker
|
||||
above (same _BATCH_MAX_PROCS / _BATCH_POOL_MIN_BYTES thresholds): an
|
||||
FFT-derived (CH1/Velocity) view is exactly the same expensive per-file
|
||||
compute Batch Compute FFT already parallelizes this way. Never mutates
|
||||
*paths* — each file is opened read-only — so unlike BatchCacheWorker
|
||||
there is no version gate.
|
||||
|
||||
Emits progress(int) (0-100 by files completed), file_done(str, str, str)
|
||||
(path, error message or "", output filename this file targeted — set
|
||||
even on most failures so the caller can flag same-stem collisions across
|
||||
the batch without any cross-process bookkeeping), and finished().
|
||||
"""
|
||||
progress = pyqtSignal(int)
|
||||
file_done = pyqtSignal(str, str, str)
|
||||
finished = pyqtSignal()
|
||||
|
||||
def __init__(self, paths: list[str], **render_kwargs):
|
||||
"""*render_kwargs* is exactly export_view_image's keyword-only
|
||||
settings (out_dir, angle_idx, ch_idx, is_fft_mode, is_velocity,
|
||||
dc_threshold_mv, apply_bg_sub, pad_factor, min_freq_mhz, grating_um,
|
||||
cmap, auto_scale, vmin, vmax, highlight_masked, mode_str,
|
||||
colorbar_label, mask_color) — bundled rather than repeated as
|
||||
positional params across __init__/_run_pooled/_run_inline."""
|
||||
super().__init__()
|
||||
self._paths = paths
|
||||
self._kw = render_kwargs
|
||||
|
||||
def _report(self, path: str, err: str, out_name: str, done: int, total: int):
|
||||
self.file_done.emit(path, err, out_name)
|
||||
self.progress.emit(int(done / max(1, total) * 100))
|
||||
|
||||
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
|
||||
"""Same contract as BatchCacheWorker._run_pooled: returns the paths
|
||||
that never got a real answer because the pool itself died, so the
|
||||
caller can retry them in-process."""
|
||||
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
|
||||
unresolved: list[str] = []
|
||||
done = 0
|
||||
|
||||
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||||
futures = {
|
||||
executor.submit(export_view_image, p,
|
||||
max_workers=per_proc_workers, **self._kw): p
|
||||
for p in paths
|
||||
}
|
||||
for fut in as_completed(futures):
|
||||
path = futures[fut]
|
||||
try:
|
||||
err, out_name = fut.result()
|
||||
except BrokenProcessPool:
|
||||
unresolved.append(path)
|
||||
continue
|
||||
except Exception as exc:
|
||||
err, out_name = str(exc), ""
|
||||
done += 1
|
||||
self._report(path, err, out_name, done, len(paths))
|
||||
|
||||
return unresolved
|
||||
|
||||
def _run_inline(self, paths: list[str], done: int, total: int):
|
||||
for path in paths:
|
||||
try:
|
||||
err, out_name = export_view_image(
|
||||
path, max_workers=compute.default_max_workers(), **self._kw)
|
||||
except Exception as exc:
|
||||
err, out_name = str(exc), ""
|
||||
done += 1
|
||||
self._report(path, err, out_name, done, total)
|
||||
|
||||
def _worth_pooling(self, paths: list[str]) -> bool:
|
||||
if len(paths) < 2:
|
||||
return False
|
||||
total = 0
|
||||
for p in paths:
|
||||
try:
|
||||
total += os.path.getsize(p)
|
||||
except OSError:
|
||||
pass # unreadable files are reported by export_view_image
|
||||
return total >= _BATCH_POOL_MIN_BYTES
|
||||
|
||||
def run(self):
|
||||
paths = self._paths
|
||||
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
|
||||
|
||||
if not self._worth_pooling(paths):
|
||||
self._run_inline(paths, 0, len(paths))
|
||||
self.finished.emit()
|
||||
return
|
||||
|
||||
try:
|
||||
unresolved = self._run_pooled(paths, n_procs)
|
||||
except Exception:
|
||||
# The pool could not be created or collapsed wholesale.
|
||||
unresolved = list(paths)
|
||||
|
||||
if unresolved:
|
||||
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
|
||||
|
||||
self.finished.emit()
|
||||
|
||||
|
||||
class Ch4MaskWorker(_PooledWorker):
|
||||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
|
||||
ManualAlignmentDialog's initial threshold-mask overlay.
|
||||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
|
||||
alignment wizard's initial threshold-mask stack.
|
||||
|
||||
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
|
||||
@@ -328,7 +596,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
|
||||
ManualAlignmentDialog._start_mask_prep).
|
||||
CorrelatePage._start_mask_prep).
|
||||
"""
|
||||
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||||
|
||||
@@ -355,8 +623,8 @@ class Ch4MaskWorker(_PooledWorker):
|
||||
|
||||
class CrossCorrelateWorker(_PooledWorker):
|
||||
"""Rigid registration (rotation + translation, never scale) of each of
|
||||
*angle_indices* against *ref_angle_idx*, for ManualAlignmentDialog's Auto
|
||||
Cross-Correlate button.
|
||||
*angle_indices* against *ref_angle_idx*, for the alignment wizard's
|
||||
Run/Re-run Correlation button.
|
||||
|
||||
Runs on a background thread — registering a real many-angle,
|
||||
high-resolution scan takes long enough that doing it on the GUI thread
|
||||
@@ -370,17 +638,18 @@ 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], *,
|
||||
sources: tuple[str, ...], dc_threshold_mv: float,
|
||||
search_deg: float):
|
||||
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."""
|
||||
super().__init__()
|
||||
self._sras = sras
|
||||
self._ref = ref_angle_idx
|
||||
self._angles = angle_indices
|
||||
self._dc4_mv = dc4_mv
|
||||
self._sources = sources
|
||||
self._threshold = dc_threshold_mv
|
||||
self._search_deg = search_deg
|
||||
self._reg_kwargs = dict(reg_kwargs or {})
|
||||
|
||||
def _plan(self) -> int:
|
||||
return compute.registration_workers(self._sras)
|
||||
@@ -390,11 +659,45 @@ 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,
|
||||
dc_threshold_mv=self._threshold, sources=self._sources,
|
||||
search_deg=self._search_deg)
|
||||
self._sras, a, self._ref, self._dc4_mv, **self._reg_kwargs)
|
||||
|
||||
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))
|
||||
|
||||
@@ -0,0 +1,646 @@
|
||||
"""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
|
||||
@@ -155,7 +155,7 @@ def test_all_angles_stack(rig):
|
||||
|
||||
|
||||
def test_downsampled_preview_lands_with_full_res(rig):
|
||||
# ManualAlignmentDialog reprojects block-mean-downsampled masks, so the
|
||||
# The wizard 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.
|
||||
|
||||
@@ -0,0 +1,719 @@
|
||||
"""Both batch image exports -- Convert -> Batch Export View as Images (one
|
||||
PNG per *file*, via sras_render.export_view_image) and Export -> Batch
|
||||
Export Images (one PNG per *angle x channel* of the open file, via
|
||||
sras_workers.BatchExportWorker). Does the exported PNG actually match what
|
||||
the live view would show, and does the batch dispatch (menu action ->
|
||||
worker -> per-file render) behave like Batch Compute's proven pattern?
|
||||
|
||||
The recurring hazard both halves guard is a dropped view setting: every
|
||||
parameter that decides what an FFT-derived pixel *is* (bg-sub, pad,
|
||||
row-averaging, the min peak frequency floor) defaults to "off" in
|
||||
compute_rf_image, so an argument the export forgets to forward does not
|
||||
degrade gracefully -- it silently renders a different image than the screen,
|
||||
and worse, makes this file's stored cache look like a match so the export
|
||||
hands back peaks from an earlier compute.
|
||||
|
||||
sras_render.export_view_image is tested directly (no Qt) for the plumbing
|
||||
that decides *what* gets rendered -- pad_factor derived per file, masked-
|
||||
pixel NaN fill, per-file auto-scale, velocity scaling -- via a
|
||||
draw_view_image spy rather than pixel-diffing PNGs, the same "spy on the
|
||||
seam, don't inspect the rendered artifact" approach test_highlight_masked_
|
||||
pixels (tests/test_gui.py) uses for the live canvas.
|
||||
|
||||
The GUI-dispatch half drives SrasViewerWindow._on_batch_export_images()
|
||||
end-to-end with patched file dialogs, the same shape as
|
||||
test_stored_cache.py's test_viewer_batch_row_average_dispatch.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PyQt6.QtCore import QEventLoop, QTimer
|
||||
from PyQt6.QtWidgets import QApplication
|
||||
|
||||
import sras_render
|
||||
import sras_workers
|
||||
from sras_compute import cache_file, compute_rf_image, dc_image_mv
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile
|
||||
from sras_render import DEFAULT_FIGSIZE, export_view_image, sanitize_figsize
|
||||
from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX
|
||||
from sras_workers import (
|
||||
BatchExportImagesWorker, BatchExportWorker, ExportChannel,
|
||||
)
|
||||
import tools.make_test_sras as gen
|
||||
|
||||
_THRESHOLD_MV = 50.0
|
||||
|
||||
|
||||
def pump(ms: int = 200):
|
||||
loop = QEventLoop()
|
||||
QTimer.singleShot(ms, loop.quit)
|
||||
loop.exec()
|
||||
|
||||
|
||||
def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
|
||||
waited = 0
|
||||
while waited < timeout_ms:
|
||||
if pred():
|
||||
return True
|
||||
pump(step)
|
||||
waited += step
|
||||
return pred()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sras_render.export_view_image -- pure function, no Qt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DEFAULT_KW = dict(
|
||||
is_fft_mode=False, is_velocity=False, dc_threshold_mv=_THRESHOLD_MV,
|
||||
apply_bg_sub=True, pad_factor=1, min_freq_mhz=0.0, grating_um=1.0,
|
||||
cmap="viridis", auto_scale=True, vmin=0.0, vmax=1.0,
|
||||
highlight_masked=False, mode_str="DC", colorbar_label="mV",
|
||||
)
|
||||
|
||||
|
||||
def _kw(**overrides):
|
||||
kw = dict(_DEFAULT_KW)
|
||||
kw.update(overrides)
|
||||
return kw
|
||||
|
||||
|
||||
def _png_size(path) -> tuple[int, int]:
|
||||
"""(width, height) in pixels, read straight out of the PNG's IHDR chunk
|
||||
(two big-endian uint32s at byte 16) -- no image library needed just to
|
||||
check the shape of an export."""
|
||||
raw = Path(path).read_bytes()
|
||||
assert raw[:8] == b"\x89PNG\r\n\x1a\n", "not a valid PNG"
|
||||
return (int.from_bytes(raw[16:20], "big"),
|
||||
int.from_bytes(raw[20:24], "big"))
|
||||
|
||||
|
||||
def _spy_draw(monkeypatch):
|
||||
"""Patches sras_render.draw_view_image to record the image array and
|
||||
vmin/vmax/bad_color it was called with, then delegates to the real
|
||||
implementation so the PNG is still written -- lets a test check *what*
|
||||
export_view_image computed without depending on rendered PNG pixels."""
|
||||
orig = sras_render.draw_view_image
|
||||
captured = {}
|
||||
|
||||
def spy(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, title,
|
||||
colorbar_label="", cb_ticks=None, norm=None, bad_color=None):
|
||||
captured["img"] = np.array(img, copy=True)
|
||||
captured["vmin"] = vmin
|
||||
captured["vmax"] = vmax
|
||||
captured["bad_color"] = bad_color
|
||||
return orig(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel,
|
||||
title, colorbar_label, cb_ticks, norm, bad_color)
|
||||
|
||||
monkeypatch.setattr(sras_render, "draw_view_image", spy)
|
||||
return captured
|
||||
|
||||
|
||||
def test_export_dc_channel_writes_png(tmp_path):
|
||||
path = tmp_path / "dc.sras"
|
||||
gen.write(path, n_angles=2, seed=1, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
err, out_name = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw())
|
||||
assert err == ""
|
||||
assert out_name == f"dc_angle0_{CH_NAMES[CH4_IDX]}.png"
|
||||
out_path = out_dir / out_name
|
||||
assert out_path.exists() and out_path.stat().st_size > 0
|
||||
assert out_path.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n", "not a valid PNG"
|
||||
|
||||
|
||||
def test_export_out_of_range_angle(tmp_path):
|
||||
path = tmp_path / "short.sras"
|
||||
gen.write(path, n_angles=2, seed=2, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
err, out_name = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=5, ch_idx=CH4_IDX, **_kw())
|
||||
assert err != "" and "2" in err, "error should mention the file's actual angle count"
|
||||
assert out_name == ""
|
||||
assert list(out_dir.iterdir()) == [], "no file written for a failed export"
|
||||
|
||||
|
||||
def test_export_fft_mode_matches_compute_rf_image(tmp_path, monkeypatch):
|
||||
path = tmp_path / "fft.sras"
|
||||
gen.write(path, n_angles=1, seed=3, samples_per_frame=128)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
captured = _spy_draw(monkeypatch)
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
|
||||
**_kw(is_fft_mode=True))
|
||||
assert err == ""
|
||||
|
||||
sras = SrasFile(str(path))
|
||||
expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
|
||||
apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0)
|
||||
assert np.array_equal(captured["img"], expected)
|
||||
|
||||
|
||||
def test_export_velocity_scales_frequency(tmp_path, monkeypatch):
|
||||
path = tmp_path / "vel.sras"
|
||||
gen.write(path, n_angles=1, seed=4, samples_per_frame=128)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
captured = _spy_draw(monkeypatch)
|
||||
grating_um = 3.5
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=VELOCITY_MODE_IDX,
|
||||
**_kw(is_fft_mode=True, is_velocity=True, grating_um=grating_um))
|
||||
assert err == ""
|
||||
|
||||
sras = SrasFile(str(path))
|
||||
freq = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
|
||||
apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0)
|
||||
assert np.array_equal(captured["img"], freq * grating_um)
|
||||
|
||||
|
||||
def test_pad_factor_uses_each_files_own_samples_per_frame(tmp_path, monkeypatch):
|
||||
"""n_fft must be derived per file from that file's own samples_per_frame,
|
||||
never a value carried over from whichever file the caller had open --
|
||||
otherwise every file but one in a batch gets silently mis-padded."""
|
||||
orig_draw = sras_render.draw_view_image
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
for i, spf in enumerate((64, 256)):
|
||||
path = tmp_path / f"pad_{spf}.sras"
|
||||
gen.write(path, n_angles=1, seed=5 + i, samples_per_frame=spf)
|
||||
|
||||
captured = {}
|
||||
|
||||
def spy(ax, fig, img, *a, __c=captured, **kw):
|
||||
__c["img"] = np.array(img, copy=True)
|
||||
return orig_draw(ax, fig, img, *a, **kw)
|
||||
|
||||
monkeypatch.setattr(sras_render, "draw_view_image", spy)
|
||||
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
|
||||
**_kw(is_fft_mode=True, pad_factor=4))
|
||||
assert err == ""
|
||||
|
||||
sras = SrasFile(str(path))
|
||||
expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
|
||||
apply_bg_sub=True, n_fft=spf * 4,
|
||||
min_freq_mhz=0.0)
|
||||
assert np.array_equal(captured["img"], expected), \
|
||||
f"samples_per_frame={spf}: n_fft must use this file's own value"
|
||||
|
||||
|
||||
def test_highlight_masked_sets_nan_and_bad_color(tmp_path, monkeypatch):
|
||||
path = tmp_path / "mask.sras"
|
||||
gen.write(path, n_angles=1, seed=7, samples_per_frame=128)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
sras = SrasFile(str(path))
|
||||
dc4 = dc_image_mv(sras, 0, CH4_IDX)
|
||||
threshold = float(np.median(dc4))
|
||||
expect_masked = dc4 < threshold
|
||||
assert expect_masked.any() and not expect_masked.all(), \
|
||||
"fixture threshold should mask some but not all pixels"
|
||||
|
||||
captured = _spy_draw(monkeypatch)
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
|
||||
**_kw(is_fft_mode=True, dc_threshold_mv=threshold,
|
||||
highlight_masked=True, mask_color="magenta"))
|
||||
assert err == ""
|
||||
assert captured["bad_color"] == "magenta"
|
||||
# The export masks by value too (0 == the "no valid peak" sentinel, same
|
||||
# rule as the viewer's _redraw_image); on this fixture every
|
||||
# above-threshold pixel has a nonzero peak, so the value mask coincides
|
||||
# with the DC mask and the NaN set is exactly expect_masked.
|
||||
assert np.array_equal(np.isnan(captured["img"]), expect_masked)
|
||||
valid_vals = captured["img"][~expect_masked]
|
||||
assert not np.isnan(valid_vals).any() and (valid_vals != 0).all(), \
|
||||
"fixture precondition: every valid pixel has a nonzero peak"
|
||||
|
||||
captured2 = _spy_draw(monkeypatch)
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
|
||||
**_kw(is_fft_mode=True, dc_threshold_mv=threshold, highlight_masked=False))
|
||||
assert err == ""
|
||||
assert captured2["bad_color"] is None
|
||||
assert not np.isnan(captured2["img"]).any()
|
||||
|
||||
|
||||
def test_auto_scale_uses_per_file_min_max(tmp_path, monkeypatch):
|
||||
path = tmp_path / "scale.sras"
|
||||
gen.write(path, n_angles=1, seed=8, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
captured = _spy_draw(monkeypatch)
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
|
||||
**_kw(auto_scale=True))
|
||||
assert err == ""
|
||||
img = captured["img"]
|
||||
assert captured["vmin"] == pytest.approx(float(np.nanmin(img)))
|
||||
assert captured["vmax"] == pytest.approx(float(np.nanmax(img)))
|
||||
|
||||
captured2 = _spy_draw(monkeypatch)
|
||||
err, _ = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
|
||||
**_kw(auto_scale=False, vmin=-5.0, vmax=5.0))
|
||||
assert err == ""
|
||||
assert captured2["vmin"] == -5.0
|
||||
assert captured2["vmax"] == 5.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("figsize", [(12.0, 4.0), (4.0, 9.0)])
|
||||
def test_export_matches_requested_figsize(tmp_path, figsize):
|
||||
"""The PNG comes out at the caller's figure size, so a view the user has
|
||||
sized wide (or tall) doesn't get squeezed into a fixed 7x5 -- the image
|
||||
is drawn with aspect="auto", so the figure box *is* the map's shape."""
|
||||
path = tmp_path / "shape.sras"
|
||||
gen.write(path, n_angles=1, seed=70, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
dpi = 100
|
||||
err, out_name = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
|
||||
figsize=figsize, dpi=dpi, **_kw())
|
||||
assert err == ""
|
||||
|
||||
w_px, h_px = _png_size(out_dir / out_name)
|
||||
# Agg truncates inches*dpi to whole pixels; a pixel of slack, not an
|
||||
# aspect-ratio tolerance, is what's being allowed for here.
|
||||
assert abs(w_px - figsize[0] * dpi) <= 1
|
||||
assert abs(h_px - figsize[1] * dpi) <= 1
|
||||
|
||||
|
||||
def test_export_default_figsize_when_unspecified(tmp_path):
|
||||
"""No figsize (a caller with no live canvas to match) still renders at
|
||||
the viewer's starting size rather than failing or guessing."""
|
||||
path = tmp_path / "default_shape.sras"
|
||||
gen.write(path, n_angles=1, seed=71, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
err, out_name = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX,
|
||||
dpi=100, **_kw())
|
||||
assert err == ""
|
||||
w_px, h_px = _png_size(out_dir / out_name)
|
||||
assert abs(w_px - DEFAULT_FIGSIZE[0] * 100) <= 1
|
||||
assert abs(h_px - DEFAULT_FIGSIZE[1] * 100) <= 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", [None, (0.0, 0.0), (-3.0, 5.0), (np.nan, 5.0),
|
||||
(float("inf"), 5.0), (7.0,), "7x5"])
|
||||
def test_sanitize_figsize_never_yields_an_unrenderable_size(bad):
|
||||
"""A degenerate canvas size (collapsed pane, minimized window) must cost
|
||||
at most a wrong-looking image, never a failed export."""
|
||||
w, h = sanitize_figsize(bad)
|
||||
assert np.isfinite(w) and np.isfinite(h)
|
||||
assert w >= 1.0 and h >= 1.0
|
||||
|
||||
|
||||
def test_sanitize_figsize_passes_through_a_normal_canvas_size():
|
||||
assert sanitize_figsize(np.array([12.8, 6.4])) == pytest.approx((12.8, 6.4))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GUI dispatch: SrasViewerWindow._on_batch_export_images end-to-end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_window(path) -> SrasViewerWindow:
|
||||
app = QApplication.instance() or QApplication([]) # noqa: F841
|
||||
win = SrasViewerWindow()
|
||||
win.show()
|
||||
win._load_file(str(path))
|
||||
assert wait_until(lambda: win._sras is not None), "file loaded"
|
||||
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
|
||||
for a in range(win._sras.n_angles))), \
|
||||
"DC precompute finished"
|
||||
return win
|
||||
|
||||
|
||||
def test_batch_export_images_writes_one_png_per_file(tmp_path):
|
||||
path = tmp_path / "src.sras"
|
||||
gen.write(path, n_angles=2, seed=10, samples_per_frame=64)
|
||||
paths = [str(path)]
|
||||
for i in range(2):
|
||||
p2 = tmp_path / f"other{i}.sras"
|
||||
gen.write(p2, n_angles=2, seed=20 + i, samples_per_frame=64)
|
||||
paths.append(str(p2))
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path)
|
||||
try:
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=(paths, "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
|
||||
angle = win.spin_angle.value()
|
||||
ch_name = CH_NAMES[win.combo_channel.currentIndex()]
|
||||
expected_names = {f"{Path(p).stem}_angle{angle}_{ch_name}.png" for p in paths}
|
||||
actual_names = {p.name for p in out_dir.iterdir()}
|
||||
assert actual_names == expected_names
|
||||
assert "Batch export: 3/3 image(s)" in win.statusBar().currentMessage()
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_out_of_range_angle_reports_error_continues(tmp_path):
|
||||
good_path = tmp_path / "good.sras"
|
||||
short_path = tmp_path / "short.sras"
|
||||
gen.write(good_path, n_angles=3, seed=30, samples_per_frame=64)
|
||||
gen.write(short_path, n_angles=1, seed=31, samples_per_frame=64)
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(good_path)
|
||||
try:
|
||||
win.spin_angle.setValue(2) # valid for good_path, out of range for short_path
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(good_path), str(short_path)], "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
|
||||
msg = win.statusBar().currentMessage()
|
||||
assert "Batch export: 1/2 image(s)" in msg, msg
|
||||
assert "1 failed" in msg, msg
|
||||
assert len(list(out_dir.iterdir())) == 1, \
|
||||
"the batch must not abort — the good file still exports"
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_busy_guard_skips_dialogs(tmp_path, monkeypatch):
|
||||
path = tmp_path / "busy.sras"
|
||||
gen.write(path, n_angles=1, seed=40, samples_per_frame=64)
|
||||
win = _make_window(path)
|
||||
try:
|
||||
monkeypatch.setattr(win, "_job_running", lambda key: True)
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames") as mock_dlg:
|
||||
win._on_batch_export_images()
|
||||
assert mock_dlg.call_count == 0, \
|
||||
"the Jobs.BATCH busy guard must return before opening any dialog"
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_ignores_aligned_view_toggle(tmp_path):
|
||||
"""Aligned View is geometry specific to whichever single file the
|
||||
Alignment Wizard last ran against and cannot be meaningfully applied
|
||||
across a batch of different files -- _on_batch_export_images must not
|
||||
read chk_aligned_view / self._alignment_result at all, regardless of
|
||||
what's checked in the live view."""
|
||||
path = tmp_path / "aligned.sras"
|
||||
gen.write(path, n_angles=1, seed=41, samples_per_frame=64)
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path)
|
||||
try:
|
||||
captured_kwargs = []
|
||||
orig_init = BatchExportImagesWorker.__init__
|
||||
|
||||
def spy_init(self, paths, **kw):
|
||||
captured_kwargs.append(kw)
|
||||
return orig_init(self, paths, **kw)
|
||||
|
||||
win.chk_aligned_view.setChecked(True)
|
||||
with patch.object(BatchExportImagesWorker, "__init__", spy_init), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(path)], "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
|
||||
assert len(captured_kwargs) == 1
|
||||
assert not any("align" in k.lower() for k in captured_kwargs[0]), \
|
||||
captured_kwargs[0].keys()
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_uses_the_live_canvas_aspect_ratio(tmp_path):
|
||||
"""The exported PNG has the shape of the view on screen, not a fixed
|
||||
7x5 -- resize the window and the export follows it."""
|
||||
path = tmp_path / "aspect.sras"
|
||||
gen.write(path, n_angles=1, seed=45, samples_per_frame=64)
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path)
|
||||
try:
|
||||
win.resize(1400, 700)
|
||||
pump(300) # let the canvas' resizeEvent reach the figure
|
||||
canvas_w, canvas_h = win.image_canvas.figure.get_size_inches()
|
||||
|
||||
captured_kwargs = []
|
||||
orig_init = BatchExportImagesWorker.__init__
|
||||
|
||||
def spy_init(self, paths, **kw):
|
||||
captured_kwargs.append(kw)
|
||||
return orig_init(self, paths, **kw)
|
||||
|
||||
with patch.object(BatchExportImagesWorker, "__init__", spy_init), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(path)], "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
|
||||
assert captured_kwargs[0]["figsize"] == pytest.approx(
|
||||
(canvas_w, canvas_h)), "the live canvas size must travel to the worker"
|
||||
|
||||
out_files = list(out_dir.iterdir())
|
||||
assert len(out_files) == 1
|
||||
w_px, h_px = _png_size(out_files[0])
|
||||
assert w_px / h_px == pytest.approx(canvas_w / canvas_h, rel=0.01)
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_filename_collision_note(tmp_path):
|
||||
dir_a, dir_b = tmp_path / "dir_a", tmp_path / "dir_b"
|
||||
dir_a.mkdir()
|
||||
dir_b.mkdir()
|
||||
path_a, path_b = dir_a / "dup.sras", dir_b / "dup.sras"
|
||||
gen.write(path_a, n_angles=1, seed=50, samples_per_frame=64)
|
||||
gen.write(path_b, n_angles=1, seed=51, samples_per_frame=64)
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path_a)
|
||||
try:
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(path_a), str(path_b)], "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
|
||||
msg = win.statusBar().currentMessage()
|
||||
assert "Batch export: 2/2 image(s)" in msg, msg
|
||||
assert "collision" in msg, msg
|
||||
assert len(list(out_dir.iterdir())) == 1, \
|
||||
"same-stem inputs silently overwrite to one output file"
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_batch_export_does_not_modify_source_files(tmp_path):
|
||||
path = tmp_path / "untouched.sras"
|
||||
gen.write(path, n_angles=1, seed=60, samples_per_frame=64)
|
||||
before = path.read_bytes()
|
||||
out_dir = tmp_path / "images"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path)
|
||||
try:
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(path)], "")), \
|
||||
patch("sras_viewer.main_window.QFileDialog.getExistingDirectory",
|
||||
return_value=str(out_dir)):
|
||||
win._on_batch_export_images()
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
assert path.read_bytes() == before, "export must never write to the source file"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Export menu: BatchExportWorker (every angle of the open file)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# High enough to sit above several of the generator's synthetic peak bins
|
||||
# (bin spacing is 6.25 GS/s / 64 = 97.66 MHz, peaks land at bins 3-19), so a
|
||||
# floor at this value genuinely changes which peak each pixel resolves to --
|
||||
# without that, a test that the floor is honored would pass either way.
|
||||
_FLOOR_MHZ = 1000.0
|
||||
_GRATING_UM = 2.0
|
||||
|
||||
|
||||
def _velocity_channel() -> ExportChannel:
|
||||
return ExportChannel(ch_idx=CH1_IDX, is_velocity=True, vmin=0.0, vmax=1e5,
|
||||
label="Velocity", unit="m/s", tag="VEL")
|
||||
|
||||
|
||||
def _run_export_worker(monkeypatch, sras, out_dir, **overrides) -> list[np.ndarray]:
|
||||
"""Runs BatchExportWorker to completion on the calling thread (its run()
|
||||
is a plain loop; the QThread in _run_worker is a GUI concern) and returns
|
||||
the image arrays it rendered, captured at the _render_map_png seam."""
|
||||
captured = []
|
||||
orig = sras_workers._render_map_png
|
||||
|
||||
def spy(img, extent, **kw):
|
||||
captured.append(np.array(img, copy=True))
|
||||
return orig(img, extent, **kw)
|
||||
|
||||
monkeypatch.setattr(sras_workers, "_render_map_png", spy)
|
||||
kw = dict(cmap="viridis", apply_bg_sub=True, dc_threshold_mv=_THRESHOLD_MV,
|
||||
n_fft=None, grating_um=_GRATING_UM, min_freq_mhz=_FLOOR_MHZ)
|
||||
kw.update(overrides)
|
||||
BatchExportWorker(sras, [_velocity_channel()], str(out_dir), "vel", **kw).run()
|
||||
return captured
|
||||
|
||||
|
||||
def test_batch_export_applies_the_min_peak_freq_floor(tmp_path, monkeypatch):
|
||||
"""The exported Velocity map is the floored one the viewer shows, not the
|
||||
unfloored peaks the floor was raised to reject."""
|
||||
path = tmp_path / "floor.sras"
|
||||
gen.write(path, n_angles=1, seed=80, samples_per_frame=64, geometry=[(6, 9)])
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
sras = SrasFile(str(path))
|
||||
|
||||
def expected(floor):
|
||||
return compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV,
|
||||
apply_bg_sub=True, n_fft=None,
|
||||
min_freq_mhz=floor) * _GRATING_UM
|
||||
|
||||
floored, unfloored = expected(_FLOOR_MHZ), expected(0.0)
|
||||
assert not np.allclose(floored, unfloored), \
|
||||
"fixture must be one where the floor changes the image"
|
||||
|
||||
exported = _run_export_worker(monkeypatch, sras, out_dir)
|
||||
assert len(exported) == 1
|
||||
assert np.allclose(exported[0], floored)
|
||||
|
||||
|
||||
def test_batch_export_does_not_serve_an_unfloored_stored_cache(tmp_path,
|
||||
monkeypatch):
|
||||
"""A file batch-computed before the floor existed has a stored cache with
|
||||
no floor recorded. Asking for that cache at floor 0 (rather than the live
|
||||
floor) makes it a match, so the export would render the stored pre-floor
|
||||
peaks -- the map the user raised the floor to get rid of."""
|
||||
path = tmp_path / "stored_floor.sras"
|
||||
gen.write(path, n_angles=1, seed=81, samples_per_frame=64, geometry=[(6, 9)])
|
||||
assert cache_file(str(path), "fft", apply_bg_sub=True) == ""
|
||||
sras = SrasFile(str(path))
|
||||
assert sras.precomputed_freq_mhz[0] is not None, "stored cache written"
|
||||
assert sras.precomputed_min_freq_mhz == 0.0
|
||||
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
exported = _run_export_worker(monkeypatch, sras, out_dir)[0]
|
||||
|
||||
# 0.0 is the "no valid peak" sentinel; anything else below the floor is a
|
||||
# peak that only an unfloored search could have reported.
|
||||
below = (exported > 0) & (exported < _FLOOR_MHZ * _GRATING_UM)
|
||||
assert not below.any(), \
|
||||
f"{below.sum()} sub-floor pixel(s) survived the export"
|
||||
|
||||
|
||||
def test_batch_export_dispatch_passes_the_live_view_settings(tmp_path):
|
||||
"""The Export menu hands the worker what the panel currently says --
|
||||
the floor in particular, which is a display control and so persists
|
||||
across file loads while the window's own FFT cache does not."""
|
||||
path = tmp_path / "dispatch.sras"
|
||||
gen.write(path, n_angles=1, seed=82, samples_per_frame=64)
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
|
||||
win = _make_window(path)
|
||||
try:
|
||||
win.spin_min_freq_mhz.setValue(_FLOOR_MHZ)
|
||||
win.spin_grating_um.setValue(_GRATING_UM)
|
||||
win.spin_threshold_mv.setValue(_THRESHOLD_MV)
|
||||
|
||||
class StubDialog:
|
||||
def __init__(self, parent, **kw):
|
||||
pass
|
||||
|
||||
def exec(self):
|
||||
from PyQt6.QtWidgets import QDialog
|
||||
return QDialog.DialogCode.Accepted
|
||||
|
||||
def get_output_dir(self):
|
||||
return str(out_dir)
|
||||
|
||||
def get_prefix(self):
|
||||
return "vel"
|
||||
|
||||
def get_selected_channels(self):
|
||||
return [_velocity_channel()]
|
||||
|
||||
captured_kwargs = []
|
||||
orig_init = BatchExportWorker.__init__
|
||||
|
||||
def spy_init(self, sras, channels, output_dir, prefix, **kw):
|
||||
captured_kwargs.append(kw)
|
||||
return orig_init(self, sras, channels, output_dir, prefix, **kw)
|
||||
|
||||
with patch.object(BatchExportWorker, "__init__", spy_init), \
|
||||
patch("sras_viewer.main_window.BatchExportDialog", StubDialog):
|
||||
win._on_batch_export()
|
||||
assert wait_until(lambda: not win._job_running("export"), 60000), "export ran"
|
||||
|
||||
kw = captured_kwargs[0]
|
||||
assert kw["min_freq_mhz"] == _FLOOR_MHZ
|
||||
assert kw["grating_um"] == _GRATING_UM
|
||||
assert kw["dc_threshold_mv"] == _THRESHOLD_MV
|
||||
finally:
|
||||
win.close()
|
||||
pump(200)
|
||||
|
||||
|
||||
def test_export_view_image_serves_a_row_averaged_stored_cache(tmp_path,
|
||||
monkeypatch):
|
||||
"""Batch Export View as Images must ask this file's cache the question the
|
||||
viewer asks it. _stored_fft_image reads at the file's own
|
||||
precomputed_row_avg_n; requesting raw per-pixel instead makes a
|
||||
row-averaged cache a mismatch, and the export silently renders a full raw
|
||||
recompute where the screen shows the smoothed stored image."""
|
||||
path = tmp_path / "rowavg_cache.sras"
|
||||
gen.write(path, n_angles=1, seed=83, samples_per_frame=64, geometry=[(6, 9)])
|
||||
|
||||
# A planted stored image rather than a real row-averaged compute: the
|
||||
# question here is *which* source the export reads, and a distinctive
|
||||
# array answers it without depending on the synthetic waveforms being
|
||||
# smooth enough for averaging to move the numbers.
|
||||
sras = SrasFile(str(path))
|
||||
shape = sras.image_shape(0)
|
||||
stored = (100.0 + 10.0 * np.arange(shape[0] * shape[1], dtype=np.float32)
|
||||
).reshape(shape)
|
||||
sras.write_v7_cache(new_freq_mhz=[stored], new_row_avg_n=5,
|
||||
new_bg_sub=True, new_pad_factor=1)
|
||||
|
||||
reread = SrasFile(str(path))
|
||||
raw = compute_rf_image(reread, 0, dc_threshold_mv=None, apply_bg_sub=True,
|
||||
n_fft=None, row_avg_n=0, use_stored=False)
|
||||
assert not np.allclose(stored, raw), "planted cache must be distinguishable"
|
||||
|
||||
out_dir = tmp_path / "out"
|
||||
out_dir.mkdir()
|
||||
captured = _spy_draw(monkeypatch)
|
||||
err, _out_name = export_view_image(
|
||||
str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX,
|
||||
**_kw(is_fft_mode=True, dc_threshold_mv=-1e9, colorbar_label="MHz",
|
||||
mode_str="RF"))
|
||||
assert err == ""
|
||||
assert np.allclose(captured["img"], stored), \
|
||||
"export rendered a recompute instead of the stored image on screen"
|
||||
+180
-58
@@ -12,6 +12,7 @@ from pathlib import Path
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import sras_average
|
||||
import sras_compute as compute
|
||||
from sras_compute import (
|
||||
cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
|
||||
@@ -122,7 +123,7 @@ def test_parallel_identity(tmp_path, monkeypatch):
|
||||
# the block size so every chunk splits into many FFT tasks — the worst
|
||||
# case for boundary bugs.
|
||||
monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
|
||||
monkeypatch.setattr(compute, "_FFT_BLOCK", 4)
|
||||
monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
|
||||
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
|
||||
assert fft_rows < n_rows, \
|
||||
f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
|
||||
@@ -155,43 +156,10 @@ def test_parallel_identity(tmp_path, monkeypatch):
|
||||
"rf image identical (masked, padded/zoom)"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spf,bps", [(64, 2), (37, 1)])
|
||||
def test_zoom_identity(tmp_path, monkeypatch, spf, bps):
|
||||
"""The zoom peak search must reproduce the full padded-rfft argmax
|
||||
bit-for-bit, across pad factors, masking, bg-sub, dtype, and backend."""
|
||||
path = tmp_path / f"zoom_{spf}.sras"
|
||||
gen.write(path, n_angles=2, seed=6, samples_per_frame=spf, bps=bps)
|
||||
sras = SrasFile(str(path))
|
||||
dc4 = dc_image_mv(sras, 0, CH4_IDX)
|
||||
thr = float(np.median(dc4))
|
||||
|
||||
backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
|
||||
for backend in backends:
|
||||
monkeypatch.setattr(compute, "_fft_backend", backend)
|
||||
for pad in (4, 8, 40):
|
||||
n_fft = spf * pad
|
||||
for thr_v in (None, thr):
|
||||
for bg in (False, True):
|
||||
ref = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
|
||||
apply_bg_sub=bg, n_fft=n_fft,
|
||||
exact=True)
|
||||
zoom = compute_rf_image(sras, 0, dc_threshold_mv=thr_v,
|
||||
apply_bg_sub=bg, n_fft=n_fft)
|
||||
diff = int((ref != zoom).sum())
|
||||
assert diff == 0, \
|
||||
(f"{diff} px differ: backend={backend} pad={pad} "
|
||||
f"thr={thr_v} bg={bg} spf={spf}")
|
||||
|
||||
# A threshold above every pixel masks everything: both paths must agree
|
||||
# on an all-zero image.
|
||||
all_masked = compute_rf_image(sras, 0, dc_threshold_mv=1e9, n_fft=spf * 8)
|
||||
assert not all_masked.any()
|
||||
|
||||
|
||||
def test_zoom_identity_fuzz():
|
||||
"""Hammer _peak_bins_zoom directly with adversarial spectra: noise,
|
||||
def test_peak_bins_fuzz():
|
||||
"""Hammer _peak_bins directly with adversarial spectra: noise,
|
||||
un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs,
|
||||
and all-zero rows."""
|
||||
and all-zero rows — against an independent scipy.fft reference."""
|
||||
import scipy.fft as scipy_fft
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
@@ -222,14 +190,102 @@ def test_zoom_identity_fuzz():
|
||||
P[:, 0] = 0.0
|
||||
ref = np.argmax(P, axis=1)
|
||||
|
||||
zp = compute._zoom_plan(spf, n_fft)
|
||||
got = compute._peak_bins_zoom(w, zp)
|
||||
got = compute._peak_bins(w, n_fft)
|
||||
bad = np.nonzero(ref != got)[0]
|
||||
assert not len(bad), \
|
||||
(f"spf={spf} pad={pad}: rows {bad.tolist()} picked "
|
||||
f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
|
||||
|
||||
|
||||
def test_peak_bins_fuzz_min_freq():
|
||||
"""Same adversarial-spectra fuzz as test_peak_bins_fuzz, but with a swept
|
||||
min_bin floor: bins below the floor must be excluded from the argmax
|
||||
exactly as the independent scipy.fft reference is, when zeroed the same
|
||||
way before argmax."""
|
||||
import scipy.fft as scipy_fft
|
||||
|
||||
rng = np.random.default_rng(43)
|
||||
for _ in range(25):
|
||||
spf = int(rng.integers(16, 220))
|
||||
pad = int(rng.choice([4, 5, 8, 16, 40]))
|
||||
n_fft = spf * pad
|
||||
n_wf = 24
|
||||
w = rng.normal(scale=20.0, size=(n_wf, spf))
|
||||
t = np.arange(spf)
|
||||
for r in range(6):
|
||||
f = rng.uniform(1.0, spf / 2 - 1)
|
||||
w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
|
||||
f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
|
||||
w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
|
||||
+ 49.9 * np.sin(2 * np.pi * f2 * t / spf)
|
||||
w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
|
||||
+ 50 * np.cos(2 * np.pi * f2 * t / spf)
|
||||
w[8] = 90 + rng.normal(scale=5.0, size=spf)
|
||||
w[9] = 0.0
|
||||
w = w.astype(np.float32)
|
||||
|
||||
S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
|
||||
P = S.real ** 2
|
||||
P += S.imag ** 2
|
||||
n_bins_fine = n_fft // 2 + 1
|
||||
min_bin = int(rng.integers(1, max(2, n_bins_fine // 3)))
|
||||
P[:, :min_bin] = 0.0
|
||||
ref = np.argmax(P, axis=1)
|
||||
|
||||
got = compute._peak_bins(w, n_fft, min_bin)
|
||||
bad = np.nonzero(ref != got)[0]
|
||||
assert not len(bad), \
|
||||
(f"spf={spf} pad={pad} min_bin={min_bin}: rows {bad.tolist()} picked "
|
||||
f"{got[bad].tolist()} instead of {ref[bad].tolist()}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spf", [64, 500, 2500])
|
||||
def test_fft_block_for(spf):
|
||||
"""Block size == _FFT_BLOCK_MAX at natural resolution, shrinks and stays
|
||||
>= _FFT_BLOCK_MIN as n_len grows, and the implied per-thread byte
|
||||
estimate respects _FFT_PLAN_BYTES_BUDGET except when the floor is
|
||||
engaged."""
|
||||
block_natural = compute._fft_block_for(spf, spf)
|
||||
assert block_natural == compute._FFT_BLOCK_MAX
|
||||
|
||||
prev = compute._FFT_BLOCK_MAX
|
||||
for pad in (2, 4, 8, 40, 500):
|
||||
n_len = spf * pad
|
||||
block = compute._fft_block_for(spf, n_len)
|
||||
assert compute._FFT_BLOCK_MIN <= block <= prev
|
||||
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
|
||||
if block > compute._FFT_BLOCK_MIN:
|
||||
assert block * bytes_per_wf <= compute._FFT_PLAN_BYTES_BUDGET
|
||||
prev = block
|
||||
|
||||
|
||||
def test_compute_rf_image_min_freq_mhz(tmp_path):
|
||||
"""min_freq_mhz threads through compute_rf_image end-to-end, for both
|
||||
the natural-resolution and padded paths: 0.0 (default) must reproduce
|
||||
the pre-existing image exactly, and a floor above every real peak must
|
||||
collapse the image to bin 0 (0 MHz) — the same fallback the low-level
|
||||
search uses when nothing survives the floor."""
|
||||
path = tmp_path / "floor_e2e.sras"
|
||||
gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
|
||||
sras = SrasFile(str(path))
|
||||
huge_floor = float(sras.freq_axis_mhz(None)[-1]) + 1.0 # above Nyquist
|
||||
|
||||
for n_fft in (None, 64 * 8):
|
||||
unfiltered = compute_rf_image(sras, 0, dc_threshold_mv=None,
|
||||
apply_bg_sub=False, n_fft=n_fft)
|
||||
same = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
|
||||
n_fft=n_fft, min_freq_mhz=0.0)
|
||||
assert np.array_equal(unfiltered, same), \
|
||||
f"n_fft={n_fft}: min_freq_mhz=0.0 changed the output"
|
||||
assert unfiltered.any(), \
|
||||
f"n_fft={n_fft}: fixture should have real signal"
|
||||
|
||||
collapsed = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
|
||||
n_fft=n_fft, min_freq_mhz=huge_floor)
|
||||
assert not collapsed.any(), \
|
||||
f"n_fft={n_fft}: floor above Nyquist should collapse to 0 MHz"
|
||||
|
||||
|
||||
def test_nomask_equals_low_threshold(tmp_path):
|
||||
"""dc_threshold_mv=None must equal a threshold below every pixel, while
|
||||
skipping the CH4 read."""
|
||||
@@ -307,45 +363,111 @@ def test_legacy_parse(tmp_path):
|
||||
|
||||
|
||||
def test_sras_average(tmp_path):
|
||||
"""The sras_average.py CLI: frame averaging with remainder handling."""
|
||||
src = tmp_path / "legacy_v4.sras"
|
||||
meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
|
||||
samples_per_frame=32, seed=4)
|
||||
dst = tmp_path / "legacy_v4_avg.sras"
|
||||
"""The sras_average.py CLI: v6 frame averaging with remainder handling,
|
||||
per-angle ragged geometry, and the laser_freq_hz X-axis correction."""
|
||||
src = tmp_path / "v6.sras"
|
||||
meta = gen.write(src, n_angles=2, seed=4, samples_per_frame=32,
|
||||
geometry=[(4, 12)])
|
||||
src_sras = SrasFile(str(src))
|
||||
|
||||
dst = tmp_path / "v6_avg.sras"
|
||||
proc = subprocess.run(
|
||||
[sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "4"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]
|
||||
|
||||
avg = SrasFile(str(dst))
|
||||
assert avg.version == 4
|
||||
assert avg.version == 6
|
||||
assert list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}"
|
||||
assert (avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
|
||||
and avg.n_channels == meta["n_channels"])
|
||||
assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \
|
||||
and avg.n_channels == src_sras.n_channels == 3)
|
||||
assert np.allclose(avg.ch_ymult_mv, src_sras.ch_ymult_mv), \
|
||||
"calibration preserved"
|
||||
assert np.array_equal(avg.background, SrasFile(str(src)).background), \
|
||||
assert np.array_equal(avg.background, src_sras.background), \
|
||||
"background preserved"
|
||||
src_data = meta["data"]
|
||||
expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
|
||||
assert avg.laser_freq_hz == pytest.approx(src_sras.laser_freq_hz / 4), \
|
||||
"laser_freq_hz divided by N keeps pixel_x_mm correct after binning"
|
||||
assert np.array_equal(avg.x_start_mm, src_sras.x_start_mm), \
|
||||
"per-angle x_start unchanged"
|
||||
|
||||
src_waves = meta["waveforms"]
|
||||
# int16 (not float32) before .mean(): matches _average_block's own
|
||||
# float64-accumulator behavior for integer input, so this doesn't
|
||||
# drift from what _average_block actually guarantees.
|
||||
expect0 = src_waves[0][:, :, 0:4, :].astype(np.int16).mean(axis=2).astype(np.int16)
|
||||
assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \
|
||||
"first averaged group equals the mean of its 4 source frames"
|
||||
|
||||
# Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
|
||||
dst2 = tmp_path / "legacy_v4_avg5.sras"
|
||||
subprocess.run([sys.executable, str(REPO / "sras_average.py"),
|
||||
str(src), str(dst2), "--n", "5"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
dst2 = tmp_path / "v6_avg5.sras"
|
||||
proc2 = subprocess.run([sys.executable, str(REPO / "sras_average.py"),
|
||||
str(src), str(dst2), "--n", "5"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
assert proc2.returncode == 0, (proc2.stderr or proc2.stdout).strip()[-200:]
|
||||
assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \
|
||||
"partial trailing group kept by default"
|
||||
dst3 = tmp_path / "legacy_v4_avg5d.sras"
|
||||
subprocess.run([sys.executable, str(REPO / "sras_average.py"),
|
||||
str(src), str(dst3), "--n", "5", "--discard-remainder"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
dst3 = tmp_path / "v6_avg5d.sras"
|
||||
proc3 = subprocess.run([sys.executable, str(REPO / "sras_average.py"),
|
||||
str(src), str(dst3), "--n", "5", "--discard-remainder"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
assert proc3.returncode == 0, (proc3.stderr or proc3.stdout).strip()[-200:]
|
||||
assert list(SrasFile(str(dst3)).n_frames) == [2, 2], \
|
||||
"--discard-remainder drops the partial group"
|
||||
|
||||
|
||||
def test_sras_average_v7_cache_dropped(tmp_path):
|
||||
"""A v7 input's cache tail is indexed by frame count, so it's invalid
|
||||
after averaging changes that count -- the output must always be plain
|
||||
v6, never a v7 carrying a stale cache."""
|
||||
src = tmp_path / "v7.sras"
|
||||
gen.write(src, n_angles=2, seed=1, samples_per_frame=32, geometry=[(3, 8)])
|
||||
src_sras = SrasFile(str(src))
|
||||
src_sras.write_v7_cache(
|
||||
new_dc3_mv=[compute_dc_image(src_sras, a, CH3_IDX) for a in range(src_sras.n_angles)],
|
||||
new_dc4_mv=[compute_dc_image(src_sras, a, CH4_IDX) for a in range(src_sras.n_angles)])
|
||||
assert SrasFile(str(src)).version == 7
|
||||
|
||||
dst = tmp_path / "v7_avg.sras"
|
||||
proc = subprocess.run(
|
||||
[sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]
|
||||
assert SrasFile(str(dst)).version == 6, "cache-bearing input still writes plain v6"
|
||||
|
||||
|
||||
def test_sras_average_rejects_legacy(tmp_path):
|
||||
"""This tool only speaks v6/v7 now; a legacy file must fail clearly
|
||||
rather than being silently misparsed."""
|
||||
src = tmp_path / "legacy_v4.sras"
|
||||
gen.write_legacy(src, version=4, n_angles=1, n_rows=2, n_frames=8,
|
||||
samples_per_frame=16, seed=0)
|
||||
dst = tmp_path / "legacy_v4_avg.sras"
|
||||
proc = subprocess.run(
|
||||
[sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"],
|
||||
capture_output=True, text=True, cwd=REPO)
|
||||
assert proc.returncode != 0
|
||||
assert "v6" in proc.stderr and "v7" in proc.stderr
|
||||
|
||||
|
||||
def test_sras_average_chunking_matches_unchunked(tmp_path):
|
||||
"""A tiny memory budget (forcing one row per chunk) must produce
|
||||
byte-identical output to a huge budget (everything in one chunk) -- the
|
||||
load-bearing correctness claim of the memory-bounded rewrite: chunk
|
||||
boundaries must never affect the averaged result."""
|
||||
src = tmp_path / "v6.sras"
|
||||
gen.write(src, n_angles=2, seed=7, samples_per_frame=48,
|
||||
geometry=[(6, 10), (5, 13)])
|
||||
sras = SrasFile(str(src))
|
||||
|
||||
dst_tiny = tmp_path / "avg_tiny.sras"
|
||||
dst_big = tmp_path / "avg_big.sras"
|
||||
sras_average.write_v6_averaged(sras, dst_tiny, 3, False, budget=1)
|
||||
sras_average.write_v6_averaged(sras, dst_big, 3, False, budget=1 << 30)
|
||||
|
||||
assert dst_tiny.read_bytes() == dst_big.read_bytes(), \
|
||||
"chunk size must not affect the averaged output"
|
||||
|
||||
|
||||
def test_unsupported_version_reported(tmp_path):
|
||||
"""cache_file must report, not raise, for a file it can't handle."""
|
||||
bogus = tmp_path / "bogus.sras"
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
"""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()
|
||||
+605
-177
@@ -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, angle alignment, manual angle alignment, aligned view, ROI
|
||||
draw/move, and CSV export.
|
||||
changes, the alignment wizard end to end (pre-rotation, correlation, manual
|
||||
nudging, crop, export), 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,11 +20,13 @@ import numpy as np
|
||||
import pytest
|
||||
from PyQt6.QtCore import QEventLoop, Qt, QTimer
|
||||
from PyQt6.QtTest import QTest
|
||||
from PyQt6.QtWidgets import QApplication, QMessageBox
|
||||
from PyQt6.QtWidgets import QApplication, QDialog, QMessageBox
|
||||
|
||||
import sras_compute as compute
|
||||
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
|
||||
from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX
|
||||
from sras_viewer import (
|
||||
FusedRoiExportDialog, RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX,
|
||||
)
|
||||
import tools.make_test_sras as gen
|
||||
|
||||
|
||||
@@ -106,6 +108,71 @@ 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.
|
||||
"""
|
||||
win, s = ctx.win, ctx.s
|
||||
assert s.n_angles >= 3, "need room to step in both directions"
|
||||
|
||||
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)
|
||||
|
||||
seen = []
|
||||
win.spin_angle.valueChanged.connect(seen.append)
|
||||
try:
|
||||
win.spin_angle.lineEdit().selectAll()
|
||||
QTest.keyClicks(win.spin_angle, str(target))
|
||||
pump(60)
|
||||
assert seen == [], f"no signal while typing, got {seen}"
|
||||
QTest.keyClick(win.spin_angle, Qt.Key.Key_Return)
|
||||
pump(60)
|
||||
assert seen == [target], f"one signal on commit, got {seen}"
|
||||
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"
|
||||
|
||||
|
||||
def test_channel_switching(ctx):
|
||||
win = ctx.win
|
||||
win.spin_angle.setValue(0)
|
||||
@@ -156,17 +223,74 @@ def test_threshold_change_recomputes(ctx):
|
||||
f"masking zeroed some pixels ({n_zero} of {win._current_image.size})"
|
||||
|
||||
|
||||
def test_highlight_masked_pixels(ctx):
|
||||
"""Masked (below-threshold) pixels are drawn as NaN - filled with a
|
||||
highlight color, separate from the normal colormap - so they can't be
|
||||
mistaken for a real, possibly-low-frequency pixel; unchecking restores
|
||||
the old behavior where both blend into the same plain 0."""
|
||||
win = ctx.win
|
||||
assert win.chk_highlight_masked.isChecked(), "on by default"
|
||||
dc4 = win._dc_cache[(0, CH4_IDX)]
|
||||
expect_masked = dc4 < win.spin_threshold_mv.value()
|
||||
assert expect_masked.any() and not expect_masked.all(), \
|
||||
"fixture threshold should mask some but not all pixels"
|
||||
|
||||
calls = []
|
||||
orig = win.image_canvas.show_image
|
||||
|
||||
def spy(img, *a, **kw):
|
||||
calls.append((np.array(img, copy=True), kw.get("bad_color")))
|
||||
return orig(img, *a, **kw)
|
||||
|
||||
with patch.object(win.image_canvas, "show_image", side_effect=spy):
|
||||
win._redraw_image(win._current_image)
|
||||
shown, bad_color = calls[-1]
|
||||
assert bad_color is not None, "highlight color set while checkbox is on"
|
||||
# The highlight masks by value too: in an FFT mode, exactly 0 is the
|
||||
# "no valid peak" sentinel (DC-masked, below the min-freq floor, or an
|
||||
# empty spectrum), so the NaN set is the union of the DC mask and the
|
||||
# zero-valued pixels. On this fixture every above-threshold pixel has a
|
||||
# nonzero peak, so the union equals the DC mask alone.
|
||||
expect_nan = expect_masked | (win._current_image == 0)
|
||||
assert np.array_equal(np.isnan(shown), expect_nan), \
|
||||
"NaN where DC4 is below threshold or the value-0 sentinel, nowhere else"
|
||||
assert np.array_equal(expect_nan, expect_masked), \
|
||||
"fixture precondition: every valid pixel has a nonzero peak"
|
||||
|
||||
win.chk_highlight_masked.setChecked(False)
|
||||
calls.clear()
|
||||
with patch.object(win.image_canvas, "show_image", side_effect=spy):
|
||||
win._redraw_image(win._current_image)
|
||||
shown2, bad_color2 = calls[-1]
|
||||
assert bad_color2 is None, "no highlight color once unchecked"
|
||||
assert not np.isnan(shown2).any(), "unchecked: no pixel pulled out to NaN"
|
||||
assert np.array_equal(shown2, win._current_image), \
|
||||
"unchecked: displayed array is the raw, unmodified image"
|
||||
|
||||
win.chk_highlight_masked.setChecked(True)
|
||||
pump(60)
|
||||
|
||||
|
||||
def test_bg_sub_toggle(ctx):
|
||||
"""bg-sub no longer gates the display: it only affects a future live
|
||||
compute for an angle with nothing cached yet, or an explicit batch
|
||||
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)
|
||||
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"), \
|
||||
"returning to bg-sub was a cache hit (no recompute)"
|
||||
"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"
|
||||
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)
|
||||
|
||||
|
||||
def test_roi_and_csv_export(ctx):
|
||||
@@ -218,44 +342,16 @@ def test_roi_survives_switches(ctx):
|
||||
"ROI still present after channel switch"
|
||||
|
||||
|
||||
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):
|
||||
def test_alignment_geometry_is_stage_independent(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
|
||||
assert win._manual_align_act.isEnabled(), "manual alignment action enabled"
|
||||
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"
|
||||
|
||||
n_rows, n_frames = s.image_shape(0)
|
||||
assert np.allclose(compute._center_idx(s, 0),
|
||||
@@ -276,7 +372,7 @@ def test_manual_alignment_geometry(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.
|
||||
# Both signs of the stage's reported angle are searched by default.
|
||||
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)}"
|
||||
|
||||
@@ -289,128 +385,295 @@ def test_manual_alignment_geometry(ctx):
|
||||
"_shift_into moves content by exactly the requested offset"
|
||||
|
||||
|
||||
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."""
|
||||
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."""
|
||||
win, s = ctx.win, ctx.s
|
||||
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"), \
|
||||
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"), \
|
||||
"mask prep needed no background worker (already DC-cached)"
|
||||
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"
|
||||
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"
|
||||
|
||||
|
||||
def test_reference_angle_is_locked(ctx):
|
||||
dlg = ctx.dlg
|
||||
dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx)
|
||||
def test_wizard_reference_angle_is_locked(ctx):
|
||||
p1, wiz = ctx.p1, ctx.wiz
|
||||
p1.combo_active.setCurrentIndex(wiz.state.ref_angle_idx)
|
||||
pump(30)
|
||||
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, \
|
||||
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, \
|
||||
"reference angle untouched by nudge attempts"
|
||||
|
||||
|
||||
def test_nudges(ctx):
|
||||
"""Nudging a real angle (fine + coarse, translate + rotate)."""
|
||||
dlg, s = ctx.dlg, ctx.s
|
||||
def test_wizard_nudges(ctx):
|
||||
"""Manual correction, which the wizard absorbed from the old dialog."""
|
||||
p1, wiz, s = ctx.p1, ctx.wiz, ctx.s
|
||||
ctx.active = active = 1 if s.n_angles > 1 else 0
|
||||
dlg.combo_active_angle.setCurrentIndex(active)
|
||||
p1.combo_active.setCurrentIndex(active)
|
||||
pump(30)
|
||||
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, \
|
||||
|
||||
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, \
|
||||
"fine translate nudge moved shift_x by exactly one fine step"
|
||||
|
||||
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, \
|
||||
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, \
|
||||
"coarse translate nudge uses the multiplier"
|
||||
|
||||
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"
|
||||
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"
|
||||
|
||||
# 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 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], \
|
||||
"a real Right-arrow key event nudged shift_x"
|
||||
|
||||
|
||||
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"
|
||||
# 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_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), \
|
||||
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), \
|
||||
f"every non-reference angle got a fit ({label})"
|
||||
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"
|
||||
|
||||
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"
|
||||
|
||||
|
||||
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"
|
||||
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"
|
||||
ctx.sidecar = sidecar
|
||||
ctx.sidecar_raw = raw = json.loads(sidecar.read_text())
|
||||
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"
|
||||
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}"
|
||||
|
||||
|
||||
def test_stale_schema_sidecar_ignored(ctx):
|
||||
@@ -424,56 +687,221 @@ 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
|
||||
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)
|
||||
|
||||
saved = json.loads(ctx.sidecar.read_text())["per_angle"][str(active)]
|
||||
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._manual_align_dialog is None, \
|
||||
"manual dialog force-closed by a reload"
|
||||
assert win._align_wizard is None, "no wizard left open across a reload"
|
||||
assert win._alignment_result is not None, \
|
||||
"reload restores the saved manual alignment automatically"
|
||||
"reload restores the saved alignment automatically"
|
||||
assert abs(win._alignment_result.per_angle[active].rotation_deg
|
||||
- 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"
|
||||
- saved["rotation_deg"]) < 1e-9, \
|
||||
"restored rotation 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)
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
"""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 CH1_IDX, 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 agree bit-for-bit with an independent
|
||||
reference that row-averages the raw waveforms and background-subtracts
|
||||
them, then runs a plain scipy rfft + argmax at the same pad factor --
|
||||
i.e. row-averaging composes correctly with the padded peak search."""
|
||||
import scipy.fft as scipy_fft
|
||||
|
||||
path = tmp_path / "padded_rowavg.sras"
|
||||
gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
|
||||
sras = SrasFile(str(path))
|
||||
spf = sras.samples_per_frame
|
||||
n_rows, n_frames = sras.image_shape(0)
|
||||
n_fft = spf * 40
|
||||
|
||||
raw_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
|
||||
avg_natural = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
|
||||
row_avg_n=3)
|
||||
avg_padded = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True,
|
||||
row_avg_n=3, n_fft=n_fft)
|
||||
|
||||
assert avg_natural.shape == raw_natural.shape == avg_padded.shape
|
||||
assert np.all(np.isfinite(avg_padded))
|
||||
|
||||
weights = compute._row_average_weights(3)
|
||||
freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
|
||||
data = sras.data[0]
|
||||
expected = np.zeros((n_rows, n_frames), dtype=np.float32)
|
||||
for r in range(n_rows):
|
||||
v = np.ones(n_frames, dtype=bool)
|
||||
avg = compute._row_average_waveforms(
|
||||
data[r, CH1_IDX].astype(np.float32), v, weights)
|
||||
avg = avg - sras.background
|
||||
S = scipy_fft.rfft(avg, n=n_fft, axis=-1, workers=1)
|
||||
power = S.real ** 2 + S.imag ** 2
|
||||
power[:, 0] = 0.0
|
||||
expected[r] = freq32[np.argmax(power, axis=1)]
|
||||
|
||||
assert np.array_equal(avg_padded, expected), \
|
||||
"row-averaged waveforms feed the padded peak search identically " \
|
||||
"to an independent reference"
|
||||
|
||||
|
||||
def test_row_average_improves_snr_recovery():
|
||||
"""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(raw, spf)
|
||||
avg_bins = compute._peak_bins(averaged, spf)
|
||||
|
||||
raw_hits = int(np.sum(raw_bins == true_bin))
|
||||
avg_hits = int(np.sum(avg_bins == true_bin))
|
||||
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_MAX", 4)
|
||||
# row_avg_n > 0 halves the effective budget before chunk planning.
|
||||
fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET // 2)
|
||||
assert fft_rows < n_rows, \
|
||||
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"
|
||||
@@ -0,0 +1,951 @@
|
||||
"""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",):
|
||||
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, 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))
|
||||
tail_offset = v2._cache_tail_offset()
|
||||
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
|
||||
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
|
||||
fmt.SFFT_FLAG_BG_SUB, len(entries))
|
||||
for a in entries:
|
||||
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
|
||||
# Windows refuses to truncate a file with a live mapping (write_bytes
|
||||
# opens 'wb'), and every SrasFile holds its waveform memmaps for life —
|
||||
# drop the instance first. The parsed freq arrays are plain copies and
|
||||
# stay usable.
|
||||
del v2
|
||||
head = path.read_bytes()[:tail_offset]
|
||||
path.write_bytes(head + payload)
|
||||
|
||||
v1 = SrasFile(str(path))
|
||||
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))
|
||||
tail_offset = v2._cache_tail_offset()
|
||||
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 1, fmt.CACH_FLAG_FFT)
|
||||
payload += struct.pack(fmt.SFFT_HDR_FMT_V1, fmt.SFFT_MAGIC,
|
||||
fmt.SFFT_FLAG_BG_SUB, len(entries))
|
||||
for a in entries:
|
||||
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
|
||||
# See test_cach_v1_reads_as_natural_resolution: release the memmaps
|
||||
# before write_bytes truncates, or Windows raises EINVAL.
|
||||
del v2
|
||||
head = path.read_bytes()[:tail_offset]
|
||||
path.write_bytes(head + payload)
|
||||
|
||||
v1 = SrasFile(str(path))
|
||||
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"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Min peak frequency floor: serve-time masking, on-disk provenance, batch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _biting_floor(stored: np.ndarray) -> float:
|
||||
"""A floor that zeroes some-but-not-all of *stored*'s positive peaks:
|
||||
the median distinct positive value, so pixels below it get masked and
|
||||
pixels at/above it survive (the mask is a strict <)."""
|
||||
positive = np.unique(stored[stored > 0])
|
||||
assert len(positive) >= 2, "fixture must have varied peak frequencies"
|
||||
return float(positive[len(positive) // 2])
|
||||
|
||||
|
||||
def test_stored_image_served_with_raised_floor_masked(rig, no_fft):
|
||||
"""The core of the 5 m/s bug fix: a floor above the stored one (here 0)
|
||||
is re-applied when the stored image is served — pixels whose stored
|
||||
peak falls below it come back as the 0.0 invalid sentinel, everything
|
||||
else passes through, and no FFT runs. Both directly through
|
||||
cached_rf_image and through compute_rf_image's fast path."""
|
||||
assert rig.sras.precomputed_min_freq_mhz == 0.0, "batched without a floor"
|
||||
stored = rig.sras.precomputed_freq_mhz[0]
|
||||
floor = _biting_floor(stored)
|
||||
expect = np.where(stored < floor, np.float32(0.0), stored)
|
||||
|
||||
img = cached_rf_image(rig.sras, 0, None, apply_bg_sub=True,
|
||||
min_freq_mhz=floor)
|
||||
assert img is not None, "an equal-or-higher floor is servable"
|
||||
assert np.array_equal(img, expect)
|
||||
assert (img == 0).any() and (img > 0).any(), \
|
||||
"the floor bites some pixels but not all"
|
||||
|
||||
via_compute = compute_rf_image(rig.sras, 0, dc_threshold_mv=None,
|
||||
apply_bg_sub=True, min_freq_mhz=floor)
|
||||
assert np.array_equal(via_compute, expect), \
|
||||
"compute_rf_image's fast path applies the same serve-time mask"
|
||||
assert not no_fft, f"serving + masking must not run an FFT: {no_fft}"
|
||||
|
||||
|
||||
def test_min_freq_floor_round_trips_and_gates_serving(tmp_path):
|
||||
"""cache_file records the floor in the SFFT header and the accept rule
|
||||
is asymmetric: an equal-or-higher request is servable, a lower one is
|
||||
refused (the stored search never looked below its floor). 20.1 pins the
|
||||
fixed-point kHz encoding — a float32 header field would read back as
|
||||
20.10000038…, above the requested 20.1, and mismatch forever."""
|
||||
path = tmp_path / "floor_roundtrip.sras"
|
||||
gen.write(path, n_angles=2, seed=31, samples_per_frame=128)
|
||||
floor = 20.1
|
||||
assert cache_file(str(path), "fft", True, min_freq_mhz=floor) == ""
|
||||
|
||||
sras = SrasFile(str(path))
|
||||
assert sras.precomputed_min_freq_mhz == floor, "exact fixed-point round-trip"
|
||||
assert all(((img == 0) | (img >= floor)).all()
|
||||
for img in sras.precomputed_freq_mhz), \
|
||||
"no stored peak below the floor"
|
||||
|
||||
def reasons(f):
|
||||
return compute.cache_mismatch_reasons(
|
||||
sras, n_fft=None, apply_bg_sub=True, row_avg_n=0, min_freq_mhz=f)
|
||||
|
||||
assert reasons(floor) == []
|
||||
assert reasons(floor + 5.0) == [], "a higher request is servable (masked)"
|
||||
low = reasons(0.0)
|
||||
assert low and "min-peak-freq floor" in low[0], \
|
||||
"a lower request cannot be answered by the stored search"
|
||||
assert cached_rf_image(sras, 0, None, apply_bg_sub=True) is None, \
|
||||
"default floor-0 request refused against a floored store"
|
||||
assert cached_rf_image(sras, 0, None, apply_bg_sub=True,
|
||||
min_freq_mhz=floor) is not None
|
||||
|
||||
|
||||
def test_batch_recompute_resolves_not_masks(tmp_path, no_fft):
|
||||
"""Re-batching an already-cached file at a raised floor must run the
|
||||
real FFT and store re-resolved peaks — never let compute_rf_image's
|
||||
fast path serve the file's own stale cache back to it and bake the
|
||||
masked copy in as if it were a recompute (silent, permanent data
|
||||
loss: a masked pixel's true above-floor peak is unrecoverable)."""
|
||||
path = tmp_path / "rebatch.sras"
|
||||
gen.write(path, n_angles=2, seed=32, samples_per_frame=256)
|
||||
assert cache_file(str(path), "fft", True) == ""
|
||||
|
||||
first = SrasFile(str(path))
|
||||
n_angles = first.n_angles
|
||||
# On the header's kHz grid up front (as the spinbox value would be), so
|
||||
# the recorded floor reads back equal — cache_file quantizes whatever it
|
||||
# is given, and this test wants that to be the identity.
|
||||
floor = round(_biting_floor(first.precomputed_freq_mhz[0]) * 1000) / 1000.0
|
||||
bites = [img < floor for img in first.precomputed_freq_mhz]
|
||||
assert bites[0].any(), "the floor must actually bite this fixture"
|
||||
# What a real floored compute gives, from a view blinded to the cache.
|
||||
first.precomputed_freq_mhz = [None] * n_angles
|
||||
expected = [compute_rf_image(first, a, dc_threshold_mv=None,
|
||||
apply_bg_sub=True, min_freq_mhz=floor)
|
||||
for a in range(n_angles)]
|
||||
del first # release memmaps before cache_file rewrites the tail
|
||||
|
||||
no_fft.clear()
|
||||
assert cache_file(str(path), "fft", True, min_freq_mhz=floor) == ""
|
||||
assert no_fft, "the re-batch ran a real FFT"
|
||||
|
||||
after = SrasFile(str(path))
|
||||
assert after.precomputed_min_freq_mhz == floor
|
||||
for a in range(n_angles):
|
||||
assert np.array_equal(after.precomputed_freq_mhz[a], expected[a]), \
|
||||
f"angle {a}: stored image is a real floored recompute"
|
||||
assert (after.precomputed_freq_mhz[a][bites[a]] >= floor).all(), \
|
||||
f"angle {a}: bitten pixels re-resolved above the floor, not zeroed"
|
||||
|
||||
|
||||
def test_min_freq_carries_forward_through_dc_write(tmp_path):
|
||||
"""A later DC-only write must leave the FFT block's recorded floor
|
||||
untouched, like row_avg_n and pad_factor."""
|
||||
path = tmp_path / "floor_carry.sras"
|
||||
gen.write(path, n_angles=2, seed=33, samples_per_frame=128)
|
||||
assert cache_file(str(path), "fft", True, min_freq_mhz=75.0) == ""
|
||||
assert cache_file(str(path), "dc", True) == ""
|
||||
after = SrasFile(str(path))
|
||||
assert after.precomputed_min_freq_mhz == 75.0, \
|
||||
"floor survives a DC-only write"
|
||||
|
||||
|
||||
def test_cach_v3_backward_compat_defaults_floor_zero(tmp_path):
|
||||
"""A v3 CACH tail predates the min peak frequency floor entirely (no
|
||||
min_freq_khz field) — readers must still parse it in full, treating it
|
||||
as floor 0: servable as-is at floor 0, and serve-maskable at any higher
|
||||
one. This is what protects existing real-world v7 caches from silently
|
||||
becoming unusable after the v4 bump ships."""
|
||||
path = tmp_path / "v3_floor.sras"
|
||||
gen.write(path, n_angles=2, seed=34, samples_per_frame=128)
|
||||
assert cache_file(str(path), "fft", True) == ""
|
||||
|
||||
# Rewrite the tail as a genuine CACH v3 block (no min_freq field).
|
||||
v4 = SrasFile(str(path))
|
||||
freq, entries = v4.precomputed_freq_mhz, list(range(v4.n_angles))
|
||||
tail_offset = v4._cache_tail_offset()
|
||||
payload = struct.pack(fmt.CACH_HDR_FMT, fmt.CACH_MAGIC, 3, fmt.CACH_FLAG_FFT)
|
||||
payload += struct.pack(fmt.SFFT_HDR_FMT_V3, fmt.SFFT_MAGIC,
|
||||
fmt.SFFT_FLAG_BG_SUB, len(entries), 0, 1)
|
||||
for a in entries:
|
||||
payload += struct.pack(">H", a) + freq[a].astype(">f4").tobytes()
|
||||
# See test_cach_v1_reads_as_natural_resolution: release the memmaps
|
||||
# before write_bytes truncates, or Windows raises EINVAL.
|
||||
del v4
|
||||
head = path.read_bytes()[:tail_offset]
|
||||
path.write_bytes(head + payload)
|
||||
|
||||
v3 = SrasFile(str(path))
|
||||
assert v3.precomputed_min_freq_mhz == 0.0
|
||||
assert v3.precomputed_pad_factor == 1 and v3.precomputed_bg_sub is True
|
||||
assert all(np.array_equal(v3.precomputed_freq_mhz[a], freq[a])
|
||||
for a in entries), "v3 images read back unchanged"
|
||||
assert cached_rf_image(v3, 0, None, apply_bg_sub=True) is not None
|
||||
floor = _biting_floor(freq[0])
|
||||
served = cached_rf_image(v3, 0, None, apply_bg_sub=True,
|
||||
min_freq_mhz=floor)
|
||||
assert served is not None
|
||||
assert np.array_equal(served,
|
||||
np.where(freq[0] < floor, np.float32(0.0), freq[0]))
|
||||
|
||||
|
||||
def test_min_freq_validation(tmp_path):
|
||||
path = tmp_path / "floor_bad.sras"
|
||||
gen.write(path, n_angles=1, seed=35, samples_per_frame=64)
|
||||
assert cache_file(str(path), "fft", True, min_freq_mhz=-1.0), \
|
||||
"negative floor must be an error, not a write"
|
||||
assert cache_file(str(path), "fft", True, min_freq_mhz=float("nan")), \
|
||||
"NaN floor must be an error, not a write"
|
||||
sras = SrasFile(str(path))
|
||||
with pytest.raises(ValueError):
|
||||
sras.write_v7_cache(new_min_freq_mhz=-0.5)
|
||||
|
||||
|
||||
def test_viewer_reapplies_floor_to_stored_images_without_computing(
|
||||
rig, no_fft, monkeypatch):
|
||||
"""The session-cache poisoning bug behind the '95 MHz peak but 5 m/s'
|
||||
report: changing 'Min peak freq' against a stored cache used to re-file
|
||||
the identical un-floored image under a key claiming the new floor — the
|
||||
UI looked updated, the pixels weren't. Now the floor really is
|
||||
re-applied on serve (masked, still no compute), and clearing it
|
||||
restores the unmasked image, still without computing."""
|
||||
app = QApplication.instance() or QApplication([]) # noqa: F841
|
||||
win = SrasViewerWindow()
|
||||
win.show()
|
||||
|
||||
dispatched = []
|
||||
original_start = type(win)._start_compute
|
||||
monkeypatch.setattr(type(win), "_start_compute",
|
||||
lambda self: (dispatched.append(self.spin_angle.value()),
|
||||
original_start(self))[1])
|
||||
try:
|
||||
win._load_file(str(rig.path))
|
||||
assert wait_until(lambda: win._sras is not None), "file loaded"
|
||||
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
|
||||
for a in range(rig.n_angles))), \
|
||||
"DC precompute finished"
|
||||
|
||||
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
|
||||
assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3)
|
||||
|
||||
# Round to the spinbox's 3-decimal granularity; the chosen bin value
|
||||
# still survives its own (strict-<) floor after rounding down.
|
||||
floor = round(_biting_floor(rig.fresh[0]), 3)
|
||||
expect = np.where(rig.fresh[0] < floor, np.float32(0.0), rig.fresh[0])
|
||||
|
||||
dispatched.clear()
|
||||
no_fft.clear()
|
||||
win.spin_min_freq_mhz.setValue(floor)
|
||||
win._on_min_freq_changed()
|
||||
pump(120)
|
||||
assert np.allclose(win._current_image, expect, atol=1e-3), \
|
||||
"raised floor re-masks the stored image on serve"
|
||||
assert (win._current_image == 0).any() and (win._current_image > 0).any()
|
||||
assert dispatched == [] and not no_fft, \
|
||||
f"re-masked serve needs no compute (jobs={dispatched}, fft={no_fft})"
|
||||
key = (0, win.spin_threshold_mv.value(), win.spin_min_freq_mhz.value())
|
||||
assert key in win._fft_cache
|
||||
assert np.allclose(win._fft_cache[key], expect, atol=1e-3), \
|
||||
"the session cache holds the value its key claims"
|
||||
|
||||
win.spin_min_freq_mhz.setValue(0.0)
|
||||
win._on_min_freq_changed()
|
||||
pump(120)
|
||||
assert np.allclose(win._current_image, rig.fresh[0], atol=1e-3), \
|
||||
"clearing the floor restores the unmasked stored image"
|
||||
assert dispatched == [] and not no_fft
|
||||
finally:
|
||||
win.close()
|
||||
pump(300)
|
||||
|
||||
|
||||
def test_viewer_batch_fft_records_the_floor(tmp_path, no_fft, monkeypatch):
|
||||
"""Convert → Batch Compute FFT with a floor set: the live spinbox value
|
||||
reaches cache_file, lands in the reloaded file's provenance and the
|
||||
info panel, and the viewer then serves the floored cache without
|
||||
recomputing — the tooltip's promised remedy, end to end."""
|
||||
path = tmp_path / "floor_gui.sras"
|
||||
gen.write(path, n_angles=2, seed=36, samples_per_frame=256)
|
||||
|
||||
app = QApplication.instance() or QApplication([]) # noqa: F841
|
||||
win = SrasViewerWindow()
|
||||
win.show()
|
||||
|
||||
dispatched = []
|
||||
original_start = type(win)._start_compute
|
||||
monkeypatch.setattr(type(win), "_start_compute",
|
||||
lambda self: (dispatched.append(self.spin_angle.value()),
|
||||
original_start(self))[1])
|
||||
try:
|
||||
win._load_file(str(path))
|
||||
assert wait_until(lambda: win._sras is not None), "file loaded"
|
||||
assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
|
||||
for a in range(win._sras.n_angles))), \
|
||||
"DC precompute finished"
|
||||
|
||||
floor = 100.0
|
||||
win.spin_min_freq_mhz.setValue(floor)
|
||||
with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames",
|
||||
return_value=([str(path)], "")):
|
||||
win._on_batch_compute("fft")
|
||||
assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran"
|
||||
assert wait_until(lambda: win._sras is not None
|
||||
and win._sras.version == 7), "file reloaded as v7"
|
||||
pump(200)
|
||||
assert win._sras.precomputed_min_freq_mhz == floor, \
|
||||
"the viewer's floor reached the stored provenance"
|
||||
assert "floor ≥ 100 MHz" in win.lbl_frame_warn.text()
|
||||
|
||||
no_fft.clear()
|
||||
dispatched.clear()
|
||||
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||
assert wait_until(lambda: win._current_ch == CH1_IDX), "CH1 displayed"
|
||||
pump(120)
|
||||
assert dispatched == [] and not no_fft, \
|
||||
f"floored cache serves the view directly (jobs={dispatched}, fft={no_fft})"
|
||||
img = win._current_image
|
||||
assert ((img == 0) | (img >= floor)).all(), \
|
||||
"no displayed peak below the floor"
|
||||
finally:
|
||||
win.close()
|
||||
pump(300)
|
||||
|
||||
|
||||
def test_viewer_batch_row_average_dispatch(tmp_path, monkeypatch, no_fft):
|
||||
"""Driving the new 'Batch Compute Row-Averaged FFT and Store' action
|
||||
end-to-end through the real menu handler: dialog values reach the
|
||||
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)
|
||||
|
||||
|
||||
def test_compute_worker_fallback_honors_row_avg_n(tmp_path, monkeypatch, no_fft):
|
||||
"""The other half of the row-averaged recompute bug: _stored_fft_image
|
||||
isn't the only path that can serve a view. When the DC4 mask isn't known
|
||||
yet without I/O (e.g. background DC precompute hasn't reached this angle),
|
||||
_stored_fft_image's allow_dc_recompute=False guard bails and
|
||||
_refresh_display falls through to _start_compute's ComputeWorker instead.
|
||||
That worker must still ask compute_rf_image for the file's own
|
||||
row_avg_n -- not silently default to 0 -- so its own internal cache
|
||||
fast path also serves the stored row-averaged image rather than running
|
||||
a real, non-averaged FFT."""
|
||||
path = tmp_path / "rowavg_fallback.sras"
|
||||
gen.write(path, n_angles=2, seed=33, samples_per_frame=128)
|
||||
# Only fft_rowavg -- deliberately no DC block, exactly what "Batch
|
||||
# Compute Row-Averaged FFT and Store" leaves on disk by itself.
|
||||
err = cache_file(str(path), "fft_rowavg", True, dc_threshold_mv=-1e9, row_avg_n=5)
|
||||
assert err == "", err
|
||||
|
||||
app = QApplication.instance() or QApplication([]) # noqa: F841
|
||||
win = SrasViewerWindow()
|
||||
win.show()
|
||||
|
||||
# Keep _dc_cache empty so _stored_fft_image can't resolve a mask without
|
||||
# I/O and _refresh_display must fall through to _start_compute.
|
||||
monkeypatch.setattr(type(win), "_start_dc_precompute", lambda self: None)
|
||||
|
||||
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 win._sras.precomputed_row_avg_n == 5
|
||||
# The initial default view (CH4, DC) has nothing cached either, so
|
||||
# it dispatches its own one-off DC compute for angle 0 on load --
|
||||
# unrelated to this bug. Let that settle, then clear it so no DC4 is
|
||||
# available for any angle, simulating "background DC precompute
|
||||
# hasn't reached this angle yet".
|
||||
assert wait_until(lambda: not win._job_running("compute")), "initial DC view settled"
|
||||
dispatched.clear()
|
||||
win._dc_cache.clear()
|
||||
|
||||
no_fft.clear()
|
||||
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||
assert wait_until(lambda: win._current_ch == CH1_IDX
|
||||
and not win._job_running("compute")), "CH1 displayed"
|
||||
|
||||
assert dispatched == [0], \
|
||||
"with no DC cache available yet, the fallback compute must run"
|
||||
assert not no_fft, \
|
||||
f"the fallback's own compute_rf_image call must still hit the " \
|
||||
f"stored row-averaged cache internally: {no_fft}"
|
||||
expected = compute.cached_rf_image(
|
||||
win._sras, 0, dc_threshold_mv=win.spin_threshold_mv.value(),
|
||||
apply_bg_sub=win.chk_bg_sub.isChecked(), row_avg_n=5)
|
||||
assert expected is not None
|
||||
assert np.allclose(win._current_image, expected, atol=1e-3), \
|
||||
"the displayed image is the stored row-averaged one, not a raw recompute"
|
||||
finally:
|
||||
win.close()
|
||||
pump(300)
|
||||
+38
-35
@@ -1,8 +1,10 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark the FFT peak-search path: exact vs zoom, serial vs pooled.
|
||||
"""Benchmark the FFT peak-search path: serial vs pooled.
|
||||
|
||||
Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
|
||||
cores), and verifies every variant against the exact reference image.
|
||||
cores), the estimated per-thread resident pyFFTW plan footprint at each pad
|
||||
factor (see sras_compute._fft_block_for), and verifies pooled output against
|
||||
the serial reference.
|
||||
|
||||
Usage:
|
||||
python tools/bench_fft.py # synthetic, pads 1/8/40
|
||||
@@ -22,10 +24,10 @@ import numpy as np
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
import sras_compute as compute # noqa: E402
|
||||
from sras_compute import compute_rf_image, set_fft_backend # noqa: E402
|
||||
from sras_format import SrasFile # noqa: E402
|
||||
import tools.make_test_sras as gen # noqa: E402
|
||||
from tools.check_equivalence import row_slice # noqa: E402
|
||||
from sras_compute import compute_rf_image # noqa: E402
|
||||
from sras_format import SrasFile # noqa: E402
|
||||
import tools.make_test_sras as gen # noqa: E402
|
||||
from tools.check_equivalence import row_slice # noqa: E402
|
||||
|
||||
|
||||
def _timed(fn):
|
||||
@@ -38,35 +40,42 @@ def _timed(fn):
|
||||
return out, wall, cpu / max(wall, 1e-9)
|
||||
|
||||
|
||||
def bench(sras, pads, backends):
|
||||
def _plan_mb(spf: int, n_len: int, n_workers: int) -> float:
|
||||
"""Estimated resident pyFFTW plan-buffer footprint across the whole
|
||||
pool at this transform length (see sras_compute._fft_block_for)."""
|
||||
block = compute._fft_block_for(spf, n_len)
|
||||
bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
|
||||
return block * bytes_per_wf * n_workers / (1024 * 1024)
|
||||
|
||||
|
||||
def bench(sras, pads):
|
||||
n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
|
||||
for a in range(sras.n_angles))
|
||||
spf = sras.samples_per_frame
|
||||
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)")
|
||||
print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} "
|
||||
f"{'wf/s':>10} {'util':>6} match")
|
||||
n_workers = compute._MAX_WORKERS
|
||||
print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), "
|
||||
f"{n_workers} workers")
|
||||
print(f"{'pad':>4} {'variant':>8} {'wall':>9} {'wf/s':>10} {'util':>6} "
|
||||
f"{'plan MB':>9} match")
|
||||
|
||||
for pad in pads:
|
||||
n_fft = spf * pad if pad > 1 else None
|
||||
for backend in backends:
|
||||
set_fft_backend(backend)
|
||||
n_len = n_fft if n_fft is not None else spf
|
||||
plan_mb = _plan_mb(spf, n_len, n_workers)
|
||||
|
||||
def run(**kw):
|
||||
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
||||
apply_bg_sub=True, n_fft=n_fft, **kw)
|
||||
for a in range(sras.n_angles)]
|
||||
return np.concatenate([i.ravel() for i in imgs])
|
||||
def run(**kw):
|
||||
imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
||||
apply_bg_sub=True, n_fft=n_fft, **kw)
|
||||
for a in range(sras.n_angles)]
|
||||
return np.concatenate([i.ravel() for i in imgs])
|
||||
|
||||
ref, wall, util = _timed(lambda: run(exact=True))
|
||||
rows = [("exact(serial)", ref, wall, util, True)]
|
||||
for label, kw in (("zoom(serial)", dict(max_workers=1)),
|
||||
("zoom(pool)", {})):
|
||||
img, wall, util = _timed(lambda: run(**kw))
|
||||
rows.append((label, img, wall, util, bool(np.array_equal(img, ref))))
|
||||
for label, img, wall, util, ok in rows:
|
||||
print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s "
|
||||
f"{n_wf / wall:>10.0f} {util:>5.1f}x "
|
||||
f"{'OK' if ok else 'MISMATCH'}")
|
||||
ref, wall, util = _timed(lambda: run(max_workers=1))
|
||||
rows = [("serial", ref, wall, util, True)]
|
||||
img, wall, util = _timed(lambda: run())
|
||||
rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref))))
|
||||
for label, img, wall, util, ok in rows:
|
||||
print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} "
|
||||
f"{util:>5.1f}x {plan_mb:>8.1f} {'OK' if ok else 'MISMATCH'}")
|
||||
|
||||
|
||||
def main():
|
||||
@@ -76,30 +85,24 @@ def main():
|
||||
p.add_argument("--spf", type=int, default=2500)
|
||||
p.add_argument("--rows", type=int, default=8)
|
||||
p.add_argument("--frames", type=int, default=1024)
|
||||
p.add_argument("--backends", default=None,
|
||||
help="comma-separated (default: scipy,pyfftw if available)")
|
||||
p.add_argument("--real", help="path to a real .sras file")
|
||||
p.add_argument("--real-rows", type=int, default=32,
|
||||
help="rows of angle 0 to use from the real file")
|
||||
args = p.parse_args()
|
||||
|
||||
pads = [int(x) for x in args.pads.split(",")]
|
||||
if args.backends:
|
||||
backends = args.backends.split(",")
|
||||
else:
|
||||
backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])
|
||||
|
||||
if args.real:
|
||||
sras = row_slice(SrasFile(args.real), 0, args.real_rows)
|
||||
sras.data = [sras.data[0]]
|
||||
sras.n_angles = 1
|
||||
bench(sras, pads, backends)
|
||||
bench(sras, pads)
|
||||
else:
|
||||
with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
|
||||
path = Path(tmp) / "bench.sras"
|
||||
gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
|
||||
geometry=[(args.rows, args.frames)])
|
||||
bench(SrasFile(str(path)), pads, backends)
|
||||
bench(SrasFile(str(path)), pads)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -261,15 +261,10 @@ def write_rotating(path: Path, n_angles: int = 5, samples_per_frame: int = 4,
|
||||
"y_starts": y_starts, "dx_mm": _ROT_DX_MM, "dy_mm": _ROT_DY_MM}
|
||||
|
||||
|
||||
|
||||
|
||||
def write_legacy(path: Path, version: int = 4, n_angles: int = 2,
|
||||
n_rows: int = 4, n_frames: int = 10,
|
||||
samples_per_frame: int = 32, seed: int = 0) -> dict:
|
||||
"""Write a v2/v3/v4 file: uniform geometry, one flat waveform block.
|
||||
|
||||
Used to exercise sras_average.py, which only handles the legacy formats.
|
||||
"""
|
||||
"""Write a v2/v3/v4 file: uniform geometry, one flat waveform block."""
|
||||
rng = np.random.default_rng(seed)
|
||||
n_ch, bps = 3, 1
|
||||
|
||||
|
||||
Reference in New Issue
Block a user