diff --git a/PtyLab/Engines/e3PIE.py b/PtyLab/Engines/e3PIE.py index 11905f9..8ad7a56 100644 --- a/PtyLab/Engines/e3PIE.py +++ b/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() diff --git a/PtyLab/Reconstruction/Reconstruction.py b/PtyLab/Reconstruction/Reconstruction.py index b490a7d..11e5d46 100644 --- a/PtyLab/Reconstruction/Reconstruction.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 5d63894..bb3e095 100644 --- a/pyproject.toml +++ b/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" }, diff --git a/tests/regression/data/e3pie_multislice.npz b/tests/regression/data/e3pie_multislice.npz new file mode 100644 index 0000000..7df4e20 Binary files /dev/null and b/tests/regression/data/e3pie_multislice.npz differ diff --git a/tests/regression/test_engine_regression.py b/tests/regression/test_engine_regression.py index 40d1b14..b83029b 100644 --- a/tests/regression/test_engine_regression.py +++ b/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), }