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