def qdht(r: np.ndarray, f: np.ndarray, order: int = 0) -> Tuple[np.ndarray, np.ndarray]:
"""Perform a quasi-discrete Hankel transform of the function ``f`` (sampled at points
``r``) and return the transformed function and its sample points in :math:`k`-space.
If you requires the transform on a frequency axis (as opposed to the :math:`k`-axis), the
frequency axis :math:`v` can be calculated using :math:`v = \\frac{k}{2\\pi}`.
.. warning::
This method is a convenience wrapper for :meth:`.HankelTransform.qdht`, but incurs a
significant overhead in calculating the :class:`.HankelTransform` object. If you
are performing multiple transforms on the same grid, it will be much quicker to
construct a single :class:`.HankelTransform` object and call
:meth:`.HankelTransform.qdht` multiple times.
:param r: The radial coordinates at which the function is sampled
:type r: :class:`numpy.ndarray`
:param f: The value of the function to be transformed.
:type f: :class:`numpy.ndarray`
:param order: The order of the Hankel Transform to perform. Defaults to 0.
:return: A tuple containing the k coordinates of the transformed function and its values
:rtype: (:class:`numpy.ndarray`, :class:`numpy.ndarray`)
"""
transformer = HankelTransform(order=order, radial_grid=r)
f_transform = transformer.to_transform_r(f)
ht = transformer.qdht(f_transform)
return transformer.kr, ht
def test_sinc(p):
"""Tests from figure 1 of
*"Computation of quasi-discrete Hankel transforms of the integer
order for propagating optical wave fields"*
Manuel Guizar-Sicairos and Julio C. Guitierrez-Vega
J. Opt. Soc. Am. A **21** (1) 53-58 (2004)
"""
transformer = HankelTransform(p, max_radius=3, n_points=256)
v = transformer.v
gamma = 5
func = sinc(2 * np.pi * gamma * transformer.r)
expected_ht = np.zeros_like(func)
expected_ht[v < gamma] = (v[v < gamma]**p * np.cos(p * np.pi / 2)
/ (2*np.pi*gamma * np.sqrt(gamma**2 - v[v < gamma]**2)
* (gamma + np.sqrt(gamma**2 - v[v < gamma]**2))**p))
expected_ht[v >= gamma] = (np.sin(p * np.arcsin(gamma/v[v >= gamma]))
/ (2*np.pi*gamma * np.sqrt(v[v >= gamma]**2 - gamma**2)))
ht = transformer.qdht(func)
# use the same error measure as the paper
dynamical_error = 20 * np.log10(np.abs(expected_ht-ht) / np.max(ht))
not_near_gamma = np.logical_or(v > gamma*1.25,
v < gamma*0.75)
assert np.all(dynamical_error < -10)
assert np.all(dynamical_error[not_near_gamma] < -35)
def iqdht(k: np.ndarray, f: np.ndarray, order: int = 0, axis: int = -2) -> Tuple[np.ndarray, np.ndarray]:
"""Perform a inverse quasi-discrete Hankel transform of the function ``f`` (sampled at points
``k``) and return the transformed function and its sample points in radial space.
If you have the transform on a frequency axis (as opposed to a :math:`k`-axis), the
:math:`k`-axis can be calculated using :math:`k = 2\\pi{}f`.
.. warning::
This method is a convenience wrapper for :meth:`.HankelTransform.iqdht`, but incurs a
significant overhead in calculating the :class:`.HankelTransform` object. If you
are performing multiple transforms on the same grid, it will be much quicker to
construct a single :class:`.HankelTransform` object and call
:meth:`.HankelTransform.iqdht` multiple times.
:param k: The :math:`k` coordinates at which the function is sampled
:type k: :class:`numpy.ndarray`
:param f: The value of the function to be transformed.
:type f: :class:`numpy.ndarray`
:param order: The order of the Hankel Transform to perform. Defaults to 0.
:parameter axis: Axis over which to compute the Hankel transform.
:type axis: :class:`int`
:return: A tuple containing the radial coordinates of the transformed function and its values
:rtype: (:class:`numpy.ndarray`, :class:`numpy.ndarray`)
"""
transformer = HankelTransform(order=order, k_grid=k)
f_transform = transformer.to_transform_k(f, axis=axis)
ht = transformer.iqdht(f_transform, axis=axis)
return transformer.r, ht
def test_top_hat_equivalence(a: float, order: int, radius: np.ndarray):
transformer = HankelTransform(order=order, radial_grid=radius)
f = generalised_top_hat(radius, a, order)
kr, one_shot_ht = qdht(radius, f, order=order)
f_t = generalised_top_hat(transformer.r, a, transformer.order)
standard_ht = transformer.qdht(f_t)
assert np.allclose(one_shot_ht, standard_ht)
def test_round_trip_3d(two_d_size: int, axis: int, radius: np.ndarray,
transformer: HankelTransform):
dims = np.ones(3, np.int) * two_d_size
dims[axis] = radius.size
func = np.random.random(dims)
ht = transformer.qdht(func, axis=axis)
reconstructed = transformer.iqdht(ht, axis=axis)
assert np.allclose(func, reconstructed)
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_inverse_gaussian(a: float, radius: np.ndarray):
# Note the definition in Guizar-Sicairos varies by 2*pi in
# both scaling of the argument (so use kr rather than v) and
# scaling of the magnitude.
transformer = HankelTransform(order=0, radial_grid=radius)
ht = 2*np.pi*(1 / (2 * a**2)) * np.exp(-transformer.kr**2 / (4 * a**2))
actual_f = transformer.iqdht(ht)
expected_f = np.exp(-a ** 2 * transformer.r ** 2)
assert np.allclose(expected_f, actual_f)
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_gaussian_equivalence(a: float, radius: np.ndarray):
# Note the definition in Guizar-Sicairos varies by 2*pi in
# both scaling of the argument (so use kr rather than v) and
# scaling of the magnitude.
transformer = HankelTransform(order=0, radial_grid=radius)
f = np.exp(-a ** 2 * radius ** 2)
kr, one_shot_ht = qdht(radius, f)
f_t = np.exp(-a ** 2 * transformer.r ** 2)
standard_ht = transformer.qdht(f_t)
assert np.allclose(one_shot_ht, standard_ht, rtol=1e-3, atol=1e-4)
def test_1_over_r2_plus_z2_equivalence(a: float):
r = np.linspace(0, 50, 1024)
f = 1 / (r ** 2 + a ** 2)
transformer = HankelTransform(order=0, radial_grid=r)
f_transformer = 1 / (transformer.r**2 + a**2)
assert np.allclose(transformer.to_transform_r(f), f_transformer, rtol=1e-2, atol=1e-6)
kr, one_shot_ht = qdht(r, f)
assert np.allclose(kr, transformer.kr)
standard_ht = transformer.qdht(f_transformer)
assert np.allclose(one_shot_ht, standard_ht, rtol=1e-3, atol=1e-2)
def test_energy_conservation(shape: Callable, transformer: HankelTransform):
transformer = HankelTransform(transformer.order, 10, transformer.n_points)
func = shape(transformer.r, 0.5, transformer.order)
intensity_before = np.abs(func)**2
energy_before = np.trapz(y=intensity_before * 2 * np.pi * transformer.r,
x=transformer.r)
ht = transformer.qdht(func)
intensity_after = np.abs(ht)**2
energy_after = np.trapz(y=intensity_after * 2 * np.pi * transformer.v,
x=transformer.v)
assert np.isclose(energy_before, energy_after, rtol=0.01)
def test_r_creation_equivalence(n: int, max_radius: float):
transformer1 = HankelTransform(order=0, n_points=1024, max_radius=50)
r = np.linspace(0, 50, 1024)
transformer2 = HankelTransform(order=0, radial_grid=r)
for key, val in transformer1.__dict__.items():
if key == '_original_radial_grid':
continue
val2 = getattr(transformer2, key)
if val is None:
assert val2 is None
else:
assert np.allclose(val, val2)
def test_1_over_r2_plus_z2(a: float):
# Note the definition in Guizar-Sicairos varies by 2*pi in
# both scaling of the argument (so use kr rather than v) and
# scaling of the magnitude.
transformer = HankelTransform(order=0, n_points=1024, max_radius=50)
f = 1 / (transformer.r**2 + a**2)
# kn cannot handle complex arguments, so a must be real
expected_ht = 2 * np.pi * scipy_bessel.kn(0, a * transformer.kr)
actual_ht = transformer.qdht(f)
# These tolerances are pretty loose, but there seems to be large
# error here
assert np.allclose(expected_ht, actual_ht, rtol=0.1, atol=0.01)
error = np.mean(np.abs(expected_ht - actual_ht))
assert error < 4e-3
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 test_qdht(engine):
r_max = 5e-3
nr = 512
h = HankelTransform(0, r_max, nr)
hm = engine.hankel_matrix(0., r_max, float(nr), nargout=1)
dr = r_max / (nr - 1) # Radial spacing
nr = np.arange(0, nr) # Radial pixels
r = nr * dr # Radial positions
er = r < 1e-3
# noinspection PyUnresolvedReferences
er_m = matlab.double(er[np.newaxis, :].transpose().tolist())
matlab_value = engine.qdht(er_m, hm, float(3), nargout=1)
python_value = h.qdht(er)
assert matlab_python_allclose(python_value, matlab_value)
def test_inverse_gaussian_2d(axis: int, radius: np.ndarray):
# Note the definition in Guizar-Sicairos varies by 2*pi in
# both scaling of the argument (so use kr rather than v) and
# scaling of the magnitude.
transformer = HankelTransform(order=0, radial_grid=radius)
a = np.linspace(2, 10)
dims_a = np.ones(2, np.int)
dims_a[1 - axis] = len(a)
dims_r = np.ones(2, np.int)
dims_r[axis] = len(transformer.r)
a_reshaped = np.reshape(a, dims_a)
r_reshaped = np.reshape(transformer.r, dims_r)
kr_reshaped = np.reshape(transformer.kr, dims_r)
ht = 2 * np.pi * (1 / (2 * a_reshaped**2)) * np.exp(-kr_reshaped**2 /
(4 * a_reshaped**2))
actual_f = transformer.iqdht(ht, axis=axis)
expected_f = np.exp(-a_reshaped**2 * r_reshaped**2)
assert np.allclose(expected_f, actual_f)
def test_original_r_k_grid():
r_1d = np.linspace(0, 1, 10)
k_1d = r_1d.copy()
transformer = HankelTransform(order=0, max_radius=1, n_points=10)
with pytest.raises(ValueError):
_ = transformer.original_radial_grid
with pytest.raises(ValueError):
_ = transformer.original_k_grid
transformer = HankelTransform(order=0, radial_grid=r_1d)
# no error
_ = transformer.original_radial_grid
with pytest.raises(ValueError):
_ = transformer.original_k_grid
transformer = HankelTransform(order=0, k_grid=k_1d)
# no error
_ = transformer.original_k_grid
with pytest.raises(ValueError):
_ = transformer.original_radial_grid
def example():
# Gaussian function
def gauss1d(x, x0, fwhm):
return np.exp(-2 * np.log(2) * ((x - x0) / fwhm) ** 2)
nr = 1024 # Number of sample points
r_max = .05 # Maximum radius (5cm)
dr = r_max / (nr - 1) # Radial spacing
ri = np.arange(0, nr) # Radial pixels
r = ri * dr # Radial positions
beam_radius = 5e-3 # 5mm
propagation_direction = 5000
Nz = 200 # Number of z positions
z_max = .25 # Maximum propagation distance
dz = z_max / (Nz - 1)
z = np.arange(0, Nz) * dz # Propagation axis
ht = HankelTransform(0, radial_grid=r)
k_max = 2 * np.pi * ht.v_max # Maximum K vector
field = gauss1d(r, 0, beam_radius) * np.exp(1j * propagation_direction * r) # Initial field
ht_field = ht.to_transform_r(field) # Resampled field
transform = ht.qdht(ht_field) # Convert from physical field to physical wavevector
propagated_intensity = np.zeros((nr, Nz + 1))
propagated_intensity[:, 0] = np.abs(field) ** 2
for n in range(1, Nz):
z_loop = z[n]
propagation_phase = (np.sqrt(k_max ** 2 - ht.kr ** 2) - k_max) * z_loop # Propagation phase
propagated_transform = transform * np.exp(1j * propagation_phase) # Apply propagation
propagated_ht_field = ht.iqdht(propagated_transform) # iQDHT
propagated_field = _spline(ht.r, propagated_ht_field, r) # Interpolate output
propagated_intensity[:, n] = np.abs(propagated_field) ** 2
return transform, propagated_intensity
def test_hankel_matrix(engine):
matlab_to_python_mappings = {'N': 'n_points',
'p': 'order',
'alpha_N1': 'alpha_n1',
'V': 'v_max'
}
r_max = 5e-3
nr = 512
h = HankelTransform(0, r_max, nr)
hm = engine.hankel_matrix(0., r_max, float(nr), nargout=1)
for key, matlab_value in hm.items():
python_key = matlab_to_python_mappings.get(key, key)
python_value = getattr(h, python_key)
assert matlab_python_allclose(python_value, matlab_value), \
f"Hankel matrix key {key} doesn't match"
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
def hankel_transform_of_sinc(v):
ht = np.zeros_like(v)
ht[v < gamma] = (v[v < gamma]**p * np.cos(p * np.pi / 2) /
(2 * np.pi * gamma * np.sqrt(gamma**2 - v[v < gamma]**2) *
(gamma + np.sqrt(gamma**2 - v[v < gamma]**2))**p))
ht[v >= gamma] = (
np.sin(p * np.arcsin(gamma / v[v >= gamma])) /
(2 * np.pi * gamma * np.sqrt(v[v >= gamma]**2 - gamma**2)))
return ht
# %%
# Now plot the values of the hankel transform and the dynamical error as in figure 1 of |Guizar| `Guizar`_
# for order 1 and 4
for p in [1, 4]:
transformer = HankelTransform(p, max_radius=3, n_points=256)
gamma = 5
func = sinc(2 * np.pi * gamma * transformer.r)
expected_ht = hankel_transform_of_sinc(transformer.v)
ht = transformer.qdht(func)
dynamical_error = 20 * np.log10(np.abs(expected_ht - ht) / np.max(ht))
not_near_gamma = np.logical_or(transformer.v > gamma * 1.25,
transformer.v < gamma * 0.75)
plt.figure()
plt.subplot(2, 1, 1)
plt.plot(transformer.v, expected_ht, label='Analytical')
plt.plot(transformer.v, ht, marker='+', linestyle='None', label='QDHT')
plt.title(f'Hankel Transform, p={p}')
plt.legend()
"""
# %%
# First import the :class:`.HankelTransform` class and other packages
from pyhank import HankelTransform
import scipy.special
import matplotlib.pyplot as plt
# %%
# Create a :class:`.HankelTransform` object which holds the grid for :math:`r` and
# :math:`k_r` points and calculate the jinc function.
#
# Note that although the calculation fails at :math:`r = 0`, ``transformer.r`` does
# not include :math:`r=0`.
transformer = HankelTransform(order=0, max_radius=100, n_points=1024)
f = scipy.special.jv(1, transformer.r) / transformer.r
plt.figure()
plt.plot(transformer.r, f)
plt.xlabel('Radius /m')
# %%
# Now take the Hankel transform using :meth:`.HankelTransform.qdht`
ht = transformer.qdht(f)
plt.figure()
plt.plot(transformer.kr, ht)
plt.xlim([0, 5])
plt.xlabel('Radial wavevector /m$^{-1}$')
def test_round_trip_2d(two_d_size: int, radius: np.ndarray, transformer: HankelTransform):
func = np.random.random((radius.size, two_d_size))
ht = transformer.qdht(func)
reconstructed = transformer.iqdht(ht)
assert np.allclose(func, reconstructed)
def test_round_trip(radius: np.ndarray, transformer: HankelTransform):
func = np.random.random(radius.shape)
ht = transformer.qdht(func)
reconstructed = transformer.iqdht(ht)
assert np.allclose(func, reconstructed)
def transformer(request, radius) -> HankelTransform:
order = request.param
return HankelTransform(order, radial_grid=radius)
# %% # Initialise :math:`z` grid Nz = 200 # Number of z positions z_max = 0.1 # Maximum propagation distance z = np.linspace(0, z_max, Nz) # Propagation axis # %% # Set up beam parameters Dr = 100e-6 # Beam radius (100um) lambda_ = 488e-9 # wavelength 488nm k0 = 2 * np.pi / lambda_ # Vacuum k vector # %% # Set up a :class:`.HankelTransform` object, telling it the order (``0``) and # the radial grid. H = HankelTransform(order=0, radial_grid=r) # %% # Set up the electric field profile at :math:`z = 0`, and resample onto the correct radial grid # (``transformer.r``) as required for the QDHT. Er = gauss1d(r, 0, Dr) # Initial field ErH = H.to_transform_r(Er) # Resampled field # %% # Perform Hankel Transform # ------------------------ # Convert from physical field to physical wavevector EkrH = H.qdht(ErH) # %%
def test_initialisation_errors():
r_1d = np.linspace(0, 1, 10)
k_1d = r_1d.copy()
r_2d = np.repeat(r_1d[:, np.newaxis], repeats=5, axis=1)
k_2d = r_2d.copy()
with pytest.raises(ValueError):
# missing any radius or k info
HankelTransform(order=0)
with pytest.raises(ValueError):
# missing n_points
HankelTransform(order=0, max_radius=1)
with pytest.raises(ValueError):
# missing max_radius
HankelTransform(order=0, n_points=10)
with pytest.raises(ValueError):
# radial_grid and n_points
HankelTransform(order=0, radial_grid=r_1d, n_points=10)
with pytest.raises(ValueError):
# radial_grid and max_radius
HankelTransform(order=0, radial_grid=r_1d, max_radius=1)
with pytest.raises(ValueError):
# k_grid and n_points
HankelTransform(order=0, k_grid=k_1d, n_points=10)
with pytest.raises(ValueError):
# k_grid and max_radius
HankelTransform(order=0, k_grid=k_1d, max_radius=1)
with pytest.raises(ValueError):
# k_grid and r_grid
HankelTransform(order=0, k_grid=k_1d, radial_grid=r_1d)
with pytest.raises(AssertionError):
HankelTransform(order=0, radial_grid=r_2d)
with pytest.raises(AssertionError):
HankelTransform(order=0, radial_grid=k_2d)
# no error
_ = HankelTransform(order=0, max_radius=1, n_points=10)
_ = HankelTransform(order=0, radial_grid=r_1d)
_ = HankelTransform(order=0, k_grid=k_1d)
def test_top_hat(transformer: HankelTransform, a: float):
f = generalised_top_hat(transformer.r, a, transformer.order)
expected_ht = generalised_jinc(transformer.v, a, transformer.order)
actual_ht = transformer.qdht(f)
error = np.mean(np.abs(expected_ht-actual_ht))
assert error < 1e-3
plt.plot(v, actual_ht, marker='x', linestyle='None', label='QDHT')
plt.title(f'Hankel transform - generalised top-hat, order = {order}')
plt.xlabel('Frequency /$v$')
plt.xlim([0, 1.5])
plt.legend()
plt.tight_layout()
error = np.mean(np.abs(expected_ht-actual_ht))
assert error < 1e-3
# %%
# Now we repeat but the other way round: the Hankel transform of the top-hat
# function should be the jinc function.
radius = np.linspace(0, 2, 1024)
for order in [0, 1, 4]:
transformer = HankelTransform(order=order, radial_grid=radius)
f = generalised_top_hat(transformer.r, a, order)
actual_ht = transformer.qdht(f)
expected_ht = generalised_jinc(transformer.v, a, order)
plt.figure()
plt.subplot(2, 1, 1)
plt.title(f'Generalised top-hat function, order = {order}')
plt.plot(radius, f)
plt.xlabel('Radius /$r$')
plt.subplot(2, 1, 2)
plt.plot(v, expected_ht, label='Analytical')
plt.plot(v, actual_ht, marker='x', linestyle='None', label='QDHT')
plt.title(f'Hankel transform - generalised jinc, order = {order}')
plt.xlabel('Frequency /$v$')
plt.xlim([0, 1.5])