From c54cce453c5bce128f44b09619d6cc229551eda5 Mon Sep 17 00:00:00 2001 From: "Thomas Ales [M S E]" Date: Mon, 10 Aug 2026 17:24:41 -0500 Subject: [PATCH] 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 --- sras_average.py | 223 ++++++++++++++++++++++++++++------------ tests/test_compute.py | 88 +++++++++++++--- tools/make_test_sras.py | 5 +- 3 files changed, 232 insertions(+), 84 deletions(-) diff --git a/sras_average.py b/sras_average.py index 42aeaca..4e7c4ae 100644 --- a/sras_average.py +++ b/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: - """The raw bytes between the header and the waveform data (angle table, - 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 _memory_budget_bytes() -> int: + return int(os.environ.get("SRAS_MEM_BUDGET_MB", _DEFAULT_BUDGET_MB)) * 1024 * 1024 -def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray: - """Average every N frames of one angle's (n_rows, n_ch, n_frames, spf) - block. Returns int16 of shape (n_rows, n_ch, n_out, spf). +def _plan_angle(n_frames_in: int, n: int, discard_remainder: bool) -> tuple[int, int, int]: + """(n_full, remainder, n_frames_out) for averaging one angle's frames.""" + 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) - return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2) + averaged = np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16) + 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, - n: int, discard_remainder: bool) -> int: - """Stream each angle through the averager, writing as we go so peak RAM - stays at one angle's block rather than the whole file.""" +def write_v6_averaged(sras: SrasFile, out_path: Path, n: int, + discard_remainder: bool, budget: int | None = None) -> list[int]: + """Write *sras* averaged every N frames to a new v6 .sras file, streaming + 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 - if n_frames_in % n and not discard_remainder: - n_out += 1 + + plans = [_plan_angle(int(sras.n_frames[a]), n, discard_remainder) + for a in range(sras.n_angles)] + 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]), - float(sras.x_start_mm[0]), float(sras.x_delta_mm), - sras.velocity_mm_s, sras.laser_freq_hz, - n_out, # updated frame count - sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels, + HDR_FMT_V6, b"SRAS", 6, sras.n_angles, + sras.x_start_nominal_mm, sras.y_start_nominal_mm, + sras.x_delta_nominal_mm, sras.y_delta_nominal_mm, + sras.row_spacing_mm, sras.velocity_mm_s, sras.laser_freq_hz / n, + spf, sras.sample_rate_hz, bps, n_ch, ) - with open(out_path, "wb") as f: - f.write(header) - f.write(mid_sections) - for a in range(sras.n_angles): - averaged = average_rows(sras.data[a], n, discard_remainder) - if bps == 1: - f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes()) - else: - f.write(averaged.astype(">i2").tobytes()) - return n_out + geo = bytearray() + for a in range(sras.n_angles): + geo += struct.pack(GEO_FMT_V6, float(sras.x_start_mm[a]), + float(sras.x_delta_mm_per_angle[a]), + int(n_out_per_angle[a]), int(sras.n_rows[a])) + + part_path = out_path.with_name(out_path.name + ".part") + 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,18 +219,16 @@ 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]) - print(f" Version : v{sras.version}") - print(f" Angles : {sras.n_angles}") - 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" Bytes/sample : {sras.bytes_per_sample}") + 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}{aborted_note}") + print(f" Channels : {sras.n_channels}") + print(f" Samples/frame: {sras.samples_per_frame}") + print(f" Bytes/sample : {sras.bytes_per_sample}") if args.n == 1: print("--n 1: no averaging needed; copying file as-is.") @@ -145,28 +236,26 @@ def main(): print(f"Wrote {out_path}") return - if args.n > n_frames_in: - print(f"Warning: --n ({args.n}) exceeds available frames ({n_frames_in}). " - "The entire dataset will be averaged into a single frame.") + 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 angle {a}'s frames " + f"({n_frames_in}); it collapses to a single frame.") - print(f"\nAveraging every {args.n} frames ...", flush=True) - print(f"\nWriting {out_path} ...", flush=True) - mid_sections = read_header_sections(sras) - n_frames_out = write_averaged(out_path, sras, mid_sections, - args.n, args.discard_remainder) + if sras.version == 7: + print("\nNote: input has a v7 cache tail; it is indexed by frame count " + "and will be dropped. The viewer will recompute DC/FFT on next open.") - n_full = n_frames_in // args.n - remainder = n_frames_in % args.n - 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}") + print(f"\nAveraging every {args.n} frames and writing {out_path} ...", flush=True) + n_out_per_angle = write_v6_averaged(sras, out_path, args.n, args.discard_remainder) 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.") diff --git a/tests/test_compute.py b/tests/test_compute.py index f49b78b..ae100e7 100644 --- a/tests/test_compute.py +++ b/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" diff --git a/tools/make_test_sras.py b/tools/make_test_sras.py index 0a3651b..d8a171a 100644 --- a/tools/make_test_sras.py +++ b/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. - - Used to exercise sras_average.py, which only handles the legacy formats. - """ + """Write a v2/v3/v4 file: uniform geometry, one flat waveform block.""" rng = np.random.default_rng(seed) n_ch, bps = 3, 1