ptychozoon integration

pyproject.toml

@@ -75,6 +75,7 @@ module = [
     "ptycho.*",
     "ptycho_torch.*",
     "ptychonn.*",
+    "ptychozoon.*",
     "pvaccess",
     "pvapy.*",
     "scipy.*",

src/ptychodus/controller/core.py

@@ -129,6 +129,7 @@ def __init__(
             model.fluorescence_core.enhancer_chooser,
             model.fluorescence_core.two_step_enhancer,
             model.fluorescence_core.vspi_enhancer,
+            model.fluorescence_core.ptychozoon_enhancer,
             model.fluorescence_core.task_monitor,
             model.fluorescence_core.visualization_engine,
             view.probe_view,

src/ptychodus/controller/probe/core.py

@@ -13,6 +13,7 @@
 from ...model.fluorescence import (
     FluorescenceAPI,
     FluorescenceTaskMonitor,
+    PtychozoonFluorescenceEnhancer,
     TwoStepFluorescenceEnhancer,
     VSPIFluorescenceEnhancer,
 )
@@ -52,6 +53,7 @@ def __init__(
         fluorescence_enhancer_chooser: PluginChooser[FluorescenceEnhancer],
         fluorescence_two_step_enhancer: TwoStepFluorescenceEnhancer,
         fluorescence_vspi_enhancer: VSPIFluorescenceEnhancer,
+        fluorescence_ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None,
         fluorescence_task_monitor: FluorescenceTaskMonitor,
         fluorescence_visualization_engine: VisualizationEngine,
         view: RepositoryTreeView,
@@ -79,6 +81,7 @@ def __init__(
             fluorescence_enhancer_chooser,
             fluorescence_two_step_enhancer,
             fluorescence_vspi_enhancer,
+            fluorescence_ptychozoon_enhancer,
             fluorescence_task_monitor,
             fluorescence_visualization_engine,
             file_dialog_factory,

src/ptychodus/controller/probe/fluorescence.py

@@ -12,12 +12,14 @@
 from ...model.fluorescence import (
     FluorescenceAPI,
     FluorescenceTaskMonitor,
+    PtychozoonFluorescenceEnhancer,
     TwoStepFluorescenceEnhancer,
     VSPIFluorescenceEnhancer,
 )
 from ...model.visualization import VisualizationEngine
 from ...view.probe import (
     FluorescenceDialog,
+    FluorescencePtychozoonParametersView,
     FluorescenceStatusView,
     FluorescenceTwoStepParametersView,
     FluorescenceVSPIParametersView,
@@ -125,6 +127,69 @@ def _update(self, observable: Observable) -> None:
             self._sync_model_to_view()
 
 
+class FluorescencePtychozoonViewController(Observer):
+    MAX_INT: Final[int] = 0x7FFFFFFF
+
+    def __init__(self, enhancer: PtychozoonFluorescenceEnhancer) -> None:
+        super().__init__()
+        self._enhancer = enhancer
+        self._view = FluorescencePtychozoonParametersView()
+
+        self._view.damping_factor_line_edit.value_changed.connect(
+            self._sync_damping_factor_to_model
+        )
+        self._view.gradient_smoothness_line_edit.value_changed.connect(
+            self._sync_gradient_smoothness_to_model
+        )
+        self._view.max_iterations_spin_box.setRange(1, self.MAX_INT)
+        self._view.max_iterations_spin_box.valueChanged.connect(enhancer.set_max_iterations)
+        self._view.atol_line_edit.value_changed.connect(self._sync_atol_to_model)
+        self._view.btol_line_edit.value_changed.connect(self._sync_btol_to_model)
+        self._view.checkpoint_interval_spin_box.setRange(1, self.MAX_INT)
+        self._view.checkpoint_interval_spin_box.valueChanged.connect(
+            enhancer.set_checkpoint_interval
+        )
+        self._view.use_gpu_check_box.toggled.connect(enhancer.set_gpu_enabled)
+        self._view.gpu_device_index_spin_box.setRange(0, self.MAX_INT)
+        self._view.gpu_device_index_spin_box.valueChanged.connect(enhancer.set_gpu_device_index)
+
+        enhancer.add_observer(self)
+        self._sync_model_to_view()
+
+    def get_widget(self) -> QWidget:
+        return self._view
+
+    def _sync_damping_factor_to_model(self, value: Decimal) -> None:
+        self._enhancer.set_damping_factor(float(value))
+
+    def _sync_gradient_smoothness_to_model(self, value: Decimal) -> None:
+        self._enhancer.set_gradient_smoothness(float(value))
+
+    def _sync_atol_to_model(self, value: Decimal) -> None:
+        self._enhancer.set_atol(float(value))
+
+    def _sync_btol_to_model(self, value: Decimal) -> None:
+        self._enhancer.set_btol(float(value))
+
+    def _sync_model_to_view(self) -> None:
+        self._view.damping_factor_line_edit.set_value(
+            Decimal(repr(self._enhancer.get_damping_factor()))
+        )
+        self._view.gradient_smoothness_line_edit.set_value(
+            Decimal(repr(self._enhancer.get_gradient_smoothness()))
+        )
+        self._view.max_iterations_spin_box.setValue(self._enhancer.get_max_iterations())
+        self._view.atol_line_edit.set_value(Decimal(repr(self._enhancer.get_atol())))
+        self._view.btol_line_edit.set_value(Decimal(repr(self._enhancer.get_btol())))
+        self._view.checkpoint_interval_spin_box.setValue(self._enhancer.get_checkpoint_interval())
+        self._view.use_gpu_check_box.setChecked(self._enhancer.is_gpu_enabled())
+        self._view.gpu_device_index_spin_box.setValue(self._enhancer.get_gpu_device_index())
+
+    def _update(self, observable: Observable) -> None:
+        if observable is self._enhancer:
+            self._sync_model_to_view()
+
+
 class FluorescenceStatusController(Observer):
     def __init__(
         self,
@@ -170,6 +235,7 @@ def __init__(
         enhancer_chooser: PluginChooser,
         two_step_enhancer: TwoStepFluorescenceEnhancer,
         vspi_enhancer: VSPIFluorescenceEnhancer,
+        ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None,
         task_monitor: FluorescenceTaskMonitor,
         engine: VisualizationEngine,
         file_dialog_factory: FileDialogFactory,
@@ -214,6 +280,23 @@ def __init__(
             vspi_view_controller.get_widget()
         )
 
+        # Registered last, matching the enhancer_chooser order, so the combo-box
+        # index selects the correct stacked page. Only present when ptychozoon is
+        # installed (enhancer is None otherwise).
+        self._ptychozoon_view_controller: FluorescencePtychozoonViewController | None = None
+
+        if ptychozoon_enhancer is not None:
+            self._ptychozoon_view_controller = FluorescencePtychozoonViewController(
+                ptychozoon_enhancer
+            )
+            self._dialog.fluorescence_parameters_view.algorithm_combo_box.addItem(
+                PtychozoonFluorescenceEnhancer.DISPLAY_NAME,
+                self._dialog.fluorescence_parameters_view.algorithm_combo_box.count(),
+            )
+            self._dialog.fluorescence_parameters_view.stacked_widget.addWidget(
+                self._ptychozoon_view_controller.get_widget()
+            )
+
         self._dialog.fluorescence_parameters_view.algorithm_combo_box.textActivated.connect(
             enhancer_chooser.set_current_plugin
         )

src/ptychodus/model/fluorescence/__init__.py

@@ -5,6 +5,7 @@
     FluorescenceDatasetEmitter,
     FluorescenceTaskMonitor,
 )
+from .ptychozoon import PtychozoonFluorescenceEnhancer
 from .two_step import TwoStepFluorescenceEnhancer
 from .vspi import VSPIFluorescenceEnhancer
 
@@ -14,6 +15,7 @@
     'FluorescenceCore',
     'FluorescenceDatasetEmitter',
     'FluorescenceTaskMonitor',
+    'PtychozoonFluorescenceEnhancer',
     'TwoStepFluorescenceEnhancer',
     'VSPIFluorescenceEnhancer',
 ]

src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py

@@ -0,0 +1,156 @@
+"""Subprocess worker that runs the ptychozoon (GPU VSPI) fluorescence enhancement.
+
+This module is executed in a freshly ``spawn``ed process so that the CuPy GPU
+context is created cleanly and fully released when the process exits. It must not
+import ``ptychozoon`` (or CuPy) at module load time; the import happens inside the
+worker function so that the parent process never pulls GPU libraries into its own
+interpreter.
+
+All data crosses the process boundary as plain numpy arrays and scalars (see
+:class:`PtychozoonPayload`). Messages are streamed back over a single queue as
+tagged tuples:
+
+- ``('result', iteration, [(element_name, counts_per_second_array), ...])`` per checkpoint
+- ``('log', levelno, message)`` for each captured log record from the ptychozoon logger
+- ``('error', traceback_str)`` on failure
+
+followed by a ``None`` sentinel. Because a single child thread produces every
+message, ordering is preserved (a log line arrives before the result it precedes).
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from multiprocessing.queues import Queue
+from typing import Any
+
+import numpy
+
+__all__ = [
+    'PtychozoonPayload',
+    'run_vspi_enhancement',
+]
+
+
+class _QueueLogHandler(logging.Handler):
+    """Logging handler that forwards formatted records to the parent over the result queue."""
+
+    def __init__(self, result_queue: 'Queue[Any]') -> None:
+        super().__init__()
+        self._result_queue = result_queue
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            self._result_queue.put(('log', record.levelno, self.format(record)))
+        except Exception:
+            # Never let logging failures break the enhancement worker.
+            pass
+
+
+@dataclass(frozen=True)
+class PtychozoonPayload:
+    """Picklable inputs for one ptychozoon enhancement run."""
+
+    # Ptychography product (plain arrays / scalars)
+    probe_positions_m: numpy.ndarray  # (N, 2) float [y, x] meters
+    probe: numpy.ndarray  # (n_opr, modes, height, width) complex
+    object_array: numpy.ndarray  # (height, width) complex
+    pixel_size_m: tuple[float, float]  # (pixel_height_m, pixel_width_m)
+    object_center_m: tuple[float, float]  # (center_y_m, center_x_m)
+    opr_mode_weights: numpy.ndarray | None  # (n_opr, N) float or None
+
+    # Fluorescence element maps
+    element_maps: list[tuple[str, numpy.ndarray]]  # [(name, counts_per_second), ...]
+
+    # Solver settings
+    damping_factor: float
+    gradient_smoothness: float
+    max_iterations: int
+    atol: float
+    btol: float
+    checkpoint_interval: int
+    use_gpu: bool
+    gpu_device_index: int
+
+    # Effective log level of the parent's ptychozoon logger, so the child only
+    # forwards records that the parent would actually surface.
+    log_level: int
+
+
+def run_vspi_enhancement(payload: PtychozoonPayload, result_queue: 'Queue[Any]') -> None:
+    """Run ptychozoon VSPI enhancement and stream checkpoint results over a queue.
+
+    Intended as the target of a ``spawn``ed ``multiprocessing.Process``. Puts
+    ``('result', iteration, [(name, cps), ...])`` tuples per checkpoint and
+    ``('log', levelno, message)`` tuples for ptychozoon log records, then ``None``
+    when complete, or ``('error', traceback_str)`` on failure.
+    """
+    # Forward ptychozoon's own log output to the parent so it appears in the
+    # Enhance Fluorescence status view (this process has no ptychodus handlers).
+    log_handler = _QueueLogHandler(result_queue)
+    log_handler.setFormatter(logging.Formatter('%(message)s'))
+    ptychozoon_logger = logging.getLogger('ptychozoon')
+    ptychozoon_logger.addHandler(log_handler)
+    ptychozoon_logger.setLevel(payload.log_level)
+
+    try:
+        from ptychozoon.data_structures import (
+            ElementMap,
+            FluorescenceDataset,
+            PtychographyProduct,
+        )
+        from ptychozoon.settings import (
+            DeconvolutionEnhancementSettings,
+            InterpolationTypes,
+        )
+        from ptychozoon.vspi_enhance import VSPIFluorescenceEnhancingAlgorithm
+
+        product = PtychographyProduct(
+            probe_positions=payload.probe_positions_m,
+            probe=payload.probe,
+            object_array=payload.object_array,
+            pixel_size_m=payload.pixel_size_m,
+            object_center_m=payload.object_center_m,
+            opr_mode_weights=payload.opr_mode_weights,
+        )
+        dataset = FluorescenceDataset(
+            element_maps=[
+                ElementMap(name=name, counts_per_second=cps) for name, cps in payload.element_maps
+            ]
+        )
+
+        settings = DeconvolutionEnhancementSettings()
+        settings.lsmr.damping_factor = payload.damping_factor
+        settings.lsmr.gradient_smoothness = payload.gradient_smoothness
+        settings.lsmr.max_iter = payload.max_iterations
+        settings.lsmr.atol = payload.atol
+        settings.lsmr.btol = payload.btol
+        settings.lsmr.checkpoint_interval = payload.checkpoint_interval
+        settings.gpu.enabled = payload.use_gpu
+        settings.gpu.index = payload.gpu_device_index
+        # Fourier interpolation requires the GPU; fall back to Barycentric on CPU.
+        settings._interpolation = (
+            InterpolationTypes.FOURIER if payload.use_gpu else InterpolationTypes.BARYCENTRIC
+        )
+
+        algorithm = VSPIFluorescenceEnhancingAlgorithm()
+
+        for enhanced_dataset, iteration in algorithm.enhance(dataset, product, settings=settings):
+            result_queue.put(
+                (
+                    'result',
+                    int(iteration),
+                    [
+                        (emap.name, numpy.asarray(emap.counts_per_second))
+                        for emap in enhanced_dataset.element_maps
+                    ],
+                )
+            )
+    except Exception:
+        import traceback
+
+        result_queue.put(('error', traceback.format_exc()))
+    finally:
+        ptychozoon_logger.removeHandler(log_handler)
+        result_queue.put(None)

src/ptychodus/model/fluorescence/core.py

@@ -15,10 +15,13 @@
 from ..visualization import VisualizationEngine
 from .api import FluorescenceAPI
 from .monitor import FluorescenceTaskMonitor
+from .ptychozoon import PtychozoonFluorescenceEnhancer
 from .settings import FluorescenceSettings
 from .two_step import TwoStepFluorescenceEnhancer
 from .vspi import VSPIFluorescenceEnhancer
 
+logger = logging.getLogger(__name__)
+
 
 class FluorescenceCore:
     def __init__(
@@ -48,6 +51,26 @@ def __init__(
             simple_name=VSPIFluorescenceEnhancer.SIMPLE_NAME,
             display_name=VSPIFluorescenceEnhancer.DISPLAY_NAME,
         )
+
+        # The GPU VSPI enhancer is optional: register it only when ptychozoon is
+        # installed so it silently disappears otherwise. Register it last so its
+        # combo-box entry and stacked GUI page stay aligned by index. Importing
+        # the top-level package is cheap and does not pull in CuPy (that happens
+        # only inside the spawned subprocess), so probing it here is safe.
+        self.ptychozoon_enhancer: PtychozoonFluorescenceEnhancer | None = None
+
+        try:
+            import ptychozoon  # noqa: F401
+        except ModuleNotFoundError:
+            logger.info('ptychozoon not found.')
+        else:
+            self.ptychozoon_enhancer = PtychozoonFluorescenceEnhancer(self._settings)
+            self.enhancer_chooser.register_plugin(
+                self.ptychozoon_enhancer,
+                simple_name=PtychozoonFluorescenceEnhancer.SIMPLE_NAME,
+                display_name=PtychozoonFluorescenceEnhancer.DISPLAY_NAME,
+            )
+
         self.enhancer_chooser.synchronize_with_parameter(self._settings.algorithm)
 
         file_reader_chooser.synchronize_with_parameter(self._settings.file_type)