Thomas Ales
6d0e30b9ce
The alignment machinery never modified a scan: it produced an
AlignmentResult and every consumer resampled on the fly. That is right
for a viewer but means the aligned stack cannot leave the process, so no
other tool can read it and reopening the scan redoes the registration.
sras_align_export.write_aligned_sras bakes an alignment into a new v6
file: each angle is gathered onto the shared canvas by nearest neighbour,
so every output angle shares one grid and the file opens already aligned.
Registering the export against itself returns identity, which is the test
that pins the whole index chain.
Three details that are easy to get wrong and are now covered:
* Rounding must be floor(x + 0.5), not np.rint. scipy's order=0 rounds
halves away from zero, and the canvas is snapped to the reference's
pixel grid, so exact halves are common rather than hypothetical.
* Out-of-bounds must be tested on the fractional coordinate against
[0, n-1], not on the rounded index, or a one-pixel rim gets real data
everywhere the Aligned View shows padding.
* ...but with a tolerance, because the mm-space affine chain lands an
exactly-integer transform a few times 1e-13 off. A bare >= 0 drops
the *reference* angle's entire first row and last column.
Padding is the per-channel ADC code nearest 0 mV, not zero: zero ADC
decodes to ~+100 mV on real calibration and would masquerade as sample.
Source rows are served from sliding in-RAM bands. A rotated angle maps
one output row to a diagonal across the source, so indexing a memmap in
output order refaults nearly the whole angle per row — terabytes of
paging for a gigabyte of data.
Also in compute: crop_alignment_result (a crop is pure index translation,
so it folds into the affine's offset rather than becoming a second
transform), overlap_stats, largest_rect_at_least for a crop that stays
inside the overlap region, and seed/sign/refine knobs on
register_angle_to_reference whose defaults leave existing behaviour and
tests byte-identical.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
452 lines
18 KiB
Python
452 lines
18 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_angle_alignment, compute_rf_image, dc_image_mv,
|
||
)
|
||
from sras_format import CH3_IDX, CH4_IDX, SrasFile
|
||
|
||
# Concurrency caps. Batch conversion runs one process per file, and each of
|
||
# those processes threads internally, so the two must be divided rather than
|
||
# both set to the core count. Files also commonly sit on one external drive,
|
||
# where a dozen concurrent readers is slower than a few — hence the low
|
||
# default, overridable from the environment.
|
||
_BATCH_MAX_PROCS = int(os.environ.get("SRAS_BATCH_PROCS", 0)) or min(
|
||
4, os.cpu_count() or 2)
|
||
|
||
# Spawning a pool costs roughly a second of interpreter startup (each child
|
||
# re-imports the entry module). That is noise against a multi-GB scan but
|
||
# dominates a batch of small files, where it would make the job *slower* —
|
||
# so below this total size the batch just runs in the worker thread.
|
||
_BATCH_POOL_MIN_BYTES = int(os.environ.get("SRAS_BATCH_POOL_MIN_MB", 512)) * 1024 * 1024
|
||
|
||
|
||
class CancellableWorker(QObject):
|
||
"""A worker whose compute polls stop() between row chunks.
|
||
|
||
Without this a shutdown has to wait out whatever is in flight, and on a
|
||
large scan a single angle is ~40 s — far too long to block closing the
|
||
window. Chunk-level polling bounds the wait to one chunk instead.
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self._stop = False
|
||
|
||
def stop(self):
|
||
self._stop = True
|
||
|
||
def _stopped(self) -> bool:
|
||
return self._stop
|
||
|
||
|
||
class _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 AngleAlignmentWorker(QObject):
|
||
"""Computes the rigid (rotation + translation, never scale) alignment for
|
||
every angle in *sras* against *ref_angle_idx*, by cross-correlating each
|
||
angle's CH4 image against the reference's. Both the rotation and the
|
||
translation are found from image content — see compute_angle_alignment.
|
||
"""
|
||
progress = pyqtSignal(int) # 0–100
|
||
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
|
||
|
||
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
|
||
super().__init__()
|
||
self._sras = sras
|
||
self._ref = ref_angle_idx
|
||
self._threshold = dc_threshold_mv
|
||
|
||
def run(self):
|
||
try:
|
||
result = compute_angle_alignment(
|
||
self._sras, self._ref, self._threshold,
|
||
progress_cb=self.progress.emit)
|
||
self.finished.emit(result, "")
|
||
except Exception as exc:
|
||
self.finished.emit(None, str(exc))
|
||
|
||
|
||
class Ch4MaskWorker(_PooledWorker):
|
||
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
|
||
ManualAlignmentDialog's initial threshold-mask overlay.
|
||
|
||
Reuses dc_image_mv, which prefers a stored v5/v7 cache over recomputing
|
||
from raw waveforms, so this only does real work for a file that hasn't
|
||
gone through the v7 "Convert" batch step and for angles the main
|
||
window's own DcPrecomputeWorker (which runs automatically right after
|
||
every file load) hasn't reached yet. In the common case — the user opens
|
||
Fusion -> Manual Alignment after DC precompute has already finished —
|
||
*angle_indices* is empty and this worker is never even constructed (see
|
||
ManualAlignmentDialog._start_mask_prep).
|
||
"""
|
||
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
|
||
|
||
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 ManualAlignmentDialog's Auto
|
||
Cross-Correlate 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))
|