diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index f3cd25e..46e8403 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -356,34 +356,31 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r propagator : torch.Tensor A phase mask which accounts for the phase change that each plane wave will undergo. """ + # Internally, the generalized propagation function is used, so we start + # by creating an appropriate basis + basis = t.zeros([3,2], **kwargs).real + basis[1,0] = -spacing[0] + basis[0,1] = -spacing[1] + # And similarly, the offset is just z along the z direction + offset = t.zeros([3],**kwargs).real + 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) - ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]).numpy() - kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1]).numpy() - Kj, Ki = np.meshgrid(kj,ki) - - # Define this as complex so the square root properly gives - # k>k0 components imaginary frequencies - k0 = np.complex128((2*np.pi/wavelength)) - - # This properly accounts for evanescent waves, by always choosing - # setting the imaginary frequencies to damp as they propagate - if z >=0: - propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z) - else: - propagator = np.exp(1j*np.conj(np.sqrt(k0**2 - Ki**2 - Kj**2)) * z) - - if remove_z_phase: - propagator *= np.exp(-1j * k0 * z) - + # Bandlimiting is not implemented in the generalized function, because it + # has a less clear meaning in that setting, so we apply it here instead if bandlimit is not None: - Rs = np.sqrt((Ki / np.max(ki))**2 + (Kj / np.max(kj))**2) + ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]) + kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1]) + Ki, Kj = t.meshgrid(ki,kj) + min_radius = min(t.max(ki),t.max(kj)) + Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2) propagator = propagator * (Rs < bandlimit) - # Take the conjugate explicitly here instead of negating - # the previous expression to ensure that complex frequencies - # get mapped to values <1 instead of >1 - - return t.as_tensor(propagator, *args, **kwargs) + return propagator diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index 3970156..abe5e9c 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -256,9 +256,15 @@ 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)) + # Finally, 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, + dtype=t.complex128, bandlimit=0.3) - - + assert asp[140,0] == 0 + assert asp[0,180] == 0 + assert asp[130,0] != 0 + assert asp[0,175] != 0 def test_generalized_near_field():