Split into four modules, deduplicate, and parallelize the compute paths

Structure
  sras_format.py   parsing/writing, calibration, axes  (numpy + struct)
  sras_compute.py  DC/FFT images, alignment, batch cache  (+ scipy)
  sras_workers.py  Qt background workers
  sras_viewer.py   ROI, canvases, dialog, main window

format+compute import in 0.59s with no Qt or matplotlib (vs 2.96s for the
full app), which is what makes a spawn-based process pool worth using.

Deduplication
  - SrasFile.cal() and compute.dc_image_mv() replace the ADC->mV
    calibration incantation that appeared at ten call sites.
  - _read_preambles/_read_background/_set_calibration are shared by the
    legacy and v6 parsers instead of duplicated.
  - v5 PREC images are normalised into the same ragged per-angle list
    v6/v7 uses, removing four isinstance() branches and shortening
    compute_rf_image's fast path.
  - _run_worker replaces five copies of the QThread setup and five
    near-identical teardown methods; the two lifetime hazards they
    guarded against are now documented once, authoritatively.
  - _on_channel_changed defers to _update_controls_enabled rather than
    re-deriving the same six enable rules.
  - _corners_in_ref_frame, _CHANNEL_DISPLAY dict, unified progress-dialog
    helper, dead _on_roi_angle_edited stub removed.
  - sras_average.py builds on SrasFile instead of carrying a second copy
    of the header format, and streams per angle instead of loading the
    whole file (was a ~3x file-size peak).
  Executable lines: 2390 -> 2254, despite adding all of the below.

Parallelism
  - compute_rf_image/compute_dc_image map row chunks over a thread pool.
    Worker count is derived from the memory budget rather than the core
    count: on a large scan chunk_rows is already floored at 1 row (~75 MB
    of float32 at 7507x2500), so only the worker count can bound peak RAM.
  - DcPrecomputeWorker computes angles on a pool, emitting each result
    from its own QThread as it lands.
  - BatchCacheWorker runs one process per file via cache_file(); only
    paths and scalars cross the boundary. Per-process thread counts are
    divided so the two levels don't oversubscribe. Drops the old pass 1,
    which fully parsed every file just to weight a progress bar.
  - dc_threshold_mv=None means "no mask" and skips the CH4 read entirely
    — the batch FFT job previously read all of CH4 to compare against a
    threshold of -1e9.
  - Alignment mask and correlation stages map over angles; fft2/ifft2 use
    workers=-1.

Fixes found on the way
  - A partially-cached v5 file showed an all-zero image for uncached
    angles: the fast-path check was file-wide, not per-angle.
  - v7 cached images were read-only big-endian views; now native float32.

Verified: tools/check_equivalence.py produces byte-identical hashes for
167 outputs (DC/FFT across channels, angles, bg-sub, pad factors and
thresholds, plus the full alignment result) against ed0eba4, on both
synthetic files and a real 496 GB 17-angle v6 scan.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Thomas Ales
2026-07-30 23:38:08 -05:00
parent 95c273dce5
commit 55a4c5e42f
5 changed files with 2218 additions and 2275 deletions
+224
View File
@@ -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)`` (0100 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) # 0100
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))