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

Refactor cache validation and optimize alignment wizard incremental updates

Browse Source 
- 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>
This commit is contained in:
Thomas Ales
2026-08-09 14:30:02 -05:00
parent 8348ad313c
commit c30c8b1815
8 changed files with 392 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_align_export.py
+29 -58
View File
@@ -81,19 +81,22 @@ class ExportPlan:
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* 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.
*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.atleast_1d(np.asarray(rows, dtype=np.float64))[:, None]
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]
if np.isscalar(rows) or np.asarray(rows).ndim == 0:
return sr[0], sc[0]
return sr, sc
@@ -104,7 +107,7 @@ def _round_idx(coord: np.ndarray) -> np.ndarray:
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
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)
@@ -181,6 +184,13 @@ def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
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 "
@@ -188,7 +198,7 @@ def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
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
frac = plan.coverage_frac(a)
if frac == 0.0:
warnings.append(
f"Angle {a} has no data inside this crop — it will be written "
@@ -202,55 +212,12 @@ def plan_export(sras: SrasFile, result: AlignmentResult) -> ExportPlan:
"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):
if sras.scan_aborted:
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()
return plan
def _fill_row(sras: SrasFile, n_cols: int, dtype) -> np.ndarray:
@@ -335,7 +302,8 @@ class _SourceReader:
def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
*, progress_cb=None, should_stop=None) -> 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.
@@ -384,7 +352,7 @@ def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
# 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"):
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)
@@ -407,7 +375,7 @@ def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
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
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
@@ -419,8 +387,8 @@ def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
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))
fout.write(sras.encoded_preambles())
fout.write(sras.encoded_background())
pad = _fill_row(sras, n_cols, dtype)
for a in range(n_angles):
@@ -461,7 +429,10 @@ def write_aligned_sras(sras: SrasFile, result: AlignmentResult, out_path,
out[:, keep, :] = band[
idx_r[i][keep] - base, :, idx_c[i][keep], :
].transpose(1, 0, 2)
fout.write(out.tobytes())
# 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))