Finish building and testing the gaussian probe guess from a dataset

CDTools/tools/initializers.py

@@ -2,7 +2,8 @@ from __future__ import division, print_function, absolute_import
 import numpy as np
 import torch as t
 
-__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian']
+__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
+           'gaussian_probe']
 
 from CDTools.tools import cmath
 from scipy.fftpack import next_fast_len
@@ -10,7 +11,7 @@ import numpy as np
 
 
 def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0):
-    """Returns an exit wave basis and a detector slice for the given detector geometry
+    """Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
     
     It takes in the parameters for a given detector - the basis defining
     the pixel pitch and the shape, as well as the wavelength and propagation
@@ -137,7 +138,8 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
     return cmath.complex_to_torch(amplitude*result)
 
 
-def gaussian_initialization(dataset, basis, shape, sigma, propagation_distance=0):
+
+def gaussian_probe(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
@@ -169,32 +171,26 @@ def gaussian_initialization(dataset, basis, shape, sigma, propagation_distance=0
     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)
+    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
-    
+    curvature_pix = curvature * 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 = 
+    avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
+    avg_intensity = t.mean(t.Tensor(avg_intensities))
+    probe_intensity = t.sum(cmath.cabssq(probe))
+    return avg_intensity / probe_intensity * probe
     
 
 

tests/tools/test_initializers.py

@@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import initializers
 from CDTools.tools import cmath
+from CDTools.datasets import Ptycho_2D_Dataset
 import numpy as np
 import torch as t
 
@@ -108,3 +109,65 @@ def test_gaussian():
                         center=center, curvature=curvature, amplitude=10))
     assert np.allclose(init_result, np_result)
     
+
+def test_gaussian_probe(ptycho_cxi_1):
+    
+    dataset = Ptycho_2D_Dataset.from_cxi(ptycho_cxi_1[0])
+
+    det_basis = t.Tensor(dataset.detector_geometry['basis'])
+    det_shape = t.Size(dataset.patterns.shape[-2:])
+    wavelength = dataset.wavelength
+    distance = dataset.detector_geometry['distance']
+    
+    basis, shape, s = initializers.exit_wave_geometry(det_basis,
+                                                      det_shape,
+                                                      wavelength,
+                                                      distance)
+    # Basis is around 60nm in the i(y) direction, 85nm in the j(x) direction
+    # Full window is therefore about 15 um in i(y) and 20 um in the j(x) dir
+
+    # Come up with a roughly matching set of probe parameters
+    sigma = 5e-7
+
+    # Build a stage explicitly with numpy to compare against
+    x = (np.arange(256) - 127.5) * (-basis[0,1]).numpy()
+    y = (np.arange(256) - 127.5) * (-basis[1,0]).numpy()
+    Xs,Ys = np.meshgrid(x,y)
+    Rs = np.sqrt(Xs**2+Ys**2)
+
+    
+
+    # Now we first test the non-propagated probe
+    np_probe = np.exp(-1/(2*sigma**2) * Rs**2)
+
+    normalization = 0
+    for params, im in dataset:
+        normalization += np.sum(im.cpu().numpy())
+    normalization /= len(dataset)
+
+    normalization_1 = normalization /  np.sum(np.abs(np_probe)**2)
+    
+    probe = initializers.gaussian_probe(dataset, basis, shape, sigma)
+    probe = cmath.torch_to_complex(probe)
+    assert np.allclose(probe, normalization_1*np_probe)
+
+    # And then a propagated probe
+    z = 1e-4 #nm
+    k = 2 * np.pi / wavelength
+    w0 = np.sqrt(2)*sigma
+    zr = np.pi * w0**2 / wavelength
+    wz = w0 * np.sqrt(1 + (z / zr)**2)
+    Rz =  z * (1 + (zr / z)**2)
+    np_probe  =  np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * Rs**2 / (2 * Rz))
+
+    normalization_2 = normalization /  np.sum(np.abs(np_probe)**2)
+    
+    probe = initializers.gaussian_probe(dataset, basis, shape, sigma,
+                                        propagation_distance=z)
+    probe = cmath.torch_to_complex(probe)
+    
+    assert np.allclose(probe, normalization_2*np_probe)
+
+  
+    
+