Rewrite sras_average.py for v6/v7, memory-bounded for multi-hundred-GB scans

The old tool only understood the legacy uniform-geometry header and had no
bound on peak RAM, making it unusable on the ragged v6/v7 files the scanner
now produces at up to ~560GB. Reads via SrasFile's memmap and writes in
row-chunks sized to a memory budget (SRAS_MEM_BUDGET_MB), stages to .part and
os.replace()s per scan_format.md's atomic-write requirement, always writes
plain v6 (a v7 cache tail is frame-indexed and invalid after averaging), and
divides laser_freq_hz by N so x_axis_mm() stays correct after the X axis gets
spatially binned.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

sras_average.py

@@ -3,7 +3,23 @@
 sras_average.py — Waveform-averaging utility for .sras files.
 
 Reduces memory footprint by coherently averaging every N consecutive frames
-along the acquisition axis, writing a new .sras file with n_frames / N frames.
+along the acquisition axis, writing a new v6 .sras file with (about)
+n_frames / N frames per angle.
+
+Handles v6/v7 only (per-angle ragged geometry). A v7 input's cache tail is
+dropped — it's indexed by frame, which this changes — so the output is always
+written as v6; the viewer recomputes DC/FFT on next open. Reads via memmap and
+writes in row-chunks sized to a memory budget (default 1024 MB, override with
+the SRAS_MEM_BUDGET_MB env var), so peak RAM stays bounded regardless of file
+size — this is what makes the tool usable on multi-hundred-GB scans.
+
+Averaging every N frames also coarsens the physical X spacing between output
+frames (each frame is a distinct stage position — see scan_format.md's
+Spatial Mapping section: x_k = x_start + k * velocity_mm_s / laser_freq_hz) —
+so the output header's laser_freq_hz is divided by N to keep x_axis_mm()
+correct on the averaged file. This means the GUI's "Laser freq" info label
+will show that adjusted value rather than the scope's real setting for an
+averaged file; v6/v7 has no separate field for effective pixel pitch.
 
 Usage:
     python sras_average.py input.sras output.sras --n 10
@@ -13,9 +29,15 @@ Options:
     --n INT              Number of frames to average into one (required).
     --discard-remainder  Drop trailing frames that don't fill a complete group.
                          Default: include a partial average for the last group.
+
+Environment:
+    SRAS_MEM_BUDGET_MB   Ceiling on one row-chunk's working memory, in MB
+                         (default 1024). Lower it on a memory-constrained
+                         machine; the tool just takes more, smaller chunks.
 """
 
 import argparse
+import os
 import shutil
 import struct
 import sys
@@ -23,14 +45,22 @@ from pathlib import Path
 
 import numpy as np
 
-from sras_format import HDR_FMT, HDR_SIZE, SrasFile
+from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile
 
-_SUPPORTED = (2, 3, 4)
+_SUPPORTED = (6, 7)
+_DEFAULT_BUDGET_MB = 1024
+
+# Throttle per-chunk progress printing to roughly this many lines per angle,
+# so a huge angle (thousands of chunks) doesn't flood stdout while a small
+# one still gets to print every chunk.
+_MAX_PROGRESS_LINES = 40
 
 
 def parse_args():
     p = argparse.ArgumentParser(
-        description="Average every N waveforms in a .sras file and write a new file."
+        description="Average every N waveforms in a .sras file and write a new file.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
     )
     p.add_argument("input", help="Input .sras file")
     p.add_argument("output", help="Output .sras file")
@@ -41,21 +71,34 @@ def parse_args():
     return p.parse_args()
 
 
-def read_header_sections(sras: SrasFile) -> bytes:
+def _memory_budget_bytes() -> int:
-    """The raw bytes between the header and the waveform data (angle table,
+    return int(os.environ.get("SRAS_MEM_BUDGET_MB", _DEFAULT_BUDGET_MB)) * 1024 * 1024
-    row table, preambles, background), copied through verbatim so nothing is
-    lost in a re-encode."""
-    with open(sras.path, "rb") as f:
-        f.seek(HDR_SIZE)
-        return f.read(sras.data_offset - HDR_SIZE)
 
 
-def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
+def _plan_angle(n_frames_in: int, n: int, discard_remainder: bool) -> tuple[int, int, int]:
-    """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf)
+    """(n_full, remainder, n_frames_out) for averaging one angle's frames."""
-    block. Returns int16 of shape (n_rows, n_ch, n_out, spf).
+    n_full = n_frames_in // n
+    remainder = n_frames_in % n
+    n_out = n_full + (1 if remainder and not discard_remainder else 0)
+    return n_full, remainder, n_out
 
-    Averaging is done in float32 and rounded on cast, matching numpy's mean
+
-    followed by an int16 cast in the original implementation.
+def _chunk_rows(n_rows: int, n_channels: int, n_frames_in: int,
+                samples_per_frame: int, budget: int) -> int:
+    """How many rows to hold in RAM at once so one chunk's working buffers
+    (source block, float32 accumulator, output block) fit the budget."""
+    bytes_per_row = max(1, n_channels * n_frames_in * samples_per_frame * 4)
+    return max(1, min(n_rows, budget // bytes_per_row))
+
+
+def _average_block(block: np.ndarray, n: int, discard_remainder: bool,
+                   bps: int) -> np.ndarray:
+    """Average every N frames of a (rows, n_ch, n_frames, spf) block along
+    the frame axis, returning the result already encoded in the on-disk
+    dtype: int8 (clipped) if bps == 1, else big-endian int16.
+
+    Averaging is done in float32 and truncated on the int16 cast, matching
+    the original implementation's numpy mean-then-cast behavior.
     """
     n_frames = block.shape[2]
     n_full = n_frames // n
@@ -71,38 +114,88 @@ def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarr
         parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16))
 
     if not parts:
-        return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
+        averaged = np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
-    return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
+    else:
+        averaged = parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
+
+    if bps == 1:
+        return np.clip(averaged, -128, 127).astype(np.int8)
+    return averaged.astype(">i2")
 
 
-def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes,
+def write_v6_averaged(sras: SrasFile, out_path: Path, n: int,
-                   n: int, discard_remainder: bool) -> int:
+                      discard_remainder: bool, budget: int | None = None) -> list[int]:
-    """Stream each angle through the averager, writing as we go so peak RAM
+    """Write *sras* averaged every N frames to a new v6 .sras file, streaming
-    stays at one angle's block rather than the whole file."""
+    row-chunks per angle so peak RAM never holds more than one chunk (bounded
+    by *budget* bytes, default from SRAS_MEM_BUDGET_MB).
+
+    Stages to a sibling '.part' file and os.replace()s it into position on
+    success, per scan_format.md's requirement that any .sras writer must
+    never leave a half-written file visible under its final name (a short
+    v6 file opens successfully with the wrong angle/frame count rather than
+    failing loudly).
+
+    Returns the per-angle output frame counts.
+    """
+    budget = _memory_budget_bytes() if budget is None else max(1, budget)
+    n_ch = sras.n_channels
+    spf = sras.samples_per_frame
     bps = sras.bytes_per_sample
-    n_frames_in = int(sras.n_frames[0])
+
-    n_out = n_frames_in // n
+    plans = [_plan_angle(int(sras.n_frames[a]), n, discard_remainder)
-    if n_frames_in % n and not discard_remainder:
+             for a in range(sras.n_angles)]
-        n_out += 1
+    n_out_per_angle = [p[2] for p in plans]
 
     header = struct.pack(
-        HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]),
+        HDR_FMT_V6, b"SRAS", 6, sras.n_angles,
-        float(sras.x_start_mm[0]), float(sras.x_delta_mm),
+        sras.x_start_nominal_mm, sras.y_start_nominal_mm,
-        sras.velocity_mm_s, sras.laser_freq_hz,
+        sras.x_delta_nominal_mm, sras.y_delta_nominal_mm,
-        n_out,                       # updated frame count
+        sras.row_spacing_mm, sras.velocity_mm_s, sras.laser_freq_hz / n,
-        sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels,
+        spf, sras.sample_rate_hz, bps, n_ch,
     )
 
-    with open(out_path, "wb") as f:
+    geo = bytearray()
-        f.write(header)
-        f.write(mid_sections)
     for a in range(sras.n_angles):
-            averaged = average_rows(sras.data[a], n, discard_remainder)
+        geo += struct.pack(GEO_FMT_V6, float(sras.x_start_mm[a]),
-            if bps == 1:
+                           float(sras.x_delta_mm_per_angle[a]),
-                f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes())
+                           int(n_out_per_angle[a]), int(sras.n_rows[a]))
-            else:
+
-                f.write(averaged.astype(">i2").tobytes())
+    part_path = out_path.with_name(out_path.name + ".part")
-    return n_out
+    try:
+        with open(part_path, "wb") as f:
+            f.write(header)
+            f.write(sras.angles_deg.astype(">f4").tobytes())
+            f.write(bytes(geo))
+            for a in range(sras.n_angles):
+                f.write(sras.y_pos_per_angle[a].astype(">f4").tobytes())
+            f.write(sras.encoded_preambles())
+            f.write(sras.encoded_background())
+
+            for a in range(sras.n_angles):
+                n_rows = int(sras.n_rows[a])
+                n_frames_in = int(sras.n_frames[a])
+                chunk_rows = _chunk_rows(n_rows, n_ch, n_frames_in, spf, budget)
+                total_chunks = max(1, -(-n_rows // chunk_rows))
+                print_every = max(1, total_chunks // _MAX_PROGRESS_LINES)
+
+                data = sras.data[a]
+                for i, r0 in enumerate(range(0, n_rows, chunk_rows)):
+                    r1 = min(r0 + chunk_rows, n_rows)
+                    block = np.asarray(data[r0:r1])
+                    out_block = _average_block(block, n, discard_remainder, bps)
+                    f.write(out_block.tobytes())
+                    if total_chunks > 1 and (i % print_every == 0 or r1 == n_rows):
+                        print(f"  angle {a + 1}/{sras.n_angles}: "
+                              f"{r1}/{n_rows} rows", flush=True)
+
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(part_path, out_path)
+    except BaseException:
+        part_path.unlink(missing_ok=True)
+        raise
+
+    return n_out_per_angle
 
 
 def main():
@@ -126,17 +219,15 @@ def main():
     sras = SrasFile(str(in_path))
     if sras.version not in _SUPPORTED:
         print(f"Error: unsupported .sras version: {sras.version} "
-              f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))})",
+              f"(this tool handles v{'/v'.join(map(str, _SUPPORTED))} only)",
               file=sys.stderr)
         sys.exit(1)
 
-    n_frames_in = int(sras.n_frames[0])
+    aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else ""
     print(f"  Version      : v{sras.version}")
-    print(f"  Angles         : {sras.n_angles}")
+    print(f"  Angles       : {sras.n_angles}{aborted_note}")
-    print(f"  Rows           : {int(sras.n_rows[0])}")
-    print(f"  Frames (actual): {n_frames_in}")
     print(f"  Channels     : {sras.n_channels}")
-    print(f"  Samples/frame  : {sras.samples_per_frame}")
+    print(f"  Samples/frame: {sras.samples_per_frame}")
     print(f"  Bytes/sample : {sras.bytes_per_sample}")
 
     if args.n == 1:
@@ -145,28 +236,26 @@ def main():
         print(f"Wrote {out_path}")
         return
 
+    print(f"\n{'idx':>4}  {'angle_deg':>10}  {'rows':>6}  {'frames_in':>10}  {'frames_out':>11}")
+    for a in range(sras.n_angles):
+        n_frames_in = int(sras.n_frames[a])
+        n_full, remainder, n_out = _plan_angle(n_frames_in, args.n, args.discard_remainder)
+        print(f"{a:>4}  {sras.angles_deg[a]:>10.4f}  {int(sras.n_rows[a]):>6}  "
+              f"{n_frames_in:>10}  {n_out:>11}")
         if args.n > n_frames_in:
-        print(f"Warning: --n ({args.n}) exceeds available frames ({n_frames_in}). "
+            print(f"      Warning: --n ({args.n}) exceeds angle {a}'s frames "
-              "The entire dataset will be averaged into a single frame.")
+                  f"({n_frames_in}); it collapses to a single frame.")
 
-    print(f"\nAveraging every {args.n} frames ...", flush=True)
+    if sras.version == 7:
-    print(f"\nWriting {out_path} ...", flush=True)
+        print("\nNote: input has a v7 cache tail; it is indexed by frame count "
-    mid_sections = read_header_sections(sras)
+              "and will be dropped. The viewer will recompute DC/FFT on next open.")
-    n_frames_out = write_averaged(out_path, sras, mid_sections,
-                                  args.n, args.discard_remainder)
 
-    n_full = n_frames_in // args.n
+    print(f"\nAveraging every {args.n} frames and writing {out_path} ...", flush=True)
-    remainder = n_frames_in % args.n
+    n_out_per_angle = write_v6_averaged(sras, out_path, args.n, args.discard_remainder)
-    if remainder and not args.discard_remainder:
-        status = f"({n_full} full groups + 1 partial group of {remainder})"
-    elif remainder:
-        status = f"({n_full} full groups, {remainder} trailing frames discarded)"
-    else:
-        status = f"({n_full} full groups)"
-    print(f"  {n_frames_in} frames -> {n_frames_out} frames {status}")
 
     in_mb = in_path.stat().st_size / 1024**2
     out_mb = out_path.stat().st_size / 1024**2
+    print(f"\n  Frames out: {n_out_per_angle}")
     print(f"  Input size : {in_mb:.1f} MB")
     print(f"  Output size: {out_mb:.1f} MB  ({out_mb / in_mb * 100:.1f}% of input)")
     print("Done.")

tests/test_compute.py

@@ -12,6 +12,7 @@ from pathlib import Path
 import numpy as np
 import pytest
 
+import sras_average
 import sras_compute as compute
 from sras_compute import (
     cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
@@ -362,38 +363,46 @@ def test_legacy_parse(tmp_path):
 
 
 def test_sras_average(tmp_path):
-    """The sras_average.py CLI: frame averaging with remainder handling."""
+    """The sras_average.py CLI: v6 frame averaging with remainder handling,
-    src = tmp_path / "legacy_v4.sras"
+    per-angle ragged geometry, and the laser_freq_hz X-axis correction."""
-    meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
+    src = tmp_path / "v6.sras"
-                            samples_per_frame=32, seed=4)
+    meta = gen.write(src, n_angles=2, seed=4, samples_per_frame=32,
-    dst = tmp_path / "legacy_v4_avg.sras"
+                     geometry=[(4, 12)])
+    src_sras = SrasFile(str(src))
+
+    dst = tmp_path / "v6_avg.sras"
     proc = subprocess.run(
         [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "4"],
         capture_output=True, text=True, cwd=REPO)
     assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]
 
     avg = SrasFile(str(dst))
-    assert avg.version == 4
+    assert avg.version == 6
     assert list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}"
     assert (avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
-            and avg.n_channels == meta["n_channels"])
+            and avg.n_channels == src_sras.n_channels == 3)
-    assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \
+    assert np.allclose(avg.ch_ymult_mv, src_sras.ch_ymult_mv), \
         "calibration preserved"
-    assert np.array_equal(avg.background, SrasFile(str(src)).background), \
+    assert np.array_equal(avg.background, src_sras.background), \
         "background preserved"
-    src_data = meta["data"]
+    assert avg.laser_freq_hz == pytest.approx(src_sras.laser_freq_hz / 4), \
-    expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
+        "laser_freq_hz divided by N keeps pixel_x_mm correct after binning"
+    assert np.array_equal(avg.x_start_mm, src_sras.x_start_mm), \
+        "per-angle x_start unchanged"
+
+    src_waves = meta["waveforms"]
+    expect0 = src_waves[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
     assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \
         "first averaged group equals the mean of its 4 source frames"
 
     # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
-    dst2 = tmp_path / "legacy_v4_avg5.sras"
+    dst2 = tmp_path / "v6_avg5.sras"
     subprocess.run([sys.executable, str(REPO / "sras_average.py"),
                     str(src), str(dst2), "--n", "5"],
                    capture_output=True, text=True, cwd=REPO)
     assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \
         "partial trailing group kept by default"
-    dst3 = tmp_path / "legacy_v4_avg5d.sras"
+    dst3 = tmp_path / "v6_avg5d.sras"
     subprocess.run([sys.executable, str(REPO / "sras_average.py"),
                     str(src), str(dst3), "--n", "5", "--discard-remainder"],
                    capture_output=True, text=True, cwd=REPO)
@@ -401,6 +410,59 @@ def test_sras_average(tmp_path):
         "--discard-remainder drops the partial group"
 
 
+def test_sras_average_v7_cache_dropped(tmp_path):
+    """A v7 input's cache tail is indexed by frame count, so it's invalid
+    after averaging changes that count -- the output must always be plain
+    v6, never a v7 carrying a stale cache."""
+    src = tmp_path / "v7.sras"
+    gen.write(src, n_angles=2, seed=1, samples_per_frame=32, geometry=[(3, 8)])
+    src_sras = SrasFile(str(src))
+    src_sras.write_v7_cache(
+        new_dc3_mv=[compute_dc_image(src_sras, a, CH3_IDX) for a in range(src_sras.n_angles)],
+        new_dc4_mv=[compute_dc_image(src_sras, a, CH4_IDX) for a in range(src_sras.n_angles)])
+    assert SrasFile(str(src)).version == 7
+
+    dst = tmp_path / "v7_avg.sras"
+    proc = subprocess.run(
+        [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"],
+        capture_output=True, text=True, cwd=REPO)
+    assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]
+    assert SrasFile(str(dst)).version == 6, "cache-bearing input still writes plain v6"
+
+
+def test_sras_average_rejects_legacy(tmp_path):
+    """This tool only speaks v6/v7 now; a legacy file must fail clearly
+    rather than being silently misparsed."""
+    src = tmp_path / "legacy_v4.sras"
+    gen.write_legacy(src, version=4, n_angles=1, n_rows=2, n_frames=8,
+                     samples_per_frame=16, seed=0)
+    dst = tmp_path / "legacy_v4_avg.sras"
+    proc = subprocess.run(
+        [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "2"],
+        capture_output=True, text=True, cwd=REPO)
+    assert proc.returncode != 0
+    assert "v6" in proc.stderr and "v7" in proc.stderr
+
+
+def test_sras_average_chunking_matches_unchunked(tmp_path):
+    """A tiny memory budget (forcing one row per chunk) must produce
+    byte-identical output to a huge budget (everything in one chunk) -- the
+    load-bearing correctness claim of the memory-bounded rewrite: chunk
+    boundaries must never affect the averaged result."""
+    src = tmp_path / "v6.sras"
+    gen.write(src, n_angles=2, seed=7, samples_per_frame=48,
+             geometry=[(6, 10), (5, 13)])
+    sras = SrasFile(str(src))
+
+    dst_tiny = tmp_path / "avg_tiny.sras"
+    dst_big = tmp_path / "avg_big.sras"
+    sras_average.write_v6_averaged(sras, dst_tiny, 3, False, budget=1)
+    sras_average.write_v6_averaged(sras, dst_big, 3, False, budget=1 << 30)
+
+    assert dst_tiny.read_bytes() == dst_big.read_bytes(), \
+        "chunk size must not affect the averaged output"
+
+
 def test_unsupported_version_reported(tmp_path):
     """cache_file must report, not raise, for a file it can't handle."""
     bogus = tmp_path / "bogus.sras"

tools/make_test_sras.py

@@ -266,10 +266,7 @@ def write_rotating(path: Path, n_angles: int = 5, samples_per_frame: int = 4,
 def write_legacy(path: Path, version: int = 4, n_angles: int = 2,
                  n_rows: int = 4, n_frames: int = 10,
                  samples_per_frame: int = 32, seed: int = 0) -> dict:
-    """Write a v2/v3/v4 file: uniform geometry, one flat waveform block.
+    """Write a v2/v3/v4 file: uniform geometry, one flat waveform block."""
-
-    Used to exercise sras_average.py, which only handles the legacy formats.
-    """
     rng = np.random.default_rng(seed)
     n_ch, bps = 3, 1