diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index dbe04e9..fbc6ba0 100644 --- a/CDTools/tools/propagators.py +++ b/CDTools/tools/propagators.py @@ -72,7 +72,7 @@ def inverse_far_field(wavefront): return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True)) -def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **kwargs): +def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, **kwargs): """Generates an angular-spectrum based near-field propagator from experimental quantities This function generates an angular-spectrum based near field @@ -96,6 +96,8 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, * The wavelength of light to simulate propagation of z : float The distance to simulate propagation over + remove_z_phase : bool + Default False, whether to remove the dominant z-direction phase dependence Returns ------- @@ -113,6 +115,9 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, * propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z) + if remove_z_phase: + propagator *= np.exp(-1j * k0 * 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 @@ -139,12 +144,15 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o is designed to work on any wavefield defined on an array of parallelograms. - In addition, if propagate_along_offset is true, there is an assumed phase + In addition, if propagate_along_offset is True, there is an assumed phase ramp applied to the wavefield before propagation, defined such that a feature with uniform phase will propagate along the direction of the - defined offset vector. This decision provides the best numerical - stability and allows for the simple setup of light fields copropagating - with the coordinate system. + defined offset vector. This will also remove the phase variation along + the propagation direction, because it makes the most physical sense to + regard this choice as removing the dominant phase variation in 3D, allowing + for the generation of a smoothly varying wavefield over 3D volumes. + This decision provides the best numerical stability and allows for the + simple setup of light fields copropagating with the coordinate system. Parameters @@ -157,6 +165,8 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o The wavelength of light to simulate propagation of propagation_vector : array The displacement to propagate the wavefield along. + propagate_along_offset : bool + Optional, whether to include an implied phase ramp to propagate uniform phase features along the offset direction Returns ------- @@ -186,7 +196,6 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o # This may have a sign error - must be checked phase_mask = np.exp(1j * np.tensordot(offset_vector,K_xyz,axes=1)) - # Next, we apply a shift to the k-space vectors which sets up # propagation such that a uniform phase object will propagate along the @@ -198,10 +207,22 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o perpendicular_dir /= np.linalg.norm(perpendicular_dir) offset_perpendicular = np.dot(perpendicular_dir, offset_vector) offset_parallel = offset_vector - perpendicular_dir * offset_perpendicular + print(offset_perpendicular) + print(offset_parallel) k0 = 2*np.pi/wavelength k_offset = offset_parallel * k0 / np.sqrt(offset_perpendicular**2 + np.linalg.norm(offset_parallel)**2) - K_xyz = K_xyz - k_offset[:,None,None] + + # Only implement the shift if the flag is set to True + if propagate_along_offset: + K_xyz = K_xyz + k_offset[:,None,None] + + # we also need to remove the z-dependence on the phase + # This time, though, the z-dependence actually has to do with + # the oput of plane component of k at the central offset. Normally + # this is 0, so the z-component is just k0, but not in this case + phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2) + * offset_perpendicular) # Redefine this as complex so the square root properly gives # k>k0 components imaginary frequencies @@ -212,6 +233,11 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o * offset_perpendicular) propagator *= phase_mask + # This removes the z-dependence on the phase: + #if propagate_along_offset: + # print(offset_perpendicular) + # propagator *= np.exp(-1j * k0 * offset_perpendicular) + # Take the conjugate explicitly here instead of negating # the previous expression to ensure that complex frequencies # get mapped to values <1 instead of >1 diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index 5707979..a6c3f7e 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -9,7 +9,7 @@ import torch as t import pytest import scipy.misc from scipy.fftpack import fftshift, ifftshift - +from matplotlib import pyplot as plt @pytest.fixture(scope='module') @@ -68,9 +68,13 @@ def test_near_field(): # 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)) + Ez_nozphase = Ez * np.exp(1j * k * z) + + # First we check it normally 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) @@ -87,52 +91,137 @@ def test_near_field(): # 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)) + # Then, we check it with the phase correction + asp = propagators.generate_angular_spectrum_propagator( + E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True, + 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_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase)) + + + Emz = np.conj(Ez_nozphase) + + 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)) + def test_generalized_near_field(): # The strategy is to compare the propagation of a gaussian beam to # the propagation in the paraxial approximation. - + + # For this one, we want to test it on a rotated coordinate system + # First, we should do a test with the phase ramp along the z direction + # explicitly included + + basis= np.array([[0,-1.5e-9],[-1e-9,0],[0,0]]) x = (np.arange(901) - 450) * 1.5e-9 y = (np.arange(1200) - 600) * 1e-9 - Ys,Xs = np.meshgrid(y,x) - Rs = np.sqrt(Xs**2+Ys**2) + Xs_0,Ys_0 = np.meshgrid(x,y) + Zs_0 = np.zeros(Xs_0.shape) + Positions = np.stack([Xs_0,Ys_0,Zs_0]) + + # assert 0 wavelength = 3e-9 #nm sigma = 20e-9 #nm z = 1000e-9 #nm + propagation_vector = np.array([0,0,z]) 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)) + def get_w(Zs): + return w0 * np.sqrt(1 + (Zs / zr)**2) - basis= np.array([[0,-1e-9],[-1.5e-9,0],[0,0]]) - propagation_vector = np.array([0,0,z]) - asp = propagators.generate_generalized_angular_spectrum_propagator( - E0.shape,basis,wavelength,propagation_vector,dtype=t.float64) - #assert False - Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) - Ez_t = cmath.torch_to_complex(Ez_t) + def get_inv_R(Zs): + return Zs / (Zs**2 + zr**2) - # 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)) + def get_E(Xs, Ys, Zs, correct=False): + # if correct is True, remove the e^(-ikz) dependence + Rs_sq = Xs**2 + Ys**2 + Wzs = get_w(Zs) + E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\ + np.exp(-1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \ + 1j * np.arctan(Zs / zr)) + if correct: + E = E * np.exp(1j * k * Zs) + return E + + + # This tests the straight ahead case + I = np.eye(3) + + # This tests a rotation about the y axis + th = np.deg2rad(5) + Ry = np.array([[np.cos(th),0,np.sin(th)], + [0,1,0], + [-np.sin(th),0,np.cos(th)]]) + + # This tests a rotation about two axes + phi = np.deg2rad(2) + Rx = np.array([[1,0,0], + [0,np.cos(phi),-np.sin(phi)], + [0,np.sin(phi),np.cos(phi)]]) + Rboth = np.matmul(Rx,Ry) + + # This tests a shearing + shear = 0.23 + Rshear = np.array([[1,shear,0], + [0,1,0], + [0,0,1]]) + + # This tests a shearing and a rotation together + Rall = np.matmul(Rboth,Rshear) + + rot_mats = [I,Ry, Rboth, Rshear, Rboth] + purposes = ['standard','y-rot','both-rot','shear','shear-rot'] + + for purpose,rot_mat in zip(purposes,rot_mats): + print('Testing', purpose) + Xs,Ys,Zs_0 = np.tensordot(rot_mat,Positions,axes=1) + new_basis = np.dot(rot_mat, basis) + Zs_prop = Zs_0 + z + + # Check that it works both with the explicit and implicit phase ramps + for prop_oo in [False, True]: + print('Propagate Along Offset =',prop_oo) + + E0 = get_E(Xs,Ys,Zs_0, correct=prop_oo) + Ez = get_E(Xs,Ys,Zs_prop, correct=prop_oo) + + asp = propagators.generate_generalized_angular_spectrum_propagator( + E0.shape,new_basis,wavelength,propagation_vector, + dtype=t.float64, propagate_along_offset=prop_oo) - Emz = np.conj(Ez) + Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp) + Ez_t = cmath.torch_to_complex(Ez_t) - Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp) - Emz_t = cmath.torch_to_complex(Emz_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)) + + + Em0_t = propagators.inverse_near_field(cmath.complex_to_torch(Ez),asp) + Em0_t = cmath.torch_to_complex(Em0_t) + + # Again, 10^-3 is about all the accuracy we can expect + assert np.max(np.abs(E0-Em0_t)) < 1e-3 * np.max(np.abs(E0)) + + print('Test Successful') - # 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)) def test_inverse_near_field():