diff --git a/sras_viewer.py b/sras_viewer.py index e28fa73..c1c232f 100644 --- a/sras_viewer.py +++ b/sras_viewer.py @@ -1478,7 +1478,6 @@ class SrasViewerWindow(QMainWindow): # same combination is free. self._dc_cache: dict[tuple[int, int], np.ndarray] = {} self._fft_cache: dict[tuple[int, bool, int | None, float], np.ndarray] = {} - self._dc_precompute_worker: DcPrecomputeWorker | None = None self._dc_generation: int = 0 # Angle alignment ("Fusion" menu) @@ -2463,8 +2462,7 @@ class SrasViewerWindow(QMainWindow): n_angles = self._sras.n_angles worker = DcPrecomputeWorker(self._sras) - self._dc_precompute_worker = worker - started = self._run_worker( + self._run_worker( "dc_precompute", worker, connect=( ("angle_done", lambda a, dc3, dc4, g=generation: @@ -2473,10 +2471,7 @@ class SrasViewerWindow(QMainWindow): f"DC precompute error: {msg}", 5000)), ), quit_on=("finished", "error"), - on_done=lambda: setattr(self, "_dc_precompute_worker", None), ) - if not started: - self._dc_precompute_worker = None def _on_dc_precompute_angle_done(self, generation: int, angle_idx: int, dc3_mv: np.ndarray, dc4_mv: np.ndarray, diff --git a/sras_workers.py b/sras_workers.py index 334e655..5919b79 100644 --- a/sras_workers.py +++ b/sras_workers.py @@ -54,6 +54,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) @@ -117,29 +145,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). @@ -149,26 +177,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): @@ -306,7 +316,7 @@ class AngleAlignmentWorker(QObject): self.finished.emit(None, str(exc)) -class Ch4MaskWorker(QObject): +class Ch4MaskWorker(_PooledWorker): """Fetches each requested angle's CH4 (Bias B) DC image in mV, for ManualAlignmentDialog's initial threshold-mask overlay. @@ -320,35 +330,29 @@ class Ch4MaskWorker(QObject): ManualAlignmentDialog._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): +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. @@ -363,8 +367,6 @@ class CrossCorrelateWorker(QObject): """ # angle_idx, rotation_deg, shift_x_mm, shift_y_mm, score, source angle_done = pyqtSignal(int, float, float, float, float, str) - finished = pyqtSignal() - error = pyqtSignal(str) def __init__(self, sras: SrasFile, ref_angle_idx: int, angle_indices: list[int], dc4_mv: dict[int, np.ndarray], *, @@ -379,25 +381,19 @@ class CrossCorrelateWorker(QObject): self._threshold = dc_threshold_mv self._search_deg = search_deg + def _plan(self) -> int: + return compute._registration_workers(self._sras, compute._DEFAULT_FINE_DIM) + + 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) - def run(self): - try: - n_workers = compute._registration_workers( - self._sras, compute._DEFAULT_FINE_DIM) - 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, fit = fut.result() - self.angle_done.emit(a, fit.rotation_deg, fit.shift_mm[0], - fit.shift_mm[1], fit.score, fit.source) - finally: - pool.shutdown(wait=True) - self.finished.emit() - except Exception as exc: - self.error.emit(str(exc)) + 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)