From b6bc7d95ccdca1779f873c5e8ca03d9e3aa7c86d Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 7 Apr 2025 09:35:59 +0200
Subject: [PATCH] Change the behavior of FancyPtycho.center_probes to center
 the probe in real space even when fourier_probe is set to True

---
 src/cdtools/models/fancy_ptycho.py | 35 +++++++++++++++++++++++----
 tests/models/test_fancy_ptycho.py  | 39 ++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 0f0323e..6220299 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -629,15 +629,40 @@ class FancyPtycho(CDIModel):
 
         
     def center_probes(self, iterations=4):
-        """Centers the probes
+        """Centers the probes in real space
+
+        Takes the current guess of the illumination function and centers it
+        using a shift with periodic boundary conditions. It uses
+        cdtools.tools.image_processing.center internally to do the centering.
+        Multiple iterations of an algorithm are run, which is helpful if the
+        illumination is reconstructed near the corners and "wraps around" the
+        probe field of view.
+
+        Note that the centering is always performed in real space, even if
+        the probe array is defined in Fourier space.
         
-        Note that this does not compensate for the centering by adjusting
+        Note also that this does not compensate for the centering by adjusting
         the object, so it's a good idea to reset the object after centering
         the probes
+
+        Parameters
+        ----------
+        iterations : int
+            Default 4, how many iterations of the centering algorithm to run
         """
-        centered_probe = tools.image_processing.center(
-            self.probe.data.cpu(), iterations=iterations)
-        self.probe.data = centered_probe.to(device=self.probe.data.device)
+        if self.fourier_probe:
+            prs = tools.propagators.inverse_far_field(self.probe.detach()).cpu()
+        else:
+            prs = self.probe.detach().cpu()
+        
+        centered_prs = tools.image_processing.center(prs, iterations=iterations)
+
+        if self.fourier_probe:
+            self.probe.data = tools.propagators.far_field(
+                centered_prs.to(device=self.probe.data.device))
+        else:
+            self.probe.data = centered_prs.to(device=self.probe.data.device)
+
 
 
     def tidy_probes(self):
diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py
index 1591bf9..7ef429f 100644
--- a/tests/models/test_fancy_ptycho.py
+++ b/tests/models/test_fancy_ptycho.py
@@ -5,6 +5,45 @@ import torch as t
 import cdtools
 from matplotlib import pyplot as plt
 
+
+def test_center_probe(lab_ptycho_cxi):
+    dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
+    model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=False
+    )
+    base_probe = model.probe.detach().clone()
+    model.center_probes()
+    centered_probe = model.probe.detach().clone()
+
+    fourier_model = cdtools.models.FancyPtycho.from_dataset(
+        dataset,
+        n_modes=3,
+        fourier_probe=True,
+    )
+
+    fourier_model.probe.data = cdtools.tools.propagators.far_field(
+        base_probe
+    )
+    
+    fourier_base_probe = fourier_model.probe.detach().clone()
+    fourier_model.center_probes()
+    fourier_centered_probe = fourier_model.probe.detach().clone()
+    ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
+        fourier_centered_probe)
+
+    # So we know the code had to do something
+    assert not t.allclose(base_probe, centered_probe)
+    # And checking that they both do the same thing, whether or not
+    # fourier_probe was set to True
+    assert t.allclose(
+        centered_probe,
+        ifft_fourier_centered_probe,
+        atol=1e-4,
+        rtol=1e-3
+    )
+    
 @pytest.mark.slow
 def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):