Change propagation to preserve norm despite the extra cost, and generalize the gaussian initializer

CDTools/tools/initializers.py

@@ -67,7 +67,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
     full_shape = t.Size([dim for dim in full_shape])
     
     return real_space_basis, full_shape, det_slice
-                      
+
 
 def calc_object_setup(probe_shape, translations, padding=0):
     """Returns an object shape and minimum pixel translation
@@ -102,29 +102,101 @@ def calc_object_setup(probe_shape, translations, padding=0):
     
     
 
-def gaussian(shape, amplitude, sigma, center = None):
-    """Returns an array with a centered gaussian
+def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
+    """Returns an array with a centered Gaussian
 
     Takes in the shape, amplitude, and standard deviation of a gaussian
-    and returns a torch tensor with values corresponding to a two-dimensional
-    gaussian function
+    and returns a complex torch tensor (trailing dimension is 2) with
+    values corresponding to a two-dimensional gaussian function
 
     Note that [0, 0] is taken to be at the upper left corner of the array.
     Default is centered at ((shape[0]-1)/2, (shape[1]-1)/2)) because x and y are zero-indexed.
+    By default, the phase is uniformly 0, however a curvature can be
+    specified to simulate a probe that has been propagated a known distance
+    from it's focal point. The curvature is implemented by adding a quadratic
+    phase phi = exp(i*curvature/2 r^2) to the Gaussian
 
     Args:
         shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (i shape, j shape)
         amplitude (float or int): The amplitude the gaussian to simulate
         sigma (array_like): A 1x2 array-like object specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
         center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (i center, j center)
+        curvature (array_like) : Optional complex part to add to the gaussian coefficient
 
     Returns:
-        torch.Tensor : The real-valued gaussian array
+        torch.Tensor : The complex-style tensor storing the Gaussian
     """
     if center is None:
         center = ((shape[0]-1)/2, (shape[1]-1)/2)
         
     i, j = np.mgrid[:shape[0], :shape[1]]
-    result = amplitude*np.exp(-( (i-center[0])**2 / (2 * sigma[0]**2) )
-                              -( (j-center[1])**2 / (2 * sigma[1]**2) ))
-    return cmath.complex_to_torch(result)
+    isq = (i - center[0])**2
+    jsq = (j - center[1])**2
+    result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \
+        (1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq)
+    return cmath.complex_to_torch(amplitude*result)
+
+
+def gaussian_initialization(dataset, basis, shape, sigma, propagation_distance=0):
+    """Initializes a gaussian probe based on experimental parameters
+
+    This function generates a gaussian probe initialization which has a
+    total fluence matching the order of magnitude of the intensity in
+    the observed dataset, provided the object function is of order 1.
+    
+    The initialization is done using parameters defined in physical units,
+    such as sigma (in meters) and the propagation distance (in meters).
+    The internal conversion to pixel space is done with a provided probe
+    basis and probe shape.
+    
+    Sigma can be provided either as a scalar for a uniform beam, or as 
+    an iterable of length 2 with [sigma_i, sigma_j] being the components
+    of sigma in the directions parallel to the i and j basis vectors of
+    the probe basis
+    
+    Args:
+        dataset (Ptycho_2D_Dataset) : The dataset whose intensity we want to match
+        basis (array_like) : The real space basis for exit waves in our experiment
+        shape (array_like): The shape of the simulated real space arrays
+        sigma (array_like): The standard deviation of the probe at it's focus
+        propagation_distance (float) : Optional, a distance to propagate the gaussian from it's focus
+    
+    Returns:
+        torch.Tensor : The complex-style tensor storing the Gaussian
+    """
+    # First, we want to generate the parameters (sigma and curvature) for the
+    # propagated gaussian. Ignore the purely z-dependent phases
+    wavelength = dataset.wavelength
+    z = propagation_distance # for shorthand
+    sigma = np.array(sigma)
+    curvature = np.array(curvature) 
+    k = 2 * np.pi / wavelength
+    zr = k * sigma**2
+    sigmaz = sigma * np.sqrt(1 + (z / zr)**2)
+    curvature = k * z / (z**2 + zr**2)
+
+    # So both sigmaz and curvature can be either scalars or tensors here
+    # We make them consistent
+    if len(sigmaz.shape) == 0:
+        sigmaz = np.array([sigmaz, sigmaz])
+    if len(curvature.shape) == 0:
+        curvature = np.array([curvature, curvature])
+    
+    # The conversion must then be done to pixel space
+    sigma_pix = sigmaz / np.array([np.linalg.norm(basis[:,0]),
+                                   np.linalg.norm(basis[:,1])])
+    curvature_pix = sigmaz * np.array([np.linalg.norm(basis[:,0]),
+                                       np.linalg.norm(basis[:,1])])**2
+    
+    # Then we can generate the gaussian array
+    probe = gaussian(shape, sigma=sigma_pix, curvature=curvature_pix)
+        
+    # Finally, we should calculate the average pattern intensity from the
+    # dataset and normalize the gaussian probe. This should be done by
+    avg_intensity = 1
+    #probe_intensity = 
+    
+
+
+
+    

CDTools/tools/propagators.py

@@ -33,7 +33,7 @@ def far_field(wavefront):
         torch.Tensor : The JxNxMx2 propagated wavefield
     """
 
-    return fftshift(t.fft(wavefront, 2))
+    return fftshift(t.fft(wavefront, 2, normalized=True))
 
 
 def inverse_far_field(wavefront):
@@ -54,7 +54,7 @@ def inverse_far_field(wavefront):
     Returns:
         torch.Tensor : The JxNxMx2 exit wavefield
     """
-    return t.ifft(ifftshift(wavefront), 2)
+    return t.ifft(ifftshift(wavefront), 2, normalized=True)
 
 
 def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **kwargs):

tests/tools/test_initializers.py

@@ -81,15 +81,30 @@ def test_gaussian():
     y, x = np.mgrid[:shape[0], :shape[1]]
     np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
                           -0.5*((y-center[0])/sigma[0])**2)
-    init_result = cmath.torch_to_complex(initializers.gaussian([10, 10], 10, [2.5, 2.5]))
+    init_result = cmath.torch_to_complex(initializers.gaussian([10, 10], [2.5, 2.5], amplitude=10))
     assert np.allclose(init_result, np_result)
 
     # Generate gaussian as a numpy array (rectangular array)
     shape = [10, 5]
-    sigma = [2.5, 2.5]
+    sigma = [2.5, 3]
     center = ((shape[0]-1)/2, (shape[1]-1)/2)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
+    np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
                           -0.5*((y-center[0])/sigma[0])**2)
-    init_result = cmath.torch_to_complex(initializers.gaussian([10, 5], 10, [2.5, 2.5]))
+    init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma))
     assert np.allclose(init_result, np_result)
+                                    
+    # Generate gaussian with curvature
+    shape = [20, 30]
+    sigma = [2.5, 5]
+    curvature = [1,0.6]
+    center = ((shape[0]-1)/2 + 3, (shape[1]-1)/2 - 1.4)
+    y, x = np.mgrid[:shape[0], :shape[1]]
+    np_result = (10+0j)*np.exp(-0.5*((x-center[1])/sigma[1])**2
+                               -0.5*((y-center[0])/sigma[0])**2)
+    np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
+                        +0.5j*curvature[0]*(y-center[0])**2)
+    init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma,
+                        center=center, curvature=curvature, amplitude=10))
+    assert np.allclose(init_result, np_result)
+    

tests/tools/test_propagators.py

@@ -21,14 +21,14 @@ def exit_waves_1():
     obj = cmath.complex_to_torch(obj)
 
     # Construct wavefront from image
-    probe = initializers.gaussian([64, 64], 1e3, [5, 5])
+    probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
     return cmath.cmult(probe,obj)
 
     
 
 def test_far_field(exit_waves_1):
     # Far field diffraction patterns calculated by numpy with zero frequency in center
-    np_result = np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1)))
+    np_result = np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1),norm='ortho'))
 
     assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
 
@@ -37,7 +37,7 @@ def test_far_field(exit_waves_1):
 def test_inverse_far_field(exit_waves_1):
     # We want the inverse far field to map back to the exit waves with no intensity corrections
     # Far field result for exit waves calculated with numpy
-    far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1))))
+    far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1),norm='ortho')))
 
     assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))