fix the numerical stability issue in the angular spectrum propagator code, but remove the ability to choose whether or not to include an explicit z dependence

src/cdtools/tools/propagators/propagators.py

@@ -314,7 +314,7 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
 
 
 
-def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, bandlimit=None, **kwargs):
+def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, bandlimit=None, **kwargs):
     """Generates an angular-spectrum based near-field propagator from experimental quantities
     
     This function generates an angular-spectrum based near field
@@ -342,8 +342,6 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
         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
     bandlimit : float
         Optional, a fraction of the full detector radius beyond which to set the propagator to zero.
 
@@ -371,9 +369,12 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
     offset[2] = z
     
     # And we call the generalized function! 
-    propagator = generate_generalized_angular_spectrum_propagator(shape, basis,
-                            wavelength, offset, 
-                            propagate_along_offset=remove_z_phase, **kwargs)
+    propagator = generate_generalized_angular_spectrum_propagator(
+        shape,
+        basis,
+        wavelength,
+        offset, 
+        **kwargs)
     if z < 0:
         propagator = t.conj(propagator)
 
@@ -516,7 +517,6 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
         perpendicular_dir = perpendicular_dir \
             * t.sign(t.dot(perpendicular_dir,propagation_vector))
     else:
-        pass
         perpendicular_dir = perpendicular_dir * \
             t.sign(t.dot(perpendicular_dir,offset_vector))
     
@@ -526,10 +526,12 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
     if propagation_vector is not None:
         prop_dir = (propagation_vector /
                     t.linalg.norm(propagation_vector))
-        K_0 = 2*np.pi / wavelength * prop_dir
+        K_0 = 2*np.pi / wavelength * (prop_dir - perpendicular_dir)
         K_0_ip = K_0 - t.dot(perpendicular_dir,K_0) * perpendicular_dir
-        K_ip += K_0_ip[:,None,None]
+        
+        K_ip = K_ip + K_0_ip[:,None,None]
     else:
+        print('no prop vector')
         K_0 = t.zeros_like(offset_vector)
 
     # Now, we have accurate in-plane values for K, so we can calculate the
@@ -541,13 +543,37 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
     # before the square root to appropriately map negative numbers to
     # complex frequencies
 
+    # Below was the old, naive approach, that had numerical stability
+    # issues when K_oop >> K_ip, and the subtle changes in K_oop were
+    # clippping because K_ip was so much larger.
+    
+    # K = K_ip + perpendicular_dir[:,None,None] \
+    #     * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
+    
     # The correct algorithm for sqrt(1-x**2) - 1 is
     # - x**2 / (sqrt(1-x**2) + 1)
     # This one doesn't have numerical stability issues near x=0.
-    # I need to figure out how to apply that knowledge here.
-    
-    K = K_ip + perpendicular_dir[:,None,None] \
-        * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
+    # We still use the complex-valued trick in the square root, which
+    # still generates the appropriate complex valued frequencies corresponding
+    # to exponential decay if the in-plane Ks get larger than the
+    # wavenumber of the light.
+
+    # So, it seems like this is a silly way to calculate K_oop (why
+    # not just take the square root of K_oop_squared?), but it is much
+    # more numerically stable. Also note that this has one big difference from
+    # the old method, which (by default) preserved the 
+
+    K_oop_squared_complex = t.complex(K_oop_squared,
+                                      t.zeros_like(K_oop_squared))
+
+    K_oop = - t.linalg.norm(K_ip, dim=0)**2 / \
+        (t.sqrt(K_oop_squared_complex) + (2*np.pi/wavelength))
+    K = K_ip + perpendicular_dir[:,None,None] * K_oop
+
+    # This corrects for the propagation direction, if set. If not, K_0 is
+    # just 0
+    K_m_K_0 = K - K_0[:,None,None]
+    # print(propagation_vector)
 
     # 
     # In this section, we take the inner product of the calcualted
@@ -558,10 +584,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
     # We need to convert to complex because K is complex
     offset_vector = t.complex(offset_vector,t.zeros_like(offset_vector))
 
-    # We then subtract off K_0, essentially setting the phase offset
-    # experienced by K_0 to 0. K_0 will already have been set to 0 if
-    # there was no propagation vector set.
-    K_m_K_0 = K - K_0[:,None,None]
+
     # We actually calculate the phase mask
     phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1)
     

tests/tools/test_propagators.py

@@ -221,25 +221,6 @@ def test_near_field():
         E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
 
     
-    Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
-    
-    # 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(t.as_tensor(E0),asp).numpy()
-
-    # 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.complex128)
-    
-    
     Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
     
     # Check for at least 10^-3 relative accuracy in this scenario
@@ -255,7 +236,7 @@ def test_near_field():
 
     # Then, we check that the bandlimiting at least does something
     asp = propagators.generate_angular_spectrum_propagator(
-        E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
+        E0.shape,(1.5e-9,1e-9),wavelength,z,
         dtype=t.complex128, bandlimit=0.3)
 
     assert asp[140,0] == 0
@@ -315,7 +296,7 @@ def test_generalized_near_field():
     def get_inv_R(Zs):
         return Zs / (Zs**2 + zr**2)
     
-    def get_E(Xs, Ys, Zs, correct=False):
+    def get_E(Xs, Ys, Zs, correct=True):
         # Again, this follows the convention opposite from Wikipedia. See
         # the note in test_angular_spectrum_propagator.
         
@@ -326,11 +307,18 @@ def test_generalized_near_field():
         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:
+        
+        # This removes the z-dependence of the phase
+        if correct: 
             E = E * np.exp(-1j * k * Zs)
         return E
     
-
+    def check_equiv(analytical, numerical):
+        phase = np.angle(np.mean(numerical.conj()*analytical))
+        comp = np.exp(1j*phase) * numerical
+        return (np.max(np.abs(analytical-comp))
+                < 1e-3 * np.max(np.abs(analytical)))
+    
     # We make some rotation matrices to test
     
     # This tests the straight ahead case
@@ -439,27 +427,35 @@ def test_generalized_near_field():
         Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
         
         # Check for at least 10^-3 relative accuracy in this scenario
-        if not np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)):
+        if not check_equiv(Ez, Ez_t):
         #if True:
             plt.close('all')
+            plt.figure()
             plt.imshow(np.angle(E0))
+            plt.title('Angle of E0')
             plt.figure()
             plt.imshow(np.angle(Ez))
+            plt.title('Angle of analytically calculated Ez')
+            plt.figure()
+            plt.imshow(np.abs(Ez))
+            plt.title('Magnitude of analytically calculated Ez')
             plt.figure()
             plt.imshow(np.abs(Ez_t))
+            plt.title('Magnitude of numerically calculated Ez')
             plt.figure()
             plt.imshow(np.angle(Ez_t))
+            plt.title('Angle of numerically calculated Ez')
             plt.figure()
-            plt.imshow(np.abs(Ez-Ez_t)/np.max(np.abs(Ez)))
+            plt.imshow(np.abs(Ez-Ez_t))#/np.max(np.abs(Ez)))
+            plt.title('Magnitude of difference')
             plt.show()
             
-        assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
+        assert check_equiv(Ez, Ez_t)
         
             
         Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
         
-        # 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))
+        assert check_equiv(E0,Em0_t)
             
         print('Test Successful')