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_compute.py
+96 -41
View File
@@ -285,6 +285,13 @@ _TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
def memory_budget_bytes() -> int:
"""The module-wide ceiling on concurrently-live working buffers, for
callers outside this module that read the same data (the aligned exporter
sizes its source-band reader against it)."""
return _TOTAL_BYTES_BUDGET
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
budget: int = _TOTAL_BYTES_BUDGET) -> int:
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
@@ -500,6 +507,36 @@ def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int:
return max(1, n_fft // spf)
def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
apply_bg_sub: bool, row_avg_n: int) -> list[str]:
"""Why this file's stored FFT images can't answer a request, as human-
readable phrases; empty means they can.
Padding, background subtraction and row-averaging are all baked into
the stored numbers, so a request that differs in any of them has to go
back through a real FFT. This is the single accept rule: cached_rf_image
serves a stored image iff this returns nothing, and callers that want to
*explain* the miss (rather than silently recompute) format these same
strings, so the two can't drift apart.
"""
reasons = []
pad = pad_factor_for(sras, n_fft)
if pad != sras.precomputed_pad_factor:
want = f"pad {pad}x" if pad else "a ragged n_fft"
reasons.append(f"stored at pad {sras.precomputed_pad_factor}x, "
f"requested at {want}")
if sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None):
reasons.append("stored with background subtraction "
f"{'on' if sras.precomputed_bg_sub else 'off'}")
if sras.precomputed_row_avg_n != row_avg_n:
have = (f"row-averaged (n={sras.precomputed_row_avg_n})"
if sras.precomputed_row_avg_n else "raw per-pixel")
want = (f"row-averaged (n={row_avg_n})"
if row_avg_n else "raw per-pixel")
reasons.append(f"stored {have}, requested {want}")
return reasons
def cached_rf_image(sras: SrasFile, angle_idx: int,
dc_threshold_mv: float | None,
apply_bg_sub: bool = True,
@@ -536,15 +573,15 @@ def cached_rf_image(sras: SrasFile, angle_idx: int,
fall through to a real compute rather than block.
"""
cached_freq = sras.precomputed_freq_mhz[angle_idx]
if (cached_freq is None
or pad_factor_for(sras, n_fft) != sras.precomputed_pad_factor
or sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None)
or sras.precomputed_row_avg_n != row_avg_n):
if cached_freq is None or cache_mismatch_reasons(
sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n):
return None
freq_img = cached_freq.copy()
dc4_img = None
if dc_threshold_mv is not None:
# DC4 mask, in priority order: already-cached DC block, caller-
# supplied image, or a fresh (cheap — no FFT) recompute.
# supplied image, or a fresh (cheap — no FFT) recompute. Resolved
# before the copy below so the allow_dc_recompute bail-out doesn't
# allocate a full image it is about to throw away.
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
if dc4_img is None:
if dc4_mv is not None:
@@ -554,6 +591,8 @@ def cached_rf_image(sras: SrasFile, angle_idx: int,
*sras.cal(CH4_IDX))
else:
return None
freq_img = cached_freq.copy()
if dc4_img is not None:
freq_img[dc4_img < dc_threshold_mv] = 0.0
return freq_img
@@ -778,6 +817,9 @@ def cache_file(path: str, mode: str, apply_bg_sub: bool,
try:
if mode not in ("dc", "fft", "fft_rowavg"):
return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')"
# write_v7_cache enforces this bound too, but only once every angle's
# FFT has already been computed. Checking up front is the difference
# between a bad argument costing nothing and costing the whole run.
if not (1 <= pad_factor <= MAX_PAD_FACTOR):
return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}"
set_fft_backend(fft_backend)
@@ -1492,24 +1534,17 @@ def register_angle_to_reference(
def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
pitch_mm: tuple[float, float],
per_angle_params: dict[int, ManualAngleParams],
*, margin_frac: float = 0.0, snap: bool = True
) -> tuple[tuple[float, float], tuple[int, int]]:
"""Shared-canvas origin (stage mm) and (n_rows, n_cols) at *pitch_mm* that
contains every angle's footprint after its own rigid transform. Angles
missing from per_angle_params default to identity (e.g. a sidecar saved
before a rescan added more angles).
snap=True aligns the canvas grid with the reference angle's own pixel grid,
so the reference lands on integer canvas pixels and is resampled by an
exact integer translation — the concrete meaning of "the canvas carries
angle 0's X/Y coordinates". It requires pitch_mm to be the reference's own
pitch; the manual-alignment preview passes a coarser pitch and snap=False.
margin_frac pads the box on every side: 0 for a final canvas, nonzero for
the wizard's preview canvas, which needs headroom so an ordinary
translation nudge never has to trigger a full canvas resize (an extreme
nudge can still push content past this padding; accepted, and cheap to
recover from by re-opening the dialog).
The canvas grid is aligned with the reference angle's own pixel grid, so
the reference lands on integer canvas pixels and is resampled by an exact
integer translation — the concrete meaning of "the canvas carries angle
0's X/Y coordinates". This requires pitch_mm to be the reference's own
pitch.
"""
dx, dy = pitch_mm
corners = np.vstack([
@@ -1520,27 +1555,17 @@ def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
for a in range(sras.n_angles)])
x_min, y_min = corners.min(axis=0)
x_max, y_max = corners.max(axis=0)
if margin_frac:
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
x_min, x_max = x_min - pad_x, x_max + pad_x
y_min, y_max = y_min - pad_y, y_max + pad_y
center = ref_center_mm(sras, ref_angle_idx)
if snap:
# Express the box in the reference's own pixel indices and grow it
# outward to whole pixels, so canvas index k lands exactly where the
# reference's own pixel (k + const) does.
cy, cx = _center_idx(sras, ref_angle_idx)
cols = sorted((x_min / dx + cx, x_max / dx + cx))
rows = sorted((y_min / dy + cy, y_max / dy + cy))
col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
shape = (row1 - row0 + 1, col1 - col0 + 1)
else:
origin_ref = np.array([x_min, y_min if dy > 0 else y_max])
shape = (int(np.ceil((y_max - y_min) / abs(dy))) + 1,
int(np.ceil((x_max - x_min) / dx)) + 1)
# Express the box in the reference's own pixel indices and grow it
# outward to whole pixels, so canvas index k lands exactly where the
# reference's own pixel (k + const) does.
cy, cx = _center_idx(sras, ref_angle_idx)
cols = sorted((x_min / dx + cx, x_max / dx + cx))
rows = sorted((y_min / dy + cy, y_max / dy + cy))
col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
shape = (row1 - row0 + 1, col1 - col0 + 1)
origin_stage = origin_ref + center
return (float(origin_stage[0]), float(origin_stage[1])), shape
@@ -1594,7 +1619,7 @@ def _result_from_params(sras: SrasFile, ref_angle_idx: int,
thread on every manual edit."""
pitch = pixel_pitch_mm(sras, ref_angle_idx)
canvas_origin_mm, canvas_shape = canvas_for_params(
sras, ref_angle_idx, pitch, params, snap=True)
sras, ref_angle_idx, pitch, params)
extra = extra or {}
per_angle: dict[int, AngleTransform] = {}
@@ -1653,6 +1678,36 @@ def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
origin, per_angle)
def coarse_rect_to_canvas(rect: tuple[int, int, int, int],
downsample: tuple[int, int],
canvas_shape: tuple[int, int],
*, inset_blocks: int = 0
) -> tuple[int, int, int, int]:
"""A rectangle in coarse (block-mean preview) indices as canvas pixels,
both as (row0, col0, n_rows, n_cols) and clamped to the canvas.
The single place the preview grid's relation to the real canvas is
written down: the coarse grid samples canvas pixels 0, f, 2f, …, so a
coarse pixel stands for a whole block and every caller has to agree on
which canvas pixels that block means, or a crop the user approved on the
preview lands a few pixels off in the export.
*inset_blocks* shrinks the rectangle by that many coarse blocks on each
side. A coarse pixel reported as fully covered stands for a block whose
far edge may not be, so "fit to full overlap" insets by 1 to stay honestly
inside the overlap region; a union rectangle insets by 0 because it wants
to contain the region rather than fit inside it.
"""
row0, col0, nr, nc = rect
fy, fx = downsample
n_rows, n_cols = canvas_shape
r0 = min(n_rows - 1, (row0 + inset_blocks) * fy)
c0 = min(n_cols - 1, (col0 + inset_blocks) * fx)
r1 = min(n_rows, (row0 + nr - inset_blocks) * fy)
c1 = min(n_cols, (col0 + nc - inset_blocks) * fx)
return r0, c0, max(1, r1 - r0), max(1, c1 - c0)
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.
@@ -1798,8 +1853,8 @@ def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
# to call synchronously on the GUI thread on every edit. The only genuinely
# expensive per-pixel operation anywhere in this flow is reproject_mask, and
# only the wizard's own downsampled mask stack calls that per keystroke
# — see that class's docstring for how it limits each nudge to reprojecting only
# the actively-edited angle.
# — see AlignmentWizard.rebuild_stack for how it limits a nudge to reprojecting
# only the actively-edited angle.
# ---------------------------------------------------------------------------
def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,