from typing import Literal, TypeVar
from array_api._2024_12 import Array
from array_api_compat import array_namespace, to_device
from scipy.special import hankel1, hankel2, jv, jvp, yv, yvp
TArray = TypeVar("TArray", bound=Array)
[docs]
def szv(
v: TArray,
d: TArray,
z: TArray,
type: Literal["j", "y", "h1", "h2"],
derivative: bool = False,
) -> TArray:
r"""
Utility function to compute hyperspherical functions.
.. math::
f_v^{(d)} (z) = \sqrt{\frac{\pi}{2}} \frac{F_{v + d/2 - 1}(z)}{z^{d/2 - 1}}
Parameters
----------
v : TArray
The degree of the hyperspherical Hankel function.
d : TArray
The dimension of the hypersphere.
z : TArray
The argument of the hyperspherical Hankel function.
type : Literal["j", "y", "h1", "h2"]
The type of the hyperspherical function.
derivative : bool, optional
Whether to compute the derivative of the
hyperspherical Hankel function, by default False
Returns
-------
TArray
The hyperspherical function.
References
----------
McLean, W. (2000). Strongly Elliptic Systems and
Boundary Integral Equations. p.279
"""
xp = array_namespace(v, d, z)
device = z.device
v = xp.asarray(v, device=device)
d = xp.asarray(d, device=device)
if xp.any((d > 2) & (v < 0)):
raise ValueError(
"The hyperspherical Bessel function of "
"the first kind is not defined for negative degrees."
)
if (d > 2 or type in ("h1, h2")) and derivative:
return v / z * szv(v, d, z, type=type) - szv(v + 1, d, z, type=type)
d_half_minus_1 = d / 2 - 1
if type == "j":
if derivative:
zv = jvp
else:
zv = jv
elif type == "y":
if derivative:
zv = yvp
else:
zv = yv
elif type == "h1":
if derivative:
raise AssertionError()
else:
zv = hankel1
elif type == "h2":
if derivative:
raise AssertionError()
else:
zv = hankel2
dtype = xp.result_type(v, d, z, xp.float32)
if type in ("h1", "h2"):
dtype = xp.result_type(dtype, xp.complex64)
return (
xp.sqrt(xp.asarray(xp.pi / 2, device=z.device, dtype=dtype))
* xp.asarray(
zv(to_device(v + d_half_minus_1, "cpu"), to_device(z, "cpu")),
device=z.device,
dtype=dtype,
)
/ (z**d_half_minus_1)
)
[docs]
def sjv(
v: TArray,
d: TArray,
z: TArray,
derivative: bool = False,
) -> TArray:
r"""
Hyperspherical Bessel function of the first kind.
.. math::
j_v^{(d)} (z) = \sqrt{\frac{\pi}{2}} \frac{J_{v + d/2 - 1}(z)}{z^{d/2 - 1}}
Parameters
----------
v : TArray
The degree of the hyperspherical Bessel function.
d : TArray
The dimension of the hypersphere.
z : TArray
The argument of the hyperspherical Bessel function.
derivative : bool, optional
Whether to compute the derivative of the
hyperspherical Bessel function, by default False
Returns
-------
Array
The hyperspherical Bessel function of the first kind.
References
----------
McLean, W. (2000). Strongly Elliptic Systems and
Boundary Integral Equations. p.279
"""
return szv(v, d, z, type="j", derivative=derivative)
[docs]
def syv(
v: TArray,
d: TArray,
z: TArray,
derivative: bool = False,
) -> TArray:
r"""
Hyperspherical Bessel function of the second kind.
.. math::
y_v^{(d)} (z) = \sqrt{\frac{\pi}{2}} \frac{Y_{v + d/2 - 1}(z)}{z^{d/2 - 1}}
Parameters
----------
v : TArray
The degree of the hyperspherical Bessel function.
d : TArray
The dimension of the hypersphere.
z : TArray
The argument of the hyperspherical Bessel function.
derivative : bool, optional
Whether to compute the derivative of the
hyperspherical Bessel function, by default False
Returns
-------
Array
The hyperspherical Bessel function of the second kind.
References
----------
McLean, W. (2000). Strongly Elliptic Systems and
Boundary Integral Equations. p.279
"""
return szv(v, d, z, type="y", derivative=derivative)
[docs]
def shn1(
v: TArray,
d: TArray,
z: TArray,
derivative: bool = False,
) -> TArray:
r"""
Hyperspherical Hankel function of the first kind.
.. math::
h_v^{(1)(d)} (z) = \sqrt{\frac{\pi}{2}} \frac{H^{(1)}_{v + d/2 - 1}(z)}{z^{d/2 - 1}}
Parameters
----------
v : TArray
The degree of the hyperspherical Hankel function.
d : TArray
The dimension of the hypersphere.
z : TArray
The argument of the hyperspherical Hankel function.
derivative : bool, optional
Whether to compute the derivative of the
hyperspherical Hankel function, by default False
Returns
-------
Array
The hyperspherical Hankel function of the first kind.
"""
return szv(v, d, z, type="h1", derivative=derivative)
[docs]
def shn2(
v: TArray,
d: TArray,
z: TArray,
derivative: bool = False,
) -> TArray:
r"""
Hyperspherical Hankel function of the second kind.
.. math::
h_v^{(2)(d)} (z) = \sqrt{\frac{\pi}{2}} \frac{H^{(2)}_{v + d/2 - 1}(z)}{z^{d/2 - 1}}
Parameters
----------
v : TArray
The degree of the hyperspherical Hankel function.
d : TArray
The dimension of the hypersphere.
z : TArray
The argument of the hyperspherical Hankel function.
derivative : bool, optional
Whether to compute the derivative of the
hyperspherical Hankel function, by default False
Returns
-------
Array
The hyperspherical Hankel function of the second kind.
"""
return szv(v, d, z, type="h2", derivative=derivative)