"""Velocity-space basis and gyroaverage utilities."""
from __future__ import annotations
import jax
import jax.numpy as jnp
from jax.scipy.special import gammaln, i0e
import numpy as np
[docs]
def hermite_physicists(x: jnp.ndarray, n_max: int) -> jnp.ndarray:
"""Physicists' Hermite polynomials H_n(x) for n=0..n_max.
Weight: exp(-x**2). Recurrence:
H_0 = 1
H_1 = 2x
H_{n+1} = 2x H_n - 2n H_{n-1}
"""
x = jnp.asarray(x)
if n_max < 0:
raise ValueError("n_max must be >= 0")
if n_max == 0:
return jnp.expand_dims(jnp.ones_like(x), axis=0)
h0 = jnp.ones_like(x)
h1 = 2.0 * x
def step(carry, n):
h_prev, h_curr = carry
h_next = 2.0 * x * h_curr - 2.0 * n * h_prev
return (h_curr, h_next), h_next
_, tail = jax.lax.scan(step, (h0, h1), jnp.arange(1, n_max))
return jnp.concatenate([h0[None, ...], h1[None, ...], tail], axis=0)
[docs]
def hermite_normed(x: jnp.ndarray, n_max: int) -> jnp.ndarray:
"""Normalized Hermite functions with weight exp(-x**2).
psi_n = H_n(x) / sqrt(2**n * n! * sqrt(pi))
"""
h = hermite_physicists(x, n_max)
n = jnp.arange(0, n_max + 1)
log_norm = 0.5 * (n * jnp.log(2.0) + gammaln(n + 1) + 0.5 * jnp.log(jnp.pi))
norm = jnp.exp(log_norm)
return h / norm[:, None]
[docs]
def laguerre(x: jnp.ndarray, l_max: int) -> jnp.ndarray:
"""Laguerre polynomials L_l(x) for l=0..l_max.
Weight: exp(-x). Recurrence:
L_0 = 1
L_1 = 1 - x
(l+1) L_{l+1} = (2l+1-x) L_l - l L_{l-1}
"""
x = jnp.asarray(x)
if l_max < 0:
raise ValueError("l_max must be >= 0")
if l_max == 0:
return jnp.expand_dims(jnp.ones_like(x), axis=0)
l0 = jnp.ones_like(x)
l1 = 1.0 - x
def step(carry, ell_idx):
l_prev, l_curr = carry
l_next = ((2.0 * ell_idx + 1.0 - x) * l_curr - ell_idx * l_prev) / (
ell_idx + 1.0
)
return (l_curr, l_next), l_next
_, tail = jax.lax.scan(step, (l0, l1), jnp.arange(1, l_max))
return jnp.concatenate([l0[None, ...], l1[None, ...], tail], axis=0)
[docs]
def hermite_ladder_coeffs(n_max: int) -> tuple[jnp.ndarray, jnp.ndarray]:
"""Return sqrt(n+1) and sqrt(n) arrays for Hermite ladder operators."""
if n_max < 0:
raise ValueError("n_max must be >= 0")
n = jnp.arange(0, n_max + 1)
return jnp.sqrt(n + 1.0), jnp.sqrt(n)
[docs]
def gamma0(b: jnp.ndarray) -> jnp.ndarray:
"""Compute Gamma_0(b) = exp(-b) I_0(b) using i0e for stability."""
b = jnp.asarray(b)
return i0e(b)
[docs]
def bessel_j0(x: jnp.ndarray) -> jnp.ndarray:
"""Return J0(x) using a Cephes-style approximation (Cephes-compatible)."""
x = jnp.asarray(x)
ax = jnp.abs(x)
y = x * x
r = 57568490574.0 + y * (
-13362590354.0
+ y * (651619640.7 + y * (-11214424.18 + y * (77392.33017 + y * -184.9052456)))
)
s = 57568490411.0 + y * (
1029532985.0 + y * (9494680.718 + y * (59272.64853 + y * (267.8532712 + y)))
)
res_small = r / s
z = 8.0 / jnp.maximum(ax, 1.0e-30)
y2 = z * z
xx = ax - 0.785398164
p = 1.0 + y2 * (
-0.1098628627e-2
+ y2 * (0.2734510407e-4 + y2 * (-0.2073370639e-5 + y2 * 0.2093887211e-6))
)
q = -0.1562499995e-1 + y2 * (
0.1430488765e-3
+ y2 * (-0.6911147651e-5 + y2 * (0.7621095161e-6 + y2 * -0.934945152e-7))
)
res_large = jnp.sqrt(0.636619772 / jnp.maximum(ax, 1.0e-30)) * (
jnp.cos(xx) * p - z * jnp.sin(xx) * q
)
out = jnp.where(ax < 8.0, res_small, res_large)
return jnp.where(jnp.isfinite(out), out, res_small)
[docs]
def bessel_j1(x: jnp.ndarray) -> jnp.ndarray:
"""Return J1(x) using a Cephes-style approximation (Cephes-compatible)."""
x = jnp.asarray(x)
ax = jnp.abs(x)
y = x * x
r = 72362614232.0 + y * (
-7895059235.0
+ y * (242396853.1 + y * (-2972611.439 + y * (15704.48260 + y * -30.16036606)))
)
s = 144725228442.0 + y * (
2300535178.0 + y * (18583304.74 + y * (99447.43394 + y * (376.9991397 + y)))
)
res_small = x * (r / s)
z = 8.0 / jnp.maximum(ax, 1.0e-30)
y2 = z * z
xx = ax - 2.356194491
p = 1.0 + y2 * (
0.183105e-2
+ y2 * (-0.3516396496e-4 + y2 * (0.2457520174e-5 + y2 * -0.240337019e-6))
)
q = 0.04687499995 + y2 * (
-0.2002690873e-3
+ y2 * (0.8449199096e-5 + y2 * (-0.88228987e-6 + y2 * 0.105787412e-6))
)
res_large = jnp.sqrt(0.636619772 / jnp.maximum(ax, 1.0e-30)) * (
jnp.cos(xx) * p - z * jnp.sin(xx) * q
)
res_large = jnp.where(x < 0.0, -res_large, res_large)
out = jnp.where(ax < 8.0, res_small, res_large)
return jnp.where(jnp.isfinite(out), out, res_small)
[docs]
def bessel_laguerre_kernels(bessel_argument: jnp.ndarray, n_max: int) -> jnp.ndarray:
r"""Return finite-Larmor Bessel--Laguerre kernels through order ``n_max``.
The coefficients
.. math::
K_n(b) = \exp(-b^2/4)\frac{(b^2/4)^n}{n!}
expand :math:`J_0(B\sqrt{x})` in ordinary Laguerre polynomials, where
:math:`B=k_\perp v_{\mathrm{th}}/\Omega`. A recurrence
avoids factorial overflow and preserves the exact ``b=0`` limit. The leading
axis indexes ``n=0, ..., n_max``.
References
----------
Frei et al., *Journal of Plasma Physics* 87, 905870501 (2021), Eq. (2.13).
"""
if n_max < 0:
raise ValueError("n_max must be >= 0")
b = jnp.asarray(bessel_argument)
if not jnp.issubdtype(b.dtype, jnp.inexact):
b = b.astype(jnp.float32)
argument = 0.25 * b * b
kernel0 = jnp.exp(-argument)
if n_max == 0:
return kernel0[None, ...]
def step(kernel, order):
next_kernel = kernel * argument / order
return next_kernel, next_kernel
_, tail = jax.lax.scan(
step,
kernel0,
jnp.arange(1, n_max + 1, dtype=argument.dtype),
)
return jnp.concatenate([kernel0[None, ...], tail], axis=0)
[docs]
def associated_bessel_laguerre_coefficients(
bessel_argument: jnp.ndarray,
bessel_order: int,
n_max: int,
) -> jnp.ndarray:
r"""Return coefficients of the associated-Laguerre expansion of ``J_m``.
For ``m = bessel_order``, the returned coefficients are
.. math::
A_n^m(b) = \frac{n!}{(n+m)!}\left(\frac{b}{2}\right)^m K_n(b),
so that :math:`J_m(B\sqrt{x}) = x^{m/2}\sum_n A_n^m L_n^m(x)`, with
:math:`B=k_\perp v_{\mathrm{th}}/\Omega`.
The leading axis indexes ``n=0, ..., n_max``.
"""
if bessel_order < 0:
raise ValueError("bessel_order must be >= 0")
if n_max < 0:
raise ValueError("n_max must be >= 0")
b = jnp.asarray(bessel_argument)
if not jnp.issubdtype(b.dtype, jnp.inexact):
b = b.astype(jnp.float32)
half_b = 0.5 * b
argument = half_b * half_b
coefficient0 = (
jnp.exp(-argument - gammaln(jnp.asarray(bessel_order + 1, dtype=b.dtype)))
* half_b**bessel_order
)
if n_max == 0:
return coefficient0[None, ...]
def step(coefficient, denominator):
next_coefficient = coefficient * argument / denominator
return next_coefficient, next_coefficient
_, tail = jax.lax.scan(
step,
coefficient0,
jnp.arange(
bessel_order + 1,
bessel_order + n_max + 1,
dtype=argument.dtype,
),
)
return jnp.concatenate([coefficient0[None, ...], tail], axis=0)
[docs]
def single_precision_factorial(m: jnp.ndarray) -> jnp.ndarray:
"""Return the single-precision factorial approximation."""
m_arr = jnp.asarray(m)
dtype = m_arr.dtype
exact = jnp.asarray([1.0, 1.0, 2.0, 6.0, 24.0, 120.0, 720.0], dtype=dtype)
m_int = m_arr.astype(jnp.int32)
m_clamped = jnp.clip(m_int, 0, exact.shape[0] - 1)
m_safe = jnp.where(m_arr > 0, m_arr, jnp.asarray(1.0, dtype=dtype))
stirling = (
jnp.sqrt(2.0 * jnp.asarray(jnp.pi, dtype=dtype) * m_safe)
* (m_safe**m_safe)
* jnp.exp(-m_safe)
* (1.0 + 1.0 / (12.0 * m_safe) + 1.0 / (288.0 * m_safe * m_safe))
)
return jnp.where(m_int <= 6, exact[m_clamped], stirling)
[docs]
def J_l_all(b: jnp.ndarray, l_max: int) -> jnp.ndarray:
"""Gyroaveraging coefficients matching the Laguerre-Hermite quadrature convention."""
if l_max < 0:
raise ValueError("l_max must be >= 0")
b = jnp.asarray(b)
ell = jnp.arange(l_max + 1, dtype=b.dtype)
l_shape = (l_max + 1,) + (1,) * b.ndim
ell = ell.reshape(l_shape)
sign = jnp.where((ell % 2) == 0, 1.0, -1.0)
half_b = 0.5 * b
half_b_safe = jnp.where(half_b > 0.0, half_b, 1.0)
log_abs = (
ell * jnp.log(half_b_safe[None, ...]) - gammaln(ell + 1.0) - half_b[None, ...]
)
Jl = sign * jnp.exp(log_abs)
zero_mask = (b == 0.0)[None, ...]
Jl = jnp.where(zero_mask & (ell == 0), 1.0, Jl)
Jl = jnp.where(zero_mask & (ell > 0), 0.0, Jl)
return Jl
[docs]
def laguerre_gyroaverage_neighbors(
coefficients: jnp.ndarray,
b: jnp.ndarray,
*,
axis: int,
) -> tuple[jnp.ndarray, jnp.ndarray]:
r"""Return the lower and upper Laguerre neighbors of gyroaverage coefficients.
The upper neighbor at the truncation boundary is known analytically even
though the corresponding distribution moment is not retained. For
:math:`\mathcal J_\ell=(-1)^\ell e^{-b/2}(b/2)^\ell/\ell!`, it is
:math:`\mathcal J_{L}=-\mathcal J_{L-1}(b/2)/L`. Zero-padding that value
drops a physical term from the highest retained diamagnetic-drive equation.
"""
values = jnp.moveaxis(jnp.asarray(coefficients), axis, 0)
if values.shape[0] < 1:
raise ValueError("Laguerre coefficient axis must be non-empty")
b_arr = jnp.asarray(b, dtype=values.dtype)
lower = jnp.concatenate([jnp.zeros_like(values[:1]), values[:-1]], axis=0)
order = jnp.asarray(values.shape[0], dtype=values.dtype)
boundary = -values[-1] * (0.5 * b_arr) / order
upper = jnp.concatenate([values[1:], boundary[None, ...]], axis=0)
return jnp.moveaxis(lower, 0, axis), jnp.moveaxis(upper, 0, axis)
[docs]
def sum_Jl2(b: jnp.ndarray, l_max: int) -> jnp.ndarray:
"""Truncated sum of J_l(b)^2, useful for Gamma_0 convergence checks."""
Jl = J_l_all(b, l_max)
return jnp.sum(Jl * Jl, axis=0)
[docs]
def laguerre_quadrature_count(nl: int) -> int:
"""Default number of Laguerre quadrature points."""
return max(1, 3 * nl // 2 - 1)
[docs]
def laguerre_scaled_functions(x: np.ndarray, n_max: int) -> np.ndarray:
r"""Scaled Laguerre functions :math:`\tilde L_n(x)=e^{-x/2}L_n(x)`.
The exponential rides along in the seed of the standard three-term
recurrence instead of being applied afterwards, so the bare :math:`L_n(x)`
is never formed -- it reaches :math:`10^{18}` at the outermost Gauss node
of even a modest grid. Szego's bound gives :math:`|\tilde L_n(x)|\le 1` for
:math:`x\ge 0`, so every intermediate here is O(1):
.. math::
\tilde L_0 = e^{-x/2},\qquad \tilde L_1 = (1-x)\,e^{-x/2},\\
(n+1)\,\tilde L_{n+1} = (2n+1-x)\,\tilde L_n - n\,\tilde L_{n-1}.
Returns an array of shape ``(n_max + 1, x.size)``.
"""
if n_max < 0:
raise ValueError("n_max must be >= 0")
x = np.asarray(x, dtype=float)
out = np.zeros((n_max + 1, x.size), dtype=float)
out[0] = np.exp(-0.5 * x)
if n_max >= 1:
out[1] = (1.0 - x) * out[0]
for n in range(1, n_max):
out[n + 1] = ((2.0 * n + 1.0 - x) * out[n] - n * out[n - 1]) / (n + 1.0)
return out
# Round-trip accuracy that still leaves headroom below the physics tolerances.
_LAGUERRE_ROUNDTRIP_TOLERANCE = 1.0e-8
def _check_laguerre_transform_conditioning(
to_grid: np.ndarray, to_spectral: np.ndarray, nl: int
) -> None:
"""Reject a Laguerre transform that has lost its round-trip identity.
A tripwire, not a limit: with the scaled recurrence and the analytic
weights the measured error is 7e-14 at ``nl = 16``, 7e-13 at 32 and 1e-11
at 96, so the ``1e-8`` tolerance is orders of magnitude away from anything
a healthy build produces. It exists because the previous unweighted
construction degraded *silently* -- 4e-10 at ``nl = 17``, 1e-3 at 20, and
total garbage by 24, with no error raised and no visible symptom in the
physics -- and that failure mode must not be able to return unnoticed.
The remaining hard limit is the seed ``exp(-x/2)`` underflowing float64
once the largest node passes ``x ~ 1490``, around ``nl ~ 250``; the same
check catches it.
The cost is one small matrix product per cache build.
"""
identity_error = float(
np.abs(to_grid @ to_spectral - np.eye(nl, dtype=float)).max()
)
if identity_error > _LAGUERRE_ROUNDTRIP_TOLERANCE:
raise ValueError(
f"Laguerre transform for nl={nl} is numerically unusable: the "
f"round-trip identity error is {identity_error:.3e}, above the "
f"{_LAGUERRE_ROUNDTRIP_TOLERANCE:.0e} tolerance. The scaled "
"recurrence holds to ~1e-11 through nl=96, so this means either "
"the exp(-x/2) seed has underflowed (expected only past nl~250) "
"or the construction has regressed."
)
# Nonlinear Laguerre quadrature kernels live with the velocity basis they transform.
def _laguerre_to_grid(G: jnp.ndarray, laguerre_to_grid: jnp.ndarray) -> jnp.ndarray:
"""Transform Laguerre moments to the muB quadrature grid."""
G = jnp.asarray(G)
laguerre_to_grid = jnp.asarray(laguerre_to_grid)
return jnp.einsum(
"slmyxz,lj->sjmyxz",
G,
laguerre_to_grid,
precision=jax.lax.Precision.HIGHEST,
)
def _laguerre_to_spectral(
g_mu: jnp.ndarray, laguerre_to_spectral: jnp.ndarray
) -> jnp.ndarray:
"""Transform muB quadrature-grid values back to Laguerre moments."""
g_mu = jnp.asarray(g_mu)
laguerre_to_spectral = jnp.asarray(laguerre_to_spectral)
return jnp.einsum(
"sjmyxz,jl->slmyxz",
g_mu,
laguerre_to_spectral,
precision=jax.lax.Precision.HIGHEST,
)
def _laguerre_j0_field(
field: jnp.ndarray,
b: jnp.ndarray,
roots: jnp.ndarray,
factor: float,
) -> jnp.ndarray:
"""Apply J0(field) on the Laguerre quadrature grid."""
b = jnp.asarray(b)
roots = jnp.asarray(roots)
field = jnp.asarray(field)
if b.ndim == 3:
b = b[None, ...]
if roots.ndim == 0:
roots = roots[None]
alpha = jnp.sqrt(
jnp.maximum(0.0, 2.0 * roots[None, :, None, None, None] * b[:, None, ...])
)
j0 = bessel_j0(alpha)
field_b = field[None, None, ...]
return j0 * field_b * jnp.asarray(factor, dtype=field.dtype)
def _laguerre_j0_field_precomputed(
field: jnp.ndarray,
j0: jnp.ndarray,
factor: float,
) -> jnp.ndarray:
field = jnp.asarray(field)
field_b = field[None, None, ...]
return j0 * field_b * jnp.asarray(factor, dtype=field.dtype)
def _laguerre_bpar_correction(
bpar: jnp.ndarray,
b: jnp.ndarray,
roots: jnp.ndarray,
tz: jnp.ndarray,
factor: float,
) -> jnp.ndarray:
"""Return the bpar correction term on the Laguerre quadrature grid."""
b = jnp.asarray(b)
roots = jnp.asarray(roots)
bpar = jnp.asarray(bpar)
if b.ndim == 3:
b = b[None, ...]
if roots.ndim == 0:
roots = roots[None]
tz_arr = jnp.asarray(tz)
if tz_arr.ndim == 0:
tz_arr = tz_arr[None]
alpha = jnp.sqrt(
jnp.maximum(0.0, 2.0 * roots[None, :, None, None, None] * b[:, None, ...])
)
j1 = bessel_j1(alpha)
j1_over_alpha = jnp.where(alpha < 1.0e-8, 0.5, j1 / alpha)
coeff = (
tz_arr[:, None, None, None, None]
* 2.0
* roots[None, :, None, None, None]
* j1_over_alpha
)
bpar_b = bpar[None, None, ...]
return coeff * bpar_b * jnp.asarray(factor, dtype=bpar.dtype)
def _laguerre_bpar_correction_precomputed(
bpar: jnp.ndarray,
j1_over_alpha: jnp.ndarray,
roots: jnp.ndarray,
tz: jnp.ndarray,
factor: float,
) -> jnp.ndarray:
bpar = jnp.asarray(bpar)
tz_arr = jnp.asarray(tz)
if tz_arr.ndim == 0:
tz_arr = tz_arr[None]
coeff = (
tz_arr[:, None, None, None, None]
* 2.0
* roots[None, :, None, None, None]
* j1_over_alpha
)
bpar_b = bpar[None, None, ...]
return coeff * bpar_b * jnp.asarray(factor, dtype=bpar.dtype)