update controller to pty-chi 1.0 api

Author: stevehenke

pyproject.toml

@@ -31,6 +31,7 @@ globus = ["gladier", "gladier-tools>=0.5.4"]
 gui = ["PyQt5"]
 ptychonn = ["ptychonn==0.3.*,>=0.3.7"]
 tike = ["tike==0.25.*,>=0.25.3"]
+ptychi = ["ptychi==1.*"]
 
 [tool.setuptools.package-data]
 "ptychodus" = ["py.typed"]

src/ptychodus/controller/ptychi/core.py

@@ -43,7 +43,7 @@ def __init__(
 
         # FIXME verify tooltips
         self._reconstructor_view_controller = PtyChiReconstructorViewController(
-            model.reconstructor_settings,
+            model.settings,
             autodiff_settings,
             dm_settings,
             lsqml_settings,
@@ -55,23 +55,24 @@ def __init__(
             dm_settings,
             lsqml_settings,
             pie_settings,
-            model.reconstructor_settings.num_epochs,
+            model.settings.num_epochs,
             model.enumerators,
         )
         self._probe_view_controller = PtyChiProbeViewController(
             model.probe_settings,
+            dm_settings,
             lsqml_settings,
             pie_settings,
-            model.reconstructor_settings.num_epochs,
+            model.settings.num_epochs,
             model.enumerators,
         )
         self._probe_positions_view_controller = PtyChiProbePositionsViewController(
             model.probe_position_settings,
-            model.reconstructor_settings.num_epochs,
+            model.settings.num_epochs,
             model.enumerators,
         )
         self._opr_view_controller = PtyChiOPRViewController(
-            model.opr_settings, model.reconstructor_settings.num_epochs, model.enumerators
+            model.opr_settings, model.settings.num_epochs, model.enumerators
         )
 
         layout = QVBoxLayout()

src/ptychodus/controller/ptychi/object.py

@@ -28,6 +28,43 @@
 __all__ = ['PtyChiObjectViewController']
 
 
+class PtyChiOptimizeSliceSpacingViewController(CheckableGroupBoxParameterViewController):
+    def __init__(
+        self,
+        optimize_slice_spacing: BooleanParameter,
+        start: IntegerParameter,
+        stop: IntegerParameter,
+        stride: IntegerParameter,
+        optimizer: StringParameter,
+        step_size: RealParameter,
+        num_epochs: IntegerParameter,
+        enumerators: PtyChiEnumerators,
+    ) -> None:
+        super().__init__(
+            optimize_slice_spacing,
+            'Optimize Slice Spacing',
+            tool_tip='Whether to optimize the slice spacing.',
+        )
+        self._plan_view_controller = PtyChiOptimizationPlanViewController(
+            start,
+            stop,
+            stride,
+            num_epochs,
+        )
+        self._optimizer_view_controller = PtyChiOptimizerParameterViewController(
+            optimizer, enumerators
+        )
+        self._step_size_view_controller = DecimalLineEditParameterViewController(
+            step_size, tool_tip='Optimizer step size'
+        )
+
+        layout = QFormLayout()
+        layout.addRow('Plan:', self._plan_view_controller.get_widget())
+        layout.addRow('Optimizer:', self._optimizer_view_controller.get_widget())
+        layout.addRow('Step Size:', self._step_size_view_controller.get_widget())
+        self.get_widget().setLayout(layout)
+
+
 class PtyChiConstrainL1NormViewController(CheckableGroupBoxParameterViewController):
     def __init__(
         self,
@@ -57,6 +94,35 @@ def __init__(
         self.get_widget().setLayout(layout)
 
 
+class PtyChiConstrainL2NormViewController(CheckableGroupBoxParameterViewController):
+    def __init__(
+        self,
+        constrain_l2_norm: BooleanParameter,
+        start: IntegerParameter,
+        stop: IntegerParameter,
+        stride: IntegerParameter,
+        weight: RealParameter,
+        num_epochs: IntegerParameter,
+    ) -> None:
+        super().__init__(
+            constrain_l2_norm,
+            'Constrain L\u2082 Norm',
+            tool_tip='Whether to constrain the L\u2082 norm.',
+        )
+        self._plan_view_controller = PtyChiOptimizationPlanViewController(
+            start, stop, stride, num_epochs
+        )
+        self._weight_view_controller = DecimalLineEditParameterViewController(
+            weight,
+            tool_tip='Weight of the L\u2082 norm constraint.',
+        )
+
+        layout = QFormLayout()
+        layout.addRow('Plan:', self._plan_view_controller.get_widget())
+        layout.addRow('Weight:', self._weight_view_controller.get_widget())
+        self.get_widget().setLayout(layout)
+
+
 class PtyChiConstrainSmoothnessViewController(CheckableGroupBoxParameterViewController):
     def __init__(
         self,
@@ -209,6 +275,29 @@ def __init__(
         self.get_widget().setLayout(layout)
 
 
+class PtyChiRemoveObjectProbeAmbiguityViewController(CheckableGroupBoxParameterViewController):
+    def __init__(
+        self,
+        remove_object_probe_ambiguity: BooleanParameter,
+        start: IntegerParameter,
+        stop: IntegerParameter,
+        stride: IntegerParameter,
+        num_epochs: IntegerParameter,
+    ) -> None:
+        super().__init__(
+            remove_object_probe_ambiguity,
+            'Remove Object Probe Ambiguity',
+            tool_tip='Whether to remove object-probe ambiguity.',
+        )
+        self._plan_view_controller = PtyChiOptimizationPlanViewController(
+            start, stop, stride, num_epochs
+        )
+
+        layout = QFormLayout()
+        layout.addRow('Plan:', self._plan_view_controller.get_widget())
+        self.get_widget().setLayout(layout)
+
+
 class PtyChiObjectViewController(CheckableGroupBoxParameterViewController):
     def __init__(
         self,
@@ -236,10 +325,15 @@ def __init__(
         self._step_size_view_controller = DecimalLineEditParameterViewController(
             settings.step_size, tool_tip='Optimizer step size'
         )
-        self._patch_interpolator_view_controller = ComboBoxParameterViewController(
-            settings.patch_interpolator,
-            enumerators.patch_interpolation_methods(),
-            tool_tip='Interpolation method used for extracting and updating patches of the object.',
+        self._optimize_slice_spacing_view_controller = PtyChiOptimizeSliceSpacingViewController(
+            settings.optimize_slice_spacing,
+            settings.optimize_slice_spacing_start,
+            settings.optimize_slice_spacing_stop,
+            settings.optimize_slice_spacing_stride,
+            settings.optimize_slice_spacing_optimizer,
+            settings.optimize_slice_spacing_step_size,
+            num_epochs,
+            enumerators,
         )
         self._constrain_l1_norm_view_controller = PtyChiConstrainL1NormViewController(
             settings.constrain_l1_norm,
@@ -249,6 +343,14 @@ def __init__(
             settings.constrain_l1_norm_weight,
             num_epochs,
         )
+        self._constrain_l2_norm_view_controller = PtyChiConstrainL2NormViewController(
+            settings.constrain_l2_norm,
+            settings.constrain_l2_norm_start,
+            settings.constrain_l2_norm_stop,
+            settings.constrain_l2_norm_stride,
+            settings.constrain_l2_norm_weight,
+            num_epochs,
+        )
         self._constrain_smoothness_view_controller = PtyChiConstrainSmoothnessViewController(
             settings.constrain_smoothness,
             settings.constrain_smoothness_start,
@@ -291,17 +393,40 @@ def __init__(
             num_epochs,
             enumerators,
         )
+        self._patch_interpolator_view_controller = ComboBoxParameterViewController(
+            settings.patch_interpolator,
+            enumerators.patch_interpolation_methods(),
+            tool_tip='Interpolation method used for extracting and updating patches of the object.',
+        )
+        self._remove_object_probe_ambiguity_view_controller = (
+            PtyChiRemoveObjectProbeAmbiguityViewController(
+                settings.remove_object_probe_ambiguity,
+                settings.remove_object_probe_ambiguity_start,
+                settings.remove_object_probe_ambiguity_stop,
+                settings.remove_object_probe_ambiguity_stride,
+                num_epochs,
+            )
+        )
+        self._build_preconditioner_with_all_modes_view_controller = CheckBoxParameterViewController(
+            settings.build_preconditioner_with_all_modes,
+            'Build Preconditioner with All Modes',
+            tool_tip='Whether to build the preconditioner using all modes.',
+        )
 
         layout = QFormLayout()
         layout.addRow('Plan:', self._optimization_plan_view_controller.get_widget())
         layout.addRow('Optimizer:', self._optimizer_view_controller.get_widget())
         layout.addRow('Step Size:', self._step_size_view_controller.get_widget())
-        layout.addRow('Patch Interpolator:', self._patch_interpolator_view_controller.get_widget())
+        layout.addRow(self._optimize_slice_spacing_view_controller.get_widget())
         layout.addRow(self._constrain_l1_norm_view_controller.get_widget())
+        layout.addRow(self._constrain_l2_norm_view_controller.get_widget())
         layout.addRow(self._constrain_smoothness_view_controller.get_widget())
         layout.addRow(self._constrain_total_variation_view_controller.get_widget())
         layout.addRow(self._remove_grid_artifacts_view_controller.get_widget())
         layout.addRow(self._regularize_multislice_view_controller.get_widget())
+        layout.addRow('Patch Interpolator:', self._patch_interpolator_view_controller.get_widget())
+        layout.addRow(self._remove_object_probe_ambiguity_view_controller.get_widget())
+        layout.addRow(self._build_preconditioner_with_all_modes_view_controller.get_widget())
 
         if dm_settings is not None:
             self._amplitude_clamp_limit_view_controller = DecimalLineEditParameterViewController(
@@ -312,6 +437,12 @@ def __init__(
                 'Amplitude Clamp Limit:', self._amplitude_clamp_limit_view_controller.get_widget()
             )
 
+            self._inertia_view_controller = DecimalLineEditParameterViewController(
+                dm_settings.object_inertia,
+                tool_tip='Inertia for the object update.',
+            )
+            layout.addRow('Inertia:', self._inertia_view_controller.get_widget())
+
         if lsqml_settings is not None:
             self._object_optimal_step_size_scaler_view_controller = (
                 DecimalLineEditParameterViewController(

src/ptychodus/controller/ptychi/opr.py

@@ -77,16 +77,16 @@ def __init__(
         self._step_size_view_controller = DecimalLineEditParameterViewController(
             settings.step_size, tool_tip='Optimizer step size'
         )
-        self._optimize_intensities_view_controller = CheckBoxParameterViewController(
-            settings.optimize_intensities,
-            'Optimize Intensities',
-            tool_tip='Whether to optimize intensity variation (i.e., the weight of the first OPR mode).',
-        )
         self._optimize_eigenmode_weights_view_controller = CheckBoxParameterViewController(
             settings.optimize_eigenmode_weights,
             'Optimize Eigenmode Weights',
             tool_tip='Whether to optimize eigenmode weights (i.e., the weights of the second and following OPR modes).',
         )
+        self._optimize_intensities_view_controller = CheckBoxParameterViewController(
+            settings.optimize_intensities,
+            'Optimize Intensities',
+            tool_tip='Whether to optimize intensity variation (i.e., the weight of the first OPR mode).',
+        )
         self._smooth_mode_weights_view_controller = PtyChiSmoothOPRModeWeightsViewController(
             settings.smooth_mode_weights,
             settings.smooth_mode_weights_start,
@@ -109,4 +109,5 @@ def __init__(
         layout.addRow(self._optimize_intensities_view_controller.get_widget())
         layout.addRow(self._optimize_eigenmode_weights_view_controller.get_widget())
         layout.addRow(self._smooth_mode_weights_view_controller.get_widget())
+        layout.addRow('Relax Update:', self._relax_update_view_controller.get_widget())
         self.get_widget().setLayout(layout)

src/ptychodus/controller/ptychi/positions.py

@@ -65,6 +65,53 @@ def _update(self, observable: Observable) -> None:
             self._sync_model_to_view()
 
 
+class PtyChiConstrainAffineTransformViewController(CheckableGroupBoxParameterViewController):
+    def __init__(
+        self,
+        is_optimizable: BooleanParameter,
+        start: IntegerParameter,
+        stop: IntegerParameter,
+        stride: IntegerParameter,
+        degrees_of_freedom: IntegerParameter,
+        position_weight_update_interval: IntegerParameter,
+        apply_constraint: BooleanParameter,
+        max_expected_error_px: RealParameter,
+        num_epochs: IntegerParameter,
+    ) -> None:
+        super().__init__(
+            is_optimizable,
+            'Constrain Affine Transform',
+            tool_tip='Constrain the affine transform during position correction.',
+        )
+        self._plan_view_controller = PtyChiOptimizationPlanViewController(
+            start, stop, stride, num_epochs
+        )
+        # FIXME affine degrees of freedom: checkable combo box?
+        self._position_weight_update_interval_view_controller = SpinBoxParameterViewController(
+            position_weight_update_interval,
+            tool_tip='Interval for updating the position weight.',
+        )
+        self._apply_constraint_view_controller = CheckBoxParameterViewController(
+            apply_constraint, 'Apply Constraint', tool_tip='Whether to apply the constraint.'
+        )
+        self._max_expected_error_px_view_controller = DecimalLineEditParameterViewController(
+            max_expected_error_px, tool_tip='Maximum expected error in pixels.'
+        )
+
+        layout = QFormLayout()
+        layout.addRow('Plan:', self._plan_view_controller.get_widget())
+        # FIXE layout.addRow('Degrees of Freedom:', self._degrees_of_freedom_view_controller.get_widget())
+        layout.addRow(
+            'Position Weight Update Interval:',
+            self._position_weight_update_interval_view_controller.get_widget(),
+        )
+        layout.addRow(self._apply_constraint_view_controller.get_widget())
+        layout.addRow(
+            'Max Expected Error [px]:', self._max_expected_error_px_view_controller.get_widget()
+        )
+        self.get_widget().setLayout(layout)
+
+
 class PtyChiProbePositionsViewController(CheckableGroupBoxParameterViewController):
     def __init__(
         self,
@@ -89,28 +136,57 @@ def __init__(
         self._step_size_view_controller = DecimalLineEditParameterViewController(
             settings.step_size, tool_tip='Optimizer step size'
         )
-        self._algorithm_view_controller = ComboBoxParameterViewController(
+        self._constrain_centroid_view_controller = CheckBoxParameterViewController(
+            settings.constrain_centroid,
+            'Constrain Centroid',
+            tool_tip='Whether to subtract the mean from positions after updating positions.',
+        )
+        self._correction_type_view_controller = ComboBoxParameterViewController(
             settings.correction_type,
             enumerators.position_correction_types(),
             tool_tip='Algorithm used to calculate the position correction update.',
         )
+        self._differentiation_method_view_controller = ComboBoxParameterViewController(
+            settings.differentiation_method,
+            enumerators.image_gradient_methods(),
+            tool_tip='Method for calculating the object gradient.',
+        )
         self._cross_correlation_view_controller = PtyChiCrossCorrelationViewController(
             settings.correction_type,
             settings.cross_correlation_scale,
             settings.cross_correlation_real_space_width,
             settings.cross_correlation_probe_threshold,
         )
-        self._constrain_centroid_view_controller = CheckBoxParameterViewController(
-            settings.constrain_centroid,
-            'Constrain Centroid',
-            tool_tip='Whether to subtract the mean from positions after updating positions.',
+        self._update_magnitude_limit_view_controller = DecimalLineEditParameterViewController(
+            settings.update_magnitude_limit,
+            tool_tip='Maximum allowed magnitude of position update in each axis.',
+        )  # FIXME verify inf is handled correctly
+        self._constrain_affine_transform_view_controller = (
+            PtyChiConstrainAffineTransformViewController(
+                settings.constrain_affine_transform,
+                settings.constrain_affine_transform_start,
+                settings.constrain_affine_transform_stop,
+                settings.constrain_affine_transform_stride,
+                settings.constrain_affine_transform_degrees_of_freedom,
+                settings.constrain_affine_transform_position_weight_update_interval,
+                settings.constrain_affine_transform_apply_constraint,
+                settings.constrain_affine_transform_max_expected_error_px,
+                num_epochs,
+            )
         )
 
         layout = QFormLayout()
         layout.addRow('Plan:', self._optimization_plan_view_controller.get_widget())
         layout.addRow('Optimizer:', self._optimizer_view_controller.get_widget())
         layout.addRow('Step Size:', self._step_size_view_controller.get_widget())
-        layout.addRow('Algorithm:', self._algorithm_view_controller.get_widget())
-        layout.addRow(self._cross_correlation_view_controller.get_widget())
         layout.addRow(self._constrain_centroid_view_controller.get_widget())
+        layout.addRow('Correction Type:', self._correction_type_view_controller.get_widget())
+        layout.addRow(
+            'Differentiation Method:', self._differentiation_method_view_controller.get_widget()
+        )
+        layout.addRow(self._cross_correlation_view_controller.get_widget())
+        layout.addRow(
+            'Update Magnitude Limit:', self._update_magnitude_limit_view_controller.get_widget()
+        )
+        layout.addRow(self._constrain_affine_transform_view_controller.get_widget())
         self.get_widget().setLayout(layout)