diff --git a/src/ptychi/api/options/lsqml.py b/src/ptychi/api/options/lsqml.py index 2a2f063..fc02a8c 100644 --- a/src/ptychi/api/options/lsqml.py +++ b/src/ptychi/api/options/lsqml.py @@ -26,12 +26,17 @@ class LSQMLReconstructorOptions(base.ReconstructorOptions): The standard deviation of the gaussian noise. Only used when `noise_model == enums.NoiseModels.GAUSSIAN`. """ - solve_obj_prb_step_size_jointly_for_first_slice_in_multislice: bool = False + single_slice_solve_obj_prb_step_size_jointly: bool = True """ - Whether to solve the simultaneous object/probe step length calculation; - in FoldSlice they use independent (non-joint) step length calculation, but - we're adding the option of using simultaneous AND non-simultaneous step - length calculation. + Whether to solve the object/probe step size jointly for single-slice objects. + For multislice objects, use `multislice_solve_obj_prb_step_size_jointly` instead. + """ + + multislice_solve_obj_prb_step_size_jointly: bool = False + """ + Whether to solve the object/probe step size jointly for multislice objects at the first slice. + Slices other than the first are always solved independently. For single-slice objects, use + `single_slice_solve_obj_prb_step_size_jointly` instead. """ solve_step_sizes_only_using_first_probe_mode: bool = True diff --git a/src/ptychi/reconstructors/lsqml.py b/src/ptychi/reconstructors/lsqml.py index 3ebc433..cefbd3c 100644 --- a/src/ptychi/reconstructors/lsqml.py +++ b/src/ptychi/reconstructors/lsqml.py @@ -494,10 +494,15 @@ def calculate_optimal_step_sizes( The slice index of the object. """ object_ = self.parameter_group.object - if (not object_.is_multislice) or ( - object_.is_multislice - and slice_index == 0 - and self.options.solve_obj_prb_step_size_jointly_for_first_slice_in_multislice + if ( + ( + not object_.is_multislice + and self.options.single_slice_solve_obj_prb_step_size_jointly + ) or ( + object_.is_multislice + and self.options.multislice_solve_obj_prb_step_size_jointly + and slice_index == 0 + ) ): (alpha_o_i, alpha_p_i) = self.calculate_object_and_probe_update_step_sizes( chi, diff --git a/tests/test_multislice_ptycho_lsqml_joint_step_size.py b/tests/test_multislice_ptycho_lsqml_joint_step_size.py index f96d26d..5f5ef9f 100644 --- a/tests/test_multislice_ptycho_lsqml_joint_step_size.py +++ b/tests/test_multislice_ptycho_lsqml_joint_step_size.py @@ -44,7 +44,7 @@ def test_multislice_ptycho_lsqml_joint_step_size(self): options.reconstructor_options.batch_size = 96 options.reconstructor_options.noise_model = api.NoiseModels.GAUSSIAN options.reconstructor_options.num_epochs = 8 - options.reconstructor_options.solve_obj_prb_step_size_jointly_for_first_slice_in_multislice = True + options.reconstructor_options.multislice_solve_obj_prb_step_size_jointly = True options.reconstructor_options.allow_nondeterministic_algorithms = False task = PtychographyTask(options)