From 23f8ec6fb3471efae201f492f73b629f0f7778e2 Mon Sep 17 00:00:00 2001 From: Steve Henke Date: Fri, 10 Jul 2026 10:10:28 -0500 Subject: [PATCH] ptychozoon integration --- pyproject.toml | 1 + src/ptychodus/controller/core.py | 1 + src/ptychodus/controller/probe/core.py | 3 + .../controller/probe/fluorescence.py | 83 +++++++ src/ptychodus/model/fluorescence/__init__.py | 2 + .../fluorescence/_ptychozoon_subprocess.py | 156 +++++++++++++ src/ptychodus/model/fluorescence/core.py | 23 ++ .../model/fluorescence/ptychozoon.py | 208 ++++++++++++++++++ src/ptychodus/model/fluorescence/settings.py | 24 ++ src/ptychodus/view/probe.py | 26 +++ 10 files changed, 527 insertions(+) create mode 100644 src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py create mode 100644 src/ptychodus/model/fluorescence/ptychozoon.py diff --git a/pyproject.toml b/pyproject.toml index b101b201..568af0be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ module = [ "ptycho.*", "ptycho_torch.*", "ptychonn.*", + "ptychozoon.*", "pvaccess", "pvapy.*", "scipy.*", diff --git a/src/ptychodus/controller/core.py b/src/ptychodus/controller/core.py index f395ca26..6e0d23ee 100644 --- a/src/ptychodus/controller/core.py +++ b/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, diff --git a/src/ptychodus/controller/probe/core.py b/src/ptychodus/controller/probe/core.py index bc36788a..12c1db02 100644 --- a/src/ptychodus/controller/probe/core.py +++ b/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, diff --git a/src/ptychodus/controller/probe/fluorescence.py b/src/ptychodus/controller/probe/fluorescence.py index febccadf..0988fd11 100644 --- a/src/ptychodus/controller/probe/fluorescence.py +++ b/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 ) diff --git a/src/ptychodus/model/fluorescence/__init__.py b/src/ptychodus/model/fluorescence/__init__.py index 7ea9e693..4ab9acfd 100644 --- a/src/ptychodus/model/fluorescence/__init__.py +++ b/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', ] diff --git a/src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py b/src/ptychodus/model/fluorescence/_ptychozoon_subprocess.py new file mode 100644 index 00000000..e045a77b --- /dev/null +++ b/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) diff --git a/src/ptychodus/model/fluorescence/core.py b/src/ptychodus/model/fluorescence/core.py index 425f9652..72a049b0 100644 --- a/src/ptychodus/model/fluorescence/core.py +++ b/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) diff --git a/src/ptychodus/model/fluorescence/ptychozoon.py b/src/ptychodus/model/fluorescence/ptychozoon.py new file mode 100644 index 00000000..9907305c --- /dev/null +++ b/src/ptychodus/model/fluorescence/ptychozoon.py @@ -0,0 +1,208 @@ +"""GPU-accelerated VSPI fluorescence enhancer backed by the ptychozoon package. + +The heavy CuPy computation runs in a freshly ``spawn``ed subprocess (see +:mod:`._ptychozoon_subprocess`) so each run gets a clean GPU context and all GPU +memory is released when the run finishes or is stopped. This module never imports +ptychozoon or CuPy directly. +""" + +from __future__ import annotations +from collections.abc import Iterator +from typing import Final +import logging +import multiprocessing + +import numpy + +from ptychodus.api.fluorescence import ( + ElementMap, + FluorescenceDataset, + FluorescenceEnhancer, + FluorescenceEnhancerInput, + FluorescenceEnhancerOutput, +) +from ptychodus.api.observer import Observable, Observer +from ptychodus.api.product import Product + +from ._ptychozoon_subprocess import PtychozoonPayload, run_vspi_enhancement +from .settings import FluorescenceSettings + +logger = logging.getLogger(__name__) + +__all__ = [ + 'PtychozoonFluorescenceEnhancer', +] + + +class PtychozoonFluorescenceEnhancer(FluorescenceEnhancer, Observable, Observer): + SIMPLE_NAME: Final[str] = 'VSPI-GPU' + DISPLAY_NAME: Final[str] = 'Virtual Single Pixel Imaging (GPU)' + + def __init__(self, settings: FluorescenceSettings) -> None: + super().__init__() + self._settings = settings + + settings.ptychozoon_damping_factor.add_observer(self) + settings.ptychozoon_gradient_smoothness.add_observer(self) + settings.ptychozoon_max_iterations.add_observer(self) + settings.ptychozoon_atol.add_observer(self) + settings.ptychozoon_btol.add_observer(self) + settings.ptychozoon_checkpoint_interval.add_observer(self) + settings.ptychozoon_use_gpu.add_observer(self) + settings.ptychozoon_gpu_device_index.add_observer(self) + + @property + def name(self) -> str: + return self.DISPLAY_NAME + + def get_progress_goal(self) -> int: + return self._settings.ptychozoon_max_iterations.get_value() + + def _build_payload(self, parameters: FluorescenceEnhancerInput) -> PtychozoonPayload: + product: Product = parameters.product + dataset = parameters.dataset + object_geometry = product.object_.get_geometry() + + probe_positions_m = numpy.array( + [[p.coordinate_y_m, p.coordinate_x_m] for p in product.probe_positions], + dtype=float, + ).reshape((-1, 2)) + + opr_weights = product.probes.get_opr_weights_or_none() + opr_mode_weights = None if opr_weights is None else numpy.ascontiguousarray(opr_weights.T) + + return PtychozoonPayload( + probe_positions_m=probe_positions_m, + probe=product.probes.get_array(), + object_array=product.object_.get_layers_flattened(), + pixel_size_m=(object_geometry.pixel_height_m, object_geometry.pixel_width_m), + object_center_m=(object_geometry.center_y_m, object_geometry.center_x_m), + opr_mode_weights=opr_mode_weights, + element_maps=[ + (emap.name, numpy.asarray(emap.counts_per_second)) for emap in dataset.element_maps + ], + damping_factor=self._settings.ptychozoon_damping_factor.get_value(), + gradient_smoothness=self._settings.ptychozoon_gradient_smoothness.get_value(), + max_iterations=self._settings.ptychozoon_max_iterations.get_value(), + atol=self._settings.ptychozoon_atol.get_value(), + btol=self._settings.ptychozoon_btol.get_value(), + checkpoint_interval=self._settings.ptychozoon_checkpoint_interval.get_value(), + use_gpu=self._settings.ptychozoon_use_gpu.get_value(), + gpu_device_index=self._settings.ptychozoon_gpu_device_index.get_value(), + log_level=logger.getEffectiveLevel(), + ) + + def enhance( + self, parameters: FluorescenceEnhancerInput + ) -> Iterator[FluorescenceEnhancerOutput]: + dataset = parameters.dataset + payload = self._build_payload(parameters) + + # A fresh spawned process gives ptychozoon/CuPy a clean GPU context and + # releases all GPU memory when it exits. + ctx = multiprocessing.get_context('spawn') + result_queue = ctx.Queue() + process = ctx.Process(target=run_vspi_enhancement, args=(payload, result_queue)) + process.start() + + try: + while True: + item = result_queue.get() + + if item is None: + break + + tag = item[0] + + if tag == 'log': + _, levelno, message = item + # Re-emit through this process's logger so it reaches the + # fluorescence status view via the registered handler. + logger.log(levelno, message) + continue + + if tag == 'error': + raise RuntimeError(f'ptychozoon enhancement failed:\n{item[1]}') + + if tag == 'result': + _, iteration, enhanced_maps = item + element_maps = [ElementMap(name, cps) for name, cps in enhanced_maps] + yield FluorescenceEnhancerOutput( + dataset=FluorescenceDataset( + element_maps=element_maps, + counts_per_second_path=dataset.counts_per_second_path, + channel_names_path=dataset.channel_names_path, + ), + progress=iteration, + ) + continue + finally: + if process.is_alive(): + process.terminate() + + process.join(timeout=10.0) + + if process.is_alive(): + process.kill() + process.join() + + def get_damping_factor(self) -> float: + return self._settings.ptychozoon_damping_factor.get_value() + + def set_damping_factor(self, factor: float) -> None: + self._settings.ptychozoon_damping_factor.set_value(factor) + + def get_gradient_smoothness(self) -> float: + return self._settings.ptychozoon_gradient_smoothness.get_value() + + def set_gradient_smoothness(self, value: float) -> None: + self._settings.ptychozoon_gradient_smoothness.set_value(value) + + def get_max_iterations(self) -> int: + return self._settings.ptychozoon_max_iterations.get_value() + + def set_max_iterations(self, number: int) -> None: + self._settings.ptychozoon_max_iterations.set_value(number) + + def get_atol(self) -> float: + return self._settings.ptychozoon_atol.get_value() + + def set_atol(self, value: float) -> None: + self._settings.ptychozoon_atol.set_value(value) + + def get_btol(self) -> float: + return self._settings.ptychozoon_btol.get_value() + + def set_btol(self, value: float) -> None: + self._settings.ptychozoon_btol.set_value(value) + + def get_checkpoint_interval(self) -> int: + return self._settings.ptychozoon_checkpoint_interval.get_value() + + def set_checkpoint_interval(self, number: int) -> None: + self._settings.ptychozoon_checkpoint_interval.set_value(number) + + def is_gpu_enabled(self) -> bool: + return self._settings.ptychozoon_use_gpu.get_value() + + def set_gpu_enabled(self, enabled: bool) -> None: + self._settings.ptychozoon_use_gpu.set_value(enabled) + + def get_gpu_device_index(self) -> int: + return self._settings.ptychozoon_gpu_device_index.get_value() + + def set_gpu_device_index(self, index: int) -> None: + self._settings.ptychozoon_gpu_device_index.set_value(index) + + def _update(self, observable: Observable) -> None: + if observable in ( + self._settings.ptychozoon_damping_factor, + self._settings.ptychozoon_gradient_smoothness, + self._settings.ptychozoon_max_iterations, + self._settings.ptychozoon_atol, + self._settings.ptychozoon_btol, + self._settings.ptychozoon_checkpoint_interval, + self._settings.ptychozoon_use_gpu, + self._settings.ptychozoon_gpu_device_index, + ): + self.notify_observers() diff --git a/src/ptychodus/model/fluorescence/settings.py b/src/ptychodus/model/fluorescence/settings.py index ae285a1b..b67249d9 100644 --- a/src/ptychodus/model/fluorescence/settings.py +++ b/src/ptychodus/model/fluorescence/settings.py @@ -24,6 +24,30 @@ def __init__(self, registry: SettingsRegistry) -> None: 'DeconvolutionStrategy', 'Richardson-Lucy' ) + # Ptychozoon (GPU VSPI) settings + self.ptychozoon_damping_factor = self._group.create_real_parameter( + 'PtychozoonDampingFactor', 0.0, minimum=0.0 + ) + self.ptychozoon_gradient_smoothness = self._group.create_real_parameter( + 'PtychozoonGradientSmoothness', 0.0, minimum=0.0 + ) + self.ptychozoon_max_iterations = self._group.create_integer_parameter( + 'PtychozoonMaxIterations', 100, minimum=1 + ) + self.ptychozoon_atol = self._group.create_real_parameter( + 'PtychozoonATol', 1e-6, minimum=0.0 + ) + self.ptychozoon_btol = self._group.create_real_parameter( + 'PtychozoonBTol', 1e-6, minimum=0.0 + ) + self.ptychozoon_checkpoint_interval = self._group.create_integer_parameter( + 'PtychozoonCheckpointInterval', 5, minimum=1 + ) + self.ptychozoon_use_gpu = self._group.create_boolean_parameter('PtychozoonUseGPU', True) + self.ptychozoon_gpu_device_index = self._group.create_integer_parameter( + 'PtychozoonGPUDeviceIndex', 0, minimum=0 + ) + def _update(self, observable: Observable) -> None: if observable is self._group: self.notify_observers() diff --git a/src/ptychodus/view/probe.py b/src/ptychodus/view/probe.py index 06f5222f..cf244150 100644 --- a/src/ptychodus/view/probe.py +++ b/src/ptychodus/view/probe.py @@ -1,5 +1,6 @@ from PyQt5.QtCore import Qt from PyQt5.QtWidgets import ( + QCheckBox, QComboBox, QDialog, QFormLayout, @@ -220,6 +221,31 @@ def __init__(self, parent: QWidget | None = None) -> None: self.setLayout(layout) +class FluorescencePtychozoonParametersView(QWidget): + def __init__(self, parent: QWidget | None = None) -> None: + super().__init__(parent) + self.damping_factor_line_edit = DecimalLineEdit.create_instance() + self.gradient_smoothness_line_edit = DecimalLineEdit.create_instance() + self.max_iterations_spin_box = QSpinBox() + self.atol_line_edit = DecimalLineEdit.create_instance() + self.btol_line_edit = DecimalLineEdit.create_instance() + self.checkpoint_interval_spin_box = QSpinBox() + self.use_gpu_check_box = QCheckBox('Use GPU') + self.gpu_device_index_spin_box = QSpinBox() + + layout = QFormLayout() + layout.setContentsMargins(0, 0, 0, 0) + layout.addRow('Damping Factor:', self.damping_factor_line_edit) + layout.addRow('Gradient Smoothness:', self.gradient_smoothness_line_edit) + layout.addRow('Max Iterations:', self.max_iterations_spin_box) + layout.addRow('A Tolerance:', self.atol_line_edit) + layout.addRow('B Tolerance:', self.btol_line_edit) + layout.addRow('Checkpoint Interval:', self.checkpoint_interval_spin_box) + layout.addRow(self.use_gpu_check_box) + layout.addRow('CUDA Device Index:', self.gpu_device_index_spin_box) + self.setLayout(layout) + + class FluorescenceTwoStepParametersView(QWidget): def __init__(self, parent: QWidget | None = None) -> None: super().__init__(parent)