add algorithm-specific options to model

AdvancedPhotonSource · Dec 12, 2024 · 36d66a9 · 36d66a9
1 parent c643306
commit 36d66a9
Show file tree

Hide file tree

Showing 9 changed files with 156 additions and 6 deletions.
diff --git a/src/ptychodus/controller/ptychi/reconstructor.py b/src/ptychodus/controller/ptychi/reconstructor.py
@@ -14,6 +14,7 @@
 
 from ...model.ptychi import PtyChiDeviceRepository, PtyChiEnumerators, PtyChiReconstructorSettings
 from ..parametric import (
+    CheckBoxParameterViewController,
     CheckableGroupBoxParameterViewController,
     ComboBoxParameterViewController,
     ParameterViewController,
@@ -114,6 +115,11 @@ def __init__(
             settings.useDoublePrecision,
             tool_tip='Floating point precision to use for computation.',
         )
+        self._useLowMemoryViewController = CheckBoxParameterViewController(
+            settings.useLowMemoryForwardModel,
+            'Use Low Memory Forward Model',
+            tool_tip='When checked, forward propagation of ptychography will be done using less vectorized code. This reduces the speed, but also lowers memory usage.',
+        )
         self._widget = QGroupBox('Reconstructor')
 
         layout = QFormLayout()
@@ -122,6 +128,7 @@ def __init__(
         layout.addRow('Batch Mode:', self._batchingModeViewController.getWidget())
         layout.addRow('Batch Stride:', self._batchStride.getWidget())
         layout.addRow('Precision:', self._precisionViewController.getWidget())
+        layout.addRow(self._useLowMemoryViewController.getWidget())
 
         if repository:
             layout.addRow(self._deviceViewController.getWidget())

diff --git a/src/ptychodus/model/ptychi/autodiff.py b/src/ptychodus/model/ptychi/autodiff.py
@@ -106,7 +106,7 @@ def _create_reconstructor_options(self) -> AutodiffPtychographyReconstructorOpti
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:

diff --git a/src/ptychodus/model/ptychi/core.py b/src/ptychodus/model/ptychi/core.py
@@ -12,8 +12,11 @@
 from .device import PtyChiDeviceRepository
 from .enums import PtyChiEnumerators
 from .settings import (
+    PtyChiDMSettings,
+    PtyChiLSQMLSettings,
     PtyChiOPRSettings,
     PtyChiObjectSettings,
+    PtyChiPIESettings,
     PtyChiProbePositionSettings,
     PtyChiProbeSettings,
     PtyChiReconstructorSettings,
@@ -32,6 +35,9 @@ def __init__(
         self.probeSettings = PtyChiProbeSettings(settingsRegistry)
         self.probePositionSettings = PtyChiProbePositionSettings(settingsRegistry)
         self.oprSettings = PtyChiOPRSettings(settingsRegistry)
+        self.dmSettings = PtyChiDMSettings(settingsRegistry)
+        self.pieSettings = PtyChiPIESettings(settingsRegistry)
+        self.lsqmlSettings = PtyChiLSQMLSettings(settingsRegistry)
         self.enumerators = PtyChiEnumerators()
         self.deviceRepository = PtyChiDeviceRepository(
             isDeveloperModeEnabled=isDeveloperModeEnabled
@@ -59,6 +65,7 @@ def __init__(
                     self.probeSettings,
                     self.probePositionSettings,
                     self.oprSettings,
+                    self.dmSettings,
                     detector,
                 )
             )
@@ -69,6 +76,7 @@ def __init__(
                     self.probeSettings,
                     self.probePositionSettings,
                     self.oprSettings,
+                    self.pieSettings,
                     detector,
                 )
             )
@@ -79,6 +87,7 @@ def __init__(
                     self.probeSettings,
                     self.probePositionSettings,
                     self.oprSettings,
+                    self.pieSettings,
                     detector,
                 )
             )
@@ -89,6 +98,7 @@ def __init__(
                     self.probeSettings,
                     self.probePositionSettings,
                     self.oprSettings,
+                    self.pieSettings,
                     detector,
                 )
             )
@@ -99,6 +109,7 @@ def __init__(
                     self.probeSettings,
                     self.probePositionSettings,
                     self.oprSettings,
+                    self.lsqmlSettings,
                     detector,
                 )
             )

diff --git a/src/ptychodus/model/ptychi/dm.py b/src/ptychodus/model/ptychi/dm.py
@@ -35,6 +35,7 @@
 
 from ..patterns import Detector
 from .settings import (
+    PtyChiDMSettings,
     PtyChiOPRSettings,
     PtyChiObjectSettings,
     PtyChiProbePositionSettings,
@@ -53,6 +54,7 @@ def __init__(
         probeSettings: PtyChiProbeSettings,
         probePositionSettings: PtyChiProbePositionSettings,
         oprSettings: PtyChiOPRSettings,
+        dmSettings: PtyChiDMSettings,
         detector: Detector,
     ) -> None:
         super().__init__()
@@ -61,6 +63,7 @@ def __init__(
         self._probeSettings = probeSettings
         self._probePositionSettings = probePositionSettings
         self._oprSettings = oprSettings
+        self._dmSettings = dmSettings
         self._detector = detector
 
     @property
@@ -106,7 +109,9 @@ def _create_reconstructor_options(self) -> DMReconstructorOptions:
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
+            exit_wave_update_relaxation=self._dmSettings.exitWaveUpdateRelaxation.getValue(),
+            chunk_length=self._dmSettings.chunkLength.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:
@@ -255,6 +260,7 @@ def _create_object_options(self, object_: Object) -> DMObjectOptions:
             multislice_regularization_unwrap_image_integration_method=multislice_regularization_unwrap_image_integration_method,
             multislice_regularization_stride=self._objectSettings.regularizeMultisliceStride.getValue(),
             patch_interpolation_method=patch_interpolation_method,
+            # FIXME amplitude_clamp_limit=self._dmSettings.objectAmplitudeClampLimit.getValue()
         )
 
     def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> DMProbeOptions:

diff --git a/src/ptychodus/model/ptychi/epie.py b/src/ptychodus/model/ptychi/epie.py
@@ -37,6 +37,7 @@
 from .settings import (
     PtyChiOPRSettings,
     PtyChiObjectSettings,
+    PtyChiPIESettings,
     PtyChiProbePositionSettings,
     PtyChiProbeSettings,
     PtyChiReconstructorSettings,
@@ -53,6 +54,7 @@ def __init__(
         probeSettings: PtyChiProbeSettings,
         probePositionSettings: PtyChiProbePositionSettings,
         oprSettings: PtyChiOPRSettings,
+        pieSettings: PtyChiPIESettings,
         detector: Detector,
     ) -> None:
         super().__init__()
@@ -61,6 +63,7 @@ def __init__(
         self._probeSettings = probeSettings
         self._probePositionSettings = probePositionSettings
         self._oprSettings = oprSettings
+        self._pieSettings = pieSettings
         self._detector = detector
 
     @property
@@ -106,7 +109,7 @@ def _create_reconstructor_options(self) -> EPIEReconstructorOptions:
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:
@@ -255,6 +258,7 @@ def _create_object_options(self, object_: Object) -> PIEObjectOptions:
             multislice_regularization_unwrap_image_integration_method=multislice_regularization_unwrap_image_integration_method,
             multislice_regularization_stride=self._objectSettings.regularizeMultisliceStride.getValue(),
             patch_interpolation_method=patch_interpolation_method,
+            alpha=self._pieSettings.objectAlpha.getValue(),
         )
 
     def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEProbeOptions:
@@ -299,6 +303,7 @@ def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEP
             support_constraint=self._probeSettings.constrainSupport.getValue(),
             support_constraint_threshold=self._probeSettings.constrainSupportThreshold.getValue(),
             support_constraint_stride=self._probeSettings.constrainSupportStride.getValue(),
+            alpha=self._pieSettings.probeAlpha.getValue(),
         )
 
     def _create_probe_position_options(

diff --git a/src/ptychodus/model/ptychi/lsqml.py b/src/ptychodus/model/ptychi/lsqml.py
@@ -16,6 +16,7 @@
     LSQMLProbeOptions,
     LSQMLProbePositionOptions,
     LSQMLReconstructorOptions,
+    NoiseModels,
     OptimizationPlan,
     Optimizers,
     OrthogonalizationMethods,
@@ -35,6 +36,7 @@
 
 from ..patterns import Detector
 from .settings import (
+    PtyChiLSQMLSettings,
     PtyChiOPRSettings,
     PtyChiObjectSettings,
     PtyChiProbePositionSettings,
@@ -53,6 +55,7 @@ def __init__(
         probeSettings: PtyChiProbeSettings,
         probePositionSettings: PtyChiProbePositionSettings,
         oprSettings: PtyChiOPRSettings,
+        lsqmlSettings: PtyChiLSQMLSettings,
         detector: Detector,
     ) -> None:
         super().__init__()
@@ -61,6 +64,7 @@ def __init__(
         self._probeSettings = probeSettings
         self._probePositionSettings = probePositionSettings
         self._oprSettings = oprSettings
+        self._lsqmlSettings = lsqmlSettings
         self._detector = detector
 
     @property
@@ -90,6 +94,18 @@ def _create_reconstructor_options(self) -> LSQMLReconstructorOptions:
             logger.warning('Failed to parse batching mode "{batching_mode_str}"!')
             batching_mode = BatchingModes.RANDOM
 
+        ####
+
+        noise_model_str = self._lsqmlSettings.noiseModel.getValue()
+
+        try:
+            noise_model = NoiseModels[noise_model_str.upper()]
+        except KeyError:
+            logger.warning('Failed to parse batching mode "{noise_model_str}"!')
+            noise_model = NoiseModels.GAUSSIAN
+
+        ####
+
         return LSQMLReconstructorOptions(
             num_epochs=self._reconstructorSettings.numEpochs.getValue(),
             batch_size=self._reconstructorSettings.batchSize.getValue(),
@@ -106,7 +122,12 @@ def _create_reconstructor_options(self) -> LSQMLReconstructorOptions:
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
+            noise_model=noise_model,
+            gaussian_noise_std=self._lsqmlSettings.gaussianNoiseDeviation.getValue(),
+            solve_obj_prb_step_size_jointly_for_first_slice_in_multislice=self._lsqmlSettings.solveObjectProbeStepSizeJointlyForFirstSliceInMultislice.getValue(),
+            solve_step_sizes_only_using_first_probe_mode=self._lsqmlSettings.solveStepSizesOnlyUsingFirstProbeMode.getValue(),
+            momentum_acceleration_gain=self._lsqmlSettings.momentumAccelerationGain.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:
@@ -255,6 +276,8 @@ def _create_object_options(self, object_: Object) -> LSQMLObjectOptions:
             multislice_regularization_unwrap_image_integration_method=multislice_regularization_unwrap_image_integration_method,
             multislice_regularization_stride=self._objectSettings.regularizeMultisliceStride.getValue(),
             patch_interpolation_method=patch_interpolation_method,
+            optimal_step_size_scaler=self._lsqmlSettings.objectOptimalStepSizeScaler.getValue(),
+            multimodal_update=self._lsqmlSettings.objectMultimodalUpdate.getValue(),
         )
 
     def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> LSQMLProbeOptions:
@@ -299,6 +322,8 @@ def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> LSQM
             support_constraint=self._probeSettings.constrainSupport.getValue(),
             support_constraint_threshold=self._probeSettings.constrainSupportThreshold.getValue(),
             support_constraint_stride=self._probeSettings.constrainSupportStride.getValue(),
+            eigenmode_update_relaxation=self._lsqmlSettings.probeEigenmodeUpdateRelaxation.getValue(),
+            optimal_step_size_scaler=self._lsqmlSettings.probeOptimalStepSizeScaler.getValue(),
         )
 
     def _create_probe_position_options(
@@ -371,6 +396,7 @@ def _create_opr_mode_weight_options(self) -> LSQMLOPRModeWeightsOptions:
             initial_weights=numpy.array([0.0]),  # FIXME
             optimize_eigenmode_weights=self._oprSettings.optimizeEigenmodeWeights.getValue(),
             optimize_intensity_variation=self._oprSettings.optimizeIntensities.getValue(),
+            update_relaxation=self._lsqmlSettings.oprModesUpdateRelaxation.getValue(),
         )
 
     def _create_task_options(self, parameters: ReconstructInput) -> LSQMLOptions:

diff --git a/src/ptychodus/model/ptychi/pie.py b/src/ptychodus/model/ptychi/pie.py
@@ -37,6 +37,7 @@
 from .settings import (
     PtyChiOPRSettings,
     PtyChiObjectSettings,
+    PtyChiPIESettings,
     PtyChiProbePositionSettings,
     PtyChiProbeSettings,
     PtyChiReconstructorSettings,
@@ -53,6 +54,7 @@ def __init__(
         probeSettings: PtyChiProbeSettings,
         probePositionSettings: PtyChiProbePositionSettings,
         oprSettings: PtyChiOPRSettings,
+        pieSettings: PtyChiPIESettings,
         detector: Detector,
     ) -> None:
         super().__init__()
@@ -61,6 +63,7 @@ def __init__(
         self._probeSettings = probeSettings
         self._probePositionSettings = probePositionSettings
         self._oprSettings = oprSettings
+        self._pieSettings = pieSettings
         self._detector = detector
 
     @property
@@ -106,7 +109,7 @@ def _create_reconstructor_options(self) -> PIEReconstructorOptions:
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:
@@ -255,6 +258,7 @@ def _create_object_options(self, object_: Object) -> PIEObjectOptions:
             multislice_regularization_unwrap_image_integration_method=multislice_regularization_unwrap_image_integration_method,
             multislice_regularization_stride=self._objectSettings.regularizeMultisliceStride.getValue(),
             patch_interpolation_method=patch_interpolation_method,
+            alpha=self._pieSettings.objectAlpha.getValue(),
         )
 
     def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEProbeOptions:
@@ -299,6 +303,7 @@ def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEP
             support_constraint=self._probeSettings.constrainSupport.getValue(),
             support_constraint_threshold=self._probeSettings.constrainSupportThreshold.getValue(),
             support_constraint_stride=self._probeSettings.constrainSupportStride.getValue(),
+            alpha=self._pieSettings.probeAlpha.getValue(),
         )
 
     def _create_probe_position_options(

diff --git a/src/ptychodus/model/ptychi/rpie.py b/src/ptychodus/model/ptychi/rpie.py
@@ -37,6 +37,7 @@
 from .settings import (
     PtyChiOPRSettings,
     PtyChiObjectSettings,
+    PtyChiPIESettings,
     PtyChiProbePositionSettings,
     PtyChiProbeSettings,
     PtyChiReconstructorSettings,
@@ -53,6 +54,7 @@ def __init__(
         probeSettings: PtyChiProbeSettings,
         probePositionSettings: PtyChiProbePositionSettings,
         oprSettings: PtyChiOPRSettings,
+        pieSettings: PtyChiPIESettings,
         detector: Detector,
     ) -> None:
         super().__init__()
@@ -61,6 +63,7 @@ def __init__(
         self._probeSettings = probeSettings
         self._probePositionSettings = probePositionSettings
         self._oprSettings = oprSettings
+        self._pieSettings = pieSettings
         self._detector = detector
 
     @property
@@ -106,7 +109,7 @@ def _create_reconstructor_options(self) -> RPIEReconstructorOptions:
             ),
             # TODO random_seed
             # TODO displayed_loss_function
-            # TODO use_low_memory_forward_model
+            use_low_memory_forward_model=self._reconstructorSettings.useLowMemoryForwardModel.getValue(),
         )
 
     def _create_optimization_plan(self, start: int, stop: int, stride: int) -> OptimizationPlan:
@@ -255,6 +258,7 @@ def _create_object_options(self, object_: Object) -> PIEObjectOptions:
             multislice_regularization_unwrap_image_integration_method=multislice_regularization_unwrap_image_integration_method,
             multislice_regularization_stride=self._objectSettings.regularizeMultisliceStride.getValue(),
             patch_interpolation_method=patch_interpolation_method,
+            alpha=self._pieSettings.objectAlpha.getValue(),
         )
 
     def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEProbeOptions:
@@ -299,6 +303,7 @@ def _create_probe_options(self, probe: Probe, metadata: ProductMetadata) -> PIEP
             support_constraint=self._probeSettings.constrainSupport.getValue(),
             support_constraint_threshold=self._probeSettings.constrainSupportThreshold.getValue(),
             support_constraint_stride=self._probeSettings.constrainSupportStride.getValue(),
+            alpha=self._pieSettings.probeAlpha.getValue(),
         )
 
     def _create_probe_position_options(