Add aligned/cropped .sras export and the compute pieces behind it
The alignment machinery never modified a scan: it produced an
AlignmentResult and every consumer resampled on the fly. That is right
for a viewer but means the aligned stack cannot leave the process, so no
other tool can read it and reopening the scan redoes the registration.
sras_align_export.write_aligned_sras bakes an alignment into a new v6
file: each angle is gathered onto the shared canvas by nearest neighbour,
so every output angle shares one grid and the file opens already aligned.
Registering the export against itself returns identity, which is the test
that pins the whole index chain.
Three details that are easy to get wrong and are now covered:
* Rounding must be floor(x + 0.5), not np.rint. scipy's order=0 rounds
halves away from zero, and the canvas is snapped to the reference's
pixel grid, so exact halves are common rather than hypothetical.
* 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
everywhere the Aligned View shows padding.
* ...but with a tolerance, because the mm-space affine chain lands an
exactly-integer transform a few times 1e-13 off. A bare >= 0 drops
the *reference* angle's entire first row and last column.
Padding is the per-channel ADC code nearest 0 mV, not zero: zero ADC
decodes to ~+100 mV on real calibration and would masquerade as sample.
Source rows are served from sliding in-RAM bands. A rotated angle maps
one output row to a diagonal across the source, so indexing a memmap in
output order refaults nearly the whole angle per row — terabytes of
paging for a gigabyte of data.
Also in compute: crop_alignment_result (a crop is pure index translation,
so it folds into the affine's offset rather than becoming a second
transform), overlap_stats, largest_rect_at_least for a crop that stays
inside the overlap region, and seed/sign/refine knobs on
register_angle_to_reference whose defaults leave existing behaviour and
tests byte-identical.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,477 @@
|
|||||||
|
#!/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 _src_coords(t, rows, n_cols: int) -> tuple[np.ndarray, np.ndarray]:
|
||||||
|
"""Fractional source (row, col) coordinates for whole output rows.
|
||||||
|
|
||||||
|
*rows* may be a scalar row index or an array of them; the result is shaped
|
||||||
|
(len(rows), n_cols), or (n_cols,) for a scalar.
|
||||||
|
"""
|
||||||
|
cols = np.arange(n_cols, dtype=np.float64)
|
||||||
|
r = np.atleast_1d(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]
|
||||||
|
if np.isscalar(rows) or np.asarray(rows).ndim == 0:
|
||||||
|
return sr[0], sc[0]
|
||||||
|
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 (canvas_for_params snap=True), 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)
|
||||||
|
|
||||||
|
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 = valid_px[a] / (n_rows * n_cols) if n_rows and n_cols else 0.0
|
||||||
|
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 getattr(sras, "scan_aborted", False):
|
||||||
|
warnings.append(
|
||||||
|
f"Input scan was aborted: only its {n_angles} complete angle(s) "
|
||||||
|
f"are exported.")
|
||||||
|
|
||||||
|
return 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)
|
||||||
|
|
||||||
|
|
||||||
|
def _encode_preambles(sras: SrasFile) -> bytes:
|
||||||
|
"""The Preamble Blocks section for the output file.
|
||||||
|
|
||||||
|
v6/v7 inputs carry the on-disk span verbatim (sras.preambles_raw), which is
|
||||||
|
both cheaper and lossless. Legacy inputs don't keep that span, and a v2 file
|
||||||
|
has no preambles at all — there, empty strings are written. That 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.
|
||||||
|
"""
|
||||||
|
raw = getattr(sras, "preambles_raw", None)
|
||||||
|
if raw is not None:
|
||||||
|
return raw
|
||||||
|
strings = sras.preambles or [""] * sras.n_channels
|
||||||
|
out = bytearray()
|
||||||
|
for s in strings:
|
||||||
|
encoded = s.encode("utf-8")
|
||||||
|
out += struct.pack(">H", len(encoded)) + encoded
|
||||||
|
return bytes(out)
|
||||||
|
|
||||||
|
|
||||||
|
def _encode_background(sras: SrasFile) -> bytes:
|
||||||
|
"""The Background Block for the output file.
|
||||||
|
|
||||||
|
When the input has none (v2/v3), write 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(sras, "background_raw", None)
|
||||||
|
if raw is not None:
|
||||||
|
return raw
|
||||||
|
if sras.background is None:
|
||||||
|
samples = np.zeros(sras.samples_per_frame, dtype=np.int8)
|
||||||
|
else:
|
||||||
|
samples = np.rint(sras.background).astype(np.int8)
|
||||||
|
return struct.pack(">I", samples.size) + samples.tobytes()
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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) -> 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 hasattr(sras, "x_start_nominal_mm"):
|
||||||
|
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._TOTAL_BYTES_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(_encode_preambles(sras))
|
||||||
|
fout.write(_encode_background(sras))
|
||||||
|
|
||||||
|
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)
|
||||||
|
fout.write(out.tobytes())
|
||||||
|
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
|
||||||
+155
-16
@@ -1235,13 +1235,22 @@ def default_max_workers() -> int:
|
|||||||
|
|
||||||
|
|
||||||
def _rotation_candidates(nominal_deg: float, search_deg: float,
|
def _rotation_candidates(nominal_deg: float, search_deg: float,
|
||||||
step_deg: float) -> list[float]:
|
step_deg: float,
|
||||||
"""Coarse rotation candidates: a window around *both* signs of the stage's
|
signs: tuple[int, ...] = (-1, 1)) -> list[float]:
|
||||||
reported angle change. Scoring both is what makes the stage's sign
|
"""Coarse rotation candidates: a window around each requested sign of the
|
||||||
convention a non-issue — the images decide which way the stage turns, and
|
stage's reported angle change. Scoring both signs (the default) is what
|
||||||
a file whose stage reports the opposite sense registers just as well."""
|
makes the stage's sign convention a non-issue — the images decide which way
|
||||||
|
the stage turns, and a file whose stage reports the opposite sense
|
||||||
|
registers just as well.
|
||||||
|
|
||||||
|
*signs* exists only so a user who has established which way their own stage
|
||||||
|
turns can halve the coarse sweep from the alignment wizard. It is an
|
||||||
|
override, not an inference: nothing in the file says which sign is right.
|
||||||
|
"""
|
||||||
out: list[float] = []
|
out: list[float] = []
|
||||||
for center in (-nominal_deg, nominal_deg):
|
step_deg = max(abs(step_deg), 1e-9) # a 0 step would divide by zero below
|
||||||
|
for sign in signs:
|
||||||
|
center = sign * nominal_deg
|
||||||
k = int(np.floor(search_deg / step_deg))
|
k = int(np.floor(search_deg / step_deg))
|
||||||
for i in range(-k, k + 1):
|
for i in range(-k, k + 1):
|
||||||
out.append(center + i * step_deg)
|
out.append(center + i * step_deg)
|
||||||
@@ -1351,15 +1360,18 @@ def register_angle_to_reference(
|
|||||||
dc_threshold_mv: float = 0.0,
|
dc_threshold_mv: float = 0.0,
|
||||||
sources: tuple[str, ...] = ("signal", "mask"),
|
sources: tuple[str, ...] = ("signal", "mask"),
|
||||||
coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
|
coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
|
||||||
search_deg: float = 6.0, coarse_step_deg: float = 2.0) -> RigidFit:
|
search_deg: float = 6.0, coarse_step_deg: float = 2.0,
|
||||||
|
seed_deg: float | None = None,
|
||||||
|
seed_signs: tuple[int, ...] = (-1, 1),
|
||||||
|
refine: bool = True) -> RigidFit:
|
||||||
"""Rigid (rotation + translation, never scale) fit of *angle_idx* onto
|
"""Rigid (rotation + translation, never scale) fit of *angle_idx* onto
|
||||||
*ref_angle_idx*, found entirely by cross-correlating image content.
|
*ref_angle_idx*, found entirely by cross-correlating image content.
|
||||||
|
|
||||||
Two stages:
|
Two stages:
|
||||||
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around both
|
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
|
||||||
signs of the stage's reported angle change (see _rotation_candidates),
|
requested signs of the stage's reported angle change (see
|
||||||
for each requested source image, scored by _overlap_ncc. Only has to
|
_rotation_candidates), for each requested source image, scored by
|
||||||
pick the right basin.
|
_overlap_ncc. Only has to pick the right basin.
|
||||||
2. Hill-climbing refinement of the winning (source, rotation) at
|
2. Hill-climbing refinement of the winning (source, rotation) at
|
||||||
*fine_dim* with sub-pixel translation folded into every score
|
*fine_dim* with sub-pixel translation folded into every score
|
||||||
(_refine_rotation, _score_rotation), down to 0.05°.
|
(_refine_rotation, _score_rotation), down to 0.05°.
|
||||||
@@ -1368,6 +1380,22 @@ def register_angle_to_reference(
|
|||||||
local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
|
local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
|
||||||
caller can surface so a bad scan is visible rather than silently fused in.
|
caller can surface so a bad scan is visible rather than silently fused in.
|
||||||
Returns identity for the reference angle itself.
|
Returns identity for the reference angle itself.
|
||||||
|
|
||||||
|
The last three arguments only exist to let the alignment wizard expose the
|
||||||
|
rotation search; every default reproduces the search this function has
|
||||||
|
always done.
|
||||||
|
|
||||||
|
*seed_deg* replaces the stage's reported angle change as the center of the
|
||||||
|
coarse sweep — the pre-rotation the search starts from. None (the default)
|
||||||
|
means the stage angle from the file's Angle Table, which is nearly always
|
||||||
|
what you want: it puts the sweep within a couple of degrees of the answer.
|
||||||
|
Pass 0.0 to search around no rotation at all, which is the honest choice for
|
||||||
|
a file whose stage angles are known to be wrong.
|
||||||
|
|
||||||
|
*refine=False* stops after the coarse sweep, leaving rotation on the
|
||||||
|
*coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
|
||||||
|
it pins rotation to exactly the seed and searches translation only — for a
|
||||||
|
scan whose stage angles are trusted more than its image content.
|
||||||
"""
|
"""
|
||||||
if angle_idx == ref_angle_idx:
|
if angle_idx == ref_angle_idx:
|
||||||
return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
|
return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
|
||||||
@@ -1377,8 +1405,10 @@ def register_angle_to_reference(
|
|||||||
if not candidates:
|
if not candidates:
|
||||||
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
|
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
|
||||||
|
|
||||||
nominal = nominal_delta_deg(sras, angle_idx, ref_angle_idx)
|
nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
|
||||||
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg)
|
if seed_deg is None else float(seed_deg))
|
||||||
|
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
|
||||||
|
seed_signs)
|
||||||
|
|
||||||
# ---- Stage 1: coarse sweep, every source ------------------------------
|
# ---- Stage 1: coarse sweep, every source ------------------------------
|
||||||
pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
|
pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
|
||||||
@@ -1407,9 +1437,10 @@ def register_angle_to_reference(
|
|||||||
theta = best[1]
|
theta = best[1]
|
||||||
score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
|
score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
|
||||||
theta, subpixel=True)
|
theta, subpixel=True)
|
||||||
theta, score, shift = _refine_rotation(
|
if refine:
|
||||||
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
|
theta, score, shift = _refine_rotation(
|
||||||
step_deg=coarse_step_deg)
|
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
|
||||||
|
step_deg=coarse_step_deg)
|
||||||
|
|
||||||
dr, dc = shift
|
dr, dc = shift
|
||||||
return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
|
return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
|
||||||
@@ -1540,6 +1571,114 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
|
|||||||
pitch[0], pitch[1], canvas_origin_mm, per_angle)
|
pitch[0], pitch[1], canvas_origin_mm, per_angle)
|
||||||
|
|
||||||
|
|
||||||
|
def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
|
||||||
|
n_rows: int, n_cols: int) -> AlignmentResult:
|
||||||
|
"""The same alignment restricted to a rectangular window of its canvas —
|
||||||
|
canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).
|
||||||
|
|
||||||
|
Cropping folds into each angle's existing affine instead of becoming a
|
||||||
|
second transform, because a canvas crop is *pure index translation*. From
|
||||||
|
_affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
|
||||||
|
rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
|
||||||
|
origin, so:
|
||||||
|
|
||||||
|
matrix @ [r', c'] + (offset + matrix @ [row0, col0])
|
||||||
|
== matrix @ [r' + row0, c' + col0] + offset
|
||||||
|
|
||||||
|
i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
|
||||||
|
mapping at the shifted indices, exactly. That equality is what lets
|
||||||
|
apply_alignment, reproject_mask and the aligned .sras exporter all keep
|
||||||
|
working on a cropped result with no special-casing — and it is why the crop
|
||||||
|
preview a user approves is guaranteed to be the same pixels the exporter
|
||||||
|
writes.
|
||||||
|
|
||||||
|
Bounds are the caller's responsibility (the wizard's ROI page clamps to the
|
||||||
|
canvas): an out-of-range window is geometrically well-defined here and
|
||||||
|
simply resamples padding.
|
||||||
|
"""
|
||||||
|
if n_rows <= 0 or n_cols <= 0:
|
||||||
|
raise ValueError(f"empty crop: {n_rows} x {n_cols}")
|
||||||
|
|
||||||
|
delta = np.array([float(row0), float(col0)])
|
||||||
|
per_angle = {
|
||||||
|
a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
|
||||||
|
t.offset + t.matrix @ delta, t.score, t.source)
|
||||||
|
for a, t in result.per_angle.items()
|
||||||
|
}
|
||||||
|
origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
|
||||||
|
result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
|
||||||
|
return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
|
||||||
|
(int(n_rows), int(n_cols)),
|
||||||
|
result.canvas_dx_mm, result.canvas_dy_mm,
|
||||||
|
origin, per_angle)
|
||||||
|
|
||||||
|
|
||||||
|
def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
|
||||||
|
"""Summarize a per-pixel "how many angles cover this pixel" image — the
|
||||||
|
number the alignment wizard's mask-stack view is colored by.
|
||||||
|
|
||||||
|
Separated from the drawing code because it is the actual judgement the user
|
||||||
|
makes on that screen ("do the angles land on top of each other?"), and a
|
||||||
|
plain array-in/dict-out function can be tested without Qt.
|
||||||
|
"""
|
||||||
|
counts = np.asarray(counts)
|
||||||
|
union = int(np.count_nonzero(counts))
|
||||||
|
full = int(np.count_nonzero(counts >= n_angles))
|
||||||
|
return {
|
||||||
|
"union_px": union,
|
||||||
|
"full_px": full,
|
||||||
|
"full_frac": (full / union) if union else 0.0,
|
||||||
|
"mean_count": float(counts[counts > 0].mean()) if union else 0.0,
|
||||||
|
"max_count": int(counts.max()) if counts.size else 0,
|
||||||
|
"empty": union == 0,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def largest_rect_at_least(counts: np.ndarray, min_count: int
|
||||||
|
) -> tuple[int, int, int, int] | None:
|
||||||
|
"""Largest axis-aligned rectangle whose every pixel has counts >= min_count,
|
||||||
|
as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.
|
||||||
|
|
||||||
|
Backs the wizard's "fit to full overlap" button. A *bounding box* of the
|
||||||
|
qualifying pixels would be the obvious thing and is wrong: the full-overlap
|
||||||
|
region of several rotated scans is roughly a disc, whose bounding box has
|
||||||
|
corners no angle covers at all. Offering that as the crop would hand the
|
||||||
|
user padding they explicitly asked to avoid, so this finds a rectangle that
|
||||||
|
is entirely inside the region.
|
||||||
|
|
||||||
|
Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
|
||||||
|
preview-sized array, so it is instant at interactive rates.
|
||||||
|
"""
|
||||||
|
good = np.asarray(counts) >= min_count
|
||||||
|
if not good.any():
|
||||||
|
return None
|
||||||
|
|
||||||
|
n_rows, n_cols = good.shape
|
||||||
|
best = (0, 0, 0, 0) # area, row0, col0, ...
|
||||||
|
best_area = 0
|
||||||
|
heights = np.zeros(n_cols, dtype=np.int64)
|
||||||
|
|
||||||
|
for r in range(n_rows):
|
||||||
|
heights = np.where(good[r], heights + 1, 0)
|
||||||
|
# Sentinel column of height 0 flushes the stack at the end of the row.
|
||||||
|
stack: list[tuple[int, int]] = [] # (start col, height)
|
||||||
|
for c in range(n_cols + 1):
|
||||||
|
h = int(heights[c]) if c < n_cols else 0
|
||||||
|
start = c
|
||||||
|
while stack and stack[-1][1] >= h:
|
||||||
|
s, sh = stack.pop()
|
||||||
|
area = sh * (c - s)
|
||||||
|
if area > best_area:
|
||||||
|
best_area = area
|
||||||
|
best = (s, sh, c - s, r)
|
||||||
|
start = s
|
||||||
|
if h:
|
||||||
|
stack.append((start, h))
|
||||||
|
|
||||||
|
col0, height, width, row_end = best
|
||||||
|
return (row_end - height + 1, col0, height, width)
|
||||||
|
|
||||||
|
|
||||||
def _parallel_map(fn, items, n_workers: int) -> list:
|
def _parallel_map(fn, items, n_workers: int) -> list:
|
||||||
"""fn over items, in order, threaded when it pays."""
|
"""fn over items, in order, threaded when it pays."""
|
||||||
items = list(items)
|
items = list(items)
|
||||||
|
|||||||
+78
-10
@@ -1,7 +1,9 @@
|
|||||||
"""Matplotlib canvases and the ROI primitive."""
|
"""Matplotlib canvases and the ROI primitive."""
|
||||||
|
|
||||||
|
import matplotlib as mpl
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
|
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
|
||||||
|
from matplotlib.colors import BoundaryNorm, ListedColormap
|
||||||
from matplotlib.figure import Figure
|
from matplotlib.figure import Figure
|
||||||
from matplotlib.patches import Polygon
|
from matplotlib.patches import Polygon
|
||||||
from matplotlib.path import Path as MplPath
|
from matplotlib.path import Path as MplPath
|
||||||
@@ -100,7 +102,15 @@ class ImageCanvas(FigureCanvasQTAgg):
|
|||||||
_HANDLE_PX = 12
|
_HANDLE_PX = 12
|
||||||
_CLICK_THRESH_PX = 4 # releases within this of press count as a click
|
_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)
|
fig = Figure(figsize=(7, 5), tight_layout=True)
|
||||||
self.ax = fig.add_subplot(111)
|
self.ax = fig.add_subplot(111)
|
||||||
super().__init__(fig)
|
super().__init__(fig)
|
||||||
@@ -108,6 +118,7 @@ class ImageCanvas(FigureCanvasQTAgg):
|
|||||||
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
|
self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
|
||||||
self._extent = None
|
self._extent = None
|
||||||
self._img_shape = None
|
self._img_shape = None
|
||||||
|
self._rect_only = rect_only
|
||||||
|
|
||||||
# ROI state
|
# ROI state
|
||||||
self._roi: RoiQuad | None = None
|
self._roi: RoiQuad | None = None
|
||||||
@@ -132,9 +143,12 @@ class ImageCanvas(FigureCanvasQTAgg):
|
|||||||
# Public API
|
# 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,
|
vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
|
||||||
colorbar_label: str = ""):
|
colorbar_label: str = "", cb_ticks=None):
|
||||||
|
"""*cmap* may be a name or a Colormap instance; *cb_ticks* pins the
|
||||||
|
colorbar's ticks, which the wizard's integer overlap-count view needs so
|
||||||
|
each band reads as a whole number of angles rather than a shade."""
|
||||||
self.figure.clf()
|
self.figure.clf()
|
||||||
self.ax = self.figure.add_subplot(111)
|
self.ax = self.figure.add_subplot(111)
|
||||||
# Patches and lines are destroyed by figure.clf(); drop stale refs.
|
# Patches and lines are destroyed by figure.clf(); drop stale refs.
|
||||||
@@ -148,7 +162,8 @@ class ImageCanvas(FigureCanvasQTAgg):
|
|||||||
extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
|
extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
|
||||||
interpolation="nearest",
|
interpolation="nearest",
|
||||||
)
|
)
|
||||||
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
|
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
|
||||||
|
ticks=cb_ticks)
|
||||||
if colorbar_label:
|
if colorbar_label:
|
||||||
cb.set_label(colorbar_label)
|
cb.set_label(colorbar_label)
|
||||||
|
|
||||||
@@ -302,10 +317,28 @@ class ImageCanvas(FigureCanvasQTAgg):
|
|||||||
self._roi._pts = self._snapshot.corners() + delta
|
self._roi._pts = self._snapshot.corners() + delta
|
||||||
elif self._state == self._DRAG_CORNER:
|
elif self._state == self._DRAG_CORNER:
|
||||||
self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]
|
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_roi()
|
||||||
self.draw_idle()
|
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):
|
def _on_release(self, event):
|
||||||
if event.button != 1 and self._press_button != 1:
|
if event.button != 1 and self._press_button != 1:
|
||||||
return
|
return
|
||||||
@@ -475,14 +508,20 @@ class WaveformCanvas(FigureCanvasQTAgg):
|
|||||||
self.draw()
|
self.draw()
|
||||||
|
|
||||||
|
|
||||||
class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
|
class AlignOverlayCanvas(FigureCanvasQTAgg):
|
||||||
"""Renders ManualAlignmentDialog's multi-angle mask overlay and turns
|
"""Renders the alignment wizard's multi-angle mask views and turns keyboard
|
||||||
keyboard input into translate/rotate nudge requests for whichever angle
|
input into translate/rotate nudge requests for whichever angle is active.
|
||||||
the dialog currently has 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
|
A pure input+render widget — it holds no alignment state and never
|
||||||
touches SrasFile itself; ManualAlignmentDialog owns all of that and
|
touches SrasFile itself; the wizard page owns all of that and decides,
|
||||||
decides, from these signals, whether a cheap single-layer refresh or a
|
from these signals, whether a cheap single-layer refresh or a
|
||||||
full preview-canvas rebuild is needed.
|
full preview-canvas rebuild is needed.
|
||||||
|
|
||||||
FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
|
FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
|
||||||
@@ -522,6 +561,35 @@ class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
|
|||||||
self.figure.clf()
|
self.figure.clf()
|
||||||
self.ax = self.figure.add_subplot(111)
|
self.ax = self.figure.add_subplot(111)
|
||||||
self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
|
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)
|
||||||
|
n = max(1, int(n_angles))
|
||||||
|
colors = [(0.0, 0.0, 0.0, 0.0)] # count 0: nothing
|
||||||
|
base = mpl.colormaps["viridis"].resampled(n)
|
||||||
|
colors += [base(i) for i in range(n)]
|
||||||
|
im = self.ax.imshow(
|
||||||
|
np.asarray(counts), extent=extent, origin="upper", aspect="auto",
|
||||||
|
interpolation="nearest", cmap=ListedColormap(colors),
|
||||||
|
norm=BoundaryNorm(np.arange(-0.5, n + 1), len(colors)))
|
||||||
|
cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04,
|
||||||
|
ticks=np.arange(0, n + 1))
|
||||||
|
cb.set_label("angles overlapping")
|
||||||
|
self._finish(title)
|
||||||
|
|
||||||
|
def _finish(self, title: str):
|
||||||
self.ax.set_xlabel("X (mm)")
|
self.ax.set_xlabel("X (mm)")
|
||||||
self.ax.set_ylabel("Y (mm)")
|
self.ax.set_ylabel("Y (mm)")
|
||||||
self.ax.set_title(title)
|
self.ax.set_title(title)
|
||||||
|
|||||||
@@ -60,9 +60,13 @@ class Jobs:
|
|||||||
COMPUTE = "compute"
|
COMPUTE = "compute"
|
||||||
DC_PRECOMPUTE = "dc_precompute"
|
DC_PRECOMPUTE = "dc_precompute"
|
||||||
BATCH = "batch"
|
BATCH = "batch"
|
||||||
ALIGN = "align"
|
# The alignment wizard's three background steps: fetching each angle's CH4
|
||||||
MANUAL_ALIGN_MASKS = "manual_align_masks"
|
# image for the mask stack, registering the angles, and writing the aligned
|
||||||
MANUAL_ALIGN_CORRELATE = "manual_align_correlate"
|
# 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 = "",
|
def _make_dspin(lo: float, hi: float, decimals: int, *, suffix: str = "",
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ from sras_compute import (
|
|||||||
from sras_format import SrasFile
|
from sras_format import SrasFile
|
||||||
from sras_workers import Ch4MaskWorker, CrossCorrelateWorker
|
from sras_workers import Ch4MaskWorker, CrossCorrelateWorker
|
||||||
|
|
||||||
from .canvases import ManualAlignOverlayCanvas
|
from .canvases import AlignOverlayCanvas
|
||||||
from .common import (
|
from .common import (
|
||||||
_CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin,
|
_CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin,
|
||||||
_scroll_panel, _wrap_label,
|
_scroll_panel, _wrap_label,
|
||||||
@@ -354,7 +354,7 @@ class ManualAlignmentDialog(QDialog):
|
|||||||
def _build_ui(self, dc_threshold_mv: float):
|
def _build_ui(self, dc_threshold_mv: float):
|
||||||
root = QHBoxLayout(self)
|
root = QHBoxLayout(self)
|
||||||
|
|
||||||
self.canvas = ManualAlignOverlayCanvas()
|
self.canvas = AlignOverlayCanvas()
|
||||||
left = QWidget()
|
left = QWidget()
|
||||||
left_l = QVBoxLayout(left)
|
left_l = QVBoxLayout(left)
|
||||||
left_l.setContentsMargins(0, 0, 0, 0)
|
left_l.setContentsMargins(0, 0, 0, 0)
|
||||||
|
|||||||
@@ -1143,16 +1143,27 @@ class SrasViewerWindow(QMainWindow):
|
|||||||
# Progress dialogs
|
# Progress dialogs
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
def _show_progress(self, key: str, message: str, maximum: int = 0):
|
def _show_progress(self, key: str, message: str, maximum: int = 0,
|
||||||
|
on_cancel=None):
|
||||||
"""Show (or relabel) the progress dialog under *key*. maximum=0 gives
|
"""Show (or relabel) the progress dialog under *key*. maximum=0 gives
|
||||||
an indeterminate busy indicator."""
|
an indeterminate busy indicator.
|
||||||
|
|
||||||
|
*on_cancel* adds a Cancel button wired to it. Almost every job here is
|
||||||
|
short enough that a cancel button would only invite a click that does
|
||||||
|
nothing, hence the no-button default; the aligned export is the
|
||||||
|
exception, since it can spend minutes writing gigabytes.
|
||||||
|
"""
|
||||||
dlg = self._progress_dlgs.get(key)
|
dlg = self._progress_dlgs.get(key)
|
||||||
if dlg is not None:
|
if dlg is not None:
|
||||||
dlg.setLabelText(message)
|
dlg.setLabelText(message)
|
||||||
return
|
return
|
||||||
dlg = QProgressDialog(message, "", 0, maximum, self)
|
dlg = QProgressDialog(message, "", 0, maximum, self)
|
||||||
dlg.setWindowTitle("Please wait…")
|
dlg.setWindowTitle("Please wait…")
|
||||||
dlg.setCancelButton(None)
|
if on_cancel is None:
|
||||||
|
dlg.setCancelButton(None)
|
||||||
|
else:
|
||||||
|
dlg.setCancelButtonText("Cancel")
|
||||||
|
dlg.canceled.connect(on_cancel)
|
||||||
dlg.setWindowModality(Qt.WindowModality.WindowModal)
|
dlg.setWindowModality(Qt.WindowModality.WindowModal)
|
||||||
dlg.setMinimumDuration(300) # only appears if it takes > 300 ms
|
dlg.setMinimumDuration(300) # only appears if it takes > 300 ms
|
||||||
dlg.show()
|
dlg.show()
|
||||||
|
|||||||
+44
-2
@@ -15,6 +15,7 @@ import numpy as np
|
|||||||
from PyQt6.QtCore import QObject, pyqtSignal
|
from PyQt6.QtCore import QObject, pyqtSignal
|
||||||
|
|
||||||
import sras_compute as compute
|
import sras_compute as compute
|
||||||
|
from sras_align_export import write_aligned_sras
|
||||||
from sras_compute import (
|
from sras_compute import (
|
||||||
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
|
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
|
||||||
)
|
)
|
||||||
@@ -381,7 +382,11 @@ class CrossCorrelateWorker(_PooledWorker):
|
|||||||
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
||||||
dc4_mv: dict[int, np.ndarray], *,
|
dc4_mv: dict[int, np.ndarray], *,
|
||||||
sources: tuple[str, ...], dc_threshold_mv: float,
|
sources: tuple[str, ...], dc_threshold_mv: float,
|
||||||
search_deg: float):
|
search_deg: float, reg_kwargs: dict | None = None):
|
||||||
|
"""*reg_kwargs* is splatted into register_angle_to_reference on top of
|
||||||
|
the named arguments — the wizard's rotation-search controls (seed,
|
||||||
|
signs, refine, grid sizes) go through here, so exposing another knob
|
||||||
|
needs no change to this class."""
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self._sras = sras
|
self._sras = sras
|
||||||
self._ref = ref_angle_idx
|
self._ref = ref_angle_idx
|
||||||
@@ -390,6 +395,7 @@ class CrossCorrelateWorker(_PooledWorker):
|
|||||||
self._sources = sources
|
self._sources = sources
|
||||||
self._threshold = dc_threshold_mv
|
self._threshold = dc_threshold_mv
|
||||||
self._search_deg = search_deg
|
self._search_deg = search_deg
|
||||||
|
self._reg_kwargs = dict(reg_kwargs or {})
|
||||||
|
|
||||||
def _plan(self) -> int:
|
def _plan(self) -> int:
|
||||||
return compute.registration_workers(self._sras)
|
return compute.registration_workers(self._sras)
|
||||||
@@ -401,9 +407,45 @@ class CrossCorrelateWorker(_PooledWorker):
|
|||||||
return a, compute.register_angle_to_reference(
|
return a, compute.register_angle_to_reference(
|
||||||
self._sras, a, self._ref, self._dc4_mv,
|
self._sras, a, self._ref, self._dc4_mv,
|
||||||
dc_threshold_mv=self._threshold, sources=self._sources,
|
dc_threshold_mv=self._threshold, sources=self._sources,
|
||||||
search_deg=self._search_deg)
|
search_deg=self._search_deg, **self._reg_kwargs)
|
||||||
|
|
||||||
def _emit(self, result):
|
def _emit(self, result):
|
||||||
a, fit = result
|
a, fit = result
|
||||||
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
|
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
|
||||||
fit.shift_mm[1], fit.score, fit.source)
|
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,612 @@
|
|||||||
|
"""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_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
|
||||||
Reference in New Issue
Block a user