FEAT: allow using total probe intensity for COM in probe center constraint

Co-authored-by: Ashish Tripathi <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: MD <[email protected]>

Author: 4 people

src/ptychi/api/options/base.py

@@ -501,6 +501,13 @@ class ProbeCenterConstraintOptions(FeatureOptions):
     enabled: bool = False
 
     optimization_plan: OptimizationPlan = dataclasses.field(default_factory=OptimizationPlan)
+    
+    use_intensity_for_com: bool = False
+    """
+    Whether to use the magnitude of the dominant shared probe 
+    mode for computing the center of mass of the probe in order 
+    to keep it centered, or to use the total probe intensity.
+    """
 
 
 @dataclasses.dataclass

src/ptychi/data_structures/probe.py

@@ -394,7 +394,13 @@ def center_probe(self):
         """
         Move the probe's center of mass to the center of the probe array.
         """
-        com = ip.find_center_of_mass(self.get_mode_and_opr_mode(0, 0))
+        
+        if self.options.center_constraint.use_intensity_for_com:
+            probe_to_be_shifted = torch.sum(torch.abs(self.data[0, ...]) ** 2, dim=0)
+        else:
+            probe_to_be_shifted = self.get_mode_and_opr_mode(0, 0)
+        
+        com = ip.find_center_of_mass(probe_to_be_shifted)
         shift = utils.to_tensor(self.shape[-2:]) // 2 - com
         shifted_probe = self.shift(shift)
         self.set_data(shifted_probe)