sras-viewer/sras_edit_scans.py

#!/usr/bin/env python3
"""
sras_edit_scans.py — Remove one or more angle scans from a .sras file.

A .sras file holds one or more "angles" (rotation positions); the viewer
cross-correlates each non-reference angle against the reference to align
them. If one angle's acquisition went wrong (stage glitch, bad trigger,
laser dropout, ...) it throws off that alignment for the whole file. This
tool drops the bad angle(s) and renumbers the rest, writing a new .sras file
with everything else — waveform samples, calibration preambles, background
waveform, row/geometry tables — carried over byte-for-byte.

Handles v2-v7. Any precomputed FFT/DC cache (v5 PREC tail, v7 CACH tail) is
dropped on write, since it's indexed by angle and would be stale/misaligned
after renumbering; the viewer just recomputes it next time the file opens.

Usage:
    python sras_edit_scans.py input.sras --list
    python sras_edit_scans.py input.sras output.sras --drop 2,5
    python sras_edit_scans.py input.sras output.sras --keep 0,1,3,4,6
"""

import argparse
import struct
import sys
from pathlib import Path

from sras_format import GEO_FMT_V6, HDR_FMT, HDR_FMT_V6, HDR_SIZE, SrasFile

_LEGACY_VERSIONS = (2, 3, 4, 5)
_V6_VERSIONS = (6, 7)


def _die(msg: str):
    print(f"Error: {msg}", file=sys.stderr)
    sys.exit(1)


def parse_args():
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("input", help="Input .sras file")
    p.add_argument("output", nargs="?", help="Output .sras file (omit with --list)")
    p.add_argument("--list", action="store_true",
                   help="Print each angle's index/degrees/geometry and exit")
    g = p.add_mutually_exclusive_group()
    g.add_argument("--drop", metavar="I,J,...",
                   help="Comma-separated angle indices to remove")
    g.add_argument("--keep", metavar="I,J,...",
                   help="Comma-separated angle indices to keep (all others dropped)")
    return p.parse_args()


def _parse_index_list(s: str, n_angles: int) -> set[int]:
    out = set()
    for piece in s.split(","):
        piece = piece.strip()
        if not piece:
            continue
        i = int(piece)
        if not (0 <= i < n_angles):
            raise ValueError(f"angle index {i} out of range [0, {n_angles - 1}]")
        out.add(i)
    return out


def print_listing(sras: SrasFile):
    print(f"\n{'idx':>4}  {'angle_deg':>10}  {'x_start_mm':>11}  {'rows':>6}  {'frames':>7}")
    for a in range(sras.n_angles):
        print(f"{a:>4}  {sras.angles_deg[a]:>10.4f}  {sras.x_start_mm[a]:>11.4f}  "
              f"{int(sras.n_rows[a]):>6}  {int(sras.n_frames[a]):>7}")


def _copy_range(fin, fout, offset: int, nbytes: int, chunk: int = 64 * 1024 * 1024):
    """Stream *nbytes* raw bytes from *fin* at *offset* into *fout*, without
    ever holding more than one chunk in memory (waveform blocks can be
    hundreds of MB to low GB each)."""
    fin.seek(offset)
    remaining = nbytes
    while remaining:
        buf = fin.read(min(chunk, remaining))
        if not buf:
            raise IOError("unexpected EOF while copying waveform data")
        fout.write(buf)
        remaining -= len(buf)


# ---------------------------------------------------------------------------
# Legacy (v2-v5): uniform geometry across angles, one flat waveform block
# ---------------------------------------------------------------------------

def _write_legacy(sras: SrasFile, keep: list[int], out_path: Path):
    n_rows = int(sras.n_rows[0])
    n_frames = int(sras.n_frames[0])       # uniform across angles for v2-v5
    n_ch = sras.n_channels
    spf = sras.samples_per_frame
    bps = sras.bytes_per_sample

    header = struct.pack(
        HDR_FMT, b"SRAS", sras.version, len(keep), n_rows,
        float(sras.x_start_mm[0]), float(sras.x_delta_mm),
        sras.velocity_mm_s, sras.laser_freq_hz,
        n_frames, spf, sras.sample_rate_hz, bps, n_ch,
    )

    # Row table + preambles + background sit right after the angle table and
    # don't vary per angle — copy that whole span through unmodified.
    angle_table_size = sras.n_angles * 4
    with open(sras.path, "rb") as f:
        f.seek(HDR_SIZE + angle_table_size)
        shared_mid = f.read(sras.data_offset - (HDR_SIZE + angle_table_size))

    angle_bytes = n_rows * n_ch * n_frames * spf * bps

    with open(sras.path, "rb") as fin, open(out_path, "wb") as fout:
        fout.write(header)
        fout.write(sras.angles_deg[keep].astype(">f4").tobytes())
        fout.write(shared_mid)
        for a in keep:
            _copy_range(fin, fout, sras.data_offset + a * angle_bytes, angle_bytes)


# ---------------------------------------------------------------------------
# v6/v7: per-angle geometry, ragged waveform blocks
# ---------------------------------------------------------------------------

def _write_v6(sras: SrasFile, keep: list[int], out_path: Path):
    header = struct.pack(
        HDR_FMT_V6, b"SRAS", sras.version, len(keep),
        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,
        sras.samples_per_frame, sras.sample_rate_hz,
        sras.bytes_per_sample, sras.n_channels,
    )

    blocks = {a: (offset, nbytes) for a, offset, nbytes in sras.iter_angle_blocks()}
    with open(sras.path, "rb") as fin, open(out_path, "wb") as fout:
        fout.write(header)
        fout.write(sras.angles_deg[keep].astype(">f4").tobytes())
        for i in keep:
            fout.write(struct.pack(
                GEO_FMT_V6, float(sras.x_start_mm[i]),
                float(sras.x_delta_mm_per_angle[i]),
                int(sras.n_frames[i]), int(sras.n_rows[i])))
        for i in keep:
            fout.write(sras.y_pos_per_angle[i].astype(">f4").tobytes())
        fout.write(sras.preambles_raw)
        fout.write(sras.background_raw)
        for i in keep:
            offset, nbytes = blocks[i]
            _copy_range(fin, fout, offset, nbytes)


def main():
    args = parse_args()
    in_path = Path(args.input)
    if not in_path.exists():
        _die(f"input file not found: {in_path}")

    print(f"Reading {in_path} ...", flush=True)
    try:
        sras = SrasFile(str(in_path))
    except ValueError as e:
        _die(str(e))

    if sras.version not in (*_LEGACY_VERSIONS, *_V6_VERSIONS):
        _die(f"unsupported .sras version: {sras.version}")

    aborted_note = " (scan aborted; trailing angle(s) already excluded)" if sras.scan_aborted else ""
    print(f"  Version : v{sras.version}", flush=True)
    print(f"  Angles  : {sras.n_angles}{aborted_note}", flush=True)

    if args.list:
        print_listing(sras)
        return

    if not args.output:
        _die("output path required unless --list is given.")
    if not (args.drop or args.keep):
        _die("specify --drop or --keep (see --list for indices).")

    out_path = Path(args.output)
    if out_path.resolve() == in_path.resolve():
        _die("output path must differ from input path.")

    try:
        if args.drop:
            drop = _parse_index_list(args.drop, sras.n_angles)
            keep = [a for a in range(sras.n_angles) if a not in drop]
        else:
            keep = sorted(_parse_index_list(args.keep, sras.n_angles))
    except ValueError as e:
        _die(str(e))

    if not keep:
        _die("at least one angle must remain.")

    dropped = [a for a in range(sras.n_angles) if a not in keep]
    print(f"\nDropping angle(s): {dropped}")
    print(f"Keeping angle(s) : {keep}  ({len(keep)} of {sras.n_angles})")
    print(f"\nWriting {out_path} ...", flush=True)

    if sras.version in _LEGACY_VERSIONS:
        _write_legacy(sras, keep, out_path)
    else:
        _write_v6(sras, keep, out_path)

    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")
    print("Done.")
    print("Note: any precomputed FFT/DC cache was dropped (it's indexed by "
          "angle); the viewer will recompute it next time this file opens.")


if __name__ == "__main__":
    main()