From 4378485e5f0eb3641ef97b7c0a8a851e40878450 Mon Sep 17 00:00:00 2001 From: Ashish Tripathi Date: Wed, 11 Mar 2026 16:21:22 -0500 Subject: [PATCH 1/4] sometimes when orthogonalization of the shared probe modes causes the 0th mode to not have the most power content; this fix ensures that the shared modes will always be sorted in descending order --- src/ptychi/api/options/base.py | 4 +++- src/ptychi/data_structures/probe.py | 7 +++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/ptychi/api/options/base.py b/src/ptychi/api/options/base.py index 4c7ed148..4c6158e8 100644 --- a/src/ptychi/api/options/base.py +++ b/src/ptychi/api/options/base.py @@ -572,7 +572,9 @@ class ProbeOrthogonalizeIncoherentModesOptions(FeatureOptions): method: enums.OrthogonalizationMethods = enums.OrthogonalizationMethods.SVD """The method to use for incoherent_mode orthogonalization.""" - + + sort_by_occupancy: bool = False + """Keep the probes sorted so that mode with highest occupancy is the 0th shared mode""" @dataclasses.dataclass class ProbeOrthogonalizeOPRModesOptions(FeatureOptions): diff --git a/src/ptychi/data_structures/probe.py b/src/ptychi/data_structures/probe.py index 9016cecf..24fbe949 100644 --- a/src/ptychi/data_structures/probe.py +++ b/src/ptychi/data_structures/probe.py @@ -221,6 +221,13 @@ def constrain_incoherent_modes_orthogonality(self): return probe = self.data + + if self.options.orthogonalize_incoherent_modes.sort_by_occupancy: + shared_occupancy = torch.sum(torch.abs(probe[0, ...]) ** 2, (-2, -1)) / torch.sum( + torch.abs(probe[0, ...]) ** 2 + ) + shared_occupancy = torch.sort(shared_occupancy, dim=0, descending=True) + probe[0, ...] = probe[0, shared_occupancy[1], ...] norm_first_mode_orig = pmath.norm(probe[0, 0], dim=(-2, -1)) From a5e198f112c7435da45ba26fcfe64990177e89f2 Mon Sep 17 00:00:00 2001 From: MD Date: Fri, 13 Mar 2026 11:22:11 -0500 Subject: [PATCH 2/4] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/ptychi/data_structures/probe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ptychi/data_structures/probe.py b/src/ptychi/data_structures/probe.py index 24fbe949..352dfd75 100644 --- a/src/ptychi/data_structures/probe.py +++ b/src/ptychi/data_structures/probe.py @@ -221,13 +221,13 @@ def constrain_incoherent_modes_orthogonality(self): return probe = self.data - if self.options.orthogonalize_incoherent_modes.sort_by_occupancy: shared_occupancy = torch.sum(torch.abs(probe[0, ...]) ** 2, (-2, -1)) / torch.sum( torch.abs(probe[0, ...]) ** 2 ) shared_occupancy = torch.sort(shared_occupancy, dim=0, descending=True) - probe[0, ...] = probe[0, shared_occupancy[1], ...] + sorted_idx = shared_occupancy[1] + probe[0] = torch.index_select(probe[0], 0, sorted_idx) norm_first_mode_orig = pmath.norm(probe[0, 0], dim=(-2, -1)) From d1df670c99fc5f0cd4c1d834a169f4d243113b7f Mon Sep 17 00:00:00 2001 From: Ming Du Date: Fri, 13 Mar 2026 11:25:10 -0500 Subject: [PATCH 3/4] CHORE: remove normalization --- src/ptychi/data_structures/probe.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ptychi/data_structures/probe.py b/src/ptychi/data_structures/probe.py index 352dfd75..1ef8200e 100644 --- a/src/ptychi/data_structures/probe.py +++ b/src/ptychi/data_structures/probe.py @@ -222,9 +222,7 @@ def constrain_incoherent_modes_orthogonality(self): probe = self.data if self.options.orthogonalize_incoherent_modes.sort_by_occupancy: - shared_occupancy = torch.sum(torch.abs(probe[0, ...]) ** 2, (-2, -1)) / torch.sum( - torch.abs(probe[0, ...]) ** 2 - ) + shared_occupancy = torch.sum(torch.abs(probe[0, ...]) ** 2, (-2, -1)) shared_occupancy = torch.sort(shared_occupancy, dim=0, descending=True) sorted_idx = shared_occupancy[1] probe[0] = torch.index_select(probe[0], 0, sorted_idx) From ad220a8df5e5bf5c17100e824ecaff9844f83234 Mon Sep 17 00:00:00 2001 From: Ming Du Date: Fri, 13 Mar 2026 11:26:49 -0500 Subject: [PATCH 4/4] DOCS: update docstring --- src/ptychi/api/options/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ptychi/api/options/base.py b/src/ptychi/api/options/base.py index 4c6158e8..f7910fcc 100644 --- a/src/ptychi/api/options/base.py +++ b/src/ptychi/api/options/base.py @@ -574,7 +574,7 @@ class ProbeOrthogonalizeIncoherentModesOptions(FeatureOptions): """The method to use for incoherent_mode orthogonalization.""" sort_by_occupancy: bool = False - """Keep the probes sorted so that mode with highest occupancy is the 0th shared mode""" + """If True, keep the probes sorted so that mode with highest occupancy is the 0th shared mode.""" @dataclasses.dataclass class ProbeOrthogonalizeOPRModesOptions(FeatureOptions):