Finish first pass of off-axis propagation, still needs to be tested off-axis

CDTools/tools/propagators.py

@@ -121,7 +121,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, *
     return propagator.to(*args, **kwargs)
 
 
-def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, propagation_vector, *args, **kwargs):
+def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset_vector, *args, propagate_along_offset=True, **kwargs):
     """Generates an angular-spectrum based near-field propagator from experimental quantities
 
     This function generates an angular-spectrum based near field
@@ -133,14 +133,17 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, p
 
     Formally, this propagator is the complex conjugate of the fourier
     transform of the convolution kernel for light propagation in free
-    space
+    space. It will map a ligh field at an input plane, with the size
+    and shape defined by the shape and basis inputs, and map it to a
+    plane of the same size and shape offset by the offset vector. It
+    is designed to work on any wavefield defined on an array of
+    parallelograms.
 
-    This function is written to work on any wavefield defined on any
-    array of parallelograms. In addition, 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 propagation vector. This decision provides the best numerical
-    stabilit and allows for the simple setup of light fields copropagating
+    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.
 
 
@@ -162,22 +165,53 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, p
     """
 
     
-    sides = basis * shape
-    fourier_basis = 2 * np.pi * np.linalg.pinv(sides).transpose()
-    print(fourier_basis)
+    # First we calculate a dual basis for the real space grid
+    inv_basis =  np.linalg.pinv(basis).transpose()
 
-    ki = 2 * np.pi * fftpack.fftfreq(shape[0],basis[1,0])
-    kj = 2 * np.pi * fftpack.fftfreq(shape[1],basis[0,1])
-    print(ki[1]-ki[0])
-    print(kj[1]-kj[0])
-    Kj, Ki = np.meshgrid(kj,ki)
+    # Then we calculate the frequencies in (i,j) space
+    ki = 2 * np.pi * fftpack.fftfreq(shape[0])
+    kj = 2 * np.pi * fftpack.fftfreq(shape[1])
+    K_ij = np.stack(np.meshgrid(ki,kj, indexing='ij'))
+
+    # Now we convert these to frequencies in reciprocal space
+    # These frequencies span the 2D plane of the input wavefield.
+    K_xyz = np.tensordot(inv_basis, K_ij, axes=1)
+
+    # Now we need to apply two corrections to the standard AS method.
+    # First, we calculate a phase mask which corresponds to the
+    # shift of the final plane away from the perpendicular direction
+    # from the input plane. We don't need to extract the perpendicular
+    # component of the shift because the K_xyz vectors are naturally in the
+    # input plane.
+
+    # This may have a sign error - must be checked
+    phase_mask = np.exp(1j * np.tensordot(offset_vector,K_xyz,axes=1))
     
-    # Define this as complex so the square root properly gives
+
+    # Next, we apply a shift to the k-space vectors which sets up
+    # propagation such that a uniform phase object will propagate along the
+    # offset axis. This is not modeling a physical effect, but simply is
+    # the clearest way to do a rigorous simulation while preventing
+    # aliasing-related challenges. If used (as is by default), be aware
+    # and prepare the input wavefields appropriately.
+    perpendicular_dir = np.cross(basis[:,1],basis[:,0])
+    perpendicular_dir /= np.linalg.norm(perpendicular_dir)
+    offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
+    offset_parallel = offset_vector - perpendicular_dir * offset_perpendicular
+    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]
+    
+    # Redefine 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)
+    k0 = np.complex128(k0)
 
+    # Finally, generate the propagator!
+    propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2)
+                        * offset_perpendicular)
+    propagator *= phase_mask
+    
     # Take the conjugate explicitly here instead of negating
     # the previous expression to ensure that complex frequencies
     # get mapped to values <1 instead of >1

tests/tools/test_propagators.py

@@ -48,7 +48,7 @@ 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
+    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)
@@ -93,7 +93,7 @@ def test_generalized_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
+    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)
@@ -114,9 +114,11 @@ def test_generalized_near_field():
     # 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)
-    
+    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)