From 1e43c095d16ff6daf5622b8effeb1b064eaea01c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 28 Mar 2019 19:30:40 -0400 Subject: [PATCH] Fix some bugs in the propagators and initializers, write the tests for the angular spectrum propagator --- CDTools/tools/initializers.py | 22 +++++---- CDTools/tools/propagators.py | 77 ++++++++++++++++++++++---------- tests/tools/test_initializers.py | 13 ++++-- tests/tools/test_projectors.py | 2 +- tests/tools/test_propagators.py | 58 ++++++++++++++++++++---- 5 files changed, 128 insertions(+), 44 deletions(-) diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index dbe6361..fbb51f0 100644 --- a/CDTools/tools/initializers.py +++ b/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) diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index 57c5e51..d3eed3c 100644 --- a/CDTools/tools/propagators.py +++ b/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) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index d4dab18..2f91fac 100644 --- a/tests/tools/test_initializers.py +++ b/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) diff --git a/tests/tools/test_projectors.py b/tests/tools/test_projectors.py index 2b3da30..060b70a 100644 --- a/tests/tools/test_projectors.py +++ b/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 diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index 3e15c80..1bfbba6 100644 --- a/tests/tools/test_propagators.py +++ b/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))