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 sys
import struct
import argparse
import numpy as np
from pathlib import Path

# ---------------------------------------------------------------------------
# Header format — must match sras_viewer.py exactly
# ---------------------------------------------------------------------------

HDR_FMT  = ">4sBHHffffIIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT)   # 43 bytes


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_sras(path: Path):
    """Read all sections of a .sras file and return them as a dict."""
    with open(path, "rb") as f:
        header_bytes = f.read(HDR_SIZE)
        fields = struct.unpack(HDR_FMT, header_bytes)

    (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
     n_frames_hdr, spf, sr, bps, n_ch) = fields

    if magic != b"SRAS":
        raise ValueError(f"Not a .sras file (bad magic: {magic!r})")
    if ver not in (2, 3, 4):
        raise ValueError(f"Unsupported .sras version: {ver}")

    with open(path, "rb") as f:
        f.seek(HDR_SIZE)

        angles = f.read(n_angles * 4)   # big-endian float32 array, raw bytes
        y_pos  = f.read(n_rows * 4)     # big-endian float32 array, raw bytes

        preambles = []          # list of raw bytes (length-prefixed strings)
        if ver >= 3:
            for _ in range(n_ch):
                (length,) = struct.unpack(">H", f.read(2))
                preambles.append(f.read(length))

        background = b""        # raw bytes for the v4 background block
        if ver >= 4:
            (n_bg,) = struct.unpack(">I", f.read(4))
            background = f.read(n_bg)

        raw = f.read()          # all waveform data

    # -----------------------------------------------------------------------
    # Determine actual frame count from file size (the header value can be
    # wrong — the viewer does the same correction)
    # -----------------------------------------------------------------------
    total_samples = len(raw) // bps
    samples_per_pixel = n_ch * spf         # samples in one (angle, row, frame) cell
    samples_per_full  = n_angles * n_rows * samples_per_pixel

    actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel)
    good_bytes      = actual_n_frames * samples_per_full * bps

    # Decode waveform data
    dtype = np.int8 if bps == 1 else ">i2"
    data  = np.frombuffer(raw[:good_bytes], dtype=dtype)
    data  = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf)
    # Work in int16 (safe intermediate for both int8 and int16 inputs)
    data  = data.astype(np.int16)

    return {
        "ver": ver, "n_angles": n_angles, "n_rows": n_rows,
        "x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq,
        "n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr,
        "bps": bps, "n_ch": n_ch,
        "angles_raw": angles, "y_pos_raw": y_pos,
        "preambles": preambles, "background": background,
        "data": data,       # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16
    }


def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
    """Average every N frames along axis 3.

    data shape: (n_angles, n_rows, n_ch, n_frames, spf)
    Returns array of shape (n_angles, n_rows, n_ch, n_out, spf).
    """
    n_frames = data.shape[3]

    n_full   = n_frames // n
    remainder = n_frames % n

    # Average full groups using reshape-trick (no Python loop)
    if n_full > 0:
        full = data[:, :, :, :n_full * n, :]           # trim to full groups
        full = full.reshape(data.shape[0], data.shape[1], data.shape[2],
                            n_full, n, data.shape[4])   # (..., n_out, n, spf)
        averaged = full.mean(axis=4).astype(np.int16)   # (..., n_out, spf)
    else:
        averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16)

    if remainder > 0 and not discard_remainder:
        tail = data[:, :, :, n_full * n:, :]           # shape (..., remainder, spf)
        tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16)
        averaged = np.concatenate([averaged, tail_avg], axis=3)

    return averaged


def write_sras(path: Path, src: dict, data_out: np.ndarray):
    """Write a new .sras file with the averaged waveform data."""
    ver    = src["ver"]
    bps    = src["bps"]
    n_out  = data_out.shape[3]

    # Pack header — update only n_frames_hdr; everything else stays the same
    header = struct.pack(
        HDR_FMT,
        b"SRAS",
        ver,
        src["n_angles"],
        src["n_rows"],
        src["x_start"],
        src["x_delta"],
        src["vel"],
        src["freq"],
        n_out,              # updated frame count
        src["spf"],
        src["sr"],
        bps,
        src["n_ch"],
    )

    # Encode waveform data back to original dtype
    if bps == 1:
        raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes()
    else:
        # big-endian int16
        raw_out = data_out.astype(">i2").tobytes()

    with open(path, "wb") as f:
        f.write(header)
        f.write(src["angles_raw"])
        f.write(src["y_pos_raw"])

        if ver >= 3:
            for preamble_bytes in src["preambles"]:
                f.write(struct.pack(">H", len(preamble_bytes)))
                f.write(preamble_bytes)

        if ver >= 4:
            bg = src["background"]
            f.write(struct.pack(">I", len(bg)))
            f.write(bg)

        f.write(raw_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)
    src = read_sras(in_path)

    n_frames_in = src["data"].shape[3]
    print(f"  Version        : v{src['ver']}")
    print(f"  Angles         : {src['n_angles']}")
    print(f"  Rows           : {src['n_rows']}")
    print(f"  Frames (actual): {n_frames_in}")
    print(f"  Channels       : {src['n_ch']}")
    print(f"  Samples/frame  : {src['spf']}")
    print(f"  Bytes/sample   : {src['bps']}")

    if args.n == 1:
        print("--n 1: no averaging needed; copying file as-is.")
        import shutil
        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)
    data_out = average_frames(src["data"], args.n, args.discard_remainder)
    n_frames_out = data_out.shape[3]

    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 and args.discard_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"\nWriting {out_path} ...", flush=True)
    write_sras(out_path, src, data_out)

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