Write a multislice and s_matrix ptychography program

CDTools/tools/cmath.py

@@ -304,3 +304,24 @@ def expi(x):
 
     """
     return t.stack((t.cos(x),t.sin(x)),dim=-1)
+
+
+def cexpi(z):
+    """Returns a complex-format tensor for exp(i* (z))
+    
+    Expects the input to be in the form of a complex-valued tensor
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        An array to be exponentiated
+    
+    Returns
+    -------
+    torch.Tensor
+        A complex-format tensor
+
+    """
+    real = t.cos(z[...,0]) * t.exp(-z[...,1])
+    imag = t.sin(z[...,0]) * t.exp(-z[...,1])
+    return t.stack((real, imag),dim=-1)

CDTools/tools/interactions.py

@@ -423,4 +423,97 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
         return t.stack(exit_waves)
 
 
+def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, padding=10):
+    """Returns a stack of exit waves accounting for subpixel shifts
+ 
+    This function returns a collection of exit waves, with the first
+    dimension as the translation index and the final dimensions
+    corresponding to the detector. The exit waves are calculated by
+    shifting the probe with each translation in turn, using sinc
+    interpolation (done via multiplication with a complex exponential
+    in Fourier space)
+
+    If shift_probe is True, it applies the subpixel shift to the probe,
+    otherwise the subpixel shift is applied to the object
+
+    This needs to be edited to use a slightly different meaning of the s-wave
+    format. Currently, each pixel in the latter two dimensions index a location
+    on the input wavefield, and the first two indexes index differences from
+    that pixel. It is easier to interpret the resulting matrix though if the
+    latter two indices index locations in the output plane.
+
+    Parameters
+    ----------
+    probe : torch.Tensor
+        An MxL probe function for the exit waves
+    s_matrix : torch.Tensor
+        The 4D S-Matrix tensor (2B+1x2B+1xObject Shape) to be probed
+    translations : torch.Tensor
+        The Nx2 array of translations to simulate
+    shift_probe : bool
+        Default True, Whether to subpixel shift the probe or object
+    padding : int
+         Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects
+
+    Returns
+    -------
+    exit_waves : torch.Tensor
+        An NxMxL tensor of the calculated exit waves
+    """
+    single_translation = False
+    if translations.dim() == 1:
+        translations = translations[None,:]
+        single_translation = True
+        
+    # Separate the translations into a part that chooses the window
+    # And a part that defines the windowing function
+    integer_translations = t.floor(translations)
+    subpixel_translations = translations - integer_translations
+    integer_translations = integer_translations.to(dtype=t.int32)
+    
+    exit_waves = []
+
+    B = s_matrix.shape[0]//2
+    
+    if shift_probe:
+        i = t.arange(probe.shape[0]) - probe.shape[0]//2
+        j = t.arange(probe.shape[1]) - probe.shape[1]//2
+        I,J = t.meshgrid(i,j)
+        I = 2 * np.pi * I.to(t.float32) / probe.shape[0]
+        J = 2 * np.pi * J.to(t.float32) / probe.shape[1]
+        I = I.to(dtype=probe.dtype,device=probe.device)
+        J = J.to(dtype=probe.dtype,device=probe.device)
+        
+        for tr, sp in zip(integer_translations,
+                          subpixel_translations):
+            fft_probe = fftshift(t.fft(probe, 2))
+            shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J))
+            shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
+
+            s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0],
+                                      tr[1]:tr[1]+probe.shape[1]]
+
+
+            output = t.zeros([s_matrix_slice.shape[2]+2*B,
+                              s_matrix_slice.shape[3]+2*B,2]).to(
+                                  device=s_matrix_slice.device,
+                                  dtype=s_matrix_slice.dtype)
+            
+            for i in range(s_matrix.shape[0]):
+                for j in range(s_matrix.shape[1]):
+                    output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
+                        cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
+
+            exit_waves.append(output)
+            #exit_waves.append(cmult(shifted_probe, obj_slice))
+        
+    else:
+        raise NotImplementedError('Object shift not yet implemented')
+
+    if single_translation:
+        return exit_waves[0]
+    else:
+        return t.stack(exit_waves)
+
+
 

CDTools/tools/losses.py

@@ -100,7 +100,7 @@ def intensity_mse(intensities, sim_intensities, mask=None):
 
 
     
-def poisson_nll(intensities, sim_intensities, mask=None):
+def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-4):
     """ Returns the Poisson negative log likelihood for a simulated dataset's intensities
 
     Calculates the overall Poisson maximum likelihood metric using
@@ -133,12 +133,12 @@ def poisson_nll(intensities, sim_intensities, mask=None):
 
     """
     if mask is None:
-        return t.sum(sim_intensities -
-                     intensities * t.log(sim_intensities)) \
+        return t.sum(sim_intensities+eps -
+                     intensities * t.log(sim_intensities+eps)) \
                      / intensities.view(-1).shape[0]
     
     else:
         masked_intensities = intensities.masked_select(mask)
         masked_sims = sim_intensities.masked_select(mask)
         return t.sum(masked_sims - masked_intensities *
-                     t.log(masked_sims)) / masked_intensities.shape[0]
+                     t.log(masked_sims+eps)) / masked_intensities.shape[0]

CDTools/tools/plotting.py

@@ -127,6 +127,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
     if basis is not None:
         if isinstance(basis,t.Tensor):
             basis = basis.detach().cpu().numpy()
+        # This fails if the 
         basis_norm = np.linalg.norm(basis, axis = 0)
         basis_norm = basis_norm * get_units_factor(units)
 
@@ -207,7 +208,7 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
         except:
             plt.imshow(phase, cmap = 'hsv', extent=extent)
     else:
-        plt.imshow(phase)#, cmap = cmap, extent=extent)
+        plt.imshow(phase, cmap = cmap, extent=extent)
 
     cbar = plt.colorbar()
     cbar.set_label('Phase (rad)')
@@ -222,6 +223,9 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
     return fig
 
 
+def plot_amplitude_surfacenorm():
+    pass
+
 def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
     """ Plots the colorized version of a complex array with dimensions NxM
 

CDTools/tools/propagators.py

@@ -74,7 +74,7 @@ def inverse_far_field(wavefront):
     return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
 
 
-def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, **kwargs):
+def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, lens=False, **kwargs):
     """Generates k-space and intensity maps to allow for high-NA far-field propagation of light
 
     At high numerical apertures or for very tilted samples, the simple
@@ -97,6 +97,16 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     The intensity map is simply an object, the shape of the detector, which
     encodes intensity corrections between 0 and 1 per pixel.
 
+    If the optional "lens" parameter is set to True, the intensity map will
+    be set to a uniform map, and the distortion of Fourier space due to the
+    flat nature of the detector (that is, the portion of the distortion
+    that exists even if the sample is not tilted) will be disabled. This is
+    to account for the fact that a good, infinity-conjugate imaging lens
+    will do it's best to correct for these abberations in the lens. Of course,
+    the lens will not be perfect, but in such a case it is a better
+    approximation to assume that the lens is perfect than to assume that it
+    is not there at all.
+
     Parameters
     ----------
     sample_basis: array
@@ -109,7 +119,8 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
         The sample-to-detector distance
     wavelength: float
         The wavelength of light being propagated
-
+    lens: bool
+        Whether the diffraction pattern is formed by a lens or not.
 
     Returns
     -------
@@ -153,15 +164,39 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     samp_det_vec = np.cross(det_basis[:,0],det_basis[:,1])
     samp_det_vec *= distance / np.linalg.norm(samp_det_vec)
 
-    Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \
-        + samp_det_vec[:,None,None]
+    # This could potentially correct for a mistake in the implied
+    # propagation direction (e.g. choosing e^ikx instead of e^-ikx)
+    #samp_det_vec *= -1
 
+    if lens == False:
+        # This correctly reproduces the sample-to-each-pixel vectors
+        # in the case where the diffraction pattern is actually formed
+        # by Fraunhoffer diffraction
+        Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \
+            + samp_det_vec[:,None,None]
+    else:
+        # This forms a distorted set of vectors designed to produce the
+        # correct Fourier space map in the case where an imaging lens is
+        # used in the 2f geometry. One should not read too much meaning
+        # into these vectors, they are simply set up to produce the
+        # correct final K-map
+        Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1)#
+        Rs += (samp_det_vec / np.linalg.norm(samp_det_vec))[:,None,None] * \
+            np.sqrt(np.sum((samp_det_vec)**2)-np.sum(Rs**2,axis=0))[None,:,:]
+        
     k0 = 2*np.pi/wavelength
     
     Ks = k0 * Rs / np.linalg.norm(Rs, axis=0)
 
+    # My attempt at seeing what happens if I flip the Ks
+    #Ks *= -1
+    
     # This is the cosine of the angle with the detector normal
     intensity_map = np.tensordot(samp_det_vec/(k0*distance),Ks,axes=1)
+    if lens:
+        # Set the intensity map to be uniform if a lens is being used
+        intensity_map = np.ones_like(intensity_map)
+    
     intensity_map = t.Tensor(intensity_map).to(*args, **kwargs)
 
 
@@ -171,6 +206,10 @@ def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance,
     # uniform phase.
     Ks -= k0 * samp_det_vec[:,None,None] / distance
 
+    # A potential alternative when Ks are flipped
+    #Ks += k0 * samp_det_vec[:,None,None] / distance
+
+    
     # Now we move on to finding the conversion into k-space
     # for the sample grid. It turns out we can do this by multiplying
     # them with the real space basis (dual of the reciprocal space
@@ -233,9 +272,13 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
     low_NA_wavefield = far_field(wavefront)
     # I'm going to need to separately interpolate the real and complex parts
     # This can be done
-
+    
     k_map = k_map[None,:,:,:]
     # Will only work for a 4D wavefile stack.
+    #plt.figure()
+    #plt.pcolormesh(k_map[0,:,:,0].cpu().numpy(),k_map[0,:,:,1].cpu().numpy(),
+    #               np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy()))
+    #plt.show()
     def process_wavefield_stack(low_NA_wavefield):
         real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
         imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)