Thomas Ales
c30c8b1815
- Extract cache mismatch logic into reusable cache_mismatch_reasons() function for cleaner validation and consistent miss reporting - Expose memory_budget_bytes() for external callers (aligned exporter) - Optimize rebuild_stack() with incremental layer updates when only one angle parameters change, avoiding full reprojection overhead - Remove unused seed_per_angle parameter from alignment wizard - Simplify cached_rf_image() DC masking with cleaner early exit logic - Clean up internal state tracking with explicit origin caching Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
457 lines
21 KiB
Python
457 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
"""Write an aligned, cropped .sras file from an AlignmentResult.
|
|
|
|
The alignment machinery in sras_compute never modifies a scan: it produces an
|
|
AlignmentResult, and every consumer resamples on the fly (apply_alignment for
|
|
the display, reproject_mask for the overlay). That is right for a viewer, but it
|
|
means the aligned stack cannot leave the process — no other tool can read it,
|
|
and re-opening the scan re-does the registration.
|
|
|
|
This module bakes an alignment into a new file. Each angle is resampled onto the
|
|
shared canvas that AlignmentResult already defines, cropped to the caller's
|
|
window, so every output angle ends up with *identical* geometry: same rows, same
|
|
frames, same X/Y coordinates. Rotation and translation are gone, absorbed into
|
|
where each waveform sits. The result is an ordinary v6 file that opens already
|
|
aligned, and registering it against itself returns identity.
|
|
|
|
Two deliberate choices, both about not inventing data:
|
|
|
|
* The resample is a nearest-neighbour **gather** of whole waveforms, never an
|
|
interpolation. Averaging two neighbouring pixels' CH1 packets would produce
|
|
a waveform the instrument never measured, whose FFT peak is not the peak of
|
|
either — meaningless for a technique whose entire output is that peak
|
|
frequency. So each output pixel gets exactly one source pixel's three
|
|
waveforms, verbatim, and the cost is that some source pixels are duplicated
|
|
and others dropped. This matches apply_alignment's order=0 for the same
|
|
reason.
|
|
* Output pixels with no source pixel (the canvas corners a rotated scan cannot
|
|
reach, and anything outside the crop's coverage) are filled with the ADC
|
|
code for 0 mV, not with zero. See _fill_row.
|
|
|
|
Depends only on numpy/sras_format/sras_compute — no Qt — so it is directly
|
|
unit-testable and importable from a worker thread.
|
|
"""
|
|
|
|
import os
|
|
import struct
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
import sras_compute as compute
|
|
from sras_compute import AlignmentResult
|
|
from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile, mv_to_adc
|
|
|
|
# GEO_FMT_V6 stores n_rows as ">H" and n_frames as ">I". A canvas that overflows
|
|
# either is not representable, and silently truncating would write a file whose
|
|
# geometry table disagrees with its waveform block.
|
|
_MAX_ROWS = 0xFFFF
|
|
_MAX_FRAMES = 0xFFFFFFFF
|
|
|
|
# Slack, in source pixels, on the in-bounds test at the very edge of a source
|
|
# array. Absorbs the ~1e-13 of float noise an exactly-integer affine picks up
|
|
# from being built in mm space; see _in_bounds.
|
|
_EDGE_TOL = 1e-6
|
|
|
|
# Output rows per write() call. One row is n_channels * n_cols * spf bytes —
|
|
# ~1.5 MB on a full-size scan — so a handful of rows keeps the peak buffer in
|
|
# the low tens of MB no matter how large the scan is.
|
|
_ROW_CHUNK = 8
|
|
|
|
|
|
@dataclass
|
|
class ExportPlan:
|
|
"""What write_aligned_sras would produce, without producing it.
|
|
|
|
Derived from the affine transforms alone — no waveform bytes are read — so
|
|
the wizard can call it on every ROI edit to keep a live size estimate and
|
|
per-angle coverage readout in front of the user *before* they commit to a
|
|
multi-gigabyte write.
|
|
"""
|
|
n_rows: int
|
|
n_frames: int
|
|
n_angles: int
|
|
bytes_per_angle: int
|
|
total_bytes: int
|
|
valid_px: dict[int, int] # output pixels with a source pixel
|
|
warnings: list[str] = field(default_factory=list)
|
|
|
|
def coverage_frac(self, angle_idx: int) -> float:
|
|
px = self.n_rows * self.n_frames
|
|
return (self.valid_px.get(angle_idx, 0) / px) if px else 0.0
|
|
|
|
def empty_angles(self) -> list[int]:
|
|
"""Angles that would be written as pure padding — no output pixel of
|
|
theirs has a source pixel."""
|
|
return [a for a in range(self.n_angles) if not self.valid_px.get(a, 0)]
|
|
|
|
|
|
def _src_coords(t, rows, n_cols: int) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Fractional source (row, col) coordinates for whole output rows.
|
|
|
|
*rows* is an array of output row indices; both results are shaped
|
|
(len(rows), n_cols).
|
|
"""
|
|
cols = np.arange(n_cols, dtype=np.float64)
|
|
r = np.asarray(rows, dtype=np.float64)[:, None]
|
|
sr = t.matrix[0, 0] * r + t.matrix[0, 1] * cols + t.offset[0]
|
|
sc = t.matrix[1, 0] * r + t.matrix[1, 1] * cols + t.offset[1]
|
|
return sr, sc
|
|
|
|
|
|
def _round_idx(coord: np.ndarray) -> np.ndarray:
|
|
"""Nearest source index, rounding halves away from zero.
|
|
|
|
floor(x + 0.5), not np.rint: scipy.ndimage's order=0 rounds halves away
|
|
from zero while np.rint rounds them to even, and these have to be the same
|
|
source pixels the apply_alignment(order=0) preview drew. Exact halves are
|
|
not a corner case here — the canvas is snapped to the reference angle's own
|
|
pixel grid (see canvas_for_params), so an unrotated angle lands on
|
|
half-integers wherever its row pitch differs from the reference's.
|
|
"""
|
|
return np.floor(coord + 0.5).astype(np.int64)
|
|
|
|
|
|
def _in_bounds(sr: np.ndarray, sc: np.ndarray,
|
|
n_rows: int, n_frames: int) -> np.ndarray:
|
|
"""Which output pixels have a source pixel, by scipy's mode="constant" rule
|
|
plus a tolerance at the edge.
|
|
|
|
Tested on the *fractional* coordinate against the range of sample centres,
|
|
[0, n-1] inclusive — deliberately not on the rounded index. The two differ
|
|
around the whole rim: a coordinate of -0.4 rounds to a perfectly valid index
|
|
0, but scipy calls it out of bounds and writes cval there, so testing the
|
|
rounded index would put a one-pixel rim of real data everywhere the Aligned
|
|
View shows padding.
|
|
|
|
_EDGE_TOL is why this is not literally scipy's test. The affine is built
|
|
from a chain of mm-space multiplications, so an exactly-integer transform
|
|
comes out a few times 1e-13 off (the reference angle's offset lands on
|
|
-20 - 7e-15 rather than -20). Bare >= 0.0 then rejects that angle's entire
|
|
first row, and <= n-1 its last column — for the *reference* angle, whose
|
|
whole role is to pass through as an exact integer crop. The tolerance is
|
|
seven orders of magnitude above that noise and seven below the half-pixel
|
|
scale at which a rounding decision is ever meaningful, so it can only ever
|
|
change pixels whose scipy answer was itself decided by rounding noise.
|
|
"""
|
|
return ((sr >= -_EDGE_TOL) & (sr <= n_rows - 1 + _EDGE_TOL)
|
|
& (sc >= -_EDGE_TOL) & (sc <= n_frames - 1 + _EDGE_TOL))
|
|
|
|
|
|
# Output rows evaluated per numpy call when counting coverage. Counting row by
|
|
# row costs one small matmul per row (hundreds of milliseconds per angle on a
|
|
# full-size scan, on every ROI edit); counting the whole canvas at once needs
|
|
# hundreds of MB of index arrays. Blocking gets both: ~10 numpy calls per angle
|
|
# against ~30 MB of live index arrays.
|
|
_COUNT_BLOCK = 128
|
|
|
|
|
|
def _count_in_bounds(t, n_rows: int, n_cols: int,
|
|
src_rows: int, src_frames: int) -> int:
|
|
"""How many of the n_rows x n_cols output pixels have a source pixel."""
|
|
total = 0
|
|
for start in range(0, n_rows, _COUNT_BLOCK):
|
|
rows = np.arange(start, min(start + _COUNT_BLOCK, n_rows))
|
|
sr, sc = _src_coords(t, rows, n_cols)
|
|
total += int(np.count_nonzero(_in_bounds(sr, sc, src_rows, src_frames)))
|
|
return total
|
|
|
|
|
|
def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
|
|
"""Geometry, size and per-angle coverage of the file *result* would export.
|
|
|
|
Coverage is counted from the actual per-pixel index arrays rather than
|
|
approximated by the footprint parallelogram's area, because the two differ
|
|
exactly where it matters — a crop that clips one angle's scan window — and
|
|
this number is what tells the user an angle will come out mostly empty. No
|
|
waveform bytes are read, so it stays fast enough to call on every ROI edit.
|
|
"""
|
|
n_rows, n_cols = result.canvas_shape
|
|
n_angles = sras.n_angles
|
|
warnings: list[str] = []
|
|
|
|
valid_px: dict[int, int] = {}
|
|
for a in range(n_angles):
|
|
t = result.per_angle.get(a)
|
|
if t is None:
|
|
valid_px[a] = 0
|
|
warnings.append(f"Angle {a} has no transform and will be all padding.")
|
|
continue
|
|
src_rows, src_frames = sras.image_shape(a)
|
|
valid_px[a] = _count_in_bounds(t, n_rows, n_cols, src_rows, src_frames)
|
|
|
|
bytes_per_angle = (n_rows * sras.n_channels * n_cols
|
|
* sras.samples_per_frame * sras.bytes_per_sample)
|
|
|
|
# Built before the remaining warnings so they can be phrased with the
|
|
# plan's own coverage_frac rather than a second copy of the same division.
|
|
plan = ExportPlan(n_rows=n_rows, n_frames=n_cols, n_angles=n_angles,
|
|
bytes_per_angle=bytes_per_angle,
|
|
total_bytes=bytes_per_angle * n_angles,
|
|
valid_px=valid_px, warnings=warnings)
|
|
|
|
if n_rows > _MAX_ROWS:
|
|
warnings.append(
|
|
f"Crop is {n_rows} rows; the .sras geometry table caps rows at "
|
|
f"{_MAX_ROWS}. Narrow the ROI in Y.")
|
|
if n_cols > _MAX_FRAMES:
|
|
warnings.append(f"Crop is {n_cols} frames; the cap is {_MAX_FRAMES}.")
|
|
for a in range(n_angles):
|
|
frac = plan.coverage_frac(a)
|
|
if frac == 0.0:
|
|
warnings.append(
|
|
f"Angle {a} has no data inside this crop — it will be written "
|
|
f"as all padding.")
|
|
elif frac < 0.10:
|
|
warnings.append(
|
|
f"Angle {a} covers only {frac * 100:.1f}% of the crop.")
|
|
if sras.background is None:
|
|
warnings.append(
|
|
"Input has no background waveform (pre-v4 scan); a zero background "
|
|
"is written, which makes background subtraction a no-op.")
|
|
if sras.version != 6:
|
|
warnings.append(f"Input is v{sras.version}; the export is written as v6.")
|
|
if sras.scan_aborted:
|
|
warnings.append(
|
|
f"Input scan was aborted: only its {n_angles} complete angle(s) "
|
|
f"are exported.")
|
|
|
|
return plan
|
|
|
|
|
|
def _fill_row(sras: SrasFile, n_cols: int, dtype) -> np.ndarray:
|
|
"""One output row of pure padding, shape (n_channels, n_cols, spf).
|
|
|
|
Filled per channel with the ADC code for 0 mV, not with 0. Zero ADC decodes
|
|
to (0 - yoff) * ymult + yzero, which for a real scope preamble is a long way
|
|
from 0 mV — often far enough to sit above the CH4 mask threshold, which
|
|
would paint a solid rectangle of "valid" pixels around the sample and make
|
|
every DC image and every ROI statistic wrong. Rounding to the integer code
|
|
lands within half an ADC step of 0 mV, which is as close as the format can
|
|
represent.
|
|
"""
|
|
info = np.iinfo(dtype)
|
|
codes = [int(np.clip(round(mv_to_adc(0.0, *sras.cal(ch))), info.min, info.max))
|
|
for ch in range(sras.n_channels)]
|
|
row = np.empty((sras.n_channels, n_cols, sras.samples_per_frame), dtype=dtype)
|
|
for ch, code in enumerate(codes):
|
|
row[ch] = code
|
|
return row
|
|
|
|
|
|
class _SourceReader:
|
|
"""Gives the gather source rows without ever reading one twice.
|
|
|
|
This is the difference between a usable export and an unusable one, and it
|
|
is entirely about read amplification. A rotated angle maps one output row to
|
|
a *diagonal* line across the source array, so the pixels of a single output
|
|
row come from hundreds of different source rows — on a full-size scan, a
|
|
~1.4 MB source row each. Indexing a memmap pixel by pixel in output order
|
|
therefore re-faults nearly the whole angle for every output row: terabytes
|
|
of paging for a gigabyte of data.
|
|
|
|
So rows are served from a contiguous *band* held in RAM. The band for a
|
|
chunk of output rows is read in one sequential slice, and because output
|
|
rows advance monotonically through the source, consecutive chunks' bands
|
|
barely overlap: each source row is read about once, and the whole job costs
|
|
roughly 2x the source size in reads rather than a thousand times it.
|
|
|
|
Small angles skip the machinery — if the whole block fits the budget it is
|
|
materialized once and every band is a view of it.
|
|
|
|
One honest caveat: a chunk whose diagonal spans more source rows than the
|
|
budget allows still gets the band it asked for, so the budget can be
|
|
overshot. The overshoot is bounded by the span of _ROW_CHUNK output rows,
|
|
and in the worst case (an extreme rotation on a huge scan) that is the whole
|
|
angle — i.e. no worse than the in-RAM path above. Accepted deliberately:
|
|
correctness of the gather is not negotiable, and the alternative is the
|
|
memmap thrashing this class exists to avoid.
|
|
"""
|
|
|
|
def __init__(self, sras: SrasFile, angle_idx: int, budget: int):
|
|
self._src = sras.data[angle_idx]
|
|
self._n_rows = self._src.shape[0]
|
|
self._row_bytes = max(1, self._src[0].nbytes)
|
|
self._whole = np.asarray(self._src) if self._src.nbytes <= budget else None
|
|
# Leave room for the output buffer and the index arrays alongside.
|
|
self._max_band = max(1, int(budget * 0.5) // self._row_bytes)
|
|
self._band = None
|
|
self._lo = self._hi = 0
|
|
|
|
def band(self, lo: int, hi: int) -> tuple[np.ndarray, int]:
|
|
"""Rows [lo, hi) as an in-RAM array, plus the index its row 0 holds."""
|
|
lo = max(0, min(lo, self._n_rows))
|
|
hi = max(lo + 1, min(hi, self._n_rows))
|
|
if self._whole is not None:
|
|
return self._whole, 0
|
|
if self._band is None or lo < self._lo or hi > self._hi:
|
|
# Read a little more than asked so a chunk whose band creeps
|
|
# forward by a few rows does not re-read the whole span.
|
|
span = min(self._max_band, max(hi - lo, self._max_band // 2))
|
|
self._lo = lo
|
|
self._hi = min(self._n_rows, lo + span)
|
|
if self._hi < hi: # band cannot cover the ask
|
|
self._hi = hi
|
|
self._band = np.asarray(self._src[self._lo:self._hi])
|
|
return self._band, self._lo
|
|
|
|
def close(self):
|
|
self._whole = None
|
|
self._band = None
|
|
|
|
|
|
def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
|
|
*, progress_cb=None, should_stop=None,
|
|
budget: int | None = None) -> Path:
|
|
"""Write *sras*, aligned per *result* and cropped to its canvas, to a new
|
|
v6 .sras file. Returns the path written.
|
|
|
|
*result* is used exactly as given: crop the canvas first with
|
|
compute.crop_alignment_result, whose offset shift makes the cropped result
|
|
resample precisely the window the user selected.
|
|
|
|
No cache tail is written. Any DC/FFT the input had cached is indexed by the
|
|
input's grid and is meaningless on the new one, so — like sras_edit_scans —
|
|
the export drops it and lets the viewer recompute.
|
|
|
|
Writes to a sibling ".part" file and os.replace()s it into position on
|
|
success, unlinking it on error or cancellation: a half-written .sras is not
|
|
detectably broken (the v6 parser treats a short file as an aborted scan and
|
|
opens it happily), so it must never be left where the user might load it.
|
|
|
|
*should_stop* is polled once per output row chunk; returning True aborts and
|
|
raises nothing — the partial file is removed and the returned path will not
|
|
exist, so callers must check.
|
|
"""
|
|
out_path = Path(out_path)
|
|
n_rows, n_cols = result.canvas_shape
|
|
n_ch, spf = sras.n_channels, sras.samples_per_frame
|
|
n_angles = sras.n_angles
|
|
|
|
if n_rows <= 0 or n_cols <= 0:
|
|
raise ValueError(f"empty canvas: {n_rows} x {n_cols}")
|
|
if n_rows > _MAX_ROWS:
|
|
raise ValueError(
|
|
f"{n_rows} rows exceeds the .sras per-angle geometry limit of "
|
|
f"{_MAX_ROWS}; crop further in Y")
|
|
if n_cols > _MAX_FRAMES:
|
|
raise ValueError(f"{n_cols} frames exceeds the limit of {_MAX_FRAMES}")
|
|
missing = [a for a in range(n_angles) if a not in result.per_angle]
|
|
if missing:
|
|
raise ValueError(f"alignment result has no transform for angle(s) {missing}")
|
|
# The source's waveform blocks are live read-only memmaps into sras.path,
|
|
# so writing over it would corrupt the very reads the gather is making.
|
|
if out_path.exists() and out_path.samefile(sras.path):
|
|
raise ValueError(
|
|
"refusing to export onto the source scan; choose another filename")
|
|
|
|
dtype = np.dtype(np.int8 if sras.bytes_per_sample == 1 else ">i2")
|
|
x0_mm, y0_mm = result.canvas_origin_mm
|
|
y_rows = (y0_mm + np.arange(n_rows) * result.canvas_dy_mm).astype(">f4")
|
|
|
|
# Reference-only header fields. v6/v7 inputs have real ones to carry over;
|
|
# for a legacy input describe the canvas we are actually writing.
|
|
if sras.x_start_nominal_mm is not None:
|
|
nominal = (sras.x_start_nominal_mm, sras.y_start_nominal_mm,
|
|
sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
|
|
sras.row_spacing_mm)
|
|
else:
|
|
nominal = (x0_mm, y0_mm,
|
|
n_cols * sras.pixel_x_mm, n_rows * result.canvas_dy_mm,
|
|
result.canvas_dy_mm)
|
|
|
|
header = struct.pack(
|
|
HDR_FMT_V6, b"SRAS", 6, n_angles,
|
|
float(nominal[0]), float(nominal[1]), float(nominal[2]),
|
|
float(nominal[3]), float(nominal[4]),
|
|
sras.velocity_mm_s, sras.laser_freq_hz, spf, sras.sample_rate_hz,
|
|
sras.bytes_per_sample, n_ch)
|
|
|
|
# Every angle now shares one grid, so the ragged v6 tables collapse to
|
|
# n_angles copies of the same record. x_delta is the reference angle's own
|
|
# pitch (the canvas is its grid extended), which is velocity/laser_freq
|
|
# exactly, so x_axis_mm() stays self-consistent on re-read.
|
|
geo = struct.pack(GEO_FMT_V6, float(x0_mm), float(sras.pixel_x_mm),
|
|
int(n_cols), int(n_rows)) * n_angles
|
|
|
|
budget = compute.memory_budget_bytes() if budget is None else max(1, budget)
|
|
total_chunks = max(1, n_angles * ((n_rows + _ROW_CHUNK - 1) // _ROW_CHUNK))
|
|
done_chunks = 0
|
|
cancelled = False
|
|
|
|
part_path = out_path.with_name(out_path.name + ".part")
|
|
try:
|
|
with open(part_path, "wb") as fout:
|
|
fout.write(header)
|
|
fout.write(sras.angles_deg.astype(">f4").tobytes())
|
|
fout.write(geo)
|
|
fout.write(y_rows.tobytes() * n_angles)
|
|
fout.write(sras.encoded_preambles())
|
|
fout.write(sras.encoded_background())
|
|
|
|
pad = _fill_row(sras, n_cols, dtype)
|
|
for a in range(n_angles):
|
|
t = result.per_angle[a]
|
|
reader = _SourceReader(sras, a, budget)
|
|
src_rows, src_frames = sras.image_shape(a)
|
|
try:
|
|
for chunk_start in range(0, n_rows, _ROW_CHUNK):
|
|
if should_stop is not None and should_stop():
|
|
cancelled = True
|
|
break
|
|
chunk = np.arange(chunk_start,
|
|
min(chunk_start + _ROW_CHUNK, n_rows))
|
|
sr, sc = _src_coords(t, chunk, n_cols)
|
|
ok = _in_bounds(sr, sc, src_rows, src_frames)
|
|
# Clip rather than trust: _EDGE_TOL admits coordinates a
|
|
# hair outside the array, and an index off the end here
|
|
# would silently read the wrong row of the band.
|
|
idx_r = np.clip(_round_idx(sr), 0, src_rows - 1)
|
|
idx_c = np.clip(_round_idx(sc), 0, src_frames - 1)
|
|
|
|
# One band read covers the whole chunk: every source row
|
|
# any of these output rows touches, in one sequential
|
|
# slice. See _SourceReader.
|
|
if ok.any():
|
|
band, base = reader.band(int(idx_r[ok].min()),
|
|
int(idx_r[ok].max()) + 1)
|
|
else:
|
|
band, base = None, 0
|
|
|
|
for i in range(len(chunk)):
|
|
out = pad.copy()
|
|
keep = ok[i]
|
|
if keep.any():
|
|
# The two advanced indices are separated by a
|
|
# slice, so numpy puts the gathered axis first:
|
|
# (n_sel, n_ch, spf). Move it behind channels.
|
|
out[:, keep, :] = band[
|
|
idx_r[i][keep] - base, :, idx_c[i][keep], :
|
|
].transpose(1, 0, 2)
|
|
# out is C-contiguous, so the buffer protocol
|
|
# writes it straight out — .tobytes() would
|
|
# copy a full row per row written.
|
|
fout.write(out)
|
|
done_chunks += 1
|
|
if progress_cb is not None:
|
|
progress_cb(int(done_chunks / total_chunks * 100))
|
|
finally:
|
|
reader.close()
|
|
if cancelled:
|
|
break
|
|
if not cancelled:
|
|
fout.flush()
|
|
os.fsync(fout.fileno())
|
|
if cancelled:
|
|
part_path.unlink(missing_ok=True)
|
|
return out_path
|
|
os.replace(part_path, out_path)
|
|
except BaseException:
|
|
part_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
if progress_cb is not None:
|
|
progress_cb(100)
|
|
return out_path
|