sras-viewer/sras_average.py

#!/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 .sras file with n_frames / N frames.

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.
"""

import argparse
import shutil
import struct
import sys
from pathlib import Path

import numpy as np

from sras_format import HDR_FMT, HDR_SIZE, SrasFile

_SUPPORTED = (2, 3, 4)


def parse_args():
    p = argparse.ArgumentParser(
        description="Average every N waveforms in a .sras file and write a new file."
    )
    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 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 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).

    Averaging is done in float32 and rounded on cast, matching numpy's mean
    followed by an int16 cast in the original implementation.
    """
    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:
        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)


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."""
    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

    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,
    )

    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


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))})",
              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}")

    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

    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"\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)

    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}")

    in_mb = in_path.stat().st_size / 1024**2
    out_mb = out_path.stat().st_size / 1024**2
    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()