Merge pull request #332 from a4894z/main

multislice rPIE ready for main branch after testing Daniel's reworkng of multiGPU workflow

Author: a4894z

src/tike/operators/cupy/fresnelspectprop.py

@@ -123,8 +123,8 @@ def _create_fresnel_spectrum_propagator(
         xgrid = ( 0.5 + self.xp.linspace( ( -0.5 * N[1] ), ( 0.5 * N[1] - 1 ), num = N[1] )) / N[1]
         ygrid = ( 0.5 + self.xp.linspace( ( -0.5 * N[0] ), ( 0.5 * N[0] - 1 ), num = N[0] )) / N[0]
 
-        kx = 2 * self.xp.pi * N[0] * xgrid / probe_FOV[ 0 ]
-        ky = 2 * self.xp.pi * N[1] * ygrid / probe_FOV[ 1 ]
+        kx = 2 * self.xp.pi * N[1] * xgrid / probe_FOV[ 1 ]
+        ky = 2 * self.xp.pi * N[0] * ygrid / probe_FOV[ 0 ]
 
         Kx, Ky = self.xp.meshgrid(kx, ky, indexing='xy')
 

src/tike/operators/cupy/multislice.py

@@ -73,8 +73,11 @@ def fwd(
         psi: npt.NDArray[np.csingle],
         **kwargs,
     ) -> npt.NDArray[np.csingle]:
+        
         """Please see help(Multislice) for more info."""
+
         assert psi.ndim == 3
+
         exitwave = self.diffraction.fwd(
             psi=psi[0],
             scan=scan,
@@ -88,6 +91,56 @@ def fwd(
             )
         return exitwave
 
+
+
+
+    def fwd_return_intermediate_probes(
+        self,
+        probe: npt.NDArray[np.csingle],
+        scan: npt.NDArray[np.single],
+        psi: npt.NDArray[np.csingle],
+        **kwargs,
+    ) -> npt.NDArray[np.csingle]:
+        
+        """Please see help(Multislice) for more info."""
+
+        assert psi.ndim == 3
+
+
+        # exitwave = self.diffraction.fwd(
+        #     psi=psi[0],
+        #     scan=scan,
+        #     probe=probe,
+        # )
+
+        # for s in range(1, len(psi)):
+        #     exitwave = self.diffraction.fwd(
+        #         psi=psi[s],
+        #         scan=scan,
+        #         probe=self.propagation.fwd(exitwave),
+        #     )
+
+        # return exitwave
+
+        multislice_probes =  self.xp.zeros( ( psi.shape[0], scan.shape[-2], *probe.shape[-3:] ), dtype=probe.dtype )
+        multislice_probes[ 0, ... ] = probe[..., 0, :, :, :]
+
+        for tt in range(0, len(psi)) :
+
+            multislice_exwv = self.diffraction.fwd(
+                    psi   = psi[ tt, ... ],               
+                    scan  = scan,
+                    probe = multislice_probes[ tt, ... ],
+                )
+            
+            if tt == ( psi.shape[0] - 1 ) :
+                break
+
+            multislice_probes[ tt + 1, ... ] = self.propagation.fwd( nearplane = multislice_exwv, )
+
+        return multislice_exwv, multislice_probes
+    
+    
     def adj(
         self,
         nearplane: npt.NDArray[np.csingle],

src/tike/operators/cupy/ptycho.py

@@ -128,6 +128,31 @@ def fwd(
             overwrite=True,
         )[..., None, :, :, :]
 
+    def fwd_return_intermediate_probes(
+        self,
+        probe: npt.NDArray[np.csingle],
+        scan: npt.NDArray[np.single],
+        psi: npt.NDArray[np.csingle],
+        **kwargs,
+    ) -> npt.NDArray[np.csingle]:
+        
+        """Please see help(Ptycho) for more info."""
+
+        # return self.propagation.fwd(
+        #     self.diffraction.fwd(
+        #         psi=psi,
+        #         scan=scan,
+        #         probe=probe[..., 0, :, :, :],
+        #     ),
+        #     overwrite=True,
+        # )[..., None, :, :, :]
+
+        multislice_exwv, multislice_probes = self.diffraction.fwd_return_intermediate_probes( psi=psi, scan=scan, probe=probe, )
+
+        multislice_farfield = self.propagation.fwd( nearplane = multislice_exwv, overwrite=True, )[..., None, :, :, :]
+
+        return multislice_farfield, multislice_probes
+    
     def adj(
         self,
         farplane: npt.NDArray[np.csingle],

src/tike/ptycho/ptycho.py

@@ -336,7 +336,7 @@ def __init__(
                     "automatic psi initialization is not synchronized "
                     "across processes.")
         else:
-            mpi = tike.communicators.NoMPIComm
+            mpi = tike.communicators.NoMPIComm  # isinstance(mpi, tike.communicators.NoMPIComm) # THIS ALWAYS RETURNS FALSE?
 
         self.data: typing.List[npt.ArrayLike] = [data]
         self.parameters: typing.List[solvers.PtychoParameters] = [
@@ -354,7 +354,8 @@ def __init__(
             probe_FOV_lengths=parameters.probe_options.probe_FOV_lengths,
             multislice_propagation_distance=parameters.object_options.multislice_propagation_distance,
         )
-        self.comm = tike.communicators.Comm(num_gpu, mpi)
+
+        self.comm = tike.communicators.Comm(num_gpu, mpi)   
 
     def __enter__(self):
         self.device.__enter__()

src/tike/ptycho/solvers/_preconditioner.py

@@ -11,22 +11,29 @@
 
 
 def _rolling_average_object(parameters: PtychoParameters, new):
-    if parameters.object_options.preconditioner is None:
-        parameters.object_options.preconditioner = new
-    else:
-        parameters.object_options.preconditioner = 0.5 * (
-            new + parameters.object_options.preconditioner
-        )
+
+    # if parameters.object_options.preconditioner is None:
+    #     parameters.object_options.preconditioner = new
+    # else:
+    #     parameters.object_options.preconditioner = 0.5 * (
+    #         new + parameters.object_options.preconditioner
+    #     )
+
+    parameters.object_options.preconditioner = new
+
     return parameters
 
 
 def _rolling_average_probe(parameters: PtychoParameters, new):
-    if parameters.probe_options.preconditioner is None:
-        parameters.probe_options.preconditioner = new
-    else:
-        parameters.probe_options.preconditioner = 0.5 * (
-            new + parameters.probe_options.preconditioner
-        )
+
+    # if parameters.probe_options.preconditioner is None:
+    #     parameters.probe_options.preconditioner = new
+    # else:
+    #     parameters.probe_options.preconditioner = 0.5 * (
+    #         new + parameters.probe_options.preconditioner
+    #     )
+
+    parameters.probe_options.preconditioner = new    
     return parameters
 
 
@@ -55,24 +62,31 @@ def make_certain_args_constant(
         lo: int,
         hi: int,
     ) -> None:
+        
         nonlocal psi_update_denominator
 
         probe_amp = _probe_amp_sum(parameters.probe)[:, 0]
+
         psi_update_denominator[0] = operator.diffraction.patch.adj(
             patches=probe_amp,
             images=psi_update_denominator[0],
             positions=parameters.scan[lo:hi],
         )
 
         probe1 = parameters.probe[:, 0]
+
         for i in range(1, len(parameters.psi)):
+
             probe1 = operator.diffraction.diffraction.fwd(
                 probe=probe1,
                 scan=parameters.scan[lo:hi],
                 psi=parameters.psi[i-1],
             )
+
             probe1 = operator.diffraction.propagation.fwd(probe1)
+
             probe_amp = _probe_amp_sum(probe1)
+
             psi_update_denominator[i] = operator.diffraction.patch.adj(
                 patches=probe_amp,
                 images=psi_update_denominator[i],
@@ -107,7 +121,7 @@ def _probe_preconditioner(
 ) -> npt.NDArray:
 
     probe_update_denominator = cp.zeros(
-        shape=parameters.probe.shape[-2:],
+        shape=( parameters.psi.shape[0], *parameters.probe.shape[-2:] ),
         dtype=parameters.probe.dtype,
     )
 
@@ -116,16 +130,31 @@ def make_certain_args_constant(
         lo: int,
         hi: int,
     ) -> None:
+        
         nonlocal probe_update_denominator
 
-        # FIXME: Only use the first slice for the probe preconditioner
-        patches = operator.diffraction.patch.fwd(
-            images=parameters.psi[0],
-            positions=parameters.scan[lo:hi],
-            patch_width=parameters.probe.shape[-1],
-        )
-        probe_update_denominator[...] += _patch_amp_sum(patches)
-        assert probe_update_denominator.ndim == 2
+        for i in range(0, len(parameters.psi)):
+
+            patches = operator.diffraction.patch.fwd(
+                images=parameters.psi[ i, ... ],
+                positions=parameters.scan[lo:hi],
+                patch_width=parameters.probe.shape[-1],
+            )
+
+            probe_update_denominator[ i, ...] += _patch_amp_sum(patches)
+        
+
+        assert probe_update_denominator.ndim == 3
+
+        # patches = operator.diffraction.patch.fwd(
+        #     images=parameters.psi[0],
+        #     positions=parameters.scan[lo:hi],
+        #     patch_width=parameters.probe.shape[-1],
+        # )
+
+        # probe_update_denominator[...] += _patch_amp_sum(patches)
+        
+        # assert probe_update_denominator.ndim == 2
 
     tike.communicators.stream.stream_and_modify2(
         f=make_certain_args_constant,