#!/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)`` (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): 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) # 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))