Refactor/parallel and dedup #1
+70
-145
@@ -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,23 +150,21 @@ 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")
|
||||||
|
|||||||
+577
@@ -0,0 +1,577 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Image computation and angle alignment for .sras scans.
|
||||||
|
|
||||||
|
Depends only on numpy/scipy (+ optional pyfftw) and sras_format, so a
|
||||||
|
multiprocessing child can import it without loading Qt or matplotlib —
|
||||||
|
which matters because Python 3.14 on macOS spawns rather than forks.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import scipy.fft as scipy_fft
|
||||||
|
import scipy.ndimage as scipy_ndimage
|
||||||
|
|
||||||
|
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# FFT backend
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
try:
|
||||||
|
import pyfftw
|
||||||
|
pyfftw.interfaces.cache.enable()
|
||||||
|
PYFFTW_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
PYFFTW_AVAILABLE = False
|
||||||
|
|
||||||
|
_fft_backend = "numpy" # "numpy" or "pyfftw"; set via set_fft_backend()
|
||||||
|
|
||||||
|
|
||||||
|
def set_fft_backend(name: str):
|
||||||
|
"""Select the rfft implementation. Module-level state, so it must be set
|
||||||
|
explicitly inside each multiprocessing child — it does not survive a
|
||||||
|
spawn."""
|
||||||
|
global _fft_backend
|
||||||
|
_fft_backend = name if (name != "pyfftw" or PYFFTW_AVAILABLE) else "numpy"
|
||||||
|
|
||||||
|
|
||||||
|
def get_fft_backend() -> str:
|
||||||
|
return _fft_backend
|
||||||
|
|
||||||
|
|
||||||
|
def _do_rfft(x: np.ndarray, n: int | None = None, axis: int = -1,
|
||||||
|
workers: int = 1) -> np.ndarray:
|
||||||
|
"""Dispatch rfft to the selected backend with optional multithreading."""
|
||||||
|
if _fft_backend == "pyfftw" and PYFFTW_AVAILABLE:
|
||||||
|
return pyfftw.interfaces.numpy_fft.rfft(x, n=n, axis=axis, threads=workers)
|
||||||
|
return scipy_fft.rfft(x, n=n, axis=axis, workers=workers)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Chunking / parallel budget
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
#
|
||||||
|
# Rows are batched so the float32 working buffers for one chunk stay under a
|
||||||
|
# memory budget. A fixed row count (the original design) works fine for small
|
||||||
|
# legacy scans but is catastrophic for a v6 scan with a large per-angle
|
||||||
|
# frame/sample count — e.g. a 7500-frame x 2500-sample angle needs ~2.4 GB for
|
||||||
|
# a single 32-row chunk.
|
||||||
|
#
|
||||||
|
# With chunks running concurrently the budget has to cover *all* live chunks at
|
||||||
|
# once. Note that on a large scan chunk_rows is already clamped to its floor of
|
||||||
|
# 1 row (one row alone is ~75 MB of float32 at 7507x2500), so shrinking the
|
||||||
|
# per-chunk size cannot buy more concurrency — the worker count must be derived
|
||||||
|
# from the budget instead. See _plan_chunks.
|
||||||
|
|
||||||
|
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1536)) * 1024 * 1024
|
||||||
|
_CHUNK_ROWS_MAX = 32 # cap for small scans (original behavior)
|
||||||
|
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)
|
||||||
|
|
||||||
|
# An rfft chunk holds, live at once: the float32 input, the complex64
|
||||||
|
# transform, and the float32 power spectrum — roughly 3x the input buffer.
|
||||||
|
_FFT_LIVE_MULTIPLIER = 3
|
||||||
|
|
||||||
|
|
||||||
|
def _chunk_rows_for(n_frames: int, samples_per_frame: int,
|
||||||
|
budget: int = _TOTAL_BYTES_BUDGET) -> int:
|
||||||
|
bytes_per_row = max(1, n_frames * samples_per_frame * 4) # float32
|
||||||
|
return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))
|
||||||
|
|
||||||
|
|
||||||
|
def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
|
||||||
|
live_multiplier: int = 1,
|
||||||
|
max_workers: int | None = None) -> tuple[int, int]:
|
||||||
|
"""(chunk_rows, n_workers) sized so live_multiplier x chunk_rows x
|
||||||
|
n_frames x samples x 4 bytes x n_workers stays within the budget."""
|
||||||
|
per_chunk_budget = max(1, _TOTAL_BYTES_BUDGET // live_multiplier)
|
||||||
|
chunk_rows = _chunk_rows_for(n_frames, samples_per_frame, per_chunk_budget)
|
||||||
|
|
||||||
|
chunk_bytes = max(1, live_multiplier * chunk_rows * n_frames * samples_per_frame * 4)
|
||||||
|
cap = max_workers if max_workers is not None else _MAX_WORKERS
|
||||||
|
n_workers = int(max(1, min(cap,
|
||||||
|
_TOTAL_BYTES_BUDGET // chunk_bytes,
|
||||||
|
-(-n_rows // chunk_rows))))
|
||||||
|
return chunk_rows, n_workers
|
||||||
|
|
||||||
|
|
||||||
|
def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn):
|
||||||
|
"""Apply fn(r0, r1) over row chunks, in parallel when it pays.
|
||||||
|
|
||||||
|
Chunks write to disjoint output slices, so no locking is needed. numpy
|
||||||
|
ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so
|
||||||
|
threads give real parallelism here — and on a very large file they also
|
||||||
|
keep many more page-fault requests in flight, which is what the I/O path
|
||||||
|
wants.
|
||||||
|
"""
|
||||||
|
bounds = [(r0, min(r0 + chunk_rows, n_rows))
|
||||||
|
for r0 in range(0, n_rows, chunk_rows)]
|
||||||
|
if n_workers <= 1 or len(bounds) == 1:
|
||||||
|
for r0, r1 in bounds:
|
||||||
|
fn(r0, r1)
|
||||||
|
return
|
||||||
|
with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
|
||||||
|
list(pool.map(lambda b: fn(*b), bounds))
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Image computation
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
|
||||||
|
"""Mean of each waveform → (n_rows, n_frames) float32, in ADC counts."""
|
||||||
|
n_rows, n_frames = sras.image_shape(angle_idx)
|
||||||
|
data = sras.data[angle_idx]
|
||||||
|
chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame)
|
||||||
|
img = np.empty((n_rows, n_frames), dtype=np.float32)
|
||||||
|
|
||||||
|
def chunk(r0: int, r1: int):
|
||||||
|
img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)
|
||||||
|
|
||||||
|
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
|
||||||
|
return img
|
||||||
|
|
||||||
|
|
||||||
|
def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int) -> np.ndarray:
|
||||||
|
"""DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
|
||||||
|
cached = sras.cached_dc_mv(angle_idx, ch_idx)
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
return adc_to_mv(compute_dc_image(sras, angle_idx, ch_idx), *sras.cal(ch_idx))
|
||||||
|
|
||||||
|
|
||||||
|
def compute_rf_image(sras: SrasFile, angle_idx: int,
|
||||||
|
dc_threshold_mv: float | None,
|
||||||
|
apply_bg_sub: bool = True,
|
||||||
|
n_fft: int | None = None,
|
||||||
|
dc4_mv: np.ndarray | None = None) -> np.ndarray:
|
||||||
|
"""FFT of each CH1 waveform; pixel = peak frequency in MHz.
|
||||||
|
|
||||||
|
Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
|
||||||
|
never run for them, since that's the expensive part. The masked-out
|
||||||
|
CH1 samples are also never *read*: the boolean mask is applied to the
|
||||||
|
raw memmap slice before any dtype conversion, so numpy only pages in
|
||||||
|
the bytes for pixels that pass the threshold (an untouched memmap page
|
||||||
|
is never read from disk).
|
||||||
|
|
||||||
|
*dc_threshold_mv* of None means "mask nothing", and skips reading and
|
||||||
|
averaging CH4 entirely — worth a third of the I/O when caching a whole
|
||||||
|
file, where the mask is applied later at display time.
|
||||||
|
|
||||||
|
If the DC4 image for this angle is already known (e.g. from the
|
||||||
|
DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
|
||||||
|
(n_rows, n_frames)) to reuse it instead of re-reading CH4 here.
|
||||||
|
|
||||||
|
Fast path: if the file has a precomputed peak-frequency image for this
|
||||||
|
angle (v5 PREC or v7 CACH), zero-padding is off, and the bg-sub flag
|
||||||
|
matches, the stored image is used directly — no FFT is run.
|
||||||
|
"""
|
||||||
|
n_rows, n_frames = sras.image_shape(angle_idx)
|
||||||
|
data = sras.data[angle_idx]
|
||||||
|
|
||||||
|
# ---- Fast path: precomputed image (v5 PREC or v7 CACH) ----------------
|
||||||
|
cached_freq = sras.precomputed_freq_mhz[angle_idx]
|
||||||
|
if (cached_freq is not None
|
||||||
|
and n_fft is None # no custom zero-padding
|
||||||
|
and sras.precomputed_bg_sub == (apply_bg_sub and sras.background is not None)):
|
||||||
|
freq_img = cached_freq.copy()
|
||||||
|
if dc_threshold_mv is not None:
|
||||||
|
# DC4 mask, in priority order: already-cached DC block, caller-
|
||||||
|
# supplied image, or a fresh (cheap — no FFT) recompute.
|
||||||
|
dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
|
||||||
|
if dc4_img is None:
|
||||||
|
dc4_img = dc4_mv if dc4_mv is not None else adc_to_mv(
|
||||||
|
compute_dc_image(sras, angle_idx, CH4_IDX), *sras.cal(CH4_IDX))
|
||||||
|
freq_img[dc4_img < dc_threshold_mv] = 0.0
|
||||||
|
return freq_img
|
||||||
|
|
||||||
|
# ---- Chunked FFT path --------------------------------------------------
|
||||||
|
freq_axis = sras.freq_axis_mhz(n_fft)
|
||||||
|
img = np.zeros((n_rows, n_frames), dtype=np.float32)
|
||||||
|
n_fft_bins = n_fft if n_fft is not None else sras.samples_per_frame
|
||||||
|
chunk_rows, n_workers = _plan_chunks(
|
||||||
|
n_rows, n_frames, max(sras.samples_per_frame, n_fft_bins),
|
||||||
|
live_multiplier=_FFT_LIVE_MULTIPLIER)
|
||||||
|
background = sras.background if (apply_bg_sub and sras.background is not None) else None
|
||||||
|
cal4 = sras.cal(CH4_IDX)
|
||||||
|
|
||||||
|
def chunk(r0: int, r1: int):
|
||||||
|
if dc_threshold_mv is None:
|
||||||
|
valid = None
|
||||||
|
else:
|
||||||
|
if dc4_mv is not None:
|
||||||
|
dc4_chunk = dc4_mv[r0:r1]
|
||||||
|
else:
|
||||||
|
dc4_chunk = adc_to_mv(
|
||||||
|
data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4)
|
||||||
|
valid = dc4_chunk >= dc_threshold_mv # True = run the FFT
|
||||||
|
if not valid.any():
|
||||||
|
return
|
||||||
|
|
||||||
|
# Index the raw memmap slice with the boolean mask *before*
|
||||||
|
# converting dtype — this is a lazy view until touched, so only the
|
||||||
|
# selected elements are actually read from disk; masked-out pixels'
|
||||||
|
# pages are never paged in at all.
|
||||||
|
raw = data[r0:r1, CH1_IDX, :, :]
|
||||||
|
waves = (raw[valid] if valid is not None else raw).astype(np.float32)
|
||||||
|
|
||||||
|
if background is not None:
|
||||||
|
waves -= background # background is 1-D (spf,)
|
||||||
|
|
||||||
|
# Parallelism comes from the outer chunk loop, so keep the inner
|
||||||
|
# transform single-threaded to avoid oversubscribing the machine.
|
||||||
|
spectrum = _do_rfft(waves, n=n_fft, axis=-1, workers=1)
|
||||||
|
del waves
|
||||||
|
# |z|^2 without np.abs()'s extra full-size temporary.
|
||||||
|
power = spectrum.real ** 2
|
||||||
|
power += spectrum.imag ** 2
|
||||||
|
del spectrum
|
||||||
|
# [..., 0] not [:, 0]: the unmasked path keeps the (rows, frames,
|
||||||
|
# bins) shape, where [:, 0] would blank a whole frame.
|
||||||
|
power[..., 0] = 0.0 # suppress DC bin
|
||||||
|
peak_bins = np.argmax(power, axis=-1)
|
||||||
|
del power
|
||||||
|
|
||||||
|
if valid is not None:
|
||||||
|
img[r0:r1][valid] = freq_axis[peak_bins]
|
||||||
|
else:
|
||||||
|
img[r0:r1] = freq_axis[peak_bins]
|
||||||
|
|
||||||
|
_map_row_chunks(n_rows, chunk_rows, n_workers, chunk)
|
||||||
|
return img
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Batch cache (Convert menu) — process-pool entry point
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
def cache_file(path: str, mode: str, apply_bg_sub: bool,
|
||||||
|
fft_backend: str = "numpy", max_workers: int = 0) -> str:
|
||||||
|
"""Compute and store DC or FFT images for every angle of one file,
|
||||||
|
converting v6 → v7 in place. Returns "" on success or an error message.
|
||||||
|
|
||||||
|
Module-level and picklable so it can run in a ProcessPoolExecutor. The
|
||||||
|
FFT backend and worker cap are passed explicitly because module globals
|
||||||
|
do not survive a spawn.
|
||||||
|
"""
|
||||||
|
global _MAX_WORKERS
|
||||||
|
try:
|
||||||
|
set_fft_backend(fft_backend)
|
||||||
|
if max_workers:
|
||||||
|
_MAX_WORKERS = max_workers
|
||||||
|
|
||||||
|
sras = SrasFile(path)
|
||||||
|
if sras.version not in (6, 7):
|
||||||
|
return (f"unsupported version {sras.version} — only v6/v7 files "
|
||||||
|
"can be batch-cached")
|
||||||
|
|
||||||
|
n = sras.n_angles
|
||||||
|
if mode == "dc":
|
||||||
|
dc3 = [adc_to_mv(compute_dc_image(sras, a, CH3_IDX), *sras.cal(CH3_IDX))
|
||||||
|
for a in range(n)]
|
||||||
|
dc4 = [adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
|
||||||
|
for a in range(n)]
|
||||||
|
sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
|
||||||
|
else:
|
||||||
|
effective_bg = apply_bg_sub and sras.background is not None
|
||||||
|
# dc_threshold_mv=None: store unmasked images and mask at display
|
||||||
|
# time (same convention as v5's PREC block). Skipping the mask
|
||||||
|
# also skips reading CH4 entirely.
|
||||||
|
freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
|
||||||
|
apply_bg_sub=effective_bg)
|
||||||
|
for a in range(n)]
|
||||||
|
sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg)
|
||||||
|
return ""
|
||||||
|
except Exception as exc:
|
||||||
|
return str(exc)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Angle alignment (Fusion menu)
|
||||||
|
#
|
||||||
|
# Puts every angle's images onto one shared, zero-padded pixel grid using a
|
||||||
|
# rigid transform only (rotation + translation, never scale). Rotation for
|
||||||
|
# angle `a` is the *known* scan-angle delta relative to a reference angle —
|
||||||
|
# never searched. Only the residual translation is found, via FFT phase
|
||||||
|
# correlation of each angle's binarized CH4 ("dc-mask") image.
|
||||||
|
#
|
||||||
|
# Rotation is done in physical mm space rather than on raw pixel indices:
|
||||||
|
# the x-pixel pitch (SrasFile.pixel_x_mm) is file-wide constant but the
|
||||||
|
# y-pixel pitch (row spacing) can differ from it, and for v6 files can even
|
||||||
|
# vary per angle. Rotating the raw index grid directly would implicitly
|
||||||
|
# assume square pixels and shear a non-square-pixel image — an unwanted
|
||||||
|
# effective anisotropic scale. Instead each angle gets one affine that maps
|
||||||
|
# shared-canvas pixel index -> mm -> undo rotation/shift -> that angle's own
|
||||||
|
# local mm -> that angle's own raw pixel index, matching the output->input
|
||||||
|
# convention scipy.ndimage.affine_transform expects.
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AngleTransform:
|
||||||
|
rotation_deg: float
|
||||||
|
shift_mm: tuple[float, float] # (dx_mm, dy_mm) found by phase correlation
|
||||||
|
matrix: np.ndarray # (2,2): canvas (row,col) -> this angle's raw (row,col)
|
||||||
|
offset: np.ndarray # (2,)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class AlignmentResult:
|
||||||
|
ref_angle_idx: int
|
||||||
|
dc_threshold_mv: float
|
||||||
|
canvas_shape: tuple[int, int] # (n_rows, n_cols)
|
||||||
|
canvas_dx_mm: float
|
||||||
|
canvas_dy_mm: float
|
||||||
|
canvas_origin_mm: tuple[float, float] # mm at canvas pixel index (0, 0)
|
||||||
|
per_angle: dict[int, AngleTransform]
|
||||||
|
|
||||||
|
|
||||||
|
def _pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
||||||
|
"""(dx, dy) mm/pixel for one angle: dx is the file-wide constant
|
||||||
|
pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same
|
||||||
|
assumption _redraw_image makes when it builds the display extent)."""
|
||||||
|
y = sras.y_positions_mm(angle_idx)
|
||||||
|
return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0)
|
||||||
|
|
||||||
|
|
||||||
|
def _bbox_center_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
|
||||||
|
x = sras.x_axis_mm(angle_idx)
|
||||||
|
y = sras.y_positions_mm(angle_idx)
|
||||||
|
return float((x[0] + x[-1]) / 2.0), float((y[0] + y[-1]) / 2.0)
|
||||||
|
|
||||||
|
|
||||||
|
def _bbox_corners_mm(sras: SrasFile, angle_idx: int) -> np.ndarray:
|
||||||
|
"""4 corners (x, y) of this angle's raw mm bounding box, shape (4, 2)."""
|
||||||
|
x = sras.x_axis_mm(angle_idx)
|
||||||
|
y = sras.y_positions_mm(angle_idx)
|
||||||
|
return np.array([[xx, yy] for xx in (x[0], x[-1]) for yy in (y[0], y[-1])])
|
||||||
|
|
||||||
|
|
||||||
|
def _rotation_matrix(theta_deg: float) -> np.ndarray:
|
||||||
|
t = np.radians(theta_deg)
|
||||||
|
c, s = np.cos(t), np.sin(t)
|
||||||
|
return np.array([[c, -s], [s, c]]) # CCW rotation acting on (x, y)
|
||||||
|
|
||||||
|
|
||||||
|
def _theta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
|
||||||
|
return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])
|
||||||
|
|
||||||
|
|
||||||
|
def _corners_in_ref_frame(sras: SrasFile, angle_idx: int, ref_idx: int,
|
||||||
|
shift_mm=(0.0, 0.0)) -> np.ndarray:
|
||||||
|
"""Angle *angle_idx*'s bbox corners, rotated about its own centroid into
|
||||||
|
the reference frame and translated by *shift_mm*. Shape (4, 2)."""
|
||||||
|
R = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx))
|
||||||
|
c_a = np.array(_bbox_center_mm(sras, angle_idx))
|
||||||
|
c_ref = np.array(_bbox_center_mm(sras, ref_idx))
|
||||||
|
shift = np.asarray(shift_mm, dtype=np.float64)
|
||||||
|
return np.array([R @ (corner - c_a) + c_ref + shift
|
||||||
|
for corner in _bbox_corners_mm(sras, angle_idx)])
|
||||||
|
|
||||||
|
|
||||||
|
def _build_affine_canvas_to_raw(sras: SrasFile, angle_idx: int, ref_idx: int,
|
||||||
|
shift_mm: tuple[float, float],
|
||||||
|
canvas_dx: float, canvas_dy: float,
|
||||||
|
canvas_origin_mm: tuple[float, float]
|
||||||
|
) -> tuple[np.ndarray, np.ndarray]:
|
||||||
|
"""matrix, offset s.t. raw_index = matrix @ [row_out, col_out] + offset,
|
||||||
|
matching scipy.ndimage.affine_transform's output->input convention.
|
||||||
|
|
||||||
|
Pipeline (all mm unless noted):
|
||||||
|
[X;Y] = A_out @ [row_out;col_out] + b_out # canvas idx -> ref-frame mm
|
||||||
|
[lx;ly] = R(theta)^T @ ([X;Y]-c_ref-shift) + c_a # undo rotation+shift -> angle a's local mm
|
||||||
|
[row;col] = D @ ([lx;ly] - [x_start_a; y0_a]) # local mm -> angle a's raw idx
|
||||||
|
|
||||||
|
where theta = angles_deg[angle_idx] - angles_deg[ref_idx], c_ref/c_a are
|
||||||
|
each angle's own raw-bbox mm centroid (the rotation pivot — this keeps
|
||||||
|
rotated content centered, minimizing required canvas padding), and
|
||||||
|
A_out/D are the index<->mm scaling matrices for the canvas pitch and
|
||||||
|
this angle's own native pitch respectively.
|
||||||
|
"""
|
||||||
|
Rinv = _rotation_matrix(_theta_deg(sras, angle_idx, ref_idx)).T
|
||||||
|
cx_a, cy_a = _bbox_center_mm(sras, angle_idx)
|
||||||
|
cx_ref, cy_ref = _bbox_center_mm(sras, ref_idx)
|
||||||
|
dx_a, dy_a = _pixel_pitch_mm(sras, angle_idx)
|
||||||
|
x0_a = float(sras.x_start_mm[angle_idx])
|
||||||
|
y0_a = float(sras.y_positions_mm(angle_idx)[0])
|
||||||
|
|
||||||
|
A_out = np.array([[0.0, canvas_dx], [canvas_dy, 0.0]]) # [row,col] -> [X,Y]
|
||||||
|
b_out = np.array(canvas_origin_mm, dtype=np.float64)
|
||||||
|
D = np.array([[0.0, 1.0 / dy_a], [1.0 / dx_a, 0.0]]) # [x,y] -> [row,col]
|
||||||
|
shift = np.array(shift_mm, dtype=np.float64)
|
||||||
|
c_ref_v = np.array([cx_ref, cy_ref])
|
||||||
|
c_a_v = np.array([cx_a, cy_a])
|
||||||
|
origin_a = np.array([x0_a, y0_a])
|
||||||
|
|
||||||
|
matrix = D @ Rinv @ A_out
|
||||||
|
offset = D @ Rinv @ (b_out - c_ref_v - shift) + D @ (c_a_v - origin_a)
|
||||||
|
return matrix, offset
|
||||||
|
|
||||||
|
|
||||||
|
def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
|
||||||
|
order: int = 0) -> np.ndarray:
|
||||||
|
"""Resample any already-computed 2D image for `angle_idx` (same shape as
|
||||||
|
that angle's raw (n_rows, n_frames)) onto the shared alignment canvas.
|
||||||
|
order=0 (nearest) avoids blending real data with zero-padding or with
|
||||||
|
masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
|
||||||
|
agnostic: the same per-angle transform (found from the CH4 mask) works
|
||||||
|
for any channel's image of that angle."""
|
||||||
|
t = result.per_angle[angle_idx]
|
||||||
|
return scipy_ndimage.affine_transform(
|
||||||
|
img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
|
||||||
|
output_shape=result.canvas_shape, order=order,
|
||||||
|
mode="constant", cval=0.0)
|
||||||
|
|
||||||
|
|
||||||
|
def _block_mean_downsample(img: np.ndarray, factor: int) -> np.ndarray:
|
||||||
|
if factor <= 1:
|
||||||
|
return img
|
||||||
|
h, w = img.shape
|
||||||
|
h2, w2 = (h // factor) * factor, (w // factor) * factor
|
||||||
|
trimmed = img[:h2, :w2]
|
||||||
|
return trimmed.reshape(h2 // factor, factor, w2 // factor, factor).mean(axis=(1, 3))
|
||||||
|
|
||||||
|
|
||||||
|
def _phase_correlate_shift(ref_img: np.ndarray, mov_img: np.ndarray) -> tuple[int, int]:
|
||||||
|
"""FFT normalized cross-power-spectrum phase correlation. Returns the
|
||||||
|
integer (dr, dc) pixel shift of mov_img relative to ref_img; both must
|
||||||
|
be the same shape. Risk: if the true shift is near +/- half the array
|
||||||
|
size, wraparound can bias the peak — mitigated by the generous
|
||||||
|
margin_frac padding in _working_canvas_for_pair, which keeps the true
|
||||||
|
residual shift small relative to the correlation canvas."""
|
||||||
|
F1 = scipy_fft.fft2(ref_img.astype(np.float64), workers=-1)
|
||||||
|
F2 = scipy_fft.fft2(mov_img.astype(np.float64), workers=-1)
|
||||||
|
R = F1 * np.conj(F2)
|
||||||
|
R /= np.maximum(np.abs(R), 1e-12)
|
||||||
|
corr = scipy_fft.ifft2(R, workers=-1).real
|
||||||
|
dr, dc = np.unravel_index(np.argmax(corr), corr.shape)
|
||||||
|
h, w = corr.shape
|
||||||
|
if dr > h // 2:
|
||||||
|
dr -= h
|
||||||
|
if dc > w // 2:
|
||||||
|
dc -= w
|
||||||
|
return int(dr), int(dc)
|
||||||
|
|
||||||
|
|
||||||
|
def _working_canvas_for_pair(sras: SrasFile, ref_idx: int, a_idx: int,
|
||||||
|
dx: float, dy: float, margin_frac: float = 0.3
|
||||||
|
) -> tuple[tuple[float, float], tuple[int, int]]:
|
||||||
|
"""Union of the reference's own raw bbox and angle a's raw bbox rotated
|
||||||
|
(about its own center) into the ref frame with zero shift, padded by
|
||||||
|
margin_frac on each side — sized generously so the true phase-
|
||||||
|
correlation shift lands well inside the canvas (see
|
||||||
|
_phase_correlate_shift's wraparound note)."""
|
||||||
|
pts = np.vstack([_bbox_corners_mm(sras, ref_idx),
|
||||||
|
_corners_in_ref_frame(sras, a_idx, ref_idx)])
|
||||||
|
x_min, y_min = pts.min(axis=0)
|
||||||
|
x_max, y_max = pts.max(axis=0)
|
||||||
|
pad_x, pad_y = (x_max - x_min) * margin_frac, (y_max - y_min) * margin_frac
|
||||||
|
x_min, x_max = x_min - pad_x, x_max + pad_x
|
||||||
|
y_min, y_max = y_min - pad_y, y_max + pad_y
|
||||||
|
n_cols = int(np.ceil((x_max - x_min) / dx)) + 1
|
||||||
|
n_rows = int(np.ceil((y_max - y_min) / abs(dy))) + 1
|
||||||
|
origin = (x_min, y_min if dy > 0 else y_max)
|
||||||
|
return origin, (n_rows, n_cols)
|
||||||
|
|
||||||
|
|
||||||
|
def _parallel_map(fn, items, n_workers: int) -> list:
|
||||||
|
"""fn over items, in order, threaded when it pays."""
|
||||||
|
items = list(items)
|
||||||
|
if n_workers <= 1 or len(items) <= 1:
|
||||||
|
return [fn(x) for x in items]
|
||||||
|
with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool:
|
||||||
|
return list(pool.map(fn, items))
|
||||||
|
|
||||||
|
|
||||||
|
def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
|
||||||
|
dc_threshold_mv: float,
|
||||||
|
progress_cb=None) -> AlignmentResult:
|
||||||
|
"""Top-level alignment driver. Runs on a background thread (see
|
||||||
|
AngleAlignmentWorker) — deliberately recomputes CH4 DC images from
|
||||||
|
scratch rather than reading the GUI-thread _dc_cache dict, since
|
||||||
|
background-thread workers must not touch GUI-thread-owned caches."""
|
||||||
|
n = sras.n_angles # already the *complete*-angle count for aborted v6 scans
|
||||||
|
dx_ref, dy_ref = _pixel_pitch_mm(sras, ref_angle_idx)
|
||||||
|
n_workers = max(1, min(_MAX_WORKERS, n))
|
||||||
|
|
||||||
|
# progress_cb fires from pool threads, so the counter behind it must be
|
||||||
|
# atomic. list.append is, and len() of a list is a consistent read.
|
||||||
|
_ticks: list[int] = []
|
||||||
|
|
||||||
|
def tick(base: int, span: int):
|
||||||
|
if progress_cb:
|
||||||
|
_ticks.append(1)
|
||||||
|
progress_cb(base + int(min(len(_ticks), n) / n * span))
|
||||||
|
|
||||||
|
# ---- Step 1: binarized CH4 mask per angle, native per-angle grid -----
|
||||||
|
def mask_for(a: int) -> np.ndarray:
|
||||||
|
dc4 = adc_to_mv(compute_dc_image(sras, a, CH4_IDX), *sras.cal(CH4_IDX))
|
||||||
|
m = (dc4 >= dc_threshold_mv).astype(np.float32)
|
||||||
|
tick(0, 25)
|
||||||
|
return m
|
||||||
|
|
||||||
|
masks = dict(enumerate(_parallel_map(mask_for, range(n), n_workers)))
|
||||||
|
|
||||||
|
# ---- Step 2: coarse correlation stage (downsample first, then rotate)
|
||||||
|
# Downsampling before affine_transform (not after) is what keeps this
|
||||||
|
# tractable for a v6 scan with thousands of rows/frames per angle.
|
||||||
|
max_dim = max(max(m.shape) for m in masks.values())
|
||||||
|
factor = max(1, int(np.ceil(max_dim / 1024)))
|
||||||
|
dx_c, dy_c = dx_ref * factor, dy_ref * factor
|
||||||
|
masks_small = {a: _block_mean_downsample(m, factor) for a, m in masks.items()}
|
||||||
|
|
||||||
|
_ticks.clear()
|
||||||
|
|
||||||
|
def shift_for(a: int) -> tuple[float, float]:
|
||||||
|
if a == ref_angle_idx:
|
||||||
|
tick(25, 50)
|
||||||
|
return (0.0, 0.0)
|
||||||
|
work_origin, work_shape = _working_canvas_for_pair(
|
||||||
|
sras, ref_angle_idx, a, dx_c, dy_c, margin_frac=0.3)
|
||||||
|
|
||||||
|
# Coarse canvas->raw affine at native pitch, then rescale by /factor
|
||||||
|
# so it maps coarse-canvas idx -> coarse (downsampled) raw idx —
|
||||||
|
# exact for block-mean downsampling (up to the trimmed remainder).
|
||||||
|
m_a, o_a = _build_affine_canvas_to_raw(
|
||||||
|
sras, a, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
|
||||||
|
m_ref, o_ref = _build_affine_canvas_to_raw(
|
||||||
|
sras, ref_angle_idx, ref_angle_idx, (0.0, 0.0), dx_c, dy_c, work_origin)
|
||||||
|
rotated_a = scipy_ndimage.affine_transform(
|
||||||
|
masks_small[a], m_a / factor, offset=o_a / factor,
|
||||||
|
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
||||||
|
embedded_ref = scipy_ndimage.affine_transform(
|
||||||
|
masks_small[ref_angle_idx], m_ref / factor, offset=o_ref / factor,
|
||||||
|
output_shape=work_shape, order=0, mode="constant", cval=0.0)
|
||||||
|
|
||||||
|
dr, dc = _phase_correlate_shift(embedded_ref, rotated_a)
|
||||||
|
tick(25, 50)
|
||||||
|
return (dc * dx_c, dr * dy_c)
|
||||||
|
|
||||||
|
shifts_mm = dict(enumerate(_parallel_map(shift_for, range(n), n_workers)))
|
||||||
|
|
||||||
|
# ---- Step 3: union bounding box over all angles (rotation+shift applied)
|
||||||
|
corners = np.vstack([_corners_in_ref_frame(sras, a, ref_angle_idx, shifts_mm[a])
|
||||||
|
for a in range(n)])
|
||||||
|
x_min, y_min = corners.min(axis=0)
|
||||||
|
x_max, y_max = corners.max(axis=0)
|
||||||
|
n_cols = int(np.ceil((x_max - x_min) / dx_ref)) + 1
|
||||||
|
n_rows = int(np.ceil((y_max - y_min) / abs(dy_ref))) + 1
|
||||||
|
canvas_origin_mm = (float(x_min), float(y_min if dy_ref > 0 else y_max))
|
||||||
|
|
||||||
|
# ---- Step 4: final per-angle full-resolution affine (canvas -> raw idx)
|
||||||
|
per_angle: dict[int, AngleTransform] = {}
|
||||||
|
for a in range(n):
|
||||||
|
matrix, offset = _build_affine_canvas_to_raw(
|
||||||
|
sras, a, ref_angle_idx, shifts_mm[a], dx_ref, dy_ref, canvas_origin_mm)
|
||||||
|
per_angle[a] = AngleTransform(
|
||||||
|
_theta_deg(sras, a, ref_angle_idx), shifts_mm[a], matrix, offset)
|
||||||
|
if progress_cb:
|
||||||
|
progress_cb(75 + int((a + 1) / n * 25))
|
||||||
|
|
||||||
|
return AlignmentResult(ref_angle_idx, dc_threshold_mv, (n_rows, n_cols),
|
||||||
|
dx_ref, dy_ref, canvas_origin_mm, per_angle)
|
||||||
|
|
||||||
|
|
||||||
|
# Back-compat alias for the pre-split private name (used by tooling).
|
||||||
|
_compute_angle_alignment = compute_angle_alignment
|
||||||
+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
|
||||||
+678
-2054
File diff suppressed because it is too large
Load Diff
+224
@@ -0,0 +1,224 @@
|
|||||||
|
#!/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
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
|
||||||
|
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(QObject):
|
||||||
|
"""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)
|
||||||
|
else:
|
||||||
|
img = dc_image_mv(self._sras, self._angle, self._ch)
|
||||||
|
self.finished.emit(img)
|
||||||
|
except Exception as exc:
|
||||||
|
self.error.emit(str(exc))
|
||||||
|
|
||||||
|
|
||||||
|
class DcPrecomputeWorker(QObject):
|
||||||
|
"""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
|
||||||
|
self._stop = False
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
self._stop = True
|
||||||
|
|
||||||
|
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
|
||||||
|
return a, dc_image_mv(self._sras, a, CH3_IDX), dc_image_mv(self._sras, a, CH4_IDX)
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
try:
|
||||||
|
n = self._sras.n_angles
|
||||||
|
n_workers = max(1, min(compute._MAX_WORKERS, n))
|
||||||
|
with ThreadPoolExecutor(max_workers=n_workers) as pool:
|
||||||
|
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
|
||||||
|
try:
|
||||||
|
for fut in as_completed(futures):
|
||||||
|
if self._stop:
|
||||||
|
break
|
||||||
|
a, dc3, dc4 = fut.result()
|
||||||
|
self.angle_done.emit(a, dc3, dc4)
|
||||||
|
finally:
|
||||||
|
if self._stop:
|
||||||
|
for fut in futures:
|
||||||
|
fut.cancel()
|
||||||
|
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 _pool(self, n_files: int):
|
||||||
|
"""A process pool, falling back to threads if the platform refuses to
|
||||||
|
spawn (still a win: the compute releases the GIL)."""
|
||||||
|
n = max(1, min(_BATCH_MAX_PROCS, n_files))
|
||||||
|
try:
|
||||||
|
return ProcessPoolExecutor(max_workers=n), n
|
||||||
|
except (OSError, ValueError):
|
||||||
|
return ThreadPoolExecutor(max_workers=n), n
|
||||||
|
|
||||||
|
def run(self):
|
||||||
|
paths = self._paths
|
||||||
|
executor, n_procs = self._pool(len(paths))
|
||||||
|
# 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)
|
||||||
|
|
||||||
|
done = 0
|
||||||
|
with 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 Exception as exc:
|
||||||
|
err = str(exc)
|
||||||
|
done += 1
|
||||||
|
self.file_done.emit(path, err)
|
||||||
|
self.progress.emit(int(done / max(1, len(paths)) * 100))
|
||||||
|
|
||||||
|
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))
|
||||||
Reference in New Issue
Block a user