"""SRAS v6 binary scan-file format — the single implementation. Full byte-level spec: scan_format.md. Summary: header >4sBHfffffffIdBB magic ver n_angles xs_nom ys_nom xd_nom yd_nom row_spacing velocity laser_freq spf sample_rate bytes_per_sample n_channels angle table n_angles × >f geometry table n_angles × >ffIH (x_start x_delta n_frames n_rows) row tables (ragged) per angle: n_rows × >f (y positions) preambles n_channels × (>H length + utf-8 WFMOutpre string) background block >I length + raw int8 CH1 average waveform data angle-major, row-minor, channel-inner: for each angle, for each row, for each channel, n_frames × samples_per_frame × bytes_per_sample Incomplete files are valid: the data block is one contiguous append-only stream, so the readable prefix defines a single frontier past which nothing has been written yet (see ``SrasFile.angle_status``). """ from __future__ import annotations import mmap import struct from dataclasses import dataclass, field from pathlib import Path from typing import BinaryIO import numpy as np from core.scan_geometry import AngleGeometry, ScanPlan MAGIC = b"SRAS" VERSION = 6 HDR_FMT = ">4sBHfffffffIdBB" HDR_SIZE = struct.calcsize(HDR_FMT) # 49 bytes GEOM_FMT = ">ffIH" GEOM_SIZE = struct.calcsize(GEOM_FMT) # 14 bytes # Oscilloscope channels recorded, in on-disk order. SCAN_CHANNELS = [1, 3, 4] STATUS_OK = "OK" STATUS_TRUNCATED = "TRUNCATED" STATUS_MISSING = "MISSING" @dataclass class ScanHeader: """The fixed v6 global header (everything but magic/version).""" n_angles: int x_start_nominal: float y_start_nominal: float x_delta_nominal: float y_delta_nominal: float row_spacing: float velocity: float laser_freq: float samples_per_frame: int sample_rate: float bytes_per_sample: int n_channels: int @dataclass class AngleStatus: """How much of one angle's declared data is actually on disk.""" index: int angle_deg: float n_rows: int # declared row_bytes: int data_offset: int n_rows_available: int status: str # STATUS_OK / STATUS_TRUNCATED / STATUS_MISSING @property def complete(self) -> bool: return self.status == STATUS_OK def create_scan_file(path: Path, plan: ScanPlan, samples_per_frame: int, sample_rate: float, preambles: list[str], background_waveform: bytes) -> BinaryIO: """Create a new .sras file and write the v6 header + tables. Returns an open binary file positioned at the start of the data block; the caller appends waveform rows and must close it (try/finally). """ path.parent.mkdir(parents=True, exist_ok=True) f = open(path, "wb") f.write(struct.pack( HDR_FMT, MAGIC, VERSION, plan.n_angles, plan.x_start_nominal, plan.y_start_nominal, plan.x_delta_nominal, plan.y_delta_nominal, plan.row_spacing, plan.velocity_mm_s, plan.laser_freq_hz, samples_per_frame, sample_rate, 1, # bytes_per_sample: int8 from scope default len(SCAN_CHANNELS), )) f.write(struct.pack(f">{plan.n_angles}f", *plan.angles)) for pa in plan.per_angle: f.write(struct.pack(GEOM_FMT, pa.x_start, pa.x_delta, pa.n_frames, pa.n_rows)) for pa in plan.per_angle: f.write(struct.pack(f">{pa.n_rows}f", *pa.y_positions)) for p in preambles: enc = p.encode("utf-8") f.write(struct.pack(">H", len(enc))) f.write(enc) f.write(struct.pack(">I", len(background_waveform))) f.write(background_waveform) return f @dataclass class SrasFile: """Parsed v6 .sras file: header, tables, and lazy (memmap) data access. Parsing reads only the header/tables — never the waveform block — so opening a multi-GB file is cheap. ``load_angle``/``load_row`` return read-only numpy views backed by a shared mmap; no data is copied until the caller computes on it. """ path: Path header: ScanHeader = field(init=False) per_angle: list[AngleGeometry] = field(init=False) preambles: list[str] = field(init=False) preambles_raw: list[bytes] = field(init=False) background: bytes = field(init=False) data_start_offset: int = field(init=False) file_size: int = field(init=False) def __post_init__(self): self.path = Path(self.path) self._mmap: mmap.mmap | None = None self._parse() def _parse(self): self.file_size = self.path.stat().st_size with open(self.path, "rb") as f: raw = f.read(HDR_SIZE) if len(raw) < HDR_SIZE: raise ValueError(f"{self.path.name}: file too short to contain a valid header") (magic, version, n_angles, x_start_nominal, y_start_nominal, x_delta_nominal, y_delta_nominal, row_spacing, velocity, laser_freq, samples_per_frame, sample_rate, bytes_per_sample, n_channels) = struct.unpack(HDR_FMT, raw) if magic != MAGIC: raise ValueError(f"{self.path.name}: not a valid SRAS file (bad magic)") if version != VERSION: raise ValueError( f"{self.path.name}: unsupported SRAS format version {version} " f"(only version {VERSION} is supported)" ) self.header = ScanHeader( n_angles=n_angles, x_start_nominal=x_start_nominal, y_start_nominal=y_start_nominal, x_delta_nominal=x_delta_nominal, y_delta_nominal=y_delta_nominal, row_spacing=row_spacing, velocity=velocity, laser_freq=laser_freq, samples_per_frame=samples_per_frame, sample_rate=sample_rate, bytes_per_sample=bytes_per_sample, n_channels=n_channels, ) angles = struct.unpack(f">{n_angles}f", f.read(4 * n_angles)) self.per_angle = [] for a in angles: x_start, x_delta, n_frames, n_rows = struct.unpack(GEOM_FMT, f.read(GEOM_SIZE)) self.per_angle.append(AngleGeometry( angle_deg=a, x_start=x_start, x_delta=x_delta, n_frames=n_frames, n_rows=n_rows, )) for pa in self.per_angle: pa.y_positions = list(struct.unpack(f">{pa.n_rows}f", f.read(4 * pa.n_rows))) self.preambles_raw = [] for _ in range(n_channels): (plen,) = struct.unpack(">H", f.read(2)) self.preambles_raw.append(f.read(plen)) self.preambles = [p.decode("utf-8", errors="replace") for p in self.preambles_raw] (n_bg,) = struct.unpack(">I", f.read(4)) self.background = f.read(n_bg) self.data_start_offset = f.tell() # ── Frontier / truncation analysis ─────────────────────────────────────── def row_bytes(self, angle_idx: int) -> int: pa = self.per_angle[angle_idx] return (self.header.n_channels * pa.n_frames * self.header.samples_per_frame * self.header.bytes_per_sample) def angle_status(self) -> list[AngleStatus]: """Walk declared per-row byte counts against the actual file size. Because the data is one contiguous append-only stream, once an angle is found short every later angle is necessarily absent too — there is a single frontier past which nothing has been written yet. """ statuses = [] cursor = self.data_start_offset frontier_seen = False for ai, pa in enumerate(self.per_angle): row_bytes = self.row_bytes(ai) data_offset = cursor if frontier_seen: n_rows_available = 0 status = STATUS_MISSING else: declared_bytes = row_bytes * pa.n_rows if row_bytes > 0 and cursor + declared_bytes <= self.file_size: n_rows_available = pa.n_rows status = STATUS_OK cursor += declared_bytes else: remaining = max(0, self.file_size - cursor) n_rows_available = remaining // row_bytes if row_bytes > 0 else 0 status = STATUS_MISSING if n_rows_available == 0 else STATUS_TRUNCATED frontier_seen = True statuses.append(AngleStatus( index=ai, angle_deg=pa.angle_deg, n_rows=pa.n_rows, row_bytes=row_bytes, data_offset=data_offset, n_rows_available=n_rows_available, status=status, )) return statuses def angle_data_offset(self, angle_idx: int) -> int: offset = self.data_start_offset for ai in range(angle_idx): offset += self.row_bytes(ai) * self.per_angle[ai].n_rows return offset # ── Lazy data access ───────────────────────────────────────────────────── def _ensure_mmap(self) -> mmap.mmap: if self._mmap is None: # The mapping stays valid after the file object is closed, so # don't hold the descriptor open for the (long) life of a viewer # session. with open(self.path, "rb") as f: self._mmap = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) return self._mmap def _dtype(self) -> np.dtype: return np.dtype(np.int16 if self.header.bytes_per_sample == 2 else np.int8) def load_angle(self, angle_idx: int, n_rows: int | None = None) -> np.ndarray: """Read-only view of one angle's data block, shape (n_rows, n_channels, n_frames, samples_per_frame). ``n_rows`` limits the view to the rows actually on disk (pass ``AngleStatus.n_rows_available`` for truncated files); default is the declared row count. """ pa = self.per_angle[angle_idx] h = self.header if n_rows is None: n_rows = pa.n_rows start = self.angle_data_offset(angle_idx) count = n_rows * h.n_channels * pa.n_frames * h.samples_per_frame arr = np.frombuffer(self._ensure_mmap(), dtype=self._dtype(), count=count, offset=start) arr = arr.reshape(n_rows, h.n_channels, pa.n_frames, h.samples_per_frame) arr.flags.writeable = False return arr def load_row(self, angle_idx: int, row: int, channel_idx: int) -> np.ndarray: """Read-only view of one row/channel, shape (n_frames, samples_per_frame).""" pa = self.per_angle[angle_idx] h = self.header ch_bytes = pa.n_frames * h.samples_per_frame * h.bytes_per_sample start = (self.angle_data_offset(angle_idx) + row * self.row_bytes(angle_idx) + channel_idx * ch_bytes) arr = np.frombuffer(self._ensure_mmap(), dtype=self._dtype(), count=pa.n_frames * h.samples_per_frame, offset=start) arr = arr.reshape(pa.n_frames, h.samples_per_frame) arr.flags.writeable = False return arr def close(self): """Release this file's hold on the mapping. Views handed out earlier stay valid — they keep the mapping alive until they are garbage-collected, at which point the OS frees it. """ if self._mmap is not None: try: self._mmap.close() except BufferError: pass # live numpy views still reference the buffer self._mmap = None def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() # ── Axes helpers (viewer conveniences, derived from header fields) ─────── def pixel_pitch_x_mm(self) -> float: """Distance between adjacent frames along X.""" return self.header.velocity / self.header.laser_freq def x_axis_mm(self, angle_idx: int) -> np.ndarray: pa = self.per_angle[angle_idx] return pa.x_start + np.arange(pa.n_frames) * self.pixel_pitch_x_mm() def time_axis_ns(self) -> np.ndarray: h = self.header return np.arange(h.samples_per_frame) / h.sample_rate * 1e9 def freq_axis_mhz(self, nfft: int) -> np.ndarray: return np.fft.rfftfreq(nfft, d=1.0 / self.header.sample_rate) / 1e6 def plan_from_header(sras: SrasFile) -> ScanPlan: """Reconstruct the ScanPlan a file was written with (for resume).""" h = sras.header return ScanPlan( x_start_nominal=h.x_start_nominal, y_start_nominal=h.y_start_nominal, x_delta_nominal=h.x_delta_nominal, y_delta_nominal=h.y_delta_nominal, row_spacing=h.row_spacing, velocity_mm_s=h.velocity, laser_freq_hz=h.laser_freq, per_angle=list(sras.per_angle), )