Thomas Ales
f40c965b74
Alignment was two disconnected menu actions. `Angle Alignment` ran the
registration with ref_angle_idx hard-coded to 0, no exposed parameters and
no way to retry; `Manual Alignment...` was a separate dialog that
deliberately refused to inherit the automatic result, so a bad fit meant
starting over by hand. Neither told you whether the alignment was any
good, and neither produced anything durable beyond an in-memory result.
Fusion -> Alignment Wizard... now covers all of it in three steps:
1. Correlate. Reference angle, DC threshold, source, rotation seed and
sign, search window, coarse step and fine grid are all on the page,
and Run/Re-run is repeatable. Angles start pre-rotated from the
stage angles in the file, so the page is informative before any
correlation runs. The verdict is a picture: every angle's DC mask
reprojected onto the shared canvas and summed, coloured by how many
angles cover each pixel, so a good alignment reads as one saturated
plateau and a bad one as a fringe of low-count halos. A per-angle
fit table flags angles that did not register or that disagree with
their stage angle by more than a degree.
2. Crop. An axis-aligned rectangle on that canvas, with numeric
canvas-pixel boxes synced both ways and a live size estimate.
"Fit to full overlap" uses a largest-rectangle sweep rather than a
bounding box: the overlap region of several rotated scans is roughly
a disc, whose bounding box has corners no angle covers.
3. Save. Writes the aligned, cropped stack to a new .sras.
Because the wizard replaces both actions it absorbs the old dialog's
by-eye nudge editor — without it, a scan the search cannot fit would
have no fallback at all. ManualAlignmentDialog is therefore deleted
rather than left orphaned, and its tests move to the wizard.
Two bugs found while driving it end to end and fixed here: QSpinBox
setRange clamps and emits valueChanged, which committed a 1x1 crop
before the default preset could run; and the mm round trip returns an
exact pixel boundary as 11.000000000000002, so a bare ceil() added a
spurious column on every rectangle edit.
Verified: 103 pass, the 6 test_stored_cache failures are pre-existing on
main, and tools/check_equivalence.py is byte-identical to main.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
425 lines
17 KiB
Python
425 lines
17 KiB
Python
#!/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_align_export import write_aligned_sras
|
||
from sras_compute import cache_file, 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 _PooledWorker(CancellableWorker):
|
||
"""Fans a per-item computation across a thread pool, emitting each result
|
||
from this worker's own thread as it lands (never from a pool thread).
|
||
|
||
Subclasses provide _plan() -> n_workers (stashing whatever per-run
|
||
context they need), _items(), _one(item) -> result, and _emit(result).
|
||
On stop(): queued items are dropped, in-flight ones are not waited for —
|
||
that is what keeps closing the window responsive on a large scan.
|
||
"""
|
||
finished = pyqtSignal()
|
||
error = pyqtSignal(str)
|
||
|
||
def run(self):
|
||
try:
|
||
pool = ThreadPoolExecutor(max_workers=max(1, self._plan()))
|
||
try:
|
||
futures = [pool.submit(self._one, it) for it in self._items()]
|
||
for fut in as_completed(futures):
|
||
if self._stop:
|
||
break
|
||
self._emit(fut.result())
|
||
finally:
|
||
pool.shutdown(wait=not self._stop, cancel_futures=True)
|
||
self.finished.emit()
|
||
except Exception as exc:
|
||
self.error.emit(str(exc))
|
||
|
||
|
||
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(_PooledWorker):
|
||
"""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.
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray, np.ndarray) # angle_idx, dc3_mv, dc4_mv
|
||
|
||
def __init__(self, sras: SrasFile):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angle_budget = 0
|
||
|
||
def _plan(self) -> int:
|
||
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
|
||
return n_workers
|
||
|
||
def _items(self):
|
||
return range(self._sras.n_angles)
|
||
|
||
def _one(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 _emit(self, result):
|
||
self.angle_done.emit(*result)
|
||
|
||
|
||
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), ``"fft"`` (CH1 peak-frequency
|
||
images, unmasked — masking is applied at display time, same as v5's PREC
|
||
convention), or ``"fft_rowavg"`` (same-row, distance-weighted CH1
|
||
averaging before the FFT — needs *dc_threshold_mv* and a positive
|
||
*row_avg_n*; see ``sras_compute.cache_file``).
|
||
|
||
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,
|
||
dc_threshold_mv: float | None = None, row_avg_n: int = 0):
|
||
super().__init__()
|
||
self._paths = paths
|
||
self._mode = mode
|
||
self._apply_bg_sub = apply_bg_sub
|
||
self._dc_threshold = dc_threshold_mv
|
||
self._row_avg_n = row_avg_n
|
||
|
||
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,
|
||
dc_threshold_mv=self._dc_threshold,
|
||
row_avg_n=self._row_avg_n): 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.default_max_workers(),
|
||
dc_threshold_mv=self._dc_threshold,
|
||
row_avg_n=self._row_avg_n)
|
||
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 Ch4MaskWorker(_PooledWorker):
|
||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for the
|
||
alignment wizard's initial threshold-mask stack.
|
||
|
||
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
|
||
CorrelatePage._start_mask_prep).
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||
|
||
def __init__(self, sras: SrasFile, angle_indices: list[int]):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._angles = angle_indices
|
||
self._budget = 0
|
||
|
||
def _plan(self) -> int:
|
||
n_workers, self._budget = compute.plan_angle_level(self._sras)
|
||
return n_workers
|
||
|
||
def _items(self):
|
||
return self._angles
|
||
|
||
def _one(self, a: int) -> tuple[int, np.ndarray]:
|
||
return a, dc_image_mv(self._sras, a, CH4_IDX,
|
||
max_workers=1, budget=self._budget)
|
||
|
||
def _emit(self, result):
|
||
self.angle_done.emit(*result)
|
||
|
||
|
||
class CrossCorrelateWorker(_PooledWorker):
|
||
"""Rigid registration (rotation + translation, never scale) of each of
|
||
*angle_indices* against *ref_angle_idx*, for the alignment wizard's
|
||
Run/Re-run Correlation button.
|
||
|
||
Runs on a background thread — registering a real many-angle,
|
||
high-resolution scan takes long enough that doing it on the GUI thread
|
||
would visibly freeze the dialog. Rotation is *searched*, not taken from the
|
||
stage's reported angle: see compute.register_angle_to_reference, which
|
||
seeds from that angle but scores both of its signs and refines from there.
|
||
dc4_mv is the dialog's own already-in-memory per-angle CH4 image — this
|
||
worker does no fetching of its own.
|
||
"""
|
||
# angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source
|
||
angle_done = pyqtSignal(int, float, float, float, float, str)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int],
|
||
dc4_mv: dict[int, np.ndarray], *,
|
||
sources: tuple[str, ...], dc_threshold_mv: float,
|
||
search_deg: float, reg_kwargs: dict | None = None):
|
||
"""*reg_kwargs* is splatted into register_angle_to_reference on top of
|
||
the named arguments — the wizard's rotation-search controls (seed,
|
||
signs, refine, grid sizes) go through here, so exposing another knob
|
||
needs no change to this class."""
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._angles = angle_indices
|
||
self._dc4_mv = dc4_mv
|
||
self._sources = sources
|
||
self._threshold = dc_threshold_mv
|
||
self._search_deg = search_deg
|
||
self._reg_kwargs = dict(reg_kwargs or {})
|
||
|
||
def _plan(self) -> int:
|
||
return compute.registration_workers(self._sras)
|
||
|
||
def _items(self):
|
||
return self._angles
|
||
|
||
def _one(self, a: int) -> tuple[int, compute.RigidFit]:
|
||
return a, compute.register_angle_to_reference(
|
||
self._sras, a, self._ref, self._dc4_mv,
|
||
dc_threshold_mv=self._threshold, sources=self._sources,
|
||
search_deg=self._search_deg, **self._reg_kwargs)
|
||
|
||
def _emit(self, result):
|
||
a, fit = result
|
||
self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0],
|
||
fit.shift_mm[1], fit.score, fit.source)
|
||
|
||
|
||
class AlignedExportWorker(CancellableWorker):
|
||
"""Writes the aligned, cropped .sras on a background thread.
|
||
|
||
Unlike every other worker here this one produces a *file*, which changes
|
||
what cancellation has to mean: write_aligned_sras stages into a ".part"
|
||
sibling and removes it when should_stop() fires, so a cancelled or crashed
|
||
export leaves nothing behind. That matters more than it sounds — a
|
||
truncated .sras is not detectably broken, since the v6 parser reads a short
|
||
file as an aborted scan and opens it happily.
|
||
|
||
Cancellation is polled per output row chunk, the same granularity
|
||
CancellableWorker's docstring justifies, so closing the window never waits
|
||
on a multi-gigabyte write.
|
||
"""
|
||
progress = pyqtSignal(int) # 0-100
|
||
finished = pyqtSignal(str, str) # written path ("" = none), error
|
||
|
||
def __init__(self, sras: SrasFile, result, out_path: str):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._result = result
|
||
self._out_path = out_path
|
||
|
||
def run(self):
|
||
try:
|
||
written = write_aligned_sras(
|
||
self._sras, self._result, self._out_path,
|
||
progress_cb=self.progress.emit, should_stop=self._stopped)
|
||
if self._stopped():
|
||
self.finished.emit("", "") # cancelled: no file, no error
|
||
else:
|
||
self.finished.emit(str(written), "")
|
||
except Exception as exc:
|
||
self.finished.emit("", str(exc))
|