tka/sras-viewer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/main' into main

Browse Source 
# Conflicts:
#	sras_compute.py
#	sras_viewer.py
#	sras_workers.py
#	tools/make_test_sras.py
#	tools/test_refactor.py
This commit is contained in:
Thomas Ales [M S E]
2026-08-10 18:58:21 -05:00
parent 4212b8c313 191d1b8946
commit 105b9514a5
34 changed files with 10768 additions and 5738 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sras_workers.py
+149 -125
View File
@@ -19,9 +19,8 @@ from matplotlib.figure import Figure
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_align_export import write_aligned_sras
from sras_compute import cache_file, compute_rf_image, dc_image_mv
from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, SrasFile
# Concurrency caps. Batch conversion runs one process per file, and each of
@@ -58,6 +57,34 @@ class CancellableWorker(QObject):
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)
@@ -121,29 +148,29 @@ class ComputeWorker(CancellableWorker):
self.error.emit(str(exc))
class DcPrecomputeWorker(CancellableWorker):
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.
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._angle_budget = 0
def _one_angle(self, a: int) -> tuple[int, np.ndarray, np.ndarray]:
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).
@@ -153,26 +180,8 @@ class DcPrecomputeWorker(CancellableWorker):
dc_image_mv(self._sras, a, CH3_IDX, **kw),
dc_image_mv(self._sras, a, CH4_IDX, **kw))
def run(self):
try:
n = self._sras.n_angles
n_workers, self._angle_budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=n_workers)
try:
futures = {pool.submit(self._one_angle, a): a for a in range(n)}
for fut in as_completed(futures):
if self._stop:
break
a, dc3, dc4 = fut.result()
self.angle_done.emit(a, dc3, dc4)
finally:
# cancel_futures drops the queued angles; should_stop lets the
# in-flight ones bail within a chunk. Not waiting here is what
# keeps closing the window responsive on a large scan.
pool.shutdown(wait=not self._stop, cancel_futures=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
def _emit(self, result):
self.angle_done.emit(*result)
class BatchCacheWorker(QObject):
@@ -181,9 +190,15 @@ class BatchCacheWorker(QObject):
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
*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).
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``).
Both FFT modes cache at *pad_factor*, which the caller sets from the
viewer's own padding — a cache stored at a pad the user is not viewing
at is one the display can never use.
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
@@ -195,11 +210,16 @@ class BatchCacheWorker(QObject):
file_done = pyqtSignal(str, str)
finished = pyqtSignal()
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool):
def __init__(self, paths: list[str], mode: str, apply_bg_sub: bool,
dc_threshold_mv: float | None = None, row_avg_n: int = 0,
pad_factor: int = 1):
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
self._pad_factor = pad_factor
def _report(self, path: str, err: str, done: int, total: int):
self.file_done.emit(path, err)
@@ -224,7 +244,10 @@ class BatchCacheWorker(QObject):
with ProcessPoolExecutor(max_workers=n_procs) as executor:
futures = {
executor.submit(cache_file, p, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), per_proc_workers): p
compute.get_fft_backend(), per_proc_workers,
pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n): p
for p in paths
}
for fut in as_completed(futures):
@@ -247,7 +270,11 @@ class BatchCacheWorker(QObject):
for path in paths:
try:
err = cache_file(path, self._mode, self._apply_bg_sub,
compute.get_fft_backend(), compute._MAX_WORKERS)
compute.get_fft_backend(),
compute.default_max_workers(),
pad_factor=self._pad_factor,
dc_threshold_mv=self._dc_threshold,
row_avg_n=self._row_avg_n)
except Exception as exc:
err = str(exc)
done += 1
@@ -419,34 +446,9 @@ class BatchExportWorker(CancellableWorker):
self.finished.emit()
class AngleAlignmentWorker(QObject):
"""Computes rotation+translation alignment for every angle in *sras*,
referenced to *ref_angle_idx*, from each angle's binarized CH4 mask.
Rotation is analytic (from sras.angles_deg); only translation is found by
phase correlation.
"""
progress = pyqtSignal(int) # 0–100
finished = pyqtSignal(object, str) # AlignmentResult|None, error ("" = success)
def __init__(self, sras: SrasFile, ref_angle_idx: int, dc_threshold_mv: float):
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._threshold = dc_threshold_mv
def run(self):
try:
result = compute_angle_alignment(
self._sras, self._ref, self._threshold,
progress_cb=self.progress.emit)
self.finished.emit(result, "")
except Exception as exc:
self.finished.emit(None, str(exc))
class Ch4MaskWorker(QObject):
"""Fetches each requested angle's CH4 (Bias B) DC image in mV, for
ManualAlignmentDialog's initial threshold-mask overlay.
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
@@ -455,86 +457,108 @@ class Ch4MaskWorker(QObject):
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).
CorrelatePage._start_mask_prep).
"""
angle_done = pyqtSignal(int, np.ndarray) # angle_idx, dc4_mv
finished = pyqtSignal()
error = pyqtSignal(str)
def __init__(self, sras: SrasFile, angle_indices: list[int]):
super().__init__()
self._sras = sras
self._angles = angle_indices
self._budget = 0
def run(self):
try:
n_workers, budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=n_workers)
try:
futures = {
pool.submit(dc_image_mv, self._sras, a, CH4_IDX,
max_workers=1, budget=budget): a
for a in self._angles
}
for fut in as_completed(futures):
a = futures[fut]
self.angle_done.emit(a, fut.result())
finally:
pool.shutdown(wait=True)
self.finished.emit()
except Exception as exc:
self.error.emit(str(exc))
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(QObject):
"""FFT phase-correlation translation for each of *angle_indices* against
*ref_angle_idx*, for ManualAlignmentDialog's Auto Cross-Correlate button.
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 — a real many-angle, high-resolution scan's
correlation (even at its downsampled working resolution) can take long
enough that doing all of them on the GUI thread would visibly freeze the
dialog. Rotation is set to the same analytic scan-angle delta Auto
De-rotate uses alongside the correlated shift, since a translation
search is only meaningful once both angles' content is already oriented
the same way. dc4_mv/pivot_mm are the dialog's own already-in-memory
per-angle images/pivots — this worker does no fetching of its own.
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_done = pyqtSignal(int, float, float, float) # angle_idx, rotation_deg, shift_x_mm, shift_y_mm
finished = pyqtSignal()
error = pyqtSignal(str)
# 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], pivot_mm: dict[int, tuple[float, float]],
*, use_mask: bool, dc_threshold_mv: float, margin_frac: float):
dc4_mv: dict[int, np.ndarray], *, reg_kwargs: dict | None = None):
"""*reg_kwargs* is splatted into register_angle_to_reference — every
registration setting the wizard exposes (sources, threshold, search
width, seed, signs, refine, grid sizes) travels in it, so this class
holds no opinion about which knobs exist and exposing another needs no
change here."""
super().__init__()
self._sras = sras
self._ref = ref_angle_idx
self._angles = angle_indices
self._dc4_mv = dc4_mv
self._pivot_mm = pivot_mm
self._use_mask = use_mask
self._threshold = dc_threshold_mv
self._margin = margin_frac
self._reg_kwargs = dict(reg_kwargs or {})
def _one(self, a: int) -> tuple[int, float, float, float]:
theta = compute._theta_deg(self._sras, a, self._ref)
dx, dy = compute.correlate_translation_mm(
self._sras, a, self._ref, self._dc4_mv, self._pivot_mm,
use_mask=self._use_mask, dc_threshold_mv=self._threshold,
margin_frac=self._margin)
return a, theta, dx, dy
def _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, **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:
n_workers, _budget = compute.plan_angle_level(self._sras)
pool = ThreadPoolExecutor(max_workers=max(1, n_workers))
try:
futures = [pool.submit(self._one, a) for a in self._angles]
for fut in as_completed(futures):
a, theta, dx, dy = fut.result()
self.angle_done.emit(a, theta, dx, dy)
finally:
pool.shutdown(wait=True)
self.finished.emit()
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.error.emit(str(exc))
self.finished.emit("", str(exc))