From f3b6ac9056799dd4501e91f2a9643757e84b918b Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 4 Dec 2023 15:29:40 +0100 Subject: [PATCH] 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 | 57 ++++++++++++++------ tests/tools/test_propagators.py | 52 +++++++++--------- 2 files changed, 64 insertions(+), 45 deletions(-) diff --git a/src/cdtools/tools/propagators/propagators.py b/src/cdtools/tools/propagators/propagators.py index 5d10fba..84ca9d1 100644 --- a/src/cdtools/tools/propagators/propagators.py +++ b/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) diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index 2e1b1df..ab1e432 100644 --- a/tests/tools/test_propagators.py +++ b/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')