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

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

Browse Source 
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>
This commit is contained in:
Thomas Ales [M S E]
2026-08-10 17:24:41 -05:00
parent 1caf6373cb
commit c54cce453c
3 changed files with 232 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_average.py
+150 -61
View File
@@ -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)
geo = bytearray()
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 += 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,15 +219,13 @@ 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" Rows : {int(sras.n_rows[0])}")
print(f" Frames (actual): {n_frames_in}")
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}")
@@ -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}). "
"The entire dataset will be averaged into a single frame.")
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.")
tests/test_compute.py
+75 -13
View File
@@ -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"
tools/make_test_sras.py
+1 -4
View File
@@ -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