#!/usr/bin/env python3 """ 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 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 python sras_average.py input.sras output.sras --n 10 --discard-remainder 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 from pathlib import Path import numpy as np from sras_format import GEO_FMT_V6, HDR_FMT_V6, SrasFile _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.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) p.add_argument("input", help="Input .sras file") p.add_argument("output", help="Output .sras file") p.add_argument("--n", type=int, required=True, metavar="N", help="Number of consecutive frames to average into one") p.add_argument("--discard-remainder", action="store_true", help="Drop trailing frames that don't fill a complete group of N") return p.parse_args() def _memory_budget_bytes() -> int: return int(os.environ.get("SRAS_MEM_BUDGET_MB", _DEFAULT_BUDGET_MB)) * 1024 * 1024 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 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 remainder = n_frames % n parts = [] if n_full: full = block[:, :, :n_full * n, :].astype(np.float32) full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3]) parts.append(full.mean(axis=3).astype(np.int16)) if remainder and not discard_remainder: tail = block[:, :, n_full * n:, :].astype(np.float32) parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16)) if not parts: 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_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 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_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, ) 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(): args = parse_args() if args.n < 1: print("Error: --n must be at least 1.", file=sys.stderr) sys.exit(1) in_path = Path(args.input) out_path = Path(args.output) if not in_path.exists(): print(f"Error: input file not found: {in_path}", file=sys.stderr) sys.exit(1) if out_path.resolve() == in_path.resolve(): print("Error: output path must differ from input path.", file=sys.stderr) sys.exit(1) print(f"Reading {in_path} ...", flush=True) 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))} only)", file=sys.stderr) sys.exit(1) 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.") shutil.copy2(in_path, out_path) 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 angle {a}'s frames " f"({n_frames_in}); it collapses to a single frame.") 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.") 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.") if __name__ == "__main__": main()