Fix some bugs in the propagators and initializers, write the tests for the angular spectrum propagator

CDTools/tools/initializers.py

@@ -4,26 +4,32 @@ import torch as t
 
 all = ['gaussian']
 
+from CDTools.tools import cmath
+
 
 def gaussian(shape, amplitude, sigma, center = None):
     """Returns an array with a centered gaussian
 
     Takes in the shape, amplitude, and standard deviation of a gaussian
-    and returns an array with values corresponding to a two-dimensional gaussian function
-        z = amplitude*exp(-(x-center[0])**2/sigma[0]**2+(y-center[1])**2/sigma[1]**2)
+    and returns a torch tensor 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.
 
     Args:
-        shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (y shape, x shape)
+        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 x- and y- standard deviation of the gaussian in the form (y stdev, y stdev)
-        center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (y center, x center)
+        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)
 
     Returns:
-        numpy.array : The real-valued gaussian array
+        torch.Tensor : The real-valued gaussian array
     """
     if center is None:
         center = ((shape[0]-1)/2, (shape[1]-1)/2)
-    y, x = np.mgrid[:shape[0], :shape[1]]
-    return amplitude*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**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)

CDTools/tools/propagators.py

@@ -1,8 +1,13 @@
 from __future__ import division, print_function, absolute_import
+
 from CDTools.tools.cmath import *
 import torch as t
+from scipy import fftpack
+import numpy as np
 
-__all__ = ['far_field', 'near_field', 'inverse_far_field', 'inverse_near_field']
+__all__ = ['far_field', 'near_field',
+           'generate_angular_spectrum_propagator',
+           'inverse_far_field', 'inverse_near_field']
 
 
 def far_field(wavefront):
@@ -30,6 +35,7 @@ def far_field(wavefront):
 
     return fftshift(t.fft(wavefront, 2))
 
+
 def inverse_far_field(wavefront):
     """Implements the inverse of the far-field propagator in torch
 
@@ -51,7 +57,7 @@ def inverse_far_field(wavefront):
     return t.ifft(ifftshift(wavefront), 2)
 
 
-def generate_angular_spectrum_propagator(shape, spacing, wavelength, z):
+def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **kwargs):
     """Generates an angular-spectrum based near-field propagator from experimental quantities
 
     This function generates an angular-spectrum based near field
@@ -61,56 +67,81 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z):
     propagator is used in a reconstruction program, then it will be best
     to calculate this mask once and close over it.
 
+    Formally, this propagator is the complex conjugate of the fourier
+    transform of the convolution kernel for light propagation in free
+    space
+
     Args:
         shape (iterable) : The shape of the arrays to be propagated
         spacing (iterable) : The pixel size in each dimension of the arrays to be propagated
         wavelength (float) : The wavelength of light to simulate propagation of
         z (float) : The distance to simulate propagation over
+
     Returns:
-        torch.Tensor : A propagation term which accounts for the phase change that each plane wave will undergo on its journey to the prediction plane.
+        torch.Tensor : A phase mask which accounts for the phase change that each plane wave will undergo.
     """
 
-    ki = fftpack.fftfreq(shape[0],spacing[0])
-    kj = fftpack.fftfreq(shape[1],spacing[1])
-    Ki, Kj = np.meshgrid(ki,kj)
-    propagator = np.exp(1j*np.sqrt((2*np.pi/wavelength)**2
-                                - Ki**2 - Kj**2) * z)
-    propagator = complex_to_float(propagator).astype(np.float32)
-    propagator = t.from_numpy(propagator).cuda()
+    ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0])
+    kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1])
+    Kj, Ki = np.meshgrid(kj,ki)
 
-    return propagator
+    # Define this as complex so the square root properly gives
+    # k>k0 components imaginary frequencies    
+    k0 = np.complex128((2*np.pi/wavelength))
+    
+    propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
+
+    # Take the conjugate explicitly here instead of negating
+    # the previous expression to ensure that complex frequencies
+    # get mapped to values <1 instead of >1
+    propagator = complex_to_torch(np.conj(propagator)) 
+
+    return propagator.to(*args, **kwargs)
 
 
 def near_field(wavefront, angular_spectrum_propagator):
-    """This function accepts an 3d torch tensor, where the
-    last dimension represents the real and imaginary components of
-    the wavefield, and returns the near-field propagated version of it.
+    """ Propagates a wavefront via the angular spectrum method
+
+    This function accepts an 3D torch tensor, where the last dimension
+    represents the real and imaginary components of the wavefield, and
+    returns the near-field propagated version of it. It does this
+    using the supplied angular spectrum propagator, which is a premade
+    phase mask.
 
 
     Args:
-        angular_spectrum_propagator (torch.Tensor) : The near field propagator
         wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
+        angular_spectrum_propagator (torch.Tensor) : The NxM phase mask to be applied during propagation
+
     Returns:
-        function : The wavefront propagated to the near field
+        torch.Tensor : The propagated wavefront 
     """
 
-    return t.ifft(angular_spectrum_propagator * t.fft(wavefront,2), 2)
+    return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
 
 
 
 def inverse_near_field(wavefront, angular_spectrum_propagator):
-    """This function accepts a 3d torch tensor, where the
-    last dimension represents the real and imaginary components of
-    the near-field propagated wavefield, and returns the exit wavefront via an inverse transformation.
+    """ Inverse ropagates a wavefront via the angular spectrum method
+
+    This function accepts an 3D torch tensor, where the last dimension
+    represents the real and imaginary components of the wavefield, and
+    returns the near-field propagated version of it. It does this
+    using the supplied angular spectrum propagator, which is a premade
+    phase mask.
+
+    It propagates the wave using the conjugate of the supplied phase mask,
+    which corresponds to the inverse propagation problem.
 
 
     Args:
-        angular_spectrum_propagator (torch.Tensor) : The pixel size in each dimension of the arrays to be propagated
         wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
+        angular_spectrum_propagator (torch.Tensor) : The NxM phase mask to be applied during propagation
+
     Returns:
-        function : A function to propagate a torch tensor.
+        torch.Tensor : The inverse propagated wavefront
     """
-    return t.ifft(t.fft(wavefront,2) * angular_spectrum_propagator**-1, 2)
+    return t.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2)
 
 
 

tests/tools/test_initializers.py

@@ -1,6 +1,7 @@
 from __future__ import division, print_function, absolute_import
 
 from CDTools.tools import initializers
+from CDTools.tools import cmath
 import numpy as np
 import torch as t
 
@@ -12,13 +13,17 @@ def test_gaussian():
     sigma = [2.5, 2.5]
     center = ((shape[0]-1)/2, (shape[1]-1)/2)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
-    assert(np.allclose(initializers.gaussian([10, 10], 10, [2.5, 2.5]), np_result))
+    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]))
+    assert np.allclose(init_result, np_result)
 
     # Generate gaussian as a numpy array (rectangular array)
     shape = [10, 5]
     sigma = [2.5, 2.5]
     center = ((shape[0]-1)/2, (shape[1]-1)/2)
     y, x = np.mgrid[:shape[0], :shape[1]]
-    np_result = 10*np.exp(-((x-center[1])/sigma[1])**2-((y-center[0])/sigma[0])**2)
-    assert(np.allclose(initializers.gaussian([10, 5], 10, [2.5, 2.5]), np_result))
+    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, 5], 10, [2.5, 2.5]))
+    assert np.allclose(init_result, np_result)

tests/tools/test_projectors.py

@@ -4,7 +4,7 @@ from CDTools.tools import cmath
 from CDTools.tools import projectors
 import numpy as np
 import torch as t
-from scipy.fftpack import fftshift, ifftshift
+
 
 def test_modulus():
     # Create a complex array with random modulus and known phase

tests/tools/test_propagators.py

@@ -14,14 +14,16 @@ from scipy.fftpack import fftshift, ifftshift
 @pytest.fixture(scope='module')
 def exit_waves_1():
     # Import scipy test image and add a random phase
-    object = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
+    obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
     arr = np.random.random_sample((64,64))
-    object *= (arr+(1-arr**2)**.5*1j)
+    obj *= (arr+(1-arr**2)**.5*1j)
+    obj = cmath.complex_to_torch(obj)
 
     # Construct wavefront from image
-    probe = initializers.gaussian([64, 64], 1e3, [5, 5])*(1+1j)
-    return cmath.complex_to_torch(probe*object)
+    probe = initializers.gaussian([64, 64], 1e3, [5, 5])
+    return cmath.cmult(probe,obj)
 
+    
 
 def test_far_field(exit_waves_1):
     # Far field diffraction patterns calculated by numpy with zero frequency in center
@@ -30,16 +32,56 @@ def test_far_field(exit_waves_1):
     assert(np.allclose(np_result, cmath.torch_to_complex(propagators.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
-    np_result = exit_waves_1
     # 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))))
 
-    assert(np.allclose(np_result, propagators.inverse_far_field(far_field_np_result)))
+    assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
 
 
 
-def test_near_field(exit_waves_1):
-    pass
+def test_near_field():
 
+    # The strategy is to compare the propagation of a gaussian beam to
+    # the propagation in the paraxial approximation.
+    
+    x = (np.arange(800) - 400) * 1.5e-9
+    y = (np.arange(1200) - 600) * 1e-9
+    Ys,Xs = np.meshgrid(y,x)
+    Rs = np.sqrt(Xs**2+Ys**2)
+
+    wavelength = 3e-9 #nm
+    sigma = 20e-9 #nm
+    z = 1000e-9 #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)
+    
+    E0 = np.exp(-Rs**2 / w0**2)
+
+    # The analytical expression for propagation of a gaussian beam in the
+    # paraxial approx
+    Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * ( z + Rs**2 / (2 * Rz)) + 1j * np.arctan(z / zr))
+
+    asp = propagators.generate_angular_spectrum_propagator(
+        E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64)
+    
+    Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
+    Ez_t = cmath.torch_to_complex(Ez_t)
+    
+    # Check for at least 10^-3 relative accuracy in this scenario
+    assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
+
+
+    Emz = np.conj(Ez)
+
+    Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
+    Emz_t = cmath.torch_to_complex(Emz_t)    
+
+    # Again, 10^-3 is about all the accuracy we can expect
+    assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))