Thomas Ales
666 lines
29 KiB
Python
666 lines
29 KiB
Python
#!/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)
|
||
# v2 added the SFFT pad factor. Writers always emit v2; readers still accept
|
||
# v1, whose FFT block is natural-resolution by construction. An older reader
|
||
# meeting a v2 tail rejects the whole section on the version check and
|
||
# recomputes — no cache, but no misreading either.
|
||
CACH_VERSION = 2
|
||
CACH_VERSION_MIN = 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_V1 = ">4sBH" # magic, flags, n_stored
|
||
SFFT_HDR_FMT = ">4sBHH" # magic, flags, n_stored, pad_factor
|
||
SFFT_HDR_SIZE = struct.calcsize(SFFT_HDR_FMT)
|
||
SFFT_FLAG_BG_SUB = 0x01
|
||
# The viewer clamps its pad factor to this, and the field is a uint16.
|
||
PAD_FACTOR_MAX = 256
|
||
|
||
# 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. The settings the FFT images were computed under travel
|
||
with them as ``precomputed_bg_sub`` / ``precomputed_pad_factor``, since a
|
||
stored image is only usable by a view asking for the same two.
|
||
"""
|
||
|
||
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
|
||
# Zero-padding the stored FFT images were computed at; 1 = natural
|
||
# resolution, which is all a v5 PREC or CACH v1 tail can hold.
|
||
self.precomputed_pad_factor: int = 1
|
||
|
||
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)
|
||
x_delta = 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], x_delta[a], n_frames[a], n_rows[a] = xs, xd, 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)
|
||
]
|
||
|
||
# Verbatim on-disk spans of the preamble and background sections,
|
||
# kept so file-rewriting tools (sras_edit_scans) can carry them
|
||
# over byte-for-byte without re-parsing.
|
||
span_start = f.tell()
|
||
self._set_calibration(_read_preambles(f, n_ch), n_ch)
|
||
span_end = f.tell()
|
||
f.seek(span_start)
|
||
self.preambles_raw = f.read(span_end - span_start)
|
||
|
||
span_start = span_end
|
||
self.background = _read_background(f)
|
||
span_end = f.tell()
|
||
f.seek(span_start)
|
||
self.background_raw = f.read(span_end - span_start)
|
||
|
||
data_offset = span_end
|
||
|
||
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.x_delta_mm_per_angle = x_delta[: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)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Byte-layout accessors (public: used by file-rewriting tools)
|
||
# ------------------------------------------------------------------
|
||
|
||
@property
|
||
def data_offset(self) -> int:
|
||
"""File offset where the waveform data begins (headers end)."""
|
||
return self._data_offset
|
||
|
||
@property
|
||
def y_pos_per_angle(self) -> list[np.ndarray]:
|
||
"""Per-angle Y row positions (mm). The list and its arrays are the
|
||
live parsed state — tools that reproject may replace entries."""
|
||
return self._y_pos_per_angle
|
||
|
||
@y_pos_per_angle.setter
|
||
def y_pos_per_angle(self, value: list[np.ndarray]):
|
||
self._y_pos_per_angle = value
|
||
|
||
def iter_angle_blocks(self):
|
||
"""Yields (angle_idx, byte_offset, byte_count) for each complete
|
||
angle's waveform block. Works for every version: legacy files have
|
||
uniform per-angle geometry, so the same walk applies."""
|
||
offset = self._data_offset
|
||
for a in range(self.n_angles):
|
||
nbytes = (int(self.n_rows[a]) * self.n_channels
|
||
* int(self.n_frames[a]) * self.samples_per_frame
|
||
* self.bytes_per_sample)
|
||
yield a, offset, nbytes
|
||
offset += nbytes
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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."""
|
||
end = self._data_offset
|
||
for _, offset, nbytes in self.iter_angle_blocks():
|
||
end = offset + nbytes
|
||
return end
|
||
|
||
def _read_cache_block(self, f, hdr_fmt: str, magic: bytes,
|
||
stores: list[list]) -> tuple | None:
|
||
"""Read one CACH sub-block header, then its per-angle image entries
|
||
into *stores* (one list per image the block stores per angle).
|
||
|
||
Every sub-block header is (magic, flags, n_stored, *extras). Returns
|
||
everything after the magic, so a caller that has extras knows how to
|
||
read them, 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, *extras = 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, *extras)
|
||
|
||
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 not CACH_VERSION_MIN <= 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:
|
||
# v1 has no pad field: those images are natural-resolution.
|
||
hdr_fmt = SFFT_HDR_FMT if cach_version >= 2 else SFFT_HDR_FMT_V1
|
||
header = self._read_cache_block(
|
||
f, hdr_fmt, SFFT_MAGIC, [self.precomputed_freq_mhz])
|
||
if header is None:
|
||
return
|
||
flags, *extras = header
|
||
self.precomputed_bg_sub = bool(flags & SFFT_FLAG_BG_SUB)
|
||
self.precomputed_pad_factor = (
|
||
min(max(1, extras[0]), PAD_FACTOR_MAX) if extras else 1)
|
||
|
||
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,
|
||
new_pad_factor: int | 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
|
||
final_pad = (new_pad_factor if new_pad_factor is not None
|
||
else self.precomputed_pad_factor)
|
||
if not 1 <= final_pad <= PAD_FACTOR_MAX:
|
||
raise ValueError(
|
||
f"pad factor {final_pad} outside 1..{PAD_FACTOR_MAX}")
|
||
|
||
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), final_pad)
|
||
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
|
||
self.precomputed_pad_factor = final_pad
|
||
|
||
# ------------------------------------------------------------------
|
||
# 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
|