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:
Thomas Ales
2026-08-07 22:34:18 -05:00
parent 3989c2a1b8
commit 6d0e30b9ce
8 changed files with 1389 additions and 36 deletions
+477
View File
@@ -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
+152 -13
View File
@@ -1235,13 +1235,22 @@ def default_max_workers() -> int:
def _rotation_candidates(nominal_deg: float, search_deg: float,
step_deg: float) -> list[float]:
"""Coarse rotation candidates: a window around *both* signs of the stage's
reported angle change. Scoring both is what makes the stage's sign
convention a non-issue — the images decide which way the stage turns, and
a file whose stage reports the opposite sense registers just as well."""
step_deg: float,
signs: tuple[int, ...] = (-1, 1)) -> list[float]:
"""Coarse rotation candidates: a window around each requested sign of the
stage's reported angle change. Scoring both signs (the default) is what
makes the stage's sign convention a non-issue — the images decide which way
the stage turns, and a file whose stage reports the opposite sense
registers just as well.
*signs* exists only so a user who has established which way their own stage
turns can halve the coarse sweep from the alignment wizard. It is an
override, not an inference: nothing in the file says which sign is right.
"""
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))
for i in range(-k, k + 1):
out.append(center + i * step_deg)
@@ -1351,15 +1360,18 @@ def register_angle_to_reference(
dc_threshold_mv: float = 0.0,
sources: tuple[str, ...] = ("signal", "mask"),
coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
search_deg: float = 6.0, coarse_step_deg: float = 2.0) -> 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
*ref_angle_idx*, found entirely by cross-correlating image content.
Two stages:
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around both
signs of the stage's reported angle change (see _rotation_candidates),
for each requested source image, scored by _overlap_ncc. Only has to
pick the right basin.
1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
requested signs of the stage's reported angle change (see
_rotation_candidates), for each requested source image, scored by
_overlap_ncc. Only has to pick the right basin.
2. Hill-climbing refinement of the winning (source, rotation) at
*fine_dim* with sub-pixel translation folded into every score
(_refine_rotation, _score_rotation), down to 0.05°.
@@ -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
caller can surface so a bad scan is visible rather than silently fused in.
Returns identity for the reference angle itself.
The last three arguments only exist to let the alignment wizard expose the
rotation search; every default reproduces the search this function has
always done.
*seed_deg* replaces the stage's reported angle change as the center of the
coarse sweep — the pre-rotation the search starts from. None (the default)
means the stage angle from the file's Angle Table, which is nearly always
what you want: it puts the sweep within a couple of degrees of the answer.
Pass 0.0 to search around no rotation at all, which is the honest choice for
a file whose stage angles are known to be wrong.
*refine=False* stops after the coarse sweep, leaving rotation on the
*coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
it pins rotation to exactly the seed and searches translation only — for a
scan whose stage angles are trusted more than its image content.
"""
if angle_idx == ref_angle_idx:
return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")
@@ -1377,8 +1405,10 @@ def register_angle_to_reference(
if not candidates:
return RigidFit(0.0, (0.0, 0.0), -1.0, "none")
nominal = nominal_delta_deg(sras, angle_idx, ref_angle_idx)
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg)
nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
if seed_deg is None else float(seed_deg))
thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
seed_signs)
# ---- Stage 1: coarse sweep, every source ------------------------------
pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
@@ -1407,6 +1437,7 @@ def register_angle_to_reference(
theta = best[1]
score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
theta, subpixel=True)
if refine:
theta, score, shift = _refine_rotation(
ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
step_deg=coarse_step_deg)
@@ -1540,6 +1571,114 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
pitch[0], pitch[1], canvas_origin_mm, per_angle)
def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
n_rows: int, n_cols: int) -> AlignmentResult:
"""The same alignment restricted to a rectangular window of its canvas —
canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).
Cropping folds into each angle's existing affine instead of becoming a
second transform, because a canvas crop is *pure index translation*. From
_affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
origin, so:
matrix @ [r', c'] + (offset + matrix @ [row0, col0])
== matrix @ [r' + row0, c' + col0] + offset
i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
mapping at the shifted indices, exactly. That equality is what lets
apply_alignment, reproject_mask and the aligned .sras exporter all keep
working on a cropped result with no special-casing — and it is why the crop
preview a user approves is guaranteed to be the same pixels the exporter
writes.
Bounds are the caller's responsibility (the wizard's ROI page clamps to the
canvas): an out-of-range window is geometrically well-defined here and
simply resamples padding.
"""
if n_rows <= 0 or n_cols <= 0:
raise ValueError(f"empty crop: {n_rows} x {n_cols}")
delta = np.array([float(row0), float(col0)])
per_angle = {
a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
t.offset + t.matrix @ delta, t.score, t.source)
for a, t in result.per_angle.items()
}
origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
(int(n_rows), int(n_cols)),
result.canvas_dx_mm, result.canvas_dy_mm,
origin, per_angle)
def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
"""Summarize a per-pixel "how many angles cover this pixel" image — the
number the alignment wizard's mask-stack view is colored by.
Separated from the drawing code because it is the actual judgement the user
makes on that screen ("do the angles land on top of each other?"), and a
plain array-in/dict-out function can be tested without Qt.
"""
counts = np.asarray(counts)
union = int(np.count_nonzero(counts))
full = int(np.count_nonzero(counts >= n_angles))
return {
"union_px": union,
"full_px": full,
"full_frac": (full / union) if union else 0.0,
"mean_count": float(counts[counts > 0].mean()) if union else 0.0,
"max_count": int(counts.max()) if counts.size else 0,
"empty": union == 0,
}
def largest_rect_at_least(counts: np.ndarray, min_count: int
) -> tuple[int, int, int, int] | None:
"""Largest axis-aligned rectangle whose every pixel has counts >= min_count,
as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.
Backs the wizard's "fit to full overlap" button. A *bounding box* of the
qualifying pixels would be the obvious thing and is wrong: the full-overlap
region of several rotated scans is roughly a disc, whose bounding box has
corners no angle covers at all. Offering that as the crop would hand the
user padding they explicitly asked to avoid, so this finds a rectangle that
is entirely inside the region.
Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
preview-sized array, so it is instant at interactive rates.
"""
good = np.asarray(counts) >= min_count
if not good.any():
return None
n_rows, n_cols = good.shape
best = (0, 0, 0, 0) # area, row0, col0, ...
best_area = 0
heights = np.zeros(n_cols, dtype=np.int64)
for r in range(n_rows):
heights = np.where(good[r], heights + 1, 0)
# Sentinel column of height 0 flushes the stack at the end of the row.
stack: list[tuple[int, int]] = [] # (start col, height)
for c in range(n_cols + 1):
h = int(heights[c]) if c < n_cols else 0
start = c
while stack and stack[-1][1] >= h:
s, sh = stack.pop()
area = sh * (c - s)
if area > best_area:
best_area = area
best = (s, sh, c - s, r)
start = s
if h:
stack.append((start, h))
col0, height, width, row_end = best
return (row_end - height + 1, col0, height, width)
def _parallel_map(fn, items, n_workers: int) -> list:
"""fn over items, in order, threaded when it pays."""
items = list(items)
+78 -10
View File
@@ -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
@@ -100,7 +102,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 +118,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 +143,12 @@ 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):
"""*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.ax = self.figure.add_subplot(111)
# 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,
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:
cb.set_label(colorbar_label)
@@ -302,10 +317,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
@@ -475,14 +508,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 +561,35 @@ 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)
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_ylabel("Y (mm)")
self.ax.set_title(title)
+7 -3
View File
@@ -60,9 +60,13 @@ class Jobs:
COMPUTE = "compute"
DC_PRECOMPUTE = "dc_precompute"
BATCH = "batch"
ALIGN = "align"
MANUAL_ALIGN_MASKS = "manual_align_masks"
MANUAL_ALIGN_CORRELATE = "manual_align_correlate"
# 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 = "",
+2 -2
View File
@@ -20,7 +20,7 @@ from sras_compute import (
from sras_format import SrasFile
from sras_workers import Ch4MaskWorker, CrossCorrelateWorker
from .canvases import ManualAlignOverlayCanvas
from .canvases import AlignOverlayCanvas
from .common import (
_CSS_HINT, _CSS_MUTED, _CSS_WARN, Jobs, _axes_extent, _form, _group, _make_dspin,
_scroll_panel, _wrap_label,
@@ -354,7 +354,7 @@ class ManualAlignmentDialog(QDialog):
def _build_ui(self, dc_threshold_mv: float):
root = QHBoxLayout(self)
self.canvas = ManualAlignOverlayCanvas()
self.canvas = AlignOverlayCanvas()
left = QWidget()
left_l = QVBoxLayout(left)
left_l.setContentsMargins(0, 0, 0, 0)
+13 -2
View File
@@ -1143,16 +1143,27 @@ class SrasViewerWindow(QMainWindow):
# 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
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)
if dlg is not None:
dlg.setLabelText(message)
return
dlg = QProgressDialog(message, "", 0, maximum, self)
dlg.setWindowTitle("Please wait…")
if on_cancel is None:
dlg.setCancelButton(None)
else:
dlg.setCancelButtonText("Cancel")
dlg.canceled.connect(on_cancel)
dlg.setWindowModality(Qt.WindowModality.WindowModal)
dlg.setMinimumDuration(300) # only appears if it takes > 300 ms
dlg.show()
+44 -2
View File
@@ -15,6 +15,7 @@ import numpy as np
from PyQt6.QtCore import QObject, pyqtSignal
import sras_compute as compute
from sras_align_export import write_aligned_sras
from sras_compute import (
cache_file, compute_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],
dc4_mv: dict[int, np.ndarray], *,
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__()
self._sras = sras
self._ref = ref_angle_idx
@@ -390,6 +395,7 @@ class CrossCorrelateWorker(_PooledWorker):
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)
@@ -401,9 +407,45 @@ class CrossCorrelateWorker(_PooledWorker):
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)
search_deg=self._search_deg, **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))
+612
View File
@@ -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