sras-viewer/sras_workers.py

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