Source code for gkx.core.velocity

"""Velocity-space basis and gyroaverage utilities."""

from __future__ import annotations

import math

import jax
import jax.numpy as jnp
from jax.scipy.special import gammaln, i0e
import numpy as np




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)





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]





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)





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)





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)





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)





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)





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)





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)





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)





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





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)





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)





def laguerre_quadrature_count(nl: int) -> int:
    """Default number of Laguerre quadrature points."""

    return max(1, 3 * nl // 2 - 1)





def laguerre_transform(nl: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return Laguerre transform matrices and roots."""

    if nl < 1:
        raise ValueError("nl must be >= 1")
    nj = laguerre_quadrature_count(nl)
    jac = np.zeros((nj, nj), dtype=float)
    for i in range(nj - 1):
        jac[i, i] = 2.0 * i + 1.0
        jac[i, i + 1] = i + 1.0
        jac[i + 1, i] = i + 1.0
    jac[nj - 1, nj - 1] = 2.0 * nj - 1.0

    evals, evecs = np.linalg.eigh(jac)
    idx = np.argsort(np.abs(evals))
    evals = evals[idx]
    evecs = evecs[:, idx]

    roots = evals.astype(float)
    poly = np.zeros((nl, nj), dtype=float)
    for i in range(nl):
        for j in range(i + 1):
            tmp = float(math.comb(i, j))
            tmp *= (-1.0) ** j / math.factorial(j) * ((-1.0) ** i)
            poly[i, j] = tmp

    to_grid = np.zeros((nl, nj), dtype=float)
    to_spectral = np.zeros((nj, nl), dtype=float)
    for j in range(nj):
        x_i = roots[j]
        wgt = float(evecs[0, j] ** 2)
        for ell in range(nl):
            coeffs = poly[ell, : ell + 1]
            Lmat = 0.0
            for c in coeffs[::-1]:
                Lmat = Lmat * x_i + c
            to_grid[ell, j] = Lmat
            to_spectral[j, ell] = Lmat * wgt

    return to_grid, to_spectral, roots



# 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)