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>
This commit is contained in:
Thomas Ales [M S E]
+156
-67
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user