From 724a4cfff2e5665057fc1c25ee0cc03a1f3b4c4f Mon Sep 17 00:00:00 2001 From: Ashish Tripathi Date: Tue, 17 Mar 2026 13:39:59 -0500 Subject: [PATCH 1/2] bugfix for lsqml where we *always* use joint obj and probe step lengths even though the flag is set to false --- src/ptychi/api/options/lsqml.py | 2 +- src/ptychi/reconstructors/lsqml.py | 7 ++----- tests/test_multislice_ptycho_lsqml_joint_step_size.py | 2 +- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/ptychi/api/options/lsqml.py b/src/ptychi/api/options/lsqml.py index 2a2f0631..86293858 100644 --- a/src/ptychi/api/options/lsqml.py +++ b/src/ptychi/api/options/lsqml.py @@ -26,7 +26,7 @@ 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 + solve_obj_prb_step_size_jointly: bool = False """ Whether to solve the simultaneous object/probe step length calculation; in FoldSlice they use independent (non-joint) step length calculation, but diff --git a/src/ptychi/reconstructors/lsqml.py b/src/ptychi/reconstructors/lsqml.py index 3ebc433e..22e9b6ca 100644 --- a/src/ptychi/reconstructors/lsqml.py +++ b/src/ptychi/reconstructors/lsqml.py @@ -494,11 +494,8 @@ 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 (self.options.solve_obj_prb_step_size_jointly + and ((not object_.is_multislice) or (object_.is_multislice and slice_index == 0))): (alpha_o_i, alpha_p_i) = self.calculate_object_and_probe_update_step_sizes( chi, obj_patches, diff --git a/tests/test_multislice_ptycho_lsqml_joint_step_size.py b/tests/test_multislice_ptycho_lsqml_joint_step_size.py index f96d26d4..2f97c341 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.solve_obj_prb_step_size_jointly = True options.reconstructor_options.allow_nondeterministic_algorithms = False task = PtychographyTask(options) From c00d6bd66d25632c71ad18b7a3831b1e6a501368 Mon Sep 17 00:00:00 2001 From: Ming Du Date: Wed, 18 Mar 2026 13:48:10 -0500 Subject: [PATCH 2/2] CHORE: separate the joint step option into single- and multi-slice, retaining default behavior --- src/ptychi/api/options/lsqml.py | 15 ++++++++++----- src/ptychi/reconstructors/lsqml.py | 12 ++++++++++-- ...est_multislice_ptycho_lsqml_joint_step_size.py | 2 +- 3 files changed, 21 insertions(+), 8 deletions(-) diff --git a/src/ptychi/api/options/lsqml.py b/src/ptychi/api/options/lsqml.py index 86293858..fc02a8ca 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: 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 22e9b6ca..cefbd3c3 100644 --- a/src/ptychi/reconstructors/lsqml.py +++ b/src/ptychi/reconstructors/lsqml.py @@ -494,8 +494,16 @@ def calculate_optimal_step_sizes( The slice index of the object. """ object_ = self.parameter_group.object - if (self.options.solve_obj_prb_step_size_jointly - and ((not object_.is_multislice) or (object_.is_multislice and slice_index == 0))): + 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, obj_patches, diff --git a/tests/test_multislice_ptycho_lsqml_joint_step_size.py b/tests/test_multislice_ptycho_lsqml_joint_step_size.py index 2f97c341..5f5ef9f1 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 = True + options.reconstructor_options.multislice_solve_obj_prb_step_size_jointly = True options.reconstructor_options.allow_nondeterministic_algorithms = False task = PtychographyTask(options)