Compare commits
9 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bbb075ff34 | |||
| 52e91d46fd | |||
| 8312668c02 | |||
| ca0c736c28 | |||
| 3c835f4592 | |||
| 013c7739a1 | |||
| 9a7cc396e7 | |||
| 55a4c5e42f | |||
| 95c273dce5 |
+71
-146
@@ -15,18 +15,17 @@ Options:
|
|||||||
Default: include a partial average for the last group.
|
Default: include a partial average for the last group.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import sys
|
|
||||||
import struct
|
|
||||||
import argparse
|
import argparse
|
||||||
import numpy as np
|
import shutil
|
||||||
|
import struct
|
||||||
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
import numpy as np
|
||||||
# Header format — must match sras_viewer.py exactly
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
HDR_FMT = ">4sBHHffffIIdBB"
|
from sras_format import HDR_FMT, HDR_SIZE, SrasFile
|
||||||
HDR_SIZE = struct.calcsize(HDR_FMT) # 43 bytes
|
|
||||||
|
_SUPPORTED = (2, 3, 4)
|
||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
@@ -42,143 +41,68 @@ def parse_args():
|
|||||||
return p.parse_args()
|
return p.parse_args()
|
||||||
|
|
||||||
|
|
||||||
def read_sras(path: Path):
|
def read_header_sections(sras: SrasFile) -> bytes:
|
||||||
"""Read all sections of a .sras file and return them as a dict."""
|
"""The raw bytes between the header and the waveform data (angle table,
|
||||||
with open(path, "rb") as f:
|
row table, preambles, background), copied through verbatim so nothing is
|
||||||
header_bytes = f.read(HDR_SIZE)
|
lost in a re-encode."""
|
||||||
fields = struct.unpack(HDR_FMT, header_bytes)
|
with open(sras.path, "rb") as f:
|
||||||
|
|
||||||
(magic, ver, n_angles, n_rows, x_start, x_delta, vel, freq,
|
|
||||||
n_frames_hdr, spf, sr, bps, n_ch) = fields
|
|
||||||
|
|
||||||
if magic != b"SRAS":
|
|
||||||
raise ValueError(f"Not a .sras file (bad magic: {magic!r})")
|
|
||||||
if ver not in (2, 3, 4):
|
|
||||||
raise ValueError(f"Unsupported .sras version: {ver}")
|
|
||||||
|
|
||||||
with open(path, "rb") as f:
|
|
||||||
f.seek(HDR_SIZE)
|
f.seek(HDR_SIZE)
|
||||||
|
return f.read(sras._data_offset - HDR_SIZE)
|
||||||
angles = f.read(n_angles * 4) # big-endian float32 array, raw bytes
|
|
||||||
y_pos = f.read(n_rows * 4) # big-endian float32 array, raw bytes
|
|
||||||
|
|
||||||
preambles = [] # list of raw bytes (length-prefixed strings)
|
|
||||||
if ver >= 3:
|
|
||||||
for _ in range(n_ch):
|
|
||||||
(length,) = struct.unpack(">H", f.read(2))
|
|
||||||
preambles.append(f.read(length))
|
|
||||||
|
|
||||||
background = b"" # raw bytes for the v4 background block
|
|
||||||
if ver >= 4:
|
|
||||||
(n_bg,) = struct.unpack(">I", f.read(4))
|
|
||||||
background = f.read(n_bg)
|
|
||||||
|
|
||||||
raw = f.read() # all waveform data
|
|
||||||
|
|
||||||
# -----------------------------------------------------------------------
|
|
||||||
# Determine actual frame count from file size (the header value can be
|
|
||||||
# wrong — the viewer does the same correction)
|
|
||||||
# -----------------------------------------------------------------------
|
|
||||||
total_samples = len(raw) // bps
|
|
||||||
samples_per_pixel = n_ch * spf # samples in one (angle, row, frame) cell
|
|
||||||
samples_per_full = n_angles * n_rows * samples_per_pixel
|
|
||||||
|
|
||||||
actual_n_frames = total_samples // (n_angles * n_rows * samples_per_pixel)
|
|
||||||
good_bytes = actual_n_frames * samples_per_full * bps
|
|
||||||
|
|
||||||
# Decode waveform data
|
|
||||||
dtype = np.int8 if bps == 1 else ">i2"
|
|
||||||
data = np.frombuffer(raw[:good_bytes], dtype=dtype)
|
|
||||||
data = data.reshape(n_angles, n_rows, n_ch, actual_n_frames, spf)
|
|
||||||
# Work in int16 (safe intermediate for both int8 and int16 inputs)
|
|
||||||
data = data.astype(np.int16)
|
|
||||||
|
|
||||||
return {
|
|
||||||
"ver": ver, "n_angles": n_angles, "n_rows": n_rows,
|
|
||||||
"x_start": x_start, "x_delta": x_delta, "vel": vel, "freq": freq,
|
|
||||||
"n_frames_hdr": n_frames_hdr, "spf": spf, "sr": sr,
|
|
||||||
"bps": bps, "n_ch": n_ch,
|
|
||||||
"angles_raw": angles, "y_pos_raw": y_pos,
|
|
||||||
"preambles": preambles, "background": background,
|
|
||||||
"data": data, # shape: (n_angles, n_rows, n_ch, n_frames, spf), int16
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def average_frames(data: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
|
def average_rows(block: np.ndarray, n: int, discard_remainder: bool) -> np.ndarray:
|
||||||
"""Average every N frames along axis 3.
|
"""Average every N frames of one angle's (n_rows, n_ch, n_frames, spf)
|
||||||
|
block. Returns int16 of shape (n_rows, n_ch, n_out, spf).
|
||||||
|
|
||||||
data shape: (n_angles, n_rows, n_ch, n_frames, spf)
|
Averaging is done in float32 and rounded on cast, matching numpy's mean
|
||||||
Returns array of shape (n_angles, n_rows, n_ch, n_out, spf).
|
followed by an int16 cast in the original implementation.
|
||||||
"""
|
"""
|
||||||
n_frames = data.shape[3]
|
n_frames = block.shape[2]
|
||||||
|
|
||||||
n_full = n_frames // n
|
n_full = n_frames // n
|
||||||
remainder = n_frames % n
|
remainder = n_frames % n
|
||||||
|
|
||||||
# Average full groups using reshape-trick (no Python loop)
|
parts = []
|
||||||
if n_full > 0:
|
if n_full:
|
||||||
full = data[:, :, :, :n_full * n, :] # trim to full groups
|
full = block[:, :, :n_full * n, :].astype(np.float32)
|
||||||
full = full.reshape(data.shape[0], data.shape[1], data.shape[2],
|
full = full.reshape(block.shape[0], block.shape[1], n_full, n, block.shape[3])
|
||||||
n_full, n, data.shape[4]) # (..., n_out, n, spf)
|
parts.append(full.mean(axis=3).astype(np.int16))
|
||||||
averaged = full.mean(axis=4).astype(np.int16) # (..., n_out, spf)
|
if remainder and not discard_remainder:
|
||||||
else:
|
tail = block[:, :, n_full * n:, :].astype(np.float32)
|
||||||
averaged = np.empty((*data.shape[:3], 0, data.shape[4]), dtype=np.int16)
|
parts.append(tail.mean(axis=2, keepdims=True).astype(np.int16))
|
||||||
|
|
||||||
if remainder > 0 and not discard_remainder:
|
if not parts:
|
||||||
tail = data[:, :, :, n_full * n:, :] # shape (..., remainder, spf)
|
return np.empty((*block.shape[:2], 0, block.shape[3]), dtype=np.int16)
|
||||||
tail_avg = tail.mean(axis=3, keepdims=True).astype(np.int16)
|
return parts[0] if len(parts) == 1 else np.concatenate(parts, axis=2)
|
||||||
averaged = np.concatenate([averaged, tail_avg], axis=3)
|
|
||||||
|
|
||||||
return averaged
|
|
||||||
|
|
||||||
|
|
||||||
def write_sras(path: Path, src: dict, data_out: np.ndarray):
|
def write_averaged(out_path: Path, sras: SrasFile, mid_sections: bytes,
|
||||||
"""Write a new .sras file with the averaged waveform data."""
|
n: int, discard_remainder: bool) -> int:
|
||||||
ver = src["ver"]
|
"""Stream each angle through the averager, writing as we go so peak RAM
|
||||||
bps = src["bps"]
|
stays at one angle's block rather than the whole file."""
|
||||||
n_out = data_out.shape[3]
|
bps = sras.bytes_per_sample
|
||||||
|
n_frames_in = int(sras.n_frames[0])
|
||||||
|
n_out = n_frames_in // n
|
||||||
|
if n_frames_in % n and not discard_remainder:
|
||||||
|
n_out += 1
|
||||||
|
|
||||||
# Pack header — update only n_frames_hdr; everything else stays the same
|
|
||||||
header = struct.pack(
|
header = struct.pack(
|
||||||
HDR_FMT,
|
HDR_FMT, b"SRAS", sras.version, sras.n_angles, int(sras.n_rows[0]),
|
||||||
b"SRAS",
|
float(sras.x_start_mm[0]), float(sras.x_delta_mm),
|
||||||
ver,
|
sras.velocity_mm_s, sras.laser_freq_hz,
|
||||||
src["n_angles"],
|
|
||||||
src["n_rows"],
|
|
||||||
src["x_start"],
|
|
||||||
src["x_delta"],
|
|
||||||
src["vel"],
|
|
||||||
src["freq"],
|
|
||||||
n_out, # updated frame count
|
n_out, # updated frame count
|
||||||
src["spf"],
|
sras.samples_per_frame, sras.sample_rate_hz, bps, sras.n_channels,
|
||||||
src["sr"],
|
|
||||||
bps,
|
|
||||||
src["n_ch"],
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Encode waveform data back to original dtype
|
with open(out_path, "wb") as f:
|
||||||
if bps == 1:
|
|
||||||
raw_out = np.clip(data_out, -128, 127).astype(np.int8).tobytes()
|
|
||||||
else:
|
|
||||||
# big-endian int16
|
|
||||||
raw_out = data_out.astype(">i2").tobytes()
|
|
||||||
|
|
||||||
with open(path, "wb") as f:
|
|
||||||
f.write(header)
|
f.write(header)
|
||||||
f.write(src["angles_raw"])
|
f.write(mid_sections)
|
||||||
f.write(src["y_pos_raw"])
|
for a in range(sras.n_angles):
|
||||||
|
averaged = average_rows(sras.data[a], n, discard_remainder)
|
||||||
if ver >= 3:
|
if bps == 1:
|
||||||
for preamble_bytes in src["preambles"]:
|
f.write(np.clip(averaged, -128, 127).astype(np.int8).tobytes())
|
||||||
f.write(struct.pack(">H", len(preamble_bytes)))
|
else:
|
||||||
f.write(preamble_bytes)
|
f.write(averaged.astype(">i2").tobytes())
|
||||||
|
return n_out
|
||||||
if ver >= 4:
|
|
||||||
bg = src["background"]
|
|
||||||
f.write(struct.pack(">I", len(bg)))
|
|
||||||
f.write(bg)
|
|
||||||
|
|
||||||
f.write(raw_out)
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
@@ -194,26 +118,29 @@ def main():
|
|||||||
if not in_path.exists():
|
if not in_path.exists():
|
||||||
print(f"Error: input file not found: {in_path}", file=sys.stderr)
|
print(f"Error: input file not found: {in_path}", file=sys.stderr)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
if out_path.resolve() == in_path.resolve():
|
if out_path.resolve() == in_path.resolve():
|
||||||
print("Error: output path must differ from input path.", file=sys.stderr)
|
print("Error: output path must differ from input path.", file=sys.stderr)
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
print(f"Reading {in_path} ...", flush=True)
|
print(f"Reading {in_path} ...", flush=True)
|
||||||
src = read_sras(in_path)
|
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))})",
|
||||||
|
file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
n_frames_in = src["data"].shape[3]
|
n_frames_in = int(sras.n_frames[0])
|
||||||
print(f" Version : v{src['ver']}")
|
print(f" Version : v{sras.version}")
|
||||||
print(f" Angles : {src['n_angles']}")
|
print(f" Angles : {sras.n_angles}")
|
||||||
print(f" Rows : {src['n_rows']}")
|
print(f" Rows : {int(sras.n_rows[0])}")
|
||||||
print(f" Frames (actual): {n_frames_in}")
|
print(f" Frames (actual): {n_frames_in}")
|
||||||
print(f" Channels : {src['n_ch']}")
|
print(f" Channels : {sras.n_channels}")
|
||||||
print(f" Samples/frame : {src['spf']}")
|
print(f" Samples/frame : {sras.samples_per_frame}")
|
||||||
print(f" Bytes/sample : {src['bps']}")
|
print(f" Bytes/sample : {sras.bytes_per_sample}")
|
||||||
|
|
||||||
if args.n == 1:
|
if args.n == 1:
|
||||||
print("--n 1: no averaging needed; copying file as-is.")
|
print("--n 1: no averaging needed; copying file as-is.")
|
||||||
import shutil
|
|
||||||
shutil.copy2(in_path, out_path)
|
shutil.copy2(in_path, out_path)
|
||||||
print(f"Wrote {out_path}")
|
print(f"Wrote {out_path}")
|
||||||
return
|
return
|
||||||
@@ -223,27 +150,25 @@ def main():
|
|||||||
"The entire dataset will be averaged into a single frame.")
|
"The entire dataset will be averaged into a single frame.")
|
||||||
|
|
||||||
print(f"\nAveraging every {args.n} frames ...", flush=True)
|
print(f"\nAveraging every {args.n} frames ...", flush=True)
|
||||||
data_out = average_frames(src["data"], args.n, args.discard_remainder)
|
print(f"\nWriting {out_path} ...", flush=True)
|
||||||
n_frames_out = data_out.shape[3]
|
mid_sections = read_header_sections(sras)
|
||||||
|
n_frames_out = write_averaged(out_path, sras, mid_sections,
|
||||||
|
args.n, args.discard_remainder)
|
||||||
|
|
||||||
n_full = n_frames_in // args.n
|
n_full = n_frames_in // args.n
|
||||||
remainder = n_frames_in % args.n
|
remainder = n_frames_in % args.n
|
||||||
if remainder and not args.discard_remainder:
|
if remainder and not args.discard_remainder:
|
||||||
status = f"({n_full} full groups + 1 partial group of {remainder})"
|
status = f"({n_full} full groups + 1 partial group of {remainder})"
|
||||||
elif remainder and args.discard_remainder:
|
elif remainder:
|
||||||
status = f"({n_full} full groups, {remainder} trailing frames discarded)"
|
status = f"({n_full} full groups, {remainder} trailing frames discarded)"
|
||||||
else:
|
else:
|
||||||
status = f"({n_full} full groups)"
|
status = f"({n_full} full groups)"
|
||||||
|
|
||||||
print(f" {n_frames_in} frames -> {n_frames_out} frames {status}")
|
print(f" {n_frames_in} frames -> {n_frames_out} frames {status}")
|
||||||
|
|
||||||
print(f"\nWriting {out_path} ...", flush=True)
|
|
||||||
write_sras(out_path, src, data_out)
|
|
||||||
|
|
||||||
in_mb = in_path.stat().st_size / 1024**2
|
in_mb = in_path.stat().st_size / 1024**2
|
||||||
out_mb = out_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" Input size : {in_mb:.1f} MB")
|
||||||
print(f" Output size: {out_mb:.1f} MB ({out_mb/in_mb*100:.1f}% of input)")
|
print(f" Output size: {out_mb:.1f} MB ({out_mb / in_mb * 100:.1f}% of input)")
|
||||||
print("Done.")
|
print("Done.")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
+1055
File diff suppressed because it is too large
Load Diff
+593
@@ -0,0 +1,593 @@
|
|||||||
|
#!/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
|
||||||
+1532
-2083
File diff suppressed because it is too large
Load Diff
+402
@@ -0,0 +1,402 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Background workers for the SRAS viewer.
|
||||||
|
|
||||||
|
Every worker is a plain QObject moved onto its own QThread by
|
||||||
|
SrasViewerWindow._run_worker, exposing signals only. Workers must never touch
|
||||||
|
GUI-thread-owned state (the display caches in particular) — they take
|
||||||
|
everything they need through their constructor and hand results back by signal.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
|
||||||
|
from concurrent.futures.process import BrokenProcessPool
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from PyQt6.QtCore import QObject, pyqtSignal
|
||||||
|
|
||||||
|
import sras_compute as compute
|
||||||
|
from sras_compute import (
|
||||||
|
cache_file, compute_angle_alignment, compute_rf_image, dc_image_mv,
|
||||||
|
)
|
||||||
|
from sras_format import CH3_IDX, CH4_IDX, SrasFile
|
||||||
|
|
||||||
|
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||||||
|
# those processes threads internally, so the two must be divided rather than
|
||||||
|
# both set to the core count. Files also commonly sit on one external drive,
|
||||||
|
# where a dozen concurrent readers is slower than a few — hence the low
|
||||||
|
# default, overridable from the environment.
|
||||||
|
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
|
||||||
|
4, os.cpu_count() or 2)
|
||||||
|
|
||||||
|
# Spawning a pool costs roughly a second of interpreter startup (each child
|
||||||
|
# re-imports the entry module). That is noise against a multi-GB scan but
|
||||||
|
# dominates a batch of small files, where it would make the job *slower* —
|
||||||
|
# so below this total size the batch just runs in the worker thread.
|
||||||
|
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
class CancellableWorker(QObject):
|
||||||
|
"""A worker whose compute polls stop() between row chunks.
|
||||||
|
|
||||||
|
Without this a shutdown has to wait out whatever is in flight, and on a
|
||||||
|
large scan a single angle is ~40 s — far too long to block closing the
|
||||||
|
window. Chunk-level polling bounds the wait to one chunk instead.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self._stop = False
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
self._stop = True
|
||||||
|
|
||||||
|
def _stopped(self) -> bool:
|
||||||
|
return self._stop
|
||||||
|
|
||||||
|
|
||||||
|
class LoadWorker(QObject):
|
||||||
|
finished = pyqtSignal(object) # SrasFile | None
|
||||||
|
error = pyqtSignal(str)
|
||||||
|
|
||||||
|
def __init__(self, path: str):
|
||||||
|
super().__init__()
|
||||||
|
self._path = path
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
self.finished.emit(SrasFile(self._path))
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
|
self.finished.emit(None)
|
||||||
|
|
||||||
|
|
||||||
|
class ComputeWorker(CancellableWorker):
|
||||||
|
"""Computes one displayable image for (angle, channel).
|
||||||
|
|
||||||
|
For CH1/Velocity (FFT-derived) channels, the FFT is only run for pixels
|
||||||
|
whose DC4 (Bias B) mean is at or above dc_threshold_mv — masked pixels are
|
||||||
|
left at 0 MHz without ever being FFT'd, since that's the expensive part of
|
||||||
|
a scan. If the DC4 image for this angle is already known, pass it in as
|
||||||
|
*dc4_mv* to skip re-reading the CH4 channel from disk entirely.
|
||||||
|
|
||||||
|
Emits a plain ``np.ndarray`` already in display units.
|
||||||
|
"""
|
||||||
|
finished = pyqtSignal(object)
|
||||||
|
error = pyqtSignal(str)
|
||||||
|
|
||||||
|
def __init__(self, sras: SrasFile, angle_idx: int, ch_idx: int,
|
||||||
|
apply_bg_sub: bool = True, n_fft: int | None = None,
|
||||||
|
dc_threshold_mv: float = 0.0,
|
||||||
|
dc4_mv: np.ndarray | None = None,
|
||||||
|
is_fft_mode: bool = False):
|
||||||
|
super().__init__()
|
||||||
|
self._sras = sras
|
||||||
|
self._angle = angle_idx
|
||||||
|
self._ch = ch_idx
|
||||||
|
self._apply_bg_sub = apply_bg_sub
|
||||||
|
self._n_fft = n_fft
|
||||||
|
self._dc_threshold = dc_threshold_mv
|
||||||
|
self._dc4_mv = dc4_mv
|
||||||
|
self._is_fft_mode = is_fft_mode
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
if self._is_fft_mode:
|
||||||
|
img = compute_rf_image(
|
||||||
|
self._sras, self._angle, dc_threshold_mv=self._dc_threshold,
|
||||||
|
apply_bg_sub=self._apply_bg_sub, n_fft=self._n_fft,
|
||||||
|
dc4_mv=self._dc4_mv, should_stop=self._stopped)
|
||||||
|
else:
|
||||||
|
img = dc_image_mv(self._sras, self._angle, self._ch,
|
||||||
|
should_stop=self._stopped)
|
||||||
|
# On cancellation the image is only partly filled, so hand back
|
||||||
|
# None rather than something that would be cached as real. The
|
||||||
|
# signal still fires either way — it is what quits the thread.
|
||||||
|
self.finished.emit(None if self._stop else img)
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
|
|
||||||
|
|
||||||
|
class DcPrecomputeWorker(CancellableWorker):
|
||||||
|
"""Computes CH3/CH4 DC images for every angle in the background.
|
||||||
|
|
||||||
|
DC images are cheap (a per-waveform mean, no FFT) compared to the
|
||||||
|
CH1/Velocity FFT, so precomputing them for the whole file right after load
|
||||||
|
makes switching angles instant while on a DC channel, and also means the
|
||||||
|
FFT masking step (which needs a DC4 image) rarely has to wait on anything.
|
||||||
|
|
||||||
|
Angles are computed on a thread pool — the work is a pure mean over the
|
||||||
|
waveform block, so it is I/O- and bandwidth-bound and embarrassingly
|
||||||
|
parallel. Results are emitted one at a time as they land (out of angle
|
||||||
|
order), and always from this worker's own thread: nothing emits a Qt
|
||||||
|
signal from a pool thread.
|
||||||
|
"""
|
||||||
|
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||||||
|
finished = pyqtSignal()
|
||||||
|
error = pyqtSignal(str)
|
||||||
|
|
||||||
|
def __init__(self, sras: SrasFile):
|
||||||
|
super().__init__()
|
||||||
|
self._sras = sras
|
||||||
|
|
||||||
|
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
|
||||||
|
# max_workers=1 *and* a budget share: this call is one of several
|
||||||
|
# concurrent angles, and both the thread count and the buffer size
|
||||||
|
# have to be divided (see compute.plan_angle_level).
|
||||||
|
kw = dict(max_workers=1, budget=self._angle_budget,
|
||||||
|
should_stop=self._stopped)
|
||||||
|
return (a,
|
||||||
|
dc_image_mv(self._sras, a, CH3_IDX, **kw),
|
||||||
|
dc_image_mv(self._sras, a, CH4_IDX, **kw))
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
n = self._sras.n_angles
|
||||||
|
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
|
||||||
|
pool = ThreadPoolExecutor(max_workers=n_workers)
|
||||||
|
try:
|
||||||
|
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
|
||||||
|
for fut in as_completed(futures):
|
||||||
|
if self._stop:
|
||||||
|
break
|
||||||
|
a, dc3, dc4 = fut.result()
|
||||||
|
self.angle_done.emit(a, dc3, dc4)
|
||||||
|
finally:
|
||||||
|
# cancel_futures drops the queued angles; should_stop lets the
|
||||||
|
# in-flight ones bail within a chunk. Not waiting here is what
|
||||||
|
# keeps closing the window responsive on a large scan.
|
||||||
|
pool.shutdown(wait=not self._stop, cancel_futures=True)
|
||||||
|
self.finished.emit()
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
|
|
||||||
|
|
||||||
|
class BatchCacheWorker(QObject):
|
||||||
|
"""Batch-computes and stores DC or FFT images into each of *paths*'s v7
|
||||||
|
CACH tail, in place — converting v6 sources to v7 on first use, or updating
|
||||||
|
an existing v7 file's cache blocks without disturbing whatever the other
|
||||||
|
block already holds.
|
||||||
|
|
||||||
|
*mode* is ``"dc"`` (CH3/CH4 mean images) or ``"fft"`` (CH1 peak-frequency
|
||||||
|
images, unmasked — masking is applied at display time, same as v5's PREC
|
||||||
|
convention).
|
||||||
|
|
||||||
|
Files are processed one per subprocess: they are fully independent, each
|
||||||
|
opens its own memmap and writes only its own bytes, and only path strings
|
||||||
|
and scalars cross the process boundary. Emits ``progress(int)`` (0–100 by
|
||||||
|
files completed), ``file_done(str, str)`` (path, error message or "") so
|
||||||
|
one file's failure doesn't abort the batch, and ``finished()``.
|
||||||
|
"""
|
||||||
|
progress = pyqtSignal(int)
|
||||||
|
file_done = pyqtSignal(str, str)
|
||||||
|
finished = pyqtSignal()
|
||||||
|
|
||||||
|
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
|
||||||
|
super().__init__()
|
||||||
|
self._paths = paths
|
||||||
|
self._mode = mode
|
||||||
|
self._apply_bg_sub = apply_bg_sub
|
||||||
|
|
||||||
|
def _report(self, path: str, err: str, done: int, total: int):
|
||||||
|
self.file_done.emit(path, err)
|
||||||
|
self.progress.emit(int(done / max(1, total) * 100))
|
||||||
|
|
||||||
|
def _run_pooled(self, paths: list[str], n_procs: int) -> list[str]:
|
||||||
|
"""Process the batch across *n_procs* subprocesses. Returns the paths
|
||||||
|
that never got a real answer because the pool itself died, so the
|
||||||
|
caller can retry them in-process.
|
||||||
|
|
||||||
|
Under spawn each child re-imports the entry module, so the batch must
|
||||||
|
survive that going wrong (an unguarded __main__, a frozen build, a
|
||||||
|
sandbox that forbids subprocesses) rather than reporting every file as
|
||||||
|
failed — hence the retry list instead of a per-file error.
|
||||||
|
"""
|
||||||
|
# Each child threads internally; divide the machine rather than
|
||||||
|
# letting every process claim every core.
|
||||||
|
per_proc_workers = max(1, (os.cpu_count() or 4) // n_procs)
|
||||||
|
unresolved: list[str] = []
|
||||||
|
done = 0
|
||||||
|
|
||||||
|
with ProcessPoolExecutor(max_workers=n_procs) as executor:
|
||||||
|
futures = {
|
||||||
|
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
|
||||||
|
compute.get_fft_backend(), per_proc_workers): p
|
||||||
|
for p in paths
|
||||||
|
}
|
||||||
|
for fut in as_completed(futures):
|
||||||
|
path = futures[fut]
|
||||||
|
try:
|
||||||
|
err = fut.result()
|
||||||
|
except BrokenProcessPool:
|
||||||
|
unresolved.append(path)
|
||||||
|
continue
|
||||||
|
except Exception as exc:
|
||||||
|
err = str(exc)
|
||||||
|
done += 1
|
||||||
|
self._report(path, err, done, len(paths))
|
||||||
|
|
||||||
|
return unresolved
|
||||||
|
|
||||||
|
def _run_inline(self, paths: list[str], done: int, total: int):
|
||||||
|
"""Fallback / single-file path: compute in this thread. Still uses the
|
||||||
|
full core count internally, since nothing else is competing."""
|
||||||
|
for path in paths:
|
||||||
|
try:
|
||||||
|
err = cache_file(path, self._mode, self._apply_bg_sub,
|
||||||
|
compute.get_fft_backend(), compute._MAX_WORKERS)
|
||||||
|
except Exception as exc:
|
||||||
|
err = str(exc)
|
||||||
|
done += 1
|
||||||
|
self._report(path, err, done, total)
|
||||||
|
|
||||||
|
def _worth_pooling(self, paths: list[str]) -> bool:
|
||||||
|
if len(paths) < 2:
|
||||||
|
return False
|
||||||
|
total = 0
|
||||||
|
for p in paths:
|
||||||
|
try:
|
||||||
|
total += os.path.getsize(p)
|
||||||
|
except OSError:
|
||||||
|
pass # unreadable files are reported by cache_file
|
||||||
|
return total >= _BATCH_POOL_MIN_BYTES
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
paths = self._paths
|
||||||
|
n_procs = max(1, min(_BATCH_MAX_PROCS, len(paths)))
|
||||||
|
|
||||||
|
if not self._worth_pooling(paths):
|
||||||
|
self._run_inline(paths, 0, len(paths))
|
||||||
|
self.finished.emit()
|
||||||
|
return
|
||||||
|
|
||||||
|
try:
|
||||||
|
unresolved = self._run_pooled(paths, n_procs)
|
||||||
|
except Exception:
|
||||||
|
# The pool could not be created or collapsed wholesale.
|
||||||
|
unresolved = list(paths)
|
||||||
|
|
||||||
|
if unresolved:
|
||||||
|
self._run_inline(unresolved, len(paths) - len(unresolved), len(paths))
|
||||||
|
|
||||||
|
self.finished.emit()
|
||||||
|
|
||||||
|
|
||||||
|
class AngleAlignmentWorker(QObject):
|
||||||
|
"""Computes rotation+translation alignment for every angle in *sras*,
|
||||||
|
referenced to *ref_angle_idx*, from each angle's binarized CH4 mask.
|
||||||
|
Rotation is analytic (from sras.angles_deg); only translation is found by
|
||||||
|
phase correlation.
|
||||||
|
"""
|
||||||
|
progress = pyqtSignal(int) # 0–100
|
||||||
|
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
|
||||||
|
|
||||||
|
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||||||
|
super().__init__()
|
||||||
|
self._sras = sras
|
||||||
|
self._ref = ref_angle_idx
|
||||||
|
self._threshold = dc_threshold_mv
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
result = compute_angle_alignment(
|
||||||
|
self._sras, self._ref, self._threshold,
|
||||||
|
progress_cb=self.progress.emit)
|
||||||
|
self.finished.emit(result, "")
|
||||||
|
except Exception as exc:
|
||||||
|
self.finished.emit(None, str(exc))
|
||||||
|
|
||||||
|
|
||||||
|
class Ch4MaskWorker(QObject):
|
||||||
|
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
|
||||||
|
ManualAlignmentDialog's initial threshold-mask overlay.
|
||||||
|
|
||||||
|
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
|
||||||
|
from raw waveforms, so this only does real work for a file that hasn't
|
||||||
|
gone through the v7 "Convert" batch step and for angles the main
|
||||||
|
window's own DcPrecomputeWorker (which runs automatically right after
|
||||||
|
every file load) hasn't reached yet. In the common case — the user opens
|
||||||
|
Fusion -> Manual Alignment after DC precompute has already finished —
|
||||||
|
*angle_indices* is empty and this worker is never even constructed (see
|
||||||
|
ManualAlignmentDialog._start_mask_prep).
|
||||||
|
"""
|
||||||
|
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||||||
|
finished = pyqtSignal()
|
||||||
|
error = pyqtSignal(str)
|
||||||
|
|
||||||
|
def __init__(self, sras: SrasFile, angle_indices: list[int]):
|
||||||
|
super().__init__()
|
||||||
|
self._sras = sras
|
||||||
|
self._angles = angle_indices
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
n_workers, budget = compute.plan_angle_level(self._sras)
|
||||||
|
pool = ThreadPoolExecutor(max_workers=n_workers)
|
||||||
|
try:
|
||||||
|
futures = {
|
||||||
|
pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
|
||||||
|
max_workers=1, budget=budget): a
|
||||||
|
for a in self._angles
|
||||||
|
}
|
||||||
|
for fut in as_completed(futures):
|
||||||
|
a = futures[fut]
|
||||||
|
self.angle_done.emit(a, fut.result())
|
||||||
|
finally:
|
||||||
|
pool.shutdown(wait=True)
|
||||||
|
self.finished.emit()
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
|
|
||||||
|
|
||||||
|
class CrossCorrelateWorker(QObject):
|
||||||
|
"""FFT phase-correlation translation for each of *angle_indices* against
|
||||||
|
*ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
|
||||||
|
|
||||||
|
Runs on a background thread — a real many-angle, high-resolution scan's
|
||||||
|
correlation (even at its downsampled working resolution) can take long
|
||||||
|
enough that doing all of them on the GUI thread would visibly freeze the
|
||||||
|
dialog. Rotation is set to the same analytic scan-angle delta Auto
|
||||||
|
De-rotate uses alongside the correlated shift, since a translation
|
||||||
|
search is only meaningful once both angles' content is already oriented
|
||||||
|
the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory
|
||||||
|
per-angle images/pivots — this worker does no fetching of its own.
|
||||||
|
"""
|
||||||
|
angle_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm
|
||||||
|
finished = pyqtSignal()
|
||||||
|
error = pyqtSignal(str)
|
||||||
|
|
||||||
|
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
||||||
|
dc4_mv: dict[int, np.ndarray], pivot_mm: dict[int, tuple[float, float]],
|
||||||
|
*, use_mask: bool, dc_threshold_mv: float, margin_frac: float):
|
||||||
|
super().__init__()
|
||||||
|
self._sras = sras
|
||||||
|
self._ref = ref_angle_idx
|
||||||
|
self._angles = angle_indices
|
||||||
|
self._dc4_mv = dc4_mv
|
||||||
|
self._pivot_mm = pivot_mm
|
||||||
|
self._use_mask = use_mask
|
||||||
|
self._threshold = dc_threshold_mv
|
||||||
|
self._margin = margin_frac
|
||||||
|
|
||||||
|
def _one(self, a: int) -> tuple[int, float, float, float]:
|
||||||
|
theta = compute._theta_deg(self._sras, a, self._ref)
|
||||||
|
dx, dy = compute.correlate_translation_mm(
|
||||||
|
self._sras, a, self._ref, self._dc4_mv, self._pivot_mm,
|
||||||
|
use_mask=self._use_mask, dc_threshold_mv=self._threshold,
|
||||||
|
margin_frac=self._margin)
|
||||||
|
return a, theta, dx, dy
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
n_workers, _budget = compute.plan_angle_level(self._sras)
|
||||||
|
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
|
||||||
|
try:
|
||||||
|
futures = [pool.submit(self._one, a) for a in self._angles]
|
||||||
|
for fut in as_completed(futures):
|
||||||
|
a, theta, dx, dy = fut.result()
|
||||||
|
self.angle_done.emit(a, theta, dx, dy)
|
||||||
|
finally:
|
||||||
|
pool.shutdown(wait=True)
|
||||||
|
self.finished.emit()
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
@@ -0,0 +1,197 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Golden-output equivalence harness for the sras-viewer refactor.
|
||||||
|
|
||||||
|
Computes a battery of DC / FFT / alignment outputs and prints a stable hash
|
||||||
|
for each. Run it on the pre-refactor commit to capture a baseline, then again
|
||||||
|
after the refactor and diff the two reports — every line must match.
|
||||||
|
|
||||||
|
Imports work against both the pre-refactor monolith (`sras_viewer`) and the
|
||||||
|
post-refactor split (`sras_format` + `sras_compute`), so the *same* script
|
||||||
|
produces both sides of the comparison.
|
||||||
|
|
||||||
|
Hashes canonicalise to native little-endian float64 before hashing, so a
|
||||||
|
deliberate dtype/byte-order change that preserves values does not show up as
|
||||||
|
a false mismatch.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python tools/check_equivalence.py --out baseline.txt
|
||||||
|
python tools/check_equivalence.py --out after.txt --real /path/to/big.sras
|
||||||
|
diff baseline.txt after.txt
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import copy
|
||||||
|
import hashlib
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
|
|
||||||
|
# --- Import shim: split modules if present, else the monolith --------------
|
||||||
|
try:
|
||||||
|
from sras_format import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
|
||||||
|
import sras_compute as C
|
||||||
|
_LAYOUT = "split"
|
||||||
|
except ImportError:
|
||||||
|
import sras_viewer as _V
|
||||||
|
from sras_viewer import SrasFile, CH1_IDX, CH3_IDX, CH4_IDX, adc_to_mv
|
||||||
|
C = _V
|
||||||
|
_LAYOUT = "monolith"
|
||||||
|
|
||||||
|
compute_dc_image = C.compute_dc_image
|
||||||
|
compute_rf_image = C.compute_rf_image
|
||||||
|
compute_alignment = C._compute_angle_alignment
|
||||||
|
apply_alignment = C.apply_alignment
|
||||||
|
|
||||||
|
import tools.make_test_sras as gen # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def h(arr) -> str:
|
||||||
|
"""Stable hash of an array's *values*, independent of dtype/byte order."""
|
||||||
|
a = np.ascontiguousarray(np.asarray(arr, dtype=np.float64))
|
||||||
|
return hashlib.sha256(a.tobytes()).hexdigest()[:16]
|
||||||
|
|
||||||
|
|
||||||
|
def row_slice(sras: SrasFile, angle_idx: int, n_rows: int) -> SrasFile:
|
||||||
|
"""A shallow view of *sras* restricted to the first *n_rows* rows of
|
||||||
|
*angle_idx*, so the huge real file can be exercised in seconds."""
|
||||||
|
view = copy.copy(sras)
|
||||||
|
n = min(int(n_rows), int(sras.n_rows[angle_idx]))
|
||||||
|
view.n_rows = np.array(sras.n_rows, copy=True)
|
||||||
|
view.n_rows[angle_idx] = n
|
||||||
|
view.data = list(sras.data)
|
||||||
|
view.data[angle_idx] = sras.data[angle_idx][:n]
|
||||||
|
view._y_pos_per_angle = list(sras._y_pos_per_angle)
|
||||||
|
view._y_pos_per_angle[angle_idx] = sras._y_pos_per_angle[angle_idx][:n]
|
||||||
|
return view
|
||||||
|
|
||||||
|
|
||||||
|
def report(lines: list[str], label: str, value: str):
|
||||||
|
lines.append(f"{label:<58} {value}")
|
||||||
|
|
||||||
|
|
||||||
|
def check_file(path: Path, lines: list[str], tag: str,
|
||||||
|
angles: list[int], n_rows: int | None):
|
||||||
|
sras = SrasFile(str(path))
|
||||||
|
report(lines, f"[{tag}] version/n_angles",
|
||||||
|
f"v{sras.version} n={sras.n_angles}")
|
||||||
|
report(lines, f"[{tag}] geometry",
|
||||||
|
f"rows={list(map(int, sras.n_rows))} frames={list(map(int, sras.n_frames))}")
|
||||||
|
report(lines, f"[{tag}] calibration",
|
||||||
|
" ".join(f"{m:.6g}/{o:.6g}/{z:.6g}" for m, o, z in
|
||||||
|
zip(sras.ch_ymult_mv, sras.ch_yoff_adc, sras.ch_yzero_mv)))
|
||||||
|
report(lines, f"[{tag}] angles_deg", h(sras.angles_deg))
|
||||||
|
if sras.background is not None:
|
||||||
|
report(lines, f"[{tag}] background", h(sras.background))
|
||||||
|
|
||||||
|
for a in angles:
|
||||||
|
if a >= sras.n_angles:
|
||||||
|
continue
|
||||||
|
s = row_slice(sras, a, n_rows) if n_rows else sras
|
||||||
|
|
||||||
|
report(lines, f"[{tag}] x_axis_mm(a={a})", h(s.x_axis_mm(a)))
|
||||||
|
report(lines, f"[{tag}] y_positions_mm(a={a})", h(s.y_positions_mm(a)))
|
||||||
|
|
||||||
|
for ch, name in ((CH3_IDX, "CH3"), (CH4_IDX, "CH4")):
|
||||||
|
raw = compute_dc_image(s, a, ch)
|
||||||
|
report(lines, f"[{tag}] dc_adc(a={a},{name})", h(raw))
|
||||||
|
mv = adc_to_mv(raw, s.ch_ymult_mv[ch], s.ch_yoff_adc[ch],
|
||||||
|
s.ch_yzero_mv[ch])
|
||||||
|
report(lines, f"[{tag}] dc_mv(a={a},{name})", h(mv))
|
||||||
|
|
||||||
|
dc4 = adc_to_mv(compute_dc_image(s, a, CH4_IDX),
|
||||||
|
s.ch_ymult_mv[CH4_IDX], s.ch_yoff_adc[CH4_IDX],
|
||||||
|
s.ch_yzero_mv[CH4_IDX])
|
||||||
|
thresholds = [-1e9, float(np.median(dc4))]
|
||||||
|
|
||||||
|
for bg in (False, True):
|
||||||
|
if bg and s.background is None:
|
||||||
|
continue
|
||||||
|
for pad in (1, 2):
|
||||||
|
n_fft = s.samples_per_frame * pad if pad > 1 else None
|
||||||
|
for ti, thr in enumerate(thresholds):
|
||||||
|
img = compute_rf_image(s, a, dc_threshold_mv=thr,
|
||||||
|
apply_bg_sub=bg, n_fft=n_fft)
|
||||||
|
report(lines,
|
||||||
|
f"[{tag}] rf(a={a},bg={int(bg)},pad={pad},thr{ti})",
|
||||||
|
h(img))
|
||||||
|
# dc4_mv passthrough must give the identical result
|
||||||
|
img2 = compute_rf_image(s, a, dc_threshold_mv=thr,
|
||||||
|
apply_bg_sub=bg, n_fft=n_fft,
|
||||||
|
dc4_mv=dc4)
|
||||||
|
same = "SAME" if h(img) == h(img2) else "DIFFER"
|
||||||
|
report(lines,
|
||||||
|
f"[{tag}] rf-dc4arg(a={a},bg={int(bg)},pad={pad},thr{ti})",
|
||||||
|
same)
|
||||||
|
|
||||||
|
|
||||||
|
def check_alignment(path: Path, lines: list[str], tag: str):
|
||||||
|
sras = SrasFile(str(path))
|
||||||
|
if sras.n_angles < 2:
|
||||||
|
return
|
||||||
|
dc4 = adc_to_mv(compute_dc_image(sras, 0, CH4_IDX),
|
||||||
|
sras.ch_ymult_mv[CH4_IDX], sras.ch_yoff_adc[CH4_IDX],
|
||||||
|
sras.ch_yzero_mv[CH4_IDX])
|
||||||
|
thr = float(np.median(dc4))
|
||||||
|
res = compute_alignment(sras, 0, thr)
|
||||||
|
report(lines, f"[{tag}] align canvas_shape", str(res.canvas_shape))
|
||||||
|
report(lines, f"[{tag}] align canvas_origin",
|
||||||
|
f"{res.canvas_origin_mm[0]:.9g},{res.canvas_origin_mm[1]:.9g}")
|
||||||
|
report(lines, f"[{tag}] align canvas_pitch",
|
||||||
|
f"{res.canvas_dx_mm:.9g},{res.canvas_dy_mm:.9g}")
|
||||||
|
for a in sorted(res.per_angle):
|
||||||
|
t = res.per_angle[a]
|
||||||
|
report(lines, f"[{tag}] align shift_mm(a={a})",
|
||||||
|
f"{t.shift_mm[0]:.9g},{t.shift_mm[1]:.9g}")
|
||||||
|
report(lines, f"[{tag}] align rot(a={a})", f"{t.rotation_deg:.9g}")
|
||||||
|
report(lines, f"[{tag}] align matrix(a={a})", h(t.matrix))
|
||||||
|
report(lines, f"[{tag}] align offset(a={a})", h(t.offset))
|
||||||
|
img = compute_dc_image(sras, a, CH4_IDX)
|
||||||
|
report(lines, f"[{tag}] align resampled(a={a})",
|
||||||
|
h(apply_alignment(res, a, img)))
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser(description=__doc__)
|
||||||
|
p.add_argument("--out", required=True, help="report file to write")
|
||||||
|
p.add_argument("--real", help="optional path to a real .sras file")
|
||||||
|
p.add_argument("--real-rows", type=int, default=2,
|
||||||
|
help="rows per angle to sample from the real file")
|
||||||
|
p.add_argument("--real-angles", type=int, default=2,
|
||||||
|
help="how many angles to sample from the real file")
|
||||||
|
p.add_argument("--scratch", default=".",
|
||||||
|
help="directory for generated synthetic files")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
lines = [f"# layout: {_LAYOUT}", f"# numpy: {np.__version__}"]
|
||||||
|
|
||||||
|
scratch = Path(args.scratch)
|
||||||
|
synth = scratch / "equiv_synth.sras"
|
||||||
|
gen.write(synth, n_angles=4, seed=0, samples_per_frame=64)
|
||||||
|
check_file(synth, lines, "synth", angles=[0, 1, 2, 3], n_rows=None)
|
||||||
|
check_alignment(synth, lines, "synth")
|
||||||
|
|
||||||
|
# A second synthetic with an odd sample count, to catch off-by-one in
|
||||||
|
# rfft bin handling and chunk-boundary arithmetic.
|
||||||
|
synth_odd = scratch / "equiv_synth_odd.sras"
|
||||||
|
gen.write(synth_odd, n_angles=2, seed=7, samples_per_frame=37)
|
||||||
|
check_file(synth_odd, lines, "odd", angles=[0, 1], n_rows=None)
|
||||||
|
|
||||||
|
if args.real:
|
||||||
|
real = Path(args.real)
|
||||||
|
if real.exists():
|
||||||
|
check_file(real, lines, "real",
|
||||||
|
angles=list(range(args.real_angles)),
|
||||||
|
n_rows=args.real_rows)
|
||||||
|
else:
|
||||||
|
lines.append(f"# real file not found: {real}")
|
||||||
|
|
||||||
|
Path(args.out).write_text("\n".join(lines) + "\n")
|
||||||
|
print("\n".join(lines))
|
||||||
|
print(f"\nWrote {args.out} ({len(lines)} lines)")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,191 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Generate small synthetic .sras files for testing.
|
||||||
|
|
||||||
|
Writes v6 files (per-angle geometry, ragged waveform blocks) matching
|
||||||
|
scan_format.md, with deterministic pseudo-random waveform content so a test
|
||||||
|
can compute expected DC/FFT images independently of the reader under test.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python tools/make_test_sras.py out.sras [--angles 3] [--seed 0]
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import struct
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
HDR_FMT_V6 = ">4sBHfffffffIdBB"
|
||||||
|
GEO_FMT_V6 = ">ffIH"
|
||||||
|
|
||||||
|
# Per-angle (n_rows, n_frames) — deliberately different per angle so ragged
|
||||||
|
# geometry handling is actually exercised.
|
||||||
|
_GEOMETRY = [(5, 7), (4, 11), (6, 9), (3, 13), (7, 6)]
|
||||||
|
|
||||||
|
_SAMPLE_RATE_HZ = 6.25e9
|
||||||
|
_VELOCITY_MM_S = 20.0
|
||||||
|
_LASER_FREQ_HZ = 1000.0
|
||||||
|
_ROW_SPACING_MM = 0.05
|
||||||
|
|
||||||
|
|
||||||
|
def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes:
|
||||||
|
"""A Tektronix WFMOutpre string in verbose (keyword) form — the reader
|
||||||
|
pulls YMULT/YOFF/YZERO out of it by name, so the keywords must be
|
||||||
|
present literally. YMULT/YZERO are in volts, as the scope reports them."""
|
||||||
|
return (
|
||||||
|
":WFMOUTPRE:BYT_NR 1;BIT_NR 8;ENCDG BIN;BN_FMT RI;BYT_OR MSB;"
|
||||||
|
'WFID "Ch1, DC coupling";NR_PT 2500;PT_FMT Y;'
|
||||||
|
"XINCR 1.6000E-10;XZERO 0.0E0;XUNIT \"s\";"
|
||||||
|
f"YMULT {ymult_v:.6E};YOFF {yoff_adc:.6E};YZERO {yzero_v:.6E};"
|
||||||
|
'YUNIT "V"'
|
||||||
|
).encode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
def build(n_angles: int, seed: int, samples_per_frame: int,
|
||||||
|
geometry: list[tuple[int, int]] | None = None) -> tuple[bytes, dict]:
|
||||||
|
rng = np.random.default_rng(seed)
|
||||||
|
src_geom = geometry or _GEOMETRY
|
||||||
|
geom = [src_geom[a % len(src_geom)] for a in range(n_angles)]
|
||||||
|
n_ch = 3
|
||||||
|
bps = 1
|
||||||
|
|
||||||
|
angles_deg = np.linspace(0.0, 60.0, n_angles, dtype=np.float32)
|
||||||
|
# Distinct calibration per channel so a swapped-channel bug is visible.
|
||||||
|
cal = [
|
||||||
|
(1.5625e-3, -87.04, 0.0),
|
||||||
|
(2.0000e-3, -60.00, 1.0e-3),
|
||||||
|
(2.5000e-3, -40.00, -2.0e-3),
|
||||||
|
]
|
||||||
|
|
||||||
|
out = bytearray()
|
||||||
|
out += struct.pack(
|
||||||
|
HDR_FMT_V6, b"SRAS", 6, n_angles,
|
||||||
|
0.0, 0.0, 1.0, 1.0, _ROW_SPACING_MM,
|
||||||
|
_VELOCITY_MM_S, _LASER_FREQ_HZ,
|
||||||
|
samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
|
||||||
|
)
|
||||||
|
out += angles_deg.astype(">f4").tobytes()
|
||||||
|
|
||||||
|
x_starts = []
|
||||||
|
for a, (n_rows, n_frames) in enumerate(geom):
|
||||||
|
x_start = -0.5 + 0.1 * a
|
||||||
|
x_starts.append(x_start)
|
||||||
|
out += struct.pack(GEO_FMT_V6, x_start, 1.0, n_frames, n_rows)
|
||||||
|
|
||||||
|
y_positions = []
|
||||||
|
for a, (n_rows, _) in enumerate(geom):
|
||||||
|
y = (0.2 * a + np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32)
|
||||||
|
y_positions.append(y)
|
||||||
|
out += y.astype(">f4").tobytes()
|
||||||
|
|
||||||
|
for ymult_v, yoff, yzero_v in cal:
|
||||||
|
p = _preamble(ymult_v, yoff, yzero_v)
|
||||||
|
out += struct.pack(">H", len(p)) + p
|
||||||
|
|
||||||
|
background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8)
|
||||||
|
out += struct.pack(">I", samples_per_frame) + background.tobytes()
|
||||||
|
|
||||||
|
# Waveform data. CH1 gets a sinusoid at a per-pixel frequency so the FFT
|
||||||
|
# peak is predictable; CH3/CH4 get per-pixel DC levels so the mean is too.
|
||||||
|
t = np.arange(samples_per_frame)
|
||||||
|
waveforms = []
|
||||||
|
for a, (n_rows, n_frames) in enumerate(geom):
|
||||||
|
block = np.empty((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8)
|
||||||
|
for r in range(n_rows):
|
||||||
|
for f in range(n_frames):
|
||||||
|
bin_idx = 3 + ((a + r + f) % 17)
|
||||||
|
phase = 2 * np.pi * bin_idx * t / samples_per_frame
|
||||||
|
block[r, 0, f] = np.clip(
|
||||||
|
np.round(60 * np.sin(phase)), -128, 127).astype(np.int8)
|
||||||
|
block[r, 1, f] = np.int8((a * 7 + r * 3 + f) % 100 - 50)
|
||||||
|
block[r, 2, f] = np.int8((a * 5 + r * 11 + f * 2) % 120 - 60)
|
||||||
|
waveforms.append(block)
|
||||||
|
out += block.tobytes()
|
||||||
|
|
||||||
|
meta = {
|
||||||
|
"n_angles": n_angles,
|
||||||
|
"geometry": geom,
|
||||||
|
"angles_deg": angles_deg,
|
||||||
|
"x_starts": x_starts,
|
||||||
|
"y_positions": y_positions,
|
||||||
|
"cal": cal,
|
||||||
|
"background": background,
|
||||||
|
"waveforms": waveforms,
|
||||||
|
"samples_per_frame": samples_per_frame,
|
||||||
|
"sample_rate_hz": _SAMPLE_RATE_HZ,
|
||||||
|
}
|
||||||
|
return bytes(out), meta
|
||||||
|
|
||||||
|
|
||||||
|
def write(path: Path, n_angles: int = 3, seed: int = 0,
|
||||||
|
samples_per_frame: int = 64,
|
||||||
|
geometry: list[tuple[int, int]] | None = None) -> dict:
|
||||||
|
payload, meta = build(n_angles, seed, samples_per_frame, geometry)
|
||||||
|
path.write_bytes(payload)
|
||||||
|
return meta
|
||||||
|
|
||||||
|
|
||||||
|
HDR_FMT_LEGACY = ">4sBHHffffIIdBB"
|
||||||
|
|
||||||
|
|
||||||
|
def write_legacy(path: Path, version: int = 4, n_angles: int = 2,
|
||||||
|
n_rows: int = 4, n_frames: int = 10,
|
||||||
|
samples_per_frame: int = 32, seed: int = 0) -> dict:
|
||||||
|
"""Write a v2/v3/v4 file: uniform geometry, one flat waveform block.
|
||||||
|
|
||||||
|
Used to exercise sras_average.py, which only handles the legacy formats.
|
||||||
|
"""
|
||||||
|
rng = np.random.default_rng(seed)
|
||||||
|
n_ch, bps = 3, 1
|
||||||
|
|
||||||
|
out = bytearray()
|
||||||
|
out += struct.pack(
|
||||||
|
HDR_FMT_LEGACY, b"SRAS", version, n_angles, n_rows,
|
||||||
|
-0.5, 1.0, _VELOCITY_MM_S, _LASER_FREQ_HZ,
|
||||||
|
n_frames, samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
|
||||||
|
)
|
||||||
|
angles = np.linspace(0.0, 45.0, n_angles, dtype=np.float32)
|
||||||
|
out += angles.astype(">f4").tobytes()
|
||||||
|
y = (np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32)
|
||||||
|
out += y.astype(">f4").tobytes()
|
||||||
|
|
||||||
|
if version >= 3:
|
||||||
|
for ymult_v, yoff, yzero_v in ((1.5625e-3, -87.04, 0.0),
|
||||||
|
(2.0e-3, -60.0, 1.0e-3),
|
||||||
|
(2.5e-3, -40.0, -2.0e-3))[:n_ch]:
|
||||||
|
p = _preamble(ymult_v, yoff, yzero_v)
|
||||||
|
out += struct.pack(">H", len(p)) + p
|
||||||
|
|
||||||
|
background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8)
|
||||||
|
if version >= 4:
|
||||||
|
out += struct.pack(">I", samples_per_frame) + background.tobytes()
|
||||||
|
|
||||||
|
data = rng.integers(-100, 101,
|
||||||
|
size=(n_angles, n_rows, n_ch, n_frames, samples_per_frame),
|
||||||
|
dtype=np.int8)
|
||||||
|
out += data.tobytes()
|
||||||
|
path.write_bytes(bytes(out))
|
||||||
|
return {"version": version, "n_angles": n_angles, "n_rows": n_rows,
|
||||||
|
"n_frames": n_frames, "samples_per_frame": samples_per_frame,
|
||||||
|
"n_channels": n_ch, "data": data, "angles_deg": angles,
|
||||||
|
"y_positions": y, "background": background}
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser(description=__doc__)
|
||||||
|
p.add_argument("output")
|
||||||
|
p.add_argument("--angles", type=int, default=3)
|
||||||
|
p.add_argument("--seed", type=int, default=0)
|
||||||
|
p.add_argument("--spf", type=int, default=64, help="samples per frame")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
out = Path(args.output)
|
||||||
|
meta = write(out, args.angles, args.seed, args.spf)
|
||||||
|
print(f"Wrote {out} ({out.stat().st_size:,} bytes)")
|
||||||
|
print(f" angles : {meta['n_angles']}")
|
||||||
|
print(f" geometry : {meta['geometry']}")
|
||||||
|
print(f" spf : {meta['samples_per_frame']}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,485 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Headless GUI test: drives SrasViewerWindow through the real Qt widgets,
|
||||||
|
signals and worker threads under the offscreen platform plugin.
|
||||||
|
|
||||||
|
Covers the interactions a manual smoke test would: load, switch angles and
|
||||||
|
channels, background DC precompute, lazy FFT compute, threshold and bg-sub
|
||||||
|
changes, angle alignment, manual angle alignment, aligned view, ROI
|
||||||
|
draw/move, and CSV export.
|
||||||
|
|
||||||
|
Usage: QT_QPA_PLATFORM=offscreen python tools/test_gui.py [file.sras]
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
|
||||||
|
|
||||||
|
import numpy as np # noqa: E402
|
||||||
|
from PyQt6.QtCore import QEventLoop, Qt, QTimer # noqa: E402
|
||||||
|
from PyQt6.QtTest import QTest # noqa: E402
|
||||||
|
from PyQt6.QtWidgets import QApplication, QMessageBox # noqa: E402
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
|
|
||||||
|
import sras_compute as compute # noqa: E402
|
||||||
|
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX # noqa: E402
|
||||||
|
from sras_viewer import RoiQuad, SrasViewerWindow, VELOCITY_MODE_IDX # noqa: E402
|
||||||
|
import tools.make_test_sras as gen # noqa: E402
|
||||||
|
|
||||||
|
_failures: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
def check(name: str, ok: bool, detail: str = ""):
|
||||||
|
print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else ""))
|
||||||
|
if not ok:
|
||||||
|
_failures.append(name)
|
||||||
|
|
||||||
|
|
||||||
|
def pump(ms: int = 250):
|
||||||
|
"""Run the event loop for a while so queued signals and worker threads
|
||||||
|
make progress."""
|
||||||
|
loop = QEventLoop()
|
||||||
|
QTimer.singleShot(ms, loop.quit)
|
||||||
|
loop.exec()
|
||||||
|
|
||||||
|
|
||||||
|
def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool:
|
||||||
|
waited = 0
|
||||||
|
while waited < timeout_ms:
|
||||||
|
if pred():
|
||||||
|
return True
|
||||||
|
pump(step)
|
||||||
|
waited += step
|
||||||
|
return pred()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
app = QApplication(sys.argv)
|
||||||
|
errors: list[str] = []
|
||||||
|
|
||||||
|
tmpdir = Path(tempfile.mkdtemp(prefix="sras_gui_"))
|
||||||
|
path = Path(sys.argv[1]) if len(sys.argv) > 1 else tmpdir / "gui.sras"
|
||||||
|
if len(sys.argv) <= 1:
|
||||||
|
gen.write(path, n_angles=4, seed=11, samples_per_frame=256)
|
||||||
|
|
||||||
|
print(f"\nloading {path.name}")
|
||||||
|
win = SrasViewerWindow()
|
||||||
|
win.show()
|
||||||
|
# Capture anything the app reports as an error via the status bar.
|
||||||
|
win.statusBar().messageChanged.connect(
|
||||||
|
lambda m: errors.append(m) if m and "error" in m.lower() else None)
|
||||||
|
|
||||||
|
win._load_file(str(path))
|
||||||
|
check("file loaded", wait_until(lambda: win._sras is not None))
|
||||||
|
s = win._sras
|
||||||
|
check("parsed as v6", s.version == 6, f"v{s.version}")
|
||||||
|
check("defaults to CH4", win.combo_channel.currentIndex() == CH4_IDX)
|
||||||
|
check("image displayed", win._current_image is not None)
|
||||||
|
check("angle spinbox ranges over all angles",
|
||||||
|
win.spin_angle.maximum() == s.n_angles - 1)
|
||||||
|
check("scan info populated",
|
||||||
|
win._info["Angles"].text() == f"Angles: {s.n_angles}",
|
||||||
|
win._info["Angles"].text())
|
||||||
|
|
||||||
|
print("\nbackground DC precompute (all angles)")
|
||||||
|
ok = wait_until(lambda: all((a, CH4_IDX) in win._dc_cache
|
||||||
|
and (a, CH3_IDX) in win._dc_cache
|
||||||
|
for a in range(s.n_angles)))
|
||||||
|
check("every angle cached for CH3 and CH4", ok,
|
||||||
|
f"{len(win._dc_cache)} entries")
|
||||||
|
check("status label reports completion",
|
||||||
|
"ready for all angles" in win.lbl_dc_precompute.text(),
|
||||||
|
win.lbl_dc_precompute.text())
|
||||||
|
|
||||||
|
print("\nangle switching (DC, should be served from cache)")
|
||||||
|
for a in range(s.n_angles):
|
||||||
|
win.spin_angle.setValue(a)
|
||||||
|
win._on_view_changed()
|
||||||
|
pump(60)
|
||||||
|
expected = win._sras.image_shape(a)
|
||||||
|
check(f"angle {a} shows its own geometry {expected}",
|
||||||
|
win._current_image.shape == expected,
|
||||||
|
str(win._current_image.shape))
|
||||||
|
check("no compute job needed for cached DC angles",
|
||||||
|
not win._job_running("compute"))
|
||||||
|
|
||||||
|
print("\nchannel switching")
|
||||||
|
win.spin_angle.setValue(0)
|
||||||
|
win._on_view_changed()
|
||||||
|
pump(60)
|
||||||
|
win.combo_channel.setCurrentIndex(CH3_IDX)
|
||||||
|
check("CH3 displayed", wait_until(lambda: win._current_ch == CH3_IDX))
|
||||||
|
|
||||||
|
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||||
|
check("CH1 (FFT) computed", wait_until(
|
||||||
|
lambda: win._current_ch == CH1_IDX and not win._job_running("compute")))
|
||||||
|
check("FFT result cached", len(win._fft_cache) > 0, f"{len(win._fft_cache)} keys")
|
||||||
|
rf_img = win._current_image
|
||||||
|
check("FFT image is non-degenerate", len(np.unique(rf_img)) > 1,
|
||||||
|
f"{len(np.unique(rf_img))} unique values")
|
||||||
|
|
||||||
|
print("\nvelocity mode (pure post-multiply, no recompute)")
|
||||||
|
n_fft_before = len(win._fft_cache)
|
||||||
|
win.combo_channel.setCurrentIndex(VELOCITY_MODE_IDX)
|
||||||
|
check("velocity displayed", wait_until(
|
||||||
|
lambda: win._current_ch == VELOCITY_MODE_IDX and not win._job_running("compute")))
|
||||||
|
grating = win.spin_grating_um.value()
|
||||||
|
check("velocity == freq x grating",
|
||||||
|
np.allclose(win._current_image, rf_img * grating, atol=1e-3))
|
||||||
|
check("velocity reused the cached FFT", len(win._fft_cache) == n_fft_before,
|
||||||
|
f"{n_fft_before} -> {len(win._fft_cache)}")
|
||||||
|
check("grating spinbox visible in velocity mode", win.grp_velocity.isVisible())
|
||||||
|
|
||||||
|
print("\nthreshold change (genuine cache-key change)")
|
||||||
|
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||||
|
wait_until(lambda: not win._job_running("compute"))
|
||||||
|
dc4 = win._dc_cache[(0, CH4_IDX)]
|
||||||
|
win.spin_threshold_mv.setValue(float(np.median(dc4)))
|
||||||
|
win._on_threshold_changed()
|
||||||
|
check("recomputed at new threshold", wait_until(
|
||||||
|
lambda: not win._job_running("compute") and len(win._fft_cache) > n_fft_before))
|
||||||
|
check("masking zeroed some pixels",
|
||||||
|
int((win._current_image == 0).sum()) > 0,
|
||||||
|
f"{int((win._current_image == 0).sum())} of {win._current_image.size}")
|
||||||
|
|
||||||
|
print("\nbackground subtraction toggle")
|
||||||
|
n_before = len(win._fft_cache)
|
||||||
|
win.chk_bg_sub.setChecked(False)
|
||||||
|
check("recomputed without bg-sub", wait_until(
|
||||||
|
lambda: not win._job_running("compute") and len(win._fft_cache) > n_before))
|
||||||
|
win.chk_bg_sub.setChecked(True)
|
||||||
|
pump(200)
|
||||||
|
check("returning to bg-sub was a cache hit (no recompute)",
|
||||||
|
not win._job_running("compute"))
|
||||||
|
|
||||||
|
print("\nROI")
|
||||||
|
x = s.x_axis_mm(0)
|
||||||
|
y = s.y_positions_mm(0)
|
||||||
|
roi = RoiQuad.from_bbox(float(x[1]), float(y[1]),
|
||||||
|
float(x[-2]), float(y[-2]))
|
||||||
|
win.image_canvas.set_roi(roi)
|
||||||
|
pump(120)
|
||||||
|
check("ROI registered", win.image_canvas.get_roi() is not None)
|
||||||
|
check("pixel count reported",
|
||||||
|
"pixels inside" in win.lbl_roi_npix.text()
|
||||||
|
and win.lbl_roi_npix.text() != "pixels inside: —",
|
||||||
|
win.lbl_roi_npix.text())
|
||||||
|
npix = int(win.lbl_roi_npix.text().split(":")[1])
|
||||||
|
check("ROI pixel count is plausible",
|
||||||
|
0 < npix <= win._current_image.size, f"{npix}")
|
||||||
|
check("Export ROI enabled", win.btn_export_roi.isEnabled())
|
||||||
|
|
||||||
|
csv_path = tmpdir / "roi.csv"
|
||||||
|
with patch("sras_viewer.QFileDialog.getSaveFileName",
|
||||||
|
return_value=(str(csv_path), "")):
|
||||||
|
win._on_export_roi_csv()
|
||||||
|
check("ROI CSV written", csv_path.exists())
|
||||||
|
if csv_path.exists():
|
||||||
|
body = [l for l in csv_path.read_text().splitlines() if not l.startswith("#")]
|
||||||
|
check("ROI CSV has header + one line per pixel",
|
||||||
|
len(body) == npix + 1, f"{len(body)} lines for {npix} pixels")
|
||||||
|
|
||||||
|
img_csv = tmpdir / "img.csv"
|
||||||
|
with patch("sras_viewer.QFileDialog.getSaveFileName",
|
||||||
|
return_value=(str(img_csv), "")):
|
||||||
|
win._on_export_csv()
|
||||||
|
check("image CSV written", img_csv.exists())
|
||||||
|
if img_csv.exists():
|
||||||
|
arr = np.loadtxt(img_csv, delimiter=",")
|
||||||
|
check("image CSV round-trips the displayed image",
|
||||||
|
arr.shape == win._current_image.shape
|
||||||
|
and np.allclose(arr, win._current_image, rtol=1e-5, atol=1e-4))
|
||||||
|
|
||||||
|
print("\nROI survives angle and channel switches")
|
||||||
|
win.spin_angle.setValue(1)
|
||||||
|
win._on_view_changed()
|
||||||
|
wait_until(lambda: not win._job_running("compute"))
|
||||||
|
check("ROI still present after angle switch",
|
||||||
|
win.image_canvas.get_roi() is not None)
|
||||||
|
win.combo_channel.setCurrentIndex(CH4_IDX)
|
||||||
|
wait_until(lambda: win._current_ch == CH4_IDX)
|
||||||
|
check("ROI still present after channel switch",
|
||||||
|
win.image_canvas.get_roi() is not None)
|
||||||
|
|
||||||
|
print("\nangle alignment (Fusion)")
|
||||||
|
win.spin_angle.setValue(0)
|
||||||
|
win._on_view_changed()
|
||||||
|
wait_until(lambda: not win._job_running("compute"))
|
||||||
|
check("alignment action enabled", win._alignment_act.isEnabled())
|
||||||
|
win._on_angle_alignment()
|
||||||
|
check("alignment completed", wait_until(
|
||||||
|
lambda: win._alignment_result is not None and not win._job_running("align"),
|
||||||
|
timeout_ms=60000))
|
||||||
|
if win._alignment_result is not None:
|
||||||
|
r = win._alignment_result
|
||||||
|
check("transform for every angle", len(r.per_angle) == s.n_angles)
|
||||||
|
check("canvas is at least as large as any single angle",
|
||||||
|
all(r.canvas_shape[0] >= int(s.n_rows[a])
|
||||||
|
and r.canvas_shape[1] >= int(s.n_frames[a])
|
||||||
|
for a in range(s.n_angles)), str(r.canvas_shape))
|
||||||
|
check("reference angle has zero shift",
|
||||||
|
r.per_angle[r.ref_angle_idx].shift_mm == (0.0, 0.0))
|
||||||
|
check("Aligned View auto-enabled and checked",
|
||||||
|
win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())
|
||||||
|
pump(200)
|
||||||
|
check("displayed image is on the alignment canvas",
|
||||||
|
win.image_canvas._img_shape == r.canvas_shape,
|
||||||
|
f"{win.image_canvas._img_shape} vs {r.canvas_shape}")
|
||||||
|
|
||||||
|
win.chk_aligned_view.setChecked(False)
|
||||||
|
pump(200)
|
||||||
|
check("unchecking returns to the raw per-angle grid",
|
||||||
|
win.image_canvas._img_shape == s.image_shape(0),
|
||||||
|
str(win.image_canvas._img_shape))
|
||||||
|
|
||||||
|
print("\nmanual alignment (Fusion)")
|
||||||
|
check("manual alignment action enabled", win._manual_align_act.isEnabled())
|
||||||
|
|
||||||
|
# --- Alignment pivot is a signal-weighted centroid, not the raw bbox --
|
||||||
|
# center, and is independent of any DC threshold (so a threshold that
|
||||||
|
# happens to leave a real angle's binary mask empty can't silently
|
||||||
|
# degrade the pivot back to the bbox center).
|
||||||
|
corner_signal = np.zeros(s.image_shape(0), dtype=np.float32)
|
||||||
|
corner_signal[0, 0] = 1.0 # single spike -> weighted centroid is exact
|
||||||
|
expected_corner = (float(s.x_axis_mm(0)[0]), float(s.y_positions_mm(0)[0]))
|
||||||
|
centroid = compute._signal_centroid_mm(s, 0, corner_signal)
|
||||||
|
check("signal-weighted centroid of a single spike pixel is that pixel exactly",
|
||||||
|
np.allclose(centroid, expected_corner), f"{centroid} vs {expected_corner}")
|
||||||
|
bbox_center = compute._bbox_center_mm(s, 0)
|
||||||
|
check("signal centroid differs from the raw scan-window bbox center",
|
||||||
|
not np.allclose(centroid, bbox_center),
|
||||||
|
f"centroid {centroid} vs bbox center {bbox_center}")
|
||||||
|
|
||||||
|
# compute_pivot_points_mm should reuse a pre-computed dc4_mv dict rather
|
||||||
|
# than recomputing from the real DC4 image (which has no such spike and
|
||||||
|
# would give a different answer if silently recomputed).
|
||||||
|
reused_pivot = compute.compute_pivot_points_mm(s, dc4_mv={0: corner_signal})[0]
|
||||||
|
check("compute_pivot_points_mm reuses a pre-computed dc4_mv dict",
|
||||||
|
np.allclose(reused_pivot, expected_corner))
|
||||||
|
|
||||||
|
# A perfectly flat signal carries no information to weight by, so it
|
||||||
|
# falls back to the bbox center rather than producing a NaN/degenerate
|
||||||
|
# centroid.
|
||||||
|
flat_signal = np.full(s.image_shape(0), 5.0, dtype=np.float32)
|
||||||
|
flat_centroid = compute._signal_centroid_mm(s, 0, flat_signal)
|
||||||
|
check("a perfectly flat signal falls back to the bbox center",
|
||||||
|
np.allclose(flat_centroid, bbox_center))
|
||||||
|
|
||||||
|
# --- Rotation sign convention: negative of the raw angles_deg delta ----
|
||||||
|
check("_theta_deg negates the raw angles_deg delta (GR stage's positive "
|
||||||
|
"angle is the opposite rotational sense from this module's CCW "
|
||||||
|
"math convention)",
|
||||||
|
all(np.isclose(compute._theta_deg(s, a, 0),
|
||||||
|
-(float(s.angles_deg[a]) - float(s.angles_deg[0])))
|
||||||
|
for a in range(s.n_angles)))
|
||||||
|
|
||||||
|
# --- FFT phase correlation recovers a known synthetic pixel shift ------
|
||||||
|
rng = np.random.default_rng(0)
|
||||||
|
corr_ref = np.zeros((40, 50), dtype=np.float32)
|
||||||
|
corr_ref[10:25, 15:35] = 1.0
|
||||||
|
corr_ref += 0.05 * rng.standard_normal(corr_ref.shape).astype(np.float32)
|
||||||
|
corr_mov = np.roll(corr_ref, shift=(4, -7), axis=(0, 1))
|
||||||
|
dr, dc = compute._phase_correlate_shift(corr_ref, corr_mov)
|
||||||
|
check("phase correlation recovers the shift that aligns mov onto ref",
|
||||||
|
(dr, dc) == (-4, 7), f"got (dr, dc)={(dr, dc)}")
|
||||||
|
|
||||||
|
# --- Open: must NOT seed from the still-live automatic AlignmentResult --
|
||||||
|
# The automatic result's translation comes from FFT phase correlation --
|
||||||
|
# the very thing manual mode exists to work around -- so manual mode
|
||||||
|
# must start from identity (centroids coincide, zero shift) regardless
|
||||||
|
# of whatever the automatic run last computed. Only a previously *saved
|
||||||
|
# manual* alignment (sidecar) should ever seed this dialog.
|
||||||
|
win._on_manual_alignment()
|
||||||
|
check("dialog opened", win._manual_align_dialog is not None)
|
||||||
|
dlg = win._manual_align_dialog
|
||||||
|
check("mask prep needed no background worker (already DC-cached)",
|
||||||
|
not win._job_running("manual_align_masks"))
|
||||||
|
check("no manual sidecar yet -> dialog starts at identity, not the "
|
||||||
|
"automatic result",
|
||||||
|
all(dlg._angle_params[a] == compute.ManualAngleParams()
|
||||||
|
for a in range(s.n_angles)))
|
||||||
|
|
||||||
|
# --- Reference angle is locked -------------------------------------------
|
||||||
|
dlg.combo_active_angle.setCurrentIndex(dlg._ref_angle_idx)
|
||||||
|
pump(30)
|
||||||
|
before_ref = dlg._angle_params[dlg._ref_angle_idx]
|
||||||
|
dlg._on_nudge_translate(1, 0, False)
|
||||||
|
dlg._on_nudge_rotate(1, False)
|
||||||
|
check("reference angle group disabled", not dlg.grp_manual_adjust.isEnabled())
|
||||||
|
check("reference angle untouched by nudge attempts",
|
||||||
|
dlg._angle_params[dlg._ref_angle_idx] == before_ref)
|
||||||
|
|
||||||
|
# --- Nudging a real angle (fine + coarse, translate + rotate) -----------
|
||||||
|
active = 1 if s.n_angles > 1 else 0
|
||||||
|
dlg.combo_active_angle.setCurrentIndex(active)
|
||||||
|
pump(30)
|
||||||
|
before = dlg._angle_params[active].shift_mm
|
||||||
|
dlg._on_nudge_translate(1, 0, False) # fine +X
|
||||||
|
fine_step = dlg.spin_step_translate_mm.value()
|
||||||
|
check("fine translate nudge moved shift_x by exactly one fine step",
|
||||||
|
abs(dlg._angle_params[active].shift_mm[0] - (before[0] + fine_step)) < 1e-9)
|
||||||
|
|
||||||
|
before = dlg._angle_params[active].shift_mm
|
||||||
|
dlg._on_nudge_translate(0, -1, True) # coarse -Y
|
||||||
|
coarse_step = fine_step * dlg.spin_step_multiplier.value()
|
||||||
|
check("coarse translate nudge uses the multiplier",
|
||||||
|
abs(dlg._angle_params[active].shift_mm[1] - (before[1] - coarse_step)) < 1e-9)
|
||||||
|
|
||||||
|
before_rot = dlg._angle_params[active].rotation_deg
|
||||||
|
dlg._on_nudge_rotate(1, False)
|
||||||
|
check("rotate nudge changed rotation_deg",
|
||||||
|
dlg._angle_params[active].rotation_deg != before_rot)
|
||||||
|
check("preview canvas rebuilt for every angle after a rotation nudge",
|
||||||
|
len(dlg._preview_layers) == s.n_angles)
|
||||||
|
|
||||||
|
# --- Real key-event wiring (proves keyPressEvent -> signal -> slot) -----
|
||||||
|
before = dlg._angle_params[active].shift_mm
|
||||||
|
QTest.keyClick(dlg.canvas, Qt.Key.Key_Right)
|
||||||
|
check("a real Right-arrow key event nudged shift_x",
|
||||||
|
dlg._angle_params[active].shift_mm[0] > before[0])
|
||||||
|
|
||||||
|
# --- Auto De-rotate: rotation only, translation untouched ---------------
|
||||||
|
shift_before_derotate = dlg._angle_params[active].shift_mm
|
||||||
|
dlg._on_auto_derotate()
|
||||||
|
expected_theta = compute._theta_deg(s, active, dlg._ref_angle_idx)
|
||||||
|
check("auto de-rotate set the known analytic angle",
|
||||||
|
abs(dlg._angle_params[active].rotation_deg - expected_theta) < 1e-6)
|
||||||
|
check("auto de-rotate left translation untouched",
|
||||||
|
dlg._angle_params[active].shift_mm == shift_before_derotate)
|
||||||
|
check("reference angle stays identity after auto de-rotate",
|
||||||
|
dlg._angle_params[dlg._ref_angle_idx].rotation_deg == 0.0)
|
||||||
|
|
||||||
|
# --- Auto Cross-Correlate: rotation + FFT-correlated shift, backgrounded -
|
||||||
|
check("cross-correlate action enabled once masks are ready",
|
||||||
|
dlg.btn_auto_correlate.isEnabled())
|
||||||
|
dlg._on_auto_correlate()
|
||||||
|
check("auto cross-correlate completed", wait_until(
|
||||||
|
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
|
||||||
|
check("auto cross-correlate set the known analytic angle for every angle",
|
||||||
|
all(abs(dlg._angle_params[a].rotation_deg
|
||||||
|
- compute._theta_deg(s, a, dlg._ref_angle_idx)) < 1e-6
|
||||||
|
for a in range(s.n_angles) if a != dlg._ref_angle_idx))
|
||||||
|
check("auto cross-correlate reference angle stays identity",
|
||||||
|
dlg._angle_params[dlg._ref_angle_idx] == compute.ManualAngleParams())
|
||||||
|
check("auto cross-correlate re-enabled controls when done",
|
||||||
|
dlg.grp_correlate.isEnabled() and dlg.btn_save.isEnabled())
|
||||||
|
check("preview canvas rebuilt after cross-correlate",
|
||||||
|
len(dlg._preview_layers) == s.n_angles)
|
||||||
|
|
||||||
|
# The thresholded-mask option should also work end to end.
|
||||||
|
dlg.combo_correlate_source.setCurrentIndex(1) # thresholded mask
|
||||||
|
dlg._on_auto_correlate()
|
||||||
|
check("auto cross-correlate (thresholded-mask option) completed", wait_until(
|
||||||
|
lambda: not win._job_running("manual_align_correlate"), timeout_ms=30000))
|
||||||
|
|
||||||
|
# --- Save -----------------------------------------------------------------
|
||||||
|
dlg._on_save()
|
||||||
|
sidecar = compute.sidecar_path(s.path)
|
||||||
|
check("sidecar file written", sidecar.exists())
|
||||||
|
raw = json.loads(sidecar.read_text()) if sidecar.exists() else {}
|
||||||
|
check("sidecar schema_version is current",
|
||||||
|
raw.get("schema_version") == compute._SIDECAR_SCHEMA_VERSION)
|
||||||
|
check("sidecar per_angle round-trips the dialog's resolved params",
|
||||||
|
all(raw.get("per_angle", {}).get(str(a), {}).get("rotation_deg")
|
||||||
|
== dlg._angle_params[a].rotation_deg for a in range(s.n_angles)))
|
||||||
|
check("main window's alignment_result replaced by the manual build",
|
||||||
|
win._alignment_result is not None
|
||||||
|
and win._alignment_result.per_angle[active].rotation_deg
|
||||||
|
== dlg._angle_params[active].rotation_deg)
|
||||||
|
check("Aligned View auto-enabled after Save",
|
||||||
|
win.chk_aligned_view.isEnabled() and win.chk_aligned_view.isChecked())
|
||||||
|
|
||||||
|
# --- An old-schema sidecar (pre-pivot/sign fix) is treated as absent ------
|
||||||
|
stale = dict(raw)
|
||||||
|
stale["schema_version"] = compute._SIDECAR_SCHEMA_VERSION - 1
|
||||||
|
sidecar.write_text(json.dumps(stale))
|
||||||
|
check("a sidecar with an old schema_version is not loaded",
|
||||||
|
compute.load_manual_alignment(s) is None)
|
||||||
|
sidecar.write_text(json.dumps(raw)) # restore for the rest of this section
|
||||||
|
|
||||||
|
# --- Clear (with confirmation) --------------------------------------------
|
||||||
|
with patch("sras_viewer.QMessageBox.question",
|
||||||
|
return_value=QMessageBox.StandardButton.Yes):
|
||||||
|
dlg._on_clear()
|
||||||
|
check("sidecar file deleted", not sidecar.exists())
|
||||||
|
check("dialog params reset to identity",
|
||||||
|
all(dlg._angle_params[a] == compute.ManualAngleParams()
|
||||||
|
for a in range(s.n_angles)))
|
||||||
|
check("main window alignment_result cleared", win._alignment_result is None)
|
||||||
|
check("Aligned View disabled after Clear",
|
||||||
|
not win.chk_aligned_view.isEnabled() and not win.chk_aligned_view.isChecked())
|
||||||
|
|
||||||
|
dlg.close()
|
||||||
|
pump(150)
|
||||||
|
check("dialog reference released on close", win._manual_align_dialog is None)
|
||||||
|
|
||||||
|
# --- Sidecar auto-restore on next load ------------------------------------
|
||||||
|
win._on_manual_alignment()
|
||||||
|
dlg = win._manual_align_dialog
|
||||||
|
dlg.combo_active_angle.setCurrentIndex(active)
|
||||||
|
pump(30)
|
||||||
|
dlg._on_auto_derotate()
|
||||||
|
dlg._on_nudge_translate(1, 1, True)
|
||||||
|
saved_rotation = dlg._angle_params[active].rotation_deg
|
||||||
|
saved_shift = dlg._angle_params[active].shift_mm
|
||||||
|
dlg._on_save()
|
||||||
|
dlg.close()
|
||||||
|
pump(150)
|
||||||
|
|
||||||
|
old_sras_id = id(win._sras)
|
||||||
|
win._load_file(str(path)) # reload the same file fresh
|
||||||
|
check("file reloaded", wait_until(
|
||||||
|
lambda: win._sras is not None and id(win._sras) != old_sras_id))
|
||||||
|
s = win._sras
|
||||||
|
check("manual dialog force-closed by a reload", win._manual_align_dialog is None)
|
||||||
|
check("reload restores the saved manual alignment automatically",
|
||||||
|
win._alignment_result is not None)
|
||||||
|
if win._alignment_result is not None:
|
||||||
|
check("restored rotation matches what was saved",
|
||||||
|
abs(win._alignment_result.per_angle[active].rotation_deg
|
||||||
|
- saved_rotation) < 1e-9)
|
||||||
|
check("restored shift matches what was saved",
|
||||||
|
win._alignment_result.per_angle[active].shift_mm == saved_shift)
|
||||||
|
check("Aligned View auto-checked after restoring a saved alignment",
|
||||||
|
win.chk_aligned_view.isChecked())
|
||||||
|
|
||||||
|
print("\npixel inspector")
|
||||||
|
win.chk_aligned_view.setChecked(False)
|
||||||
|
pump(100)
|
||||||
|
win._on_pixel_clicked(0, 0)
|
||||||
|
pump(150)
|
||||||
|
check("waveform hint hidden after a click", win.lbl_wave_hint.isHidden())
|
||||||
|
win.combo_channel.setCurrentIndex(CH1_IDX)
|
||||||
|
wait_until(lambda: not win._job_running("compute"))
|
||||||
|
win._on_pixel_clicked(1, 1)
|
||||||
|
pump(150)
|
||||||
|
check("RF waveform panel rendered",
|
||||||
|
len(win.wave_canvas.ax_wave.lines) > 0,
|
||||||
|
f"{len(win.wave_canvas.ax_wave.lines)} lines")
|
||||||
|
|
||||||
|
print("\nshutdown")
|
||||||
|
win.close()
|
||||||
|
pump(400)
|
||||||
|
check("all background jobs released", len(win._jobs) == 0,
|
||||||
|
f"{list(win._jobs)}")
|
||||||
|
|
||||||
|
print()
|
||||||
|
unexpected = [e for e in errors if e]
|
||||||
|
if unexpected:
|
||||||
|
print(f"status-bar errors seen: {unexpected}")
|
||||||
|
_failures.append("status-bar errors")
|
||||||
|
|
||||||
|
if _failures:
|
||||||
|
print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
|
||||||
|
return 1
|
||||||
|
print("All GUI checks passed.")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
sys.exit(main())
|
||||||
@@ -0,0 +1,358 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Behavioural tests for the sras-viewer refactor.
|
||||||
|
|
||||||
|
Covers what the golden-hash harness can't: the v6->v7 cache round-trip
|
||||||
|
(including block carry-forward), parallel-vs-serial identity, the no-mask
|
||||||
|
fast path, and the ROI bounding-box mask optimisation.
|
||||||
|
|
||||||
|
Usage: python tools/test_refactor.py [--scratch DIR]
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import shutil
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||||
|
|
||||||
|
import sras_compute as compute # noqa: E402
|
||||||
|
from sras_compute import ( # noqa: E402
|
||||||
|
cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
|
||||||
|
)
|
||||||
|
from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv # noqa: E402
|
||||||
|
import tools.make_test_sras as gen # noqa: E402
|
||||||
|
|
||||||
|
_failures: list[str] = []
|
||||||
|
|
||||||
|
|
||||||
|
def check(name: str, ok: bool, detail: str = ""):
|
||||||
|
print(f" {'PASS' if ok else 'FAIL'} {name}" + (f" — {detail}" if detail else ""))
|
||||||
|
if not ok:
|
||||||
|
_failures.append(name)
|
||||||
|
|
||||||
|
|
||||||
|
def test_cache_roundtrip(scratch: Path):
|
||||||
|
"""v6 -> v7 for DC, then FFT, asserting the first block survives the
|
||||||
|
second write (the carry-forward path in write_v7_cache)."""
|
||||||
|
print("\ncache round-trip (v6 -> v7, both blocks)")
|
||||||
|
path = scratch / "roundtrip.sras"
|
||||||
|
gen.write(path, n_angles=3, seed=1, samples_per_frame=64)
|
||||||
|
|
||||||
|
src = SrasFile(str(path))
|
||||||
|
check("source is v6", src.version == 6, f"got v{src.version}")
|
||||||
|
expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
|
||||||
|
expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
|
||||||
|
expect_fft = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
for a in range(src.n_angles)]
|
||||||
|
|
||||||
|
err = cache_file(str(path), "dc", True)
|
||||||
|
check("dc cache_file succeeded", err == "", err)
|
||||||
|
|
||||||
|
after_dc = SrasFile(str(path))
|
||||||
|
check("version flipped to 7", after_dc.version == 7, f"got v{after_dc.version}")
|
||||||
|
check("dc3 stored for every angle",
|
||||||
|
all(x is not None for x in after_dc.precomputed_dc3_mv))
|
||||||
|
check("dc3 values round-trip",
|
||||||
|
all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
|
||||||
|
for a in range(after_dc.n_angles)))
|
||||||
|
check("dc4 values round-trip",
|
||||||
|
all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
|
||||||
|
for a in range(after_dc.n_angles)))
|
||||||
|
check("no fft block yet",
|
||||||
|
all(x is None for x in after_dc.precomputed_freq_mhz))
|
||||||
|
check("cached images are native float32",
|
||||||
|
after_dc.precomputed_dc3_mv[0].dtype == np.float32
|
||||||
|
and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|"),
|
||||||
|
str(after_dc.precomputed_dc3_mv[0].dtype.byteorder))
|
||||||
|
check("cached images are writable",
|
||||||
|
after_dc.precomputed_dc3_mv[0].flags.writeable)
|
||||||
|
|
||||||
|
err = cache_file(str(path), "fft", True)
|
||||||
|
check("fft cache_file succeeded", err == "", err)
|
||||||
|
|
||||||
|
both = SrasFile(str(path))
|
||||||
|
check("fft stored for every angle",
|
||||||
|
all(x is not None for x in both.precomputed_freq_mhz))
|
||||||
|
check("fft values round-trip",
|
||||||
|
all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3)
|
||||||
|
for a in range(both.n_angles)))
|
||||||
|
check("DC block carried forward through the FFT write",
|
||||||
|
all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
|
||||||
|
for a in range(both.n_angles)))
|
||||||
|
check("bg_sub flag persisted", both.precomputed_bg_sub is True)
|
||||||
|
|
||||||
|
# The fast path must reproduce a fresh compute, and masking must still
|
||||||
|
# apply on top of a cached (unmasked) image.
|
||||||
|
fresh = SrasFile(str(path))
|
||||||
|
fresh.precomputed_freq_mhz = [None] * fresh.n_angles
|
||||||
|
dc4 = dc_image_mv(both, 0, CH4_IDX)
|
||||||
|
thr = float(np.median(dc4))
|
||||||
|
check("cached fast path == fresh compute (unmasked)",
|
||||||
|
np.allclose(compute_rf_image(both, 0, dc_threshold_mv=None, apply_bg_sub=True),
|
||||||
|
compute_rf_image(fresh, 0, dc_threshold_mv=None, apply_bg_sub=True),
|
||||||
|
atol=1e-3))
|
||||||
|
check("cached fast path == fresh compute (masked)",
|
||||||
|
np.allclose(compute_rf_image(both, 0, dc_threshold_mv=thr, apply_bg_sub=True),
|
||||||
|
compute_rf_image(fresh, 0, dc_threshold_mv=thr, apply_bg_sub=True),
|
||||||
|
atol=1e-3))
|
||||||
|
|
||||||
|
# Waveform data must be byte-identical to the pre-cache file.
|
||||||
|
orig = scratch / "roundtrip_orig.sras"
|
||||||
|
gen.write(orig, n_angles=3, seed=1, samples_per_frame=64)
|
||||||
|
o, n = SrasFile(str(orig)), SrasFile(str(path))
|
||||||
|
check("waveform data untouched by the cache write",
|
||||||
|
all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
|
||||||
|
for a in range(o.n_angles)))
|
||||||
|
|
||||||
|
|
||||||
|
def test_partial_v7_cache(scratch: Path):
|
||||||
|
"""Only some angles cached: uncached angles must compute, not read zeros.
|
||||||
|
This is the v5 bug the ragged normalisation fixed, checked via v7."""
|
||||||
|
print("\npartial cache (only some angles stored)")
|
||||||
|
path = scratch / "partial.sras"
|
||||||
|
gen.write(path, n_angles=3, seed=2, samples_per_frame=64)
|
||||||
|
|
||||||
|
src = SrasFile(str(path))
|
||||||
|
expected = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
for a in range(src.n_angles)]
|
||||||
|
partial = [expected[0], None, expected[2]] # angle 1 deliberately absent
|
||||||
|
src.write_v7_cache(new_freq_mhz=partial, new_bg_sub=True)
|
||||||
|
|
||||||
|
reread = SrasFile(str(path))
|
||||||
|
check("angle 1 is not cached", reread.precomputed_freq_mhz[1] is None)
|
||||||
|
check("angles 0 and 2 are cached",
|
||||||
|
reread.precomputed_freq_mhz[0] is not None
|
||||||
|
and reread.precomputed_freq_mhz[2] is not None)
|
||||||
|
img1 = compute_rf_image(reread, 1, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
check("uncached angle computes rather than returning zeros",
|
||||||
|
np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3))
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_identity(scratch: Path):
|
||||||
|
"""Forcing 1 worker vs many must give identical output — catches
|
||||||
|
chunk-boundary and race bugs."""
|
||||||
|
print("\nparallel vs serial identity")
|
||||||
|
path = scratch / "parallel.sras"
|
||||||
|
# Many rows, so the row loop actually splits into several chunks.
|
||||||
|
n_rows, n_frames, spf = 48, 9, 256
|
||||||
|
gen.write(path, n_angles=1, seed=3, samples_per_frame=spf,
|
||||||
|
geometry=[(n_rows, n_frames)])
|
||||||
|
sras = SrasFile(str(path))
|
||||||
|
|
||||||
|
saved_budget, saved_workers = compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS
|
||||||
|
try:
|
||||||
|
# Shrink the budget so chunk_rows collapses to 1 and every row is
|
||||||
|
# its own chunk — the worst case for boundary bugs.
|
||||||
|
compute._TOTAL_BYTES_BUDGET = 8 * n_frames * spf * 4
|
||||||
|
|
||||||
|
compute._MAX_WORKERS = 1
|
||||||
|
dc_serial = compute_dc_image(sras, 0, CH4_IDX)
|
||||||
|
rf_serial = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
dc4 = adc_to_mv(dc_serial, *sras.cal(CH4_IDX))
|
||||||
|
thr = float(np.median(dc4))
|
||||||
|
rf_masked_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
|
||||||
|
apply_bg_sub=True)
|
||||||
|
|
||||||
|
chunk_rows, n_workers = compute._plan_chunks(
|
||||||
|
n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
|
||||||
|
check("serial plan uses 1 worker", n_workers == 1, f"chunk_rows={chunk_rows}")
|
||||||
|
check("work actually splits into multiple chunks", chunk_rows < n_rows,
|
||||||
|
f"chunk_rows={chunk_rows} of {n_rows} rows")
|
||||||
|
|
||||||
|
compute._MAX_WORKERS = 8
|
||||||
|
chunk_rows, n_workers = compute._plan_chunks(
|
||||||
|
n_rows, n_frames, spf, live_multiplier=compute._FFT_LIVE_MULTIPLIER)
|
||||||
|
check("parallel plan uses >1 worker", n_workers > 1,
|
||||||
|
f"chunk_rows={chunk_rows} workers={n_workers}")
|
||||||
|
|
||||||
|
dc_par = compute_dc_image(sras, 0, CH4_IDX)
|
||||||
|
rf_par = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
rf_masked_par = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True)
|
||||||
|
|
||||||
|
check("dc image identical", np.array_equal(dc_serial, dc_par))
|
||||||
|
check("rf image identical (unmasked)", np.array_equal(rf_serial, rf_par))
|
||||||
|
check("rf image identical (masked)",
|
||||||
|
np.array_equal(rf_masked_serial, rf_masked_par))
|
||||||
|
finally:
|
||||||
|
compute._TOTAL_BYTES_BUDGET, compute._MAX_WORKERS = saved_budget, saved_workers
|
||||||
|
|
||||||
|
|
||||||
|
def test_nomask_equals_low_threshold(scratch: Path):
|
||||||
|
"""dc_threshold_mv=None must equal a threshold below every pixel, while
|
||||||
|
skipping the CH4 read."""
|
||||||
|
print("\nno-mask path")
|
||||||
|
path = scratch / "nomask.sras"
|
||||||
|
gen.write(path, n_angles=2, seed=4, samples_per_frame=128)
|
||||||
|
sras = SrasFile(str(path))
|
||||||
|
for a in range(sras.n_angles):
|
||||||
|
none_img = compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True)
|
||||||
|
low_img = compute_rf_image(sras, a, dc_threshold_mv=-1e9, apply_bg_sub=True)
|
||||||
|
check(f"angle {a}: None == -1e9 threshold",
|
||||||
|
np.array_equal(none_img, low_img))
|
||||||
|
check(f"angle {a}: image is non-degenerate",
|
||||||
|
len(np.unique(none_img)) > 1, f"{len(np.unique(none_img))} unique")
|
||||||
|
|
||||||
|
|
||||||
|
def test_roi_mask():
|
||||||
|
"""The bbox-restricted mask must equal a full-grid point-in-polygon test."""
|
||||||
|
print("\nROI mask (bbox fast path vs full grid)")
|
||||||
|
from matplotlib.path import Path as MplPath
|
||||||
|
from sras_viewer import RoiQuad
|
||||||
|
|
||||||
|
rng = np.random.default_rng(0)
|
||||||
|
x = np.linspace(-2.0, 3.0, 137)
|
||||||
|
y = np.linspace(1.0, 4.0, 91)
|
||||||
|
|
||||||
|
cases = {
|
||||||
|
"axis-aligned rect": np.array([[0.0, 1.5], [1.0, 1.5], [1.0, 3.0], [0.0, 3.0]]),
|
||||||
|
"skewed quad": np.array([[-0.5, 1.2], [1.7, 1.9], [1.2, 3.4], [-1.0, 2.6]]),
|
||||||
|
"entirely outside": np.array([[8.0, 8.0], [9.0, 8.0], [9.0, 9.0], [8.0, 9.0]]),
|
||||||
|
"covers whole grid": np.array([[-9.0, -9.0], [9.0, -9.0], [9.0, 9.0], [-9.0, 9.0]]),
|
||||||
|
"straddles left edge": np.array([[-4.0, 2.0], [0.5, 2.0], [0.5, 3.0], [-4.0, 3.0]]),
|
||||||
|
}
|
||||||
|
for _ in range(5):
|
||||||
|
cases[f"random {_}"] = rng.uniform([-2.5, 0.5], [3.5, 4.5], size=(4, 2))
|
||||||
|
|
||||||
|
for name, pts in cases.items():
|
||||||
|
roi = RoiQuad(pts)
|
||||||
|
fast = roi.mask_for_grid(x, y)
|
||||||
|
X, Y = np.meshgrid(x.astype(np.float64), y.astype(np.float64))
|
||||||
|
slow = MplPath(pts).contains_points(
|
||||||
|
np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
|
||||||
|
check(f"{name} ({int(slow.sum())} px inside)", np.array_equal(fast, slow))
|
||||||
|
|
||||||
|
# Descending y axis (images are stored top-down in some scans).
|
||||||
|
roi = RoiQuad(cases["skewed quad"])
|
||||||
|
y_desc = y[::-1]
|
||||||
|
fast = roi.mask_for_grid(x, y_desc)
|
||||||
|
X, Y = np.meshgrid(x.astype(np.float64), y_desc.astype(np.float64))
|
||||||
|
slow = MplPath(cases["skewed quad"]).contains_points(
|
||||||
|
np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
|
||||||
|
check("descending y axis", np.array_equal(fast, slow))
|
||||||
|
|
||||||
|
|
||||||
|
def test_legacy_parse_and_average(scratch: Path):
|
||||||
|
"""v2-v4 parsing plus the sras_average.py rewrite (which now streams via
|
||||||
|
SrasFile rather than slurping the whole file)."""
|
||||||
|
import subprocess
|
||||||
|
print("\nlegacy formats (v2-v4) and sras_average")
|
||||||
|
repo = Path(__file__).resolve().parent.parent
|
||||||
|
|
||||||
|
for version in (2, 3, 4):
|
||||||
|
path = scratch / f"legacy_v{version}.sras"
|
||||||
|
meta = gen.write_legacy(path, version=version, n_angles=2, n_rows=4,
|
||||||
|
n_frames=12, samples_per_frame=32, seed=version)
|
||||||
|
s = SrasFile(str(path))
|
||||||
|
check(f"v{version} parses", s.version == version, f"got v{s.version}")
|
||||||
|
check(f"v{version} geometry uniform across angles",
|
||||||
|
list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12],
|
||||||
|
f"rows={list(s.n_rows)} frames={list(s.n_frames)}")
|
||||||
|
check(f"v{version} waveform data matches what was written",
|
||||||
|
all(np.array_equal(np.asarray(s.data[a]), meta["data"][a])
|
||||||
|
for a in range(s.n_angles)))
|
||||||
|
check(f"v{version} background {'present' if version >= 4 else 'absent'}",
|
||||||
|
(s.background is not None) == (version >= 4))
|
||||||
|
check(f"v{version} precomputed stores are ragged lists",
|
||||||
|
isinstance(s.precomputed_freq_mhz, list)
|
||||||
|
and len(s.precomputed_freq_mhz) == s.n_angles)
|
||||||
|
# DC image must equal a direct mean of the known input.
|
||||||
|
expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1)
|
||||||
|
check(f"v{version} DC image equals a direct mean",
|
||||||
|
np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3))
|
||||||
|
|
||||||
|
src = scratch / "legacy_v4.sras"
|
||||||
|
meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
|
||||||
|
samples_per_frame=32, seed=4)
|
||||||
|
dst = scratch / "legacy_v4_avg.sras"
|
||||||
|
if dst.exists():
|
||||||
|
dst.unlink()
|
||||||
|
proc = subprocess.run(
|
||||||
|
[sys.executable, str(repo / "sras_average.py"), str(src), str(dst), "--n", "4"],
|
||||||
|
capture_output=True, text=True, cwd=repo)
|
||||||
|
check("sras_average ran", proc.returncode == 0,
|
||||||
|
(proc.stderr or proc.stdout).strip()[-200:])
|
||||||
|
|
||||||
|
if dst.exists():
|
||||||
|
avg = SrasFile(str(dst))
|
||||||
|
check("averaged file parses", avg.version == 4)
|
||||||
|
check("frame count divided by 4",
|
||||||
|
list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}")
|
||||||
|
check("angles/rows/channels unchanged",
|
||||||
|
avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
|
||||||
|
and avg.n_channels == meta["n_channels"])
|
||||||
|
check("calibration preserved",
|
||||||
|
np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv))
|
||||||
|
check("background preserved",
|
||||||
|
np.array_equal(avg.background, SrasFile(str(src)).background))
|
||||||
|
src_data = meta["data"]
|
||||||
|
expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
|
||||||
|
check("first averaged group equals the mean of its 4 source frames",
|
||||||
|
np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0))
|
||||||
|
|
||||||
|
# Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
|
||||||
|
dst2 = scratch / "legacy_v4_avg5.sras"
|
||||||
|
subprocess.run([sys.executable, str(repo / "sras_average.py"),
|
||||||
|
str(src), str(dst2), "--n", "5"],
|
||||||
|
capture_output=True, text=True, cwd=repo)
|
||||||
|
if dst2.exists():
|
||||||
|
check("partial trailing group kept by default",
|
||||||
|
list(SrasFile(str(dst2)).n_frames) == [3, 3],
|
||||||
|
f"{list(SrasFile(str(dst2)).n_frames)}")
|
||||||
|
dst3 = scratch / "legacy_v4_avg5d.sras"
|
||||||
|
subprocess.run([sys.executable, str(repo / "sras_average.py"),
|
||||||
|
str(src), str(dst3), "--n", "5", "--discard-remainder"],
|
||||||
|
capture_output=True, text=True, cwd=repo)
|
||||||
|
if dst3.exists():
|
||||||
|
check("--discard-remainder drops the partial group",
|
||||||
|
list(SrasFile(str(dst3)).n_frames) == [2, 2],
|
||||||
|
f"{list(SrasFile(str(dst3)).n_frames)}")
|
||||||
|
|
||||||
|
|
||||||
|
def test_unsupported_version_reported(scratch: Path):
|
||||||
|
"""cache_file must report, not raise, for a file it can't handle."""
|
||||||
|
print("\nerror reporting")
|
||||||
|
bogus = scratch / "bogus.sras"
|
||||||
|
bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200)
|
||||||
|
err = cache_file(str(bogus), "dc", True)
|
||||||
|
check("bad version returns an error string", bool(err), err)
|
||||||
|
missing = cache_file(str(scratch / "does_not_exist.sras"), "dc", True)
|
||||||
|
check("missing file returns an error string", bool(missing), missing)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p = argparse.ArgumentParser(description=__doc__)
|
||||||
|
p.add_argument("--scratch")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
tmp = None
|
||||||
|
if args.scratch:
|
||||||
|
scratch = Path(args.scratch)
|
||||||
|
scratch.mkdir(parents=True, exist_ok=True)
|
||||||
|
else:
|
||||||
|
tmp = tempfile.mkdtemp(prefix="sras_test_")
|
||||||
|
scratch = Path(tmp)
|
||||||
|
|
||||||
|
try:
|
||||||
|
test_cache_roundtrip(scratch)
|
||||||
|
test_partial_v7_cache(scratch)
|
||||||
|
test_parallel_identity(scratch)
|
||||||
|
test_nomask_equals_low_threshold(scratch)
|
||||||
|
test_roi_mask()
|
||||||
|
test_legacy_parse_and_average(scratch)
|
||||||
|
test_unsupported_version_reported(scratch)
|
||||||
|
finally:
|
||||||
|
if tmp:
|
||||||
|
shutil.rmtree(tmp, ignore_errors=True)
|
||||||
|
|
||||||
|
print()
|
||||||
|
if _failures:
|
||||||
|
print(f"{len(_failures)} FAILURE(S): " + ", ".join(_failures))
|
||||||
|
sys.exit(1)
|
||||||
|
print("All checks passed.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user