sras-viewer/sras_format.py

#!/usr/bin/env python3
"""SRAS binary scan file format — parsing and writing.

Reads v2–v7 .sras files. Depends only on numpy + struct, so compute workers
(including multiprocessing children) can import it without pulling in Qt or
matplotlib. See scan_format.md for the full v6/v7 spec.

Channel semantics (fixed by sc3_aui_app.py acquisition settings):
  CH1 — RF Acoustic Packet (AC-coupled, 100 mV/div): FFT → peak frequency
  CH3 — Bias A (DC-coupled, 50 mV/div): waveform mean
  CH4 — Bias B (DC-coupled, 50 mV/div): waveform mean

Frame-count correction: the scanner writes the *configured* frame count in the
header before acquisition, but the scope may acquire fewer frames.  The actual
count is computed from the file size and used for the reshape so channels are
correctly aligned.

Scan geometry: v6/v7 files scan a different bounding box per angle (x_start,
x_delta, n_frames, n_rows all vary by angle), so geometry is exposed per-angle
via SrasFile.n_rows / n_frames / x_start_mm arrays and the x_axis_mm() /
y_positions_mm() methods. v2–v5 files have uniform geometry across angles, so
those arrays simply repeat the same value n_angles times.

v7 files are v6 files with an optional trailing cache section holding
precomputed per-angle DC and/or FFT images, so display never has to recompute
them after the "Convert" menu's batch actions have stored them once.
"""

import os
import re
import struct
from pathlib import Path

import numpy as np

# ---------------------------------------------------------------------------
# Format constants
# ---------------------------------------------------------------------------

# v2–v5: fixed header, uniform geometry across angles (43 bytes)
HDR_FMT = ">4sBHHffffIIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT)

# v6/v7: fixed header, per-angle geometry in a separate table (49 bytes)
HDR_FMT_V6 = ">4sBHfffffffIdBB"
HDR_SIZE_V6 = struct.calcsize(HDR_FMT_V6)

# v6/v7: per-angle geometry record (x_start, x_delta, n_frames, n_rows)
GEO_FMT_V6 = ">ffIH"
GEO_SIZE_V6 = struct.calcsize(GEO_FMT_V6)

# v5 precomputed-image tail
PREC_MAGIC = b"PREC"
PREC_FLAG_BG_SUB = 0x01

# v7 cache tail. v7 is byte-identical to v6 (the version byte is the only
# header difference) plus this optional trailing section. Sub-block sizes are
# per-angle (n_rows[a] * n_frames[a]), taken from the Per-Angle Geometry Table
# already parsed for v6 — no new geometry fields are needed.
CACH_MAGIC = b"CACH"
CACH_HDR_FMT = ">4sBB"                        # magic, cach_version, block_flags
CACH_HDR_SIZE = struct.calcsize(CACH_HDR_FMT)
CACH_VERSION = 1
CACH_FLAG_DC = 0x01
CACH_FLAG_FFT = 0x02

SDCB_MAGIC = b"SDCB"
SDCB_HDR_FMT = ">4sBH"                        # magic, reserved, n_stored
SDCB_HDR_SIZE = struct.calcsize(SDCB_HDR_FMT)

SFFT_MAGIC = b"SFFT"
SFFT_HDR_FMT = ">4sBH"                        # magic, flags, n_stored
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
SFFT_FLAG_BG_SUB = 0x01

# Fixed channel indices into the .sras data array (CH1=RF, CH3/CH4=Bias DC)
CH1_IDX, CH3_IDX, CH4_IDX = 0, 1, 2
CH_NAMES = ["CH1", "CH3", "CH4", "VEL"]

# Fallback scope calibration used only when reading v2 files without embedded
# preambles.  v3+ files carry the WFMOutpre string so these are not used.
#   50 mV/div, 8 div full-scale, int8 ADC, position = -2.72 div
#   ymult  = 50 mV × 8 / 256  = 1.5625 mV/count
#   yoff   = position × (256/8) = -2.72 × 32 = -87.04  (ADC count for 0 V)
_FALLBACK_YMULT_MV = 1.5625   # mV per ADC count
_FALLBACK_YOFF_ADC = -87.04   # ADC count that represents 0 V


# ---------------------------------------------------------------------------
# Calibration
# ---------------------------------------------------------------------------

def _parse_preamble(preamble: str) -> dict[str, float]:
    """Extract YMULT, YOFF, YZERO from a Tektronix WFMOutpre string.

    Returns a dict with float values for whichever keys are present.
    YMULT is left in V/count as the scope reports it.
    """
    result = {}
    for key in ("YMULT", "YOFF", "YZERO"):
        m = re.search(rf'\b{key}\s+([-+]?\d*\.?\d+(?:[Ee][+-]?\d+)?)', preamble)
        if m:
            result[key] = float(m.group(1))
    return result


def mv_to_adc(mv: float, ymult_mv: float = _FALLBACK_YMULT_MV,
              yoff_adc: float = _FALLBACK_YOFF_ADC,
              yzero_mv: float = 0.0) -> float:
    return (mv - yzero_mv) / ymult_mv + yoff_adc


def adc_to_mv(adc, ymult_mv: float = _FALLBACK_YMULT_MV,
              yoff_adc: float = _FALLBACK_YOFF_ADC,
              yzero_mv: float = 0.0):
    return (adc - yoff_adc) * ymult_mv + yzero_mv


# ---------------------------------------------------------------------------
# Binary read helpers
# ---------------------------------------------------------------------------

def _read_struct(f, fmt: str) -> tuple:
    return struct.unpack(fmt, f.read(struct.calcsize(fmt)))


def _read_preambles(f, n_ch: int) -> list[str]:
    """n_ch length-prefixed UTF-8 WFMOutpre strings."""
    out = []
    for _ in range(n_ch):
        (length,) = _read_struct(f, ">H")
        out.append(f.read(length).decode("utf-8"))
    return out


def _read_background(f) -> np.ndarray:
    """uint32 sample count followed by that many int8 samples."""
    (n_bg,) = _read_struct(f, ">I")
    return np.frombuffer(f.read(n_bg), dtype=np.int8).astype(np.float32)


def _read_f32_image(f, shape: tuple[int, int]) -> np.ndarray:
    """One big-endian float32 image, converted to native float32.

    The conversion matters: np.frombuffer hands back a read-only big-endian
    view, and these arrays flow straight into the display caches and every
    downstream arithmetic op.
    """
    n_bytes = shape[0] * shape[1] * 4
    return np.frombuffer(f.read(n_bytes), dtype=">f4").reshape(shape).astype(np.float32)


# ---------------------------------------------------------------------------
# File parser
# ---------------------------------------------------------------------------

class SrasFile:
    """Parsed in-memory representation of a v2–v7 .sras file.

    Scan geometry (rows, frames, x_start) is exposed per-angle via the
    ``n_rows`` / ``n_frames`` / ``x_start_mm`` arrays and the ``x_axis_mm()``
    / ``y_positions_mm()`` methods, since v6/v7 files scan a different
    bounding box per angle. v2–v5 files have uniform geometry, so these
    arrays just repeat the same value ``n_angles`` times. Waveform data is
    likewise exposed as ``data[angle_idx]``, an array of shape
    ``(n_rows[a], n_channels, n_frames[a], samples_per_frame)``.

    Precomputed images (v5's PREC tail or v7's CACH tail) are exposed as
    ``precomputed_dc3_mv`` / ``precomputed_dc4_mv`` / ``precomputed_freq_mhz``,
    always as ragged per-angle lists (``list[np.ndarray | None]``, one entry
    per angle, ``None`` where that angle was never stored) regardless of
    source version.
    """

    def __init__(self, path: str):
        self.path = Path(path)
        self._parse()

    def _parse(self):
        with open(self.path, "rb") as f:
            magic = f.read(4)
            if magic != b"SRAS":
                raise ValueError(f"Bad magic bytes: {magic!r}")
            (version,) = struct.unpack(">B", f.read(1))

        self.version = version
        if version in (2, 3, 4, 5):
            self._parse_legacy()
        elif version in (6, 7):
            self._parse_v6()
        else:
            raise ValueError(f"Unsupported version: {version}")

    # ------------------------------------------------------------------
    # Calibration
    # ------------------------------------------------------------------

    def _set_calibration(self, preambles: list[str] | None, n_ch: int):
        """Populate the per-channel ymult/yoff/yzero lists from preamble
        strings, falling back to the hardcoded scope constants for v2 files
        that carry no preambles."""
        if preambles is None:
            self.preambles = None
            self.ch_ymult_mv = [_FALLBACK_YMULT_MV] * n_ch
            self.ch_yoff_adc = [_FALLBACK_YOFF_ADC] * n_ch
            self.ch_yzero_mv = [0.0] * n_ch
            return

        self.preambles = preambles
        self.ch_ymult_mv, self.ch_yoff_adc, self.ch_yzero_mv = [], [], []
        for p in preambles:
            cal = _parse_preamble(p)
            # The scope reports YMULT and YZERO in volts; store both as mV.
            self.ch_ymult_mv.append(cal.get("YMULT", _FALLBACK_YMULT_MV / 1000) * 1000)
            self.ch_yoff_adc.append(cal.get("YOFF", _FALLBACK_YOFF_ADC))
            self.ch_yzero_mv.append(cal.get("YZERO", 0.0) * 1000)

    def cal(self, ch_idx: int) -> tuple[float, float, float]:
        """(ymult_mv, yoff_adc, yzero_mv) for one channel — splat straight
        into adc_to_mv / mv_to_adc."""
        return (self.ch_ymult_mv[ch_idx], self.ch_yoff_adc[ch_idx],
                self.ch_yzero_mv[ch_idx])

    def _init_precomputed(self, n_angles: int):
        self.precomputed_freq_mhz: list[np.ndarray | None] = [None] * n_angles
        self.precomputed_dc4_mv: list[np.ndarray | None] = [None] * n_angles
        self.precomputed_dc3_mv: list[np.ndarray | None] = [None] * n_angles
        self.precomputed_bg_sub: bool = False

    def cached_dc_mv(self, angle_idx: int, ch_idx: int) -> np.ndarray | None:
        """A stored DC image (already in mV) for (angle, channel), or None."""
        store = self.precomputed_dc3_mv if ch_idx == CH3_IDX else self.precomputed_dc4_mv
        return store[angle_idx] if angle_idx < len(store) else None

    def image_shape(self, angle_idx: int) -> tuple[int, int]:
        return int(self.n_rows[angle_idx]), int(self.n_frames[angle_idx])

    # ------------------------------------------------------------------
    # v2–v5 parsing (uniform geometry, flat waveform block)
    # ------------------------------------------------------------------

    def _parse_legacy(self):
        with open(self.path, "rb") as f:
            (magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
             n_frames_hdr, spf, sr, bps, n_ch) = _read_struct(f, HDR_FMT)

            self.n_angles = n_angles
            self.velocity_mm_s = float(vel)
            self.laser_freq_hz = float(freq)
            self.n_frames_header = n_frames_hdr   # configured count (may be wrong)
            self.samples_per_frame = spf
            self.sample_rate_hz = float(sr)
            self.bytes_per_sample = bps
            self.n_channels = n_ch

            self._init_precomputed(n_angles)
            self.scan_aborted = False
            self.n_angles_declared = n_angles

            angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)
            y_pos = np.frombuffer(f.read(n_rows * 4), dtype=">f4").astype(np.float32)

            self._set_calibration(_read_preambles(f, n_ch) if ver >= 3 else None, n_ch)
            self.background = _read_background(f) if ver >= 4 else None

            # Record where raw waveform data begins; np.memmap maps from here.
            data_offset = f.tell()

        # ---- Determine actual frame count from file size ---------------
        # For v4 and earlier the header n_frames may be the *configured*
        # count before acquisition; the actual count is derived from the
        # bytes on disk.  For v5 files a PREC tail follows the waveform
        # data, so we must not include those extra bytes in the frame count.
        file_size = self.path.stat().st_size
        samples_per_row_per_ch = n_ch * spf
        available_bytes = file_size - data_offset

        if ver == 5:
            actual_n_frames = n_frames_hdr
            remainder = 0
        else:
            total_samples = available_bytes // bps
            per_frame = n_angles * n_rows * samples_per_row_per_ch
            actual_n_frames = total_samples // per_frame
            remainder = total_samples % per_frame

        self.frame_count_mismatch = (actual_n_frames != n_frames_hdr)
        self.n_frames_remainder = remainder

        # ---- Memory-map the waveform data (zero RAM cost) --------------
        # Instead of f.read() → astype() (which peaks at 2× file size),
        # memmap lets the OS page only the bytes that are actually touched.
        data5d = np.memmap(
            str(self.path),
            dtype=np.int8 if bps == 1 else ">i2",
            mode="r",
            offset=data_offset,
            shape=(n_angles, n_rows, n_ch, actual_n_frames, spf),
        )
        # Expose as a list of per-angle views so downstream code shares one
        # indexing convention with v6: sras.data[a][row, ch, frame, sample]
        self.data = [data5d[a] for a in range(n_angles)]

        # Uniform per-angle geometry, repeated so callers don't need to
        # special-case legacy vs. v6 files.
        self.n_rows = np.full(n_angles, n_rows, dtype=np.int64)
        self.n_frames = np.full(n_angles, actual_n_frames, dtype=np.int64)
        self.x_start_mm = np.full(n_angles, float(x_start), dtype=np.float64)
        self.x_delta_mm = float(x_delta)   # reference only; kept for re-encode
        self._y_pos_per_angle = [y_pos] * n_angles
        self.angles_deg = angles
        self._data_offset = data_offset

        if ver >= 5:
            waveform_bytes = actual_n_frames * n_angles * n_rows * n_ch * spf * bps
            prec_offset = data_offset + waveform_bytes
            if file_size > prec_offset:
                self._parse_prec_section(prec_offset)

    def _parse_prec_section(self, offset: int):
        """Parse the v5 PREC tail that holds precomputed images.

        Stored per-angle as (freq, dc4, dc3), each a full-image float32
        block prefixed by its uint16 angle index.
        """
        with open(self.path, "rb") as f:
            f.seek(offset)
            header_raw = f.read(6)      # magic(4) + fmt_ver(1) + flags(1)
            if len(header_raw) < 6 or header_raw[:4] != PREC_MAGIC:
                return
            self.precomputed_bg_sub = bool(header_raw[5] & PREC_FLAG_BG_SUB)

            (n_stored,) = _read_struct(f, ">H")
            for _ in range(n_stored):
                (aidx,) = _read_struct(f, ">H")
                if aidx >= self.n_angles:
                    break
                shape = self.image_shape(aidx)
                self.precomputed_freq_mhz[aidx] = _read_f32_image(f, shape)
                self.precomputed_dc4_mv[aidx] = _read_f32_image(f, shape)
                self.precomputed_dc3_mv[aidx] = _read_f32_image(f, shape)

    # ------------------------------------------------------------------
    # v6/v7 parsing (per-angle geometry, ragged waveform blocks)
    # ------------------------------------------------------------------

    def _parse_v6(self):
        with open(self.path, "rb") as f:
            (magic, ver, n_angles, x_start_nom, y_start_nom, x_delta_nom,
             y_delta_nom, row_spacing, vel, freq, spf, sr, bps,
             n_ch) = _read_struct(f, HDR_FMT_V6)

            n_angles_declared = n_angles

            self.velocity_mm_s = float(vel)
            self.laser_freq_hz = float(freq)
            self.samples_per_frame = spf
            self.sample_rate_hz = float(sr)
            self.bytes_per_sample = bps
            self.n_channels = n_ch

            # Reference-only fields: the ROI as entered before per-angle
            # bounding-box expansion. Actual per-angle geometry used for
            # rendering comes from the Per-Angle Geometry Table below.
            self.x_start_nominal_mm = float(x_start_nom)
            self.y_start_nominal_mm = float(y_start_nom)
            self.x_delta_nominal_mm = float(x_delta_nom)
            self.y_delta_nominal_mm = float(y_delta_nom)
            self.row_spacing_mm = float(row_spacing)

            self.n_frames_header = None
            self.frame_count_mismatch = False
            self.n_frames_remainder = 0

            angles = np.frombuffer(f.read(n_angles * 4), dtype=">f4").astype(np.float32)

            x_start = np.empty(n_angles, dtype=np.float64)
            n_frames = np.empty(n_angles, dtype=np.int64)
            n_rows = np.empty(n_angles, dtype=np.int64)
            for a in range(n_angles):
                xs, xd, nf, nr = _read_struct(f, GEO_FMT_V6)
                x_start[a], n_frames[a], n_rows[a] = xs, nf, nr

            y_pos_per_angle = [
                np.frombuffer(f.read(int(n_rows[a]) * 4), dtype=">f4").astype(np.float32)
                for a in range(n_angles)
            ]

            self._set_calibration(_read_preambles(f, n_ch), n_ch)
            self.background = _read_background(f)

            data_offset = f.tell()

        self._data_offset = data_offset

        # ---- Memory-map each angle's ragged waveform block -------------
        # v6 gives each angle its own row/frame count, so waveform data is
        # no longer one uniform (n_angles, n_rows, ...) block — each angle's
        # block sits at a different offset with its own shape. An aborted
        # scan truncates the file mid-angle; per the format spec we keep
        # whatever complete angles are present rather than refusing to open
        # the file.
        file_size = self.path.stat().st_size
        waveform_dtype = np.int8 if bps == 1 else ">i2"

        data = []
        offset = data_offset
        for a in range(n_angles):
            nr, nf = int(n_rows[a]), int(n_frames[a])
            nbytes = nr * n_ch * nf * spf * bps
            if offset + nbytes > file_size:
                break
            data.append(np.memmap(
                str(self.path), dtype=waveform_dtype, mode="r",
                offset=offset, shape=(nr, n_ch, nf, spf),
            ))
            offset += nbytes

        n_complete = len(data)
        if n_complete == 0:
            raise ValueError(
                "v6 file has no complete angle blocks — scan was aborted "
                "before the first angle finished.")

        self.data = data
        self.n_angles = n_complete
        self.n_angles_declared = n_angles_declared
        self.scan_aborted = n_complete < n_angles_declared
        self.angles_deg = angles[:n_complete]
        self.x_start_mm = x_start[:n_complete]
        self.n_frames = n_frames[:n_complete]
        self.n_rows = n_rows[:n_complete]
        self._y_pos_per_angle = y_pos_per_angle[:n_complete]

        self._init_precomputed(n_complete)
        if self.version == 7 and offset < file_size:
            self._parse_cach_section(offset)

    # ------------------------------------------------------------------
    # v7 cache tail (CACH section: precomputed DC / FFT images)
    # ------------------------------------------------------------------

    def _cache_tail_offset(self) -> int:
        """Deterministic file offset where the CACH tail starts (or would
        start), derived purely from the header + Per-Angle Geometry Table —
        independent of whether a cache tail is actually present. Used by
        both the parser and the in-place writer."""
        waveform_bytes = sum(
            int(self.n_rows[a]) * self.n_channels * int(self.n_frames[a])
            * self.samples_per_frame * self.bytes_per_sample
            for a in range(self.n_angles)
        )
        return self._data_offset + int(waveform_bytes)

    def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
                          stores: list[list]) -> int | None:
        """Read one CACH sub-block header, then its per-angle image entries
        into *stores* (one list per image the block stores per angle).

        Returns the header's flags byte, or None if the block is malformed.
        """
        raw = f.read(struct.calcsize(hdr_fmt))
        if len(raw) < struct.calcsize(hdr_fmt):
            return None
        block_magic, flags, n_stored = struct.unpack(hdr_fmt, raw)
        if block_magic != magic:
            return None
        for _ in range(n_stored):
            (angle_idx,) = _read_struct(f, ">H")
            if angle_idx >= self.n_angles:
                break
            shape = self.image_shape(angle_idx)
            for store in stores:
                store[angle_idx] = _read_f32_image(f, shape)
        return flags

    def _parse_cach_section(self, offset: int):
        """Parse the v7 CACH tail that holds precomputed DC/FFT images."""
        with open(self.path, "rb") as f:
            f.seek(offset)
            header_raw = f.read(CACH_HDR_SIZE)
            if len(header_raw) < CACH_HDR_SIZE:
                return
            magic, cach_version, block_flags = struct.unpack(CACH_HDR_FMT, header_raw)
            if magic != CACH_MAGIC or cach_version != CACH_VERSION:
                return

            if block_flags & CACH_FLAG_DC:
                if self._read_cache_block(
                        f, SDCB_HDR_FMT, SDCB_MAGIC,
                        [self.precomputed_dc3_mv, self.precomputed_dc4_mv]) is None:
                    return

            if block_flags & CACH_FLAG_FFT:
                flags = self._read_cache_block(
                    f, SFFT_HDR_FMT, SFFT_MAGIC, [self.precomputed_freq_mhz])
                if flags is None:
                    return
                self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)

    def write_v7_cache(self, *,
                       new_dc3_mv: list[np.ndarray | None] | None = None,
                       new_dc4_mv: list[np.ndarray | None] | None = None,
                       new_freq_mhz: list[np.ndarray | None] | None = None,
                       new_bg_sub: bool | None = None):
        """Store computed DC and/or FFT images into this file's CACH tail,
        in place, converting a v6 source to v7 (or updating an existing v7
        file). Only the block(s) passed in are recomputed; whichever block
        isn't passed is carried forward unchanged from whatever this
        ``SrasFile`` already has in memory (from parsing, or a prior write
        in this same session) — its bytes are never re-read from disk.

        The waveform data itself is never touched: the cache tail always
        starts at ``_cache_tail_offset()``, a fixed offset derived from the
        header and geometry table alone.
        """
        if self.version not in (6, 7):
            raise ValueError(
                f"write_v7_cache only supports v6/v7 source files, got v{self.version}")

        final_dc3 = new_dc3_mv if new_dc3_mv is not None else self.precomputed_dc3_mv
        final_dc4 = new_dc4_mv if new_dc4_mv is not None else self.precomputed_dc4_mv
        final_freq = new_freq_mhz if new_freq_mhz is not None else self.precomputed_freq_mhz
        final_bg_sub = new_bg_sub if new_bg_sub is not None else self.precomputed_bg_sub

        dc_entries = [a for a in range(self.n_angles) if final_dc3[a] is not None]
        fft_entries = [a for a in range(self.n_angles) if final_freq[a] is not None]

        block_flags = ((CACH_FLAG_DC if dc_entries else 0)
                       | (CACH_FLAG_FFT if fft_entries else 0))

        payload = bytearray()
        payload += struct.pack(CACH_HDR_FMT, CACH_MAGIC, CACH_VERSION, block_flags)

        if dc_entries:
            payload += struct.pack(SDCB_HDR_FMT, SDCB_MAGIC, 0, len(dc_entries))
            for a in dc_entries:
                payload += struct.pack(">H", a)
                payload += final_dc3[a].astype(">f4").tobytes()
                payload += final_dc4[a].astype(">f4").tobytes()

        if fft_entries:
            fft_flags = SFFT_FLAG_BG_SUB if final_bg_sub else 0
            payload += struct.pack(SFFT_HDR_FMT, SFFT_MAGIC, fft_flags, len(fft_entries))
            for a in fft_entries:
                payload += struct.pack(">H", a)
                payload += final_freq[a].astype(">f4").tobytes()

        with open(self.path, "r+b") as f:
            f.seek(self._cache_tail_offset())
            f.write(payload)
            f.truncate()
            f.flush()
            os.fsync(f.fileno())
            # Version-byte flip last: if the process dies before this point,
            # the file is still readable as plain v6 (v6 parsing only
            # bounds-checks per-angle offset+nbytes <= file_size, it never
            # asserts exactly how many bytes follow the last angle) — so an
            # interrupted write can never corrupt the file, only leave
            # harmless trailing bytes that the next successful write
            # overwrites via this same deterministic cache offset.
            f.seek(4)
            f.write(struct.pack("B", 7))
            f.flush()
            os.fsync(f.fileno())

        self.version = 7
        self.precomputed_dc3_mv = final_dc3
        self.precomputed_dc4_mv = final_dc4
        self.precomputed_freq_mhz = final_freq
        self.precomputed_bg_sub = final_bg_sub

    # ------------------------------------------------------------------
    # Axes helpers
    # ------------------------------------------------------------------

    @property
    def pixel_x_mm(self) -> float:
        return self.velocity_mm_s / self.laser_freq_hz

    def x_axis_mm(self, angle_idx: int) -> np.ndarray:
        n = int(self.n_frames[angle_idx])
        return self.x_start_mm[angle_idx] + np.arange(n) * self.pixel_x_mm

    def y_positions_mm(self, angle_idx: int) -> np.ndarray:
        return self._y_pos_per_angle[angle_idx]

    def time_axis_ns(self) -> np.ndarray:
        return np.arange(self.samples_per_frame) / self.sample_rate_hz * 1e9

    def freq_axis_mhz(self, n_fft: int | None = None) -> np.ndarray:
        n = n_fft if n_fft is not None else self.samples_per_frame
        return np.fft.rfftfreq(n, d=1.0 / self.sample_rate_hz) / 1e6