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>

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."""
-    src = tmp_path / "legacy_v4.sras"
-    meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
-                            samples_per_frame=32, seed=4)
-    dst = tmp_path / "legacy_v4_avg.sras"
+    """The sras_average.py CLI: v6 frame averaging with remainder handling,
+    per-angle ragged geometry, and the laser_freq_hz X-axis correction."""
+    src = tmp_path / "v6.sras"
+    meta = gen.write(src, n_angles=2, seed=4, samples_per_frame=32,
+                     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"])
-    assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \
+            and avg.n_channels == src_sras.n_channels == 3)
+    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"]
-    expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
+    assert avg.laser_freq_hz == pytest.approx(src_sras.laser_freq_hz / 4), \
+        "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"