Intermediate state in optimization, improved the round interaction but not sinc yet

CDTools/models/multislice_2d_ptycho.py

@@ -277,46 +277,55 @@ class Multislice2DPtycho(CDIModel):
         prs = tools.propagators.inverse_far_field(self.probe*self.probe_fourier_support[None,:,:])
         # Here is where the mixing would happen, if it happened
 
+        all_exit_waves = []
         for i in range(self.probe.shape[0]):
             pr = prs[i]
             
             #exit_waves =  pr
             #print(self.probe_norm)
             #for i in range(self.nz):
-            exit_waves = []
                 
-            for trans in pix_trans:
-                exit_wave = self.probe_norm * pr
-                for i in range(self.nz):
-                    # If only one object slice
-                    if self.obj.dim() == 3:
-                        exit_wave = tools.interactions.ptycho_2D_sinc(
-                            exit_wave,
-                            self.obj_support*cmath.cexpi(self.obj/self.nz),
-                            trans, shift_probe=True)
-                    # If separate slices
-                    elif self.obj.dim() == 4:
-                        exit_wave = tools.interactions.ptycho_2D_sinc(
-                            exit_wave,
-                            self.obj_support*cmath.cexpi(self.obj[i]/self.nz),
-                            trans, shift_probe=True)
-                        
-                    exit_wave = tools.propagators.near_field(
-                        exit_wave,self.as_prop)
+            exit_waves = self.probe_norm * pr
+            for i in range(self.nz):
+                # If only one object slice
+                if self.obj.dim() == 3:
                     
-                    #tools.plotting.plot_amplitude(exit_wave)
-                    #plt.show()
-                
-                exit_waves.append(exit_wave)
-            exit_waves = t.stack(exit_waves)
-            
-            if exit_waves.dim() == 4:
-                exit_waves =  self.weights[index][:,None,None,None] * exit_waves
-            else:
-                exit_waves =  self.weights[index] * exit_waves
+                    #exit_wave = tools.interactions.ptycho_2D_sinc(
+                    #    exit_wave,
+                    #    self.obj_support*cmath.cexpi(self.obj/self.nz),
+                    #    pix_trans, shift_probe=True)
+                    exit_waves = tools.interactions.ptycho_2D_round(
+                        exit_waves,
+                        self.obj_support*cmath.cexpi(self.obj/self.nz),
+                        trans)
+                    
+                elif self.obj.dim() == 4:
+                    # If separate slices
+                    #exit_wave = tools.interactions.ptycho_2D_sinc(
+                    #    exit_wave,
+                    #    self.obj_support*cmath.cexpi(self.obj[i]/self.nz),
+                    #    trans, shift_probe=True)
+                    # I see, it needs to know that if exit_wave is the
+                    # same shape as the translations, it should be broadcast
+                    # along that dimension
+                    exit_waves = tools.interactions.ptycho_2D_round(
+                        exit_waves,
+                        self.obj_support*cmath.cexpi(self.obj[i]/self.nz),
+                        pix_trans)
 
-            if strip_first_index:
-                exit_waves = exit_waves[0,...]
+                exit_waves = tools.propagators.near_field(
+                    exit_waves,self.as_prop)
+                    
+
+                if exit_waves.dim() == 4:
+                    # If the index is a list and not a single index
+                    exit_waves =  self.weights[index][:,None,None,None] * exit_waves
+                else:
+                    # If the index a single index
+                    exit_waves =  self.weights[index] * exit_waves
+
+                if strip_first_index:
+                    exit_waves = exit_waves[0,...]
                 
             all_exit_waves.append(exit_waves)
 

CDTools/tools/interactions/interactions.py

@@ -237,30 +237,33 @@ def ptycho_2D_round(probe, obj, translations):
     Parameters
     ----------
     probe : torch.Tensor
-        An MxL probe function for the exit waves
+        A (P)xMxL probe function for the exit waves
     object : torch.Tensor
         The object function to be probed
     translations : torch.Tensor
-        The Nx2 array of (i,j) translations to simulate
+        The (N)x2 array of (i,j) translations to simulate
     
     Returns
     -------
     exit_waves : torch.Tensor 
-        An NxMxL tensor of the calculated exit waves
+        An (N)x(P)xMxL tensor of the calculated exit waves
     """
+
     single_translation = False
     if translations.dim() == 1:
         translations = translations[None,:]
         single_translation = True
+
         
     integer_translations = t.round(translations).to(dtype=t.int32)
-    selections = [obj[tr[0]:tr[0]+probe.shape[0],
-                      tr[1]:tr[1]+probe.shape[1]]
-                  for tr in integer_translations]
+    selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-3],
+                              tr[1]:tr[1]+probe.shape[-2]]
+                          for tr in integer_translations])
+
     if single_translation:
-        return [cmult(probe,selection) for selection in selections][0]
+        return cmult(probe,selection)[0]
     else:
-        return t.stack([cmult(probe,selection) for selection in selections])
+        return cmult(probe,selections)
 
 
 

CDTools/tools/propagators/propagators.py

@@ -601,7 +601,7 @@ def near_field(wavefront, angular_spectrum_propagator):
     Parameters
     ----------
     wavefront : torch.Tensor
-        The JxNxMx2 stack of complex wavefronts to be propagated
+        The (Leading Dims)xNxMx2 stack of complex wavefronts to be propagated
     angular_spectrum_propagator : torch.Tensor
         The NxM phase mask to be applied during propagation