def test_round_trip_r_interpolation(radius: np.ndarray, order: int, shape: Callable):
transformer = HankelTransform(order=order, radial_grid=radius)
# the function must be smoothish for interpolation
# to work. Random every point doesn't work
func = shape(radius)
transform_func = transformer.to_transform_r(func)
reconstructed_func = transformer.to_original_r(transform_func)
assert np.allclose(func, reconstructed_func, rtol=1e-4)
def test_round_trip_with_interpolation(shape: Callable, radius: np.ndarray,
transformer: HankelTransform):
# the function must be smoothish for interpolation
# to work. Random every point doesn't work
func = shape(radius)
func_hr = transformer.to_transform_r(func)
ht = transformer.qdht(func_hr)
reconstructed_hr = transformer.iqdht(ht)
reconstructed = transformer.to_original_r(reconstructed_hr)
assert np.allclose(func, reconstructed, rtol=2e-4)
def test_round_trip_r_interpolation_2d(radius: np.ndarray, order: int,
shape: Callable, axis: int):
transformer = HankelTransform(order=order, radial_grid=radius)
# the function must be smoothish for interpolation
# to work. Random every point doesn't work
dims_amplitude = np.ones(2, np.int)
dims_amplitude[1 - axis] = 10
amplitude = np.random.random(dims_amplitude)
dims_radius = np.ones(2, np.int)
dims_radius[axis] = len(radius)
func = np.reshape(shape(radius), dims_radius) * np.reshape(
amplitude, dims_amplitude)
transform_func = transformer.to_transform_r(func, axis=axis)
reconstructed_func = transformer.to_original_r(transform_func, axis=axis)
assert np.allclose(func, reconstructed_func, rtol=1e-4)
def propagate_using_object(r: np.ndarray, field: np.ndarray) -> np.ndarray:
transformer = HankelTransform(order=0, radial_grid=r)
field_for_transform = transformer.to_transform_r(field) # Resampled field
hankel_transform = transformer.qdht(field_for_transform)
propagated_field = np.zeros((nr, Nz), dtype=complex)
kz = np.sqrt(k0**2 - transformer.kr**2)
for n, z_loop in enumerate(z):
phi_z = kz * z_loop # Propagation phase
hankel_transform_at_z = hankel_transform * np.exp(
1j * phi_z) # Apply propagation
field_at_z_transform_grid = transformer.iqdht(
hankel_transform_at_z) # iQDHT
propagated_field[:, n] = transformer.to_original_r(
field_at_z_transform_grid) # Interpolate output
intensity = np.abs(propagated_field)**2
return intensity
# Convert from physical field to physical wavevector
EkrH = H.qdht(ErH)
# %%
# Propagate the beam - loop
# -------------------------
# Do the propagation in a loop over :math:`z`
# Pre-allocate an array for field as a function of r and z
Erz = np.zeros((nr, Nz), dtype=complex)
kz = np.sqrt(k0**2 - H.kr**2)
for i, z_loop in enumerate(z):
phi_z = kz * z_loop # Propagation phase
EkrHz = EkrH * np.exp(1j * phi_z) # Apply propagation
ErHz = H.iqdht(EkrHz) # iQDHT
Erz[:, i] = H.to_original_r(ErHz) # Interpolate output
Irz = np.abs(Erz)**2
# %%
# Plotting
# --------
# Plot the initial field and radial wavevector distribution (given by the
# Hankel transform)
plt.figure()
plt.plot(r * 1e3,
np.abs(Er)**2, r * 1e3, np.unwrap(np.angle(Er)), H.r * 1e3,
np.abs(ErH)**2, H.r * 1e3, np.unwrap(np.angle(ErH)), '+')
plt.title('Initial electric field distribution')
plt.xlabel('Radial co-ordinate (r) /mm')
plt.ylabel('Field intensity /arb.')
plt.legend(['$|E(r)|^2$', '$\\phi(r)$', '$|E(H.r)|^2$', '$\\phi(H.r)$'])