Make e3PIE dispatch on the array module instead of forcing cupy

PtyLab/Engines/e3PIE.py

@@ -48,35 +48,25 @@ def initializeReconstructionParams(self):
         Set parameters that are specific to the e3PIE settings.
         :return:
         """
-        self.params.betaProbe = 0.25
-        self.params.betaObject = 0.25
+        # these are read back as self.betaProbe / self.betaObject in reconstruct()
+        # and objectPatchUpdate(), matching every other engine (cf. ePIE.py)
+        self.betaProbe = 0.25
+        self.betaObject = 0.25
         self.numIterations = 50
 
-        if False:
-            # preallocate transfer function
-            self.reconstruction.H = aspw(
-                np.squeeze(self.reconstruction.probe[0, 0, 0, 0, ...]),
-                self.reconstruction.dz,
-                self.reconstruction.wavelength / self.reconstruction.refrIndex,
-                self.reconstruction.Lp,
-            )[1]
-            # shift transfer function to avoid fftshifts for FFTS
-            # self.reconstruction.H = np.fft.ifftshift(self.optimizableH)
-            self.reconstruction.H = np.fft.ifftshift(self.reconstruction.H)
-
-        if True:
-            import cupy as xp
-
-            # preallocate transfer function
-            self.reconstruction.H = aspw(
-                xp.squeeze(self.reconstruction.probe[0, 0, 0, 0, ...]),
-                self.reconstruction.dz,
-                self.reconstruction.wavelength / self.reconstruction.refrIndex,
-                self.reconstruction.Lp,
-            )[1]
-            # shift transfer function to avoid fftshifts for FFTS
-            # self.reconstruction.H = np.fft.ifftshift(self.optimizableH)
-            self.reconstruction.H = xp.fft.ifftshift(self.reconstruction.H)
+        # preallocate transfer function. This runs from __init__, before
+        # _checkGPU has moved anything, so the probe is still on the host here;
+        # H is listed in Reconstruction.possible_GPU_fields and travels with the
+        # rest of the state when the engine switches to the GPU.
+        xp = getArrayModule(self.reconstruction.probe)
+        self.reconstruction.H = aspw(
+            xp.squeeze(self.reconstruction.probe[0, 0, 0, 0, ...]),
+            self.reconstruction.dz,
+            self.reconstruction.wavelength / self.reconstruction.refrIndex,
+            self.reconstruction.Lp,
+        )[1]
+        # shift transfer function to avoid fftshifts for FFTS
+        self.reconstruction.H = xp.fft.ifftshift(self.reconstruction.H)
 
     def reconstruct(self):
         self._prepareReconstruction()

PtyLab/Reconstruction/Reconstruction.py

@@ -92,7 +92,10 @@ def __init__(self, data: ExperimentalData, params: Params):
             "purityProbe",
             "purityObject",
             "reference",
-            "intensity_mask"
+            "intensity_mask",
+            # multislice (e3PIE) transfer function, built on the host in
+            # e3PIE.initializeReconstructionParams and used inside the position loop
+            "H",
         ]
 
     # @property

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ptylab"
-version = "0.2.6"
+version = "0.2.7"
 description = "A cross-platform, open-source inverse modeling toolbox for conventional and Fourier ptychography"
 authors = [
     { name = "Lars Loetgering", email = "lars.loetgering@fulbrightmail.org" },

tests/regression/test_engine_regression.py

@@ -42,9 +42,8 @@
 
 # name -> engine, propagator, (nlambda, nosm, npsm, nslice), iterations
 #
-# Engines are added here as they become testable. e3PIE is still absent: it
-# raises on its own betaProbe, so pinning it today would record a meaningless
-# baseline. It joins this table in the PR that repairs it.
+# Engines are added here as they become testable. Still absent: zPIE, aPIE and
+# mPIE_tv, which do not run on this branch at all -- see the PRs that repair them.
 CONFIGS = {
     "epie_fraunhofer": ("ePIE", "Fraunhofer", (1, 1, 1, 1), 3),
     "epie_asp": ("ePIE", "ASP", (1, 1, 1, 1), 3),
@@ -54,6 +53,7 @@
     "mpie_single": ("mPIE", "Fraunhofer", (1, 1, 1, 1), 3),
     "mpie_mixed_state": ("mPIE", "Fraunhofer", (1, 2, 3, 1), 3),
     "qnewton_single": ("qNewton", "Fraunhofer", (1, 1, 1, 1), 3),
+    "e3pie_multislice": ("e3PIE", "Fraunhofer", (1, 1, 1, 3), 2),
 }