Pspecial Tiers 3-4: Bessel K + associated Legendre + Gegenbauer (#917)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: jobovy

galpy/backend/special/__init__.py

@@ -13,6 +13,7 @@
 #   capability test asserts native-vs-fallback agreement so it can be deleted).
 ###############################################################################
 from ._router import (
+    assoc_legendre,
     ellipe,
     ellipk,
     erf,
@@ -21,10 +22,14 @@
     gammainc,
     gammaincc,
     gammaln,
+    gegenbauer,
     hyp1f1,
     hyp2f1,
     i0,
     i1,
+    k0,
+    k1,
+    kn,
     xlogy,
 )
 
@@ -42,4 +47,9 @@
     "hyp1f1",
     "ellipk",
     "ellipe",
+    "k0",
+    "k1",
+    "kn",
+    "assoc_legendre",
+    "gegenbauer",
 ]

galpy/backend/special/_fallback/assoc_legendre.py

@@ -0,0 +1,75 @@
+###############################################################################
+#   Backend-agnostic associated Legendre functions P_l^m(x) for all degrees
+#   l < L and orders 0 <= m < M, with the Condon-Shortley phase (matching
+#   scipy.special.assoc_legendre_p_all(..., branch_cut=2)). Replaces
+#   galpy.util.special.compute_legendre on the SCF / MultipoleExpansion path so
+#   those potentials run and differentiate under every backend.
+#
+#   P is built by the standard forward (Bonnet) recurrences:
+#     P_m^m   = (-1)^m (2m-1)!! (1-x^2)^{m/2}
+#     P_{m+1}^m = x (2m+1) P_m^m
+#     (l-m) P_l^m = x (2l-1) P_{l-1}^m - (l+m-1) P_{l-2}^m
+#   The optional first/second x-derivatives use
+#     (x^2-1) dP/dx = l x P_l^m - (l+m) P_{l-1}^m
+#     (1-x^2) d2P/dx^2 = 2x dP/dx - l(l+1) P + m^2/(1-x^2) P   (Legendre ODE)
+#   (these diverge at the poles x=+-1 for m>=1, exactly as scipy returns, and are
+#   multiplied by sin^2(theta) in the physical theta-derivatives downstream).
+#
+#   Everything is pure arithmetic built with lists + xp.stack (no in-place
+#   mutation), so it differentiates cleanly under jax and torch -- and the
+#   x-derivatives are also available straight from autodiff.
+###############################################################################
+
+
+def assoc_legendre(xp, L, M, x, deriv=0):
+    """P_l^m(x), shape ``x.shape + (L, M)`` (Condon-Shortley phase).
+
+    deriv: 0 -> P; 1 -> (P, dP/dx); 2 -> (P, dP/dx, d2P/dx2).
+    L, M are static ints; x is a backend array (or scalar) with |x| <= 1.
+    """
+    x = xp.asarray(x) * 1.0
+    one = xp.ones_like(x)
+    zero = xp.zeros_like(x)
+    # (1-x^2)^{1/2}; clip keeps it real at |x|=1 (interior x is unaffected).
+    somx2 = xp.sqrt(xp.where(x * x < 1.0, 1.0 - x * x, zero))
+
+    # P[l][m] as a list-of-lists of backend arrays (functional, no mutation).
+    P = [[zero for _ in range(M)] for _ in range(L)]
+    pmm = one  # running P_m^m diagonal
+    for m in range(M):
+        if m > 0:
+            pmm = pmm * (-(2 * m - 1)) * somx2
+        if m < L:
+            P[m][m] = pmm
+        if m + 1 < L:
+            P[m + 1][m] = x * (2 * m + 1) * pmm
+        for l in range(m + 2, L):
+            P[l][m] = (x * (2 * l - 1) * P[l - 1][m] - (l + m - 1) * P[l - 2][m]) / (
+                l - m
+            )
+
+    def _stack(grid):
+        return xp.stack([xp.stack(row, axis=-1) for row in grid], axis=-2)
+
+    Parr = _stack(P)
+    if deriv == 0:
+        return Parr
+
+    den = x * x - 1.0  # (x^2-1); singular only at the poles
+    dP = [[zero for _ in range(M)] for _ in range(L)]
+    for m in range(M):
+        for l in range(m, L):
+            plm1 = P[l - 1][m] if l - 1 >= m else zero
+            dP[l][m] = (l * x * P[l][m] - (l + m) * plm1) / den
+    dParr = _stack(dP)
+    if deriv == 1:
+        return Parr, dParr
+
+    om = 1.0 - x * x
+    d2 = [[zero for _ in range(M)] for _ in range(L)]
+    for m in range(M):
+        for l in range(m, L):
+            d2[l][m] = (
+                2.0 * x * dP[l][m] - l * (l + 1) * P[l][m] + (m * m) / om * P[l][m]
+            ) / om
+    return Parr, dParr, _stack(d2)

galpy/backend/special/_fallback/bessel_k.py

@@ -0,0 +1,99 @@
+###############################################################################
+#   Fallbacks for the modified Bessel functions of the second kind K0, K1, Kn
+#   on real x > 0. Needed on BOTH jax and torch:
+#     - jax.scipy.special has no k0/k1/kn at all;
+#     - torch.special has modified_bessel_k0/k1 but they are NOT differentiable
+#       (no autograd backward) and lack kn entirely, so we use the fallback there
+#       too (the router sees no torch.special.k0 attribute -> treats it missing).
+#
+#   K0, K1 (``_k01``) use two regimes, each ~1e-15 vs scipy and AD-friendly:
+#     - x <= 2: the Abramowitz & Stegun ascending series (9.6.13/9.6.11), built
+#       on the native i0/i1 (Tier 1) plus elementary terms;
+#     - x  > 2: the trapezoidal rule on K_nu(x) = int_0^inf e^{-x cosh t}
+#       cosh(nu t) dt. The integrand is double-exponentially decaying, so the
+#       trapezoidal rule converges geometrically; its e^{-x(cosh t-1)} peak has
+#       width ~1/sqrt(x), so the nodes are scaled by 1/sqrt(x) to resolve it
+#       uniformly for all large x.
+#   Each branch's argument is clamped into its valid region wherever the OTHER
+#   branch is selected, so the unused branch cannot overflow (i0 at large x) or
+#   NaN-poison reverse-mode gradients.
+#
+#   Kn (``kn_fallback``) uses the upward recurrence K_{m+1}=K_{m-1}+(2m/x)K_m
+#   from K0, K1 -- the stable direction for K.
+###############################################################################
+import numpy
+
+_GAMMA = 0.5772156649015328606  # Euler-Mascheroni
+_NSERIES = 30  # ascending-series terms (x <= 2)
+_TRAP_H = 0.25  # trapezoidal step (in the 1/sqrt(x)-scaled variable)
+_TRAP_N = 64  # trapezoidal nodes
+# node positions i*h and weights (h/2 at the endpoint i=0), as numpy constants
+_TRAP_NODES = numpy.arange(_TRAP_N + 1) * _TRAP_H
+_TRAP_W = numpy.full(_TRAP_N + 1, _TRAP_H)
+_TRAP_W[0] = _TRAP_H / 2.0
+
+
+def _k01(xp, x):
+    """Return (K0(x), K1(x)) for real x > 0, ~1e-15 vs scipy, AD-friendly."""
+    x = xp.asarray(x) * 1.0
+    inside = x <= 2.0
+    # Clamp the dead region of each branch into its valid domain.
+    xs = xp.where(inside, x, xp.ones_like(x))  # series branch (x<=2)
+    xt = xp.where(inside, 2.0 * xp.ones_like(x), x)  # trapezoid branch (x>2)
+
+    # --- ascending series (x <= 2), via native i0/i1 ---
+    from .._router import i0, i1
+
+    x2 = xs * xs / 4.0
+    K0s = -(xp.log(xs / 2.0) + _GAMMA) * i0(xs)
+    term = xp.ones_like(xs)
+    harm = 0.0
+    for k in range(1, _NSERIES):
+        harm += 1.0 / k
+        term = term * x2 / (k * k)
+        K0s = K0s + term * harm
+    s1 = xp.zeros_like(xs)
+    term = xp.ones_like(xs)
+    hk = 0.0
+    for k in range(0, _NSERIES):
+        hk1 = hk + 1.0 / (k + 1)
+        s1 = s1 + term * ((hk + hk1) / 2.0 - _GAMMA)
+        term = term * x2 / ((k + 1) * (k + 2))
+        hk = hk1
+    K1s = 1.0 / xs + xp.log(xs / 2.0) * i1(xs) - (xs / 2.0) * s1
+
+    # --- peak-resolving scaled trapezoidal (x > 2) ---
+    nodes = xp.asarray(_TRAP_NODES)
+    weights = xp.asarray(_TRAP_W)
+    sc = 1.0 / xp.sqrt(xt)
+    t = sc[..., None] * nodes  # (..., N+1)
+    cosh_t = xp.cosh(t)
+    e = xp.exp(-xt[..., None] * cosh_t) * weights
+    K0t = xp.sum(e, axis=-1) * sc
+    K1t = xp.sum(e * cosh_t, axis=-1) * sc
+
+    return xp.where(inside, K0s, K0t), xp.where(inside, K1s, K1t)
+
+
+def k0_fallback(xp, x):
+    """Modified Bessel function of the second kind, order 0."""
+    return _k01(xp, x)[0]
+
+
+def k1_fallback(xp, x):
+    """Modified Bessel function of the second kind, order 1."""
+    return _k01(xp, x)[1]
+
+
+def kn_fallback(xp, n, x):
+    """Integer-order K_n(x) via the stable upward recurrence from K0, K1."""
+    n = int(n)
+    km1, k = _k01(xp, x)  # K0, K1
+    if n == 0:
+        return km1
+    if n == 1:
+        return k
+    x = xp.asarray(x) * 1.0
+    for m in range(1, n):
+        km1, k = k, km1 + (2.0 * m / x) * k
+    return k

galpy/backend/special/_fallback/gegenbauer.py

@@ -0,0 +1,30 @@
+###############################################################################
+#   Backend-agnostic Gegenbauer (ultraspherical) polynomials C_n^alpha(x) for
+#   0 <= n < N, via the standard three-term recurrence
+#       C_0 = 1,  C_1 = 2 alpha x,
+#       (n+1) C_{n+1} = 2(n+alpha) x C_n - (n+2 alpha-1) C_{n-1}.
+#   This is the SCFPotential radial basis (galpy.potential.SCFPotential._C uses
+#   the same recurrence with alpha = 2l + 3/2). Built with lists + xp.stack (no
+#   in-place mutation), so it differentiates under jax and torch; the numpy path
+#   reproduces SCF's existing recurrence value-for-value.
+###############################################################################
+
+
+def gegenbauer(xp, N, alpha, x):
+    """C_n^alpha(x) for 0 <= n < N, shape ``x.shape + (N,)``.
+
+    N is a static int, alpha a scalar, x a backend array (or scalar).
+    """
+    x = xp.asarray(x) * 1.0
+    cols = [xp.ones_like(x)]  # C_0 = 1
+    if N > 1:
+        cnm1 = cols[0]
+        cn = 2.0 * alpha * x  # C_1 = 2 alpha x
+        cols.append(cn)
+        for n in range(1, N - 1):
+            cnp1 = (2.0 * (n + alpha) * x * cn - (n + 2.0 * alpha - 1.0) * cnm1) / (
+                n + 1.0
+            )
+            cols.append(cnp1)
+            cnm1, cn = cn, cnp1
+    return xp.stack(cols, axis=-1)

galpy/backend/special/_router.py

@@ -12,10 +12,13 @@
 # (hasattr on the backend's special module), so entries are removed as backends
 # add the native version. (numpy always has the full scipy.special.)
 _NATIVE_MISSING = {
-    "jax": frozenset(("ellipk", "ellipe")),
+    "jax": frozenset(("ellipk", "ellipe", "k0", "k1", "kn")),
+    # torch.special lacks all of these. (It does have modified_bessel_k0/k1, but
+    # they are NOT differentiable -- no autograd backward -- and there is no kn,
+    # so the k0/k1/kn fallbacks are used; the router sees no torch.special.k0.)
     "torch": frozenset(
-        ("gamma", "ellipk", "ellipe", "hyp2f1", "hyp1f1")
-    ),  # torch.special lacks all of these
+        ("gamma", "ellipk", "ellipe", "hyp2f1", "hyp1f1", "k0", "k1", "kn")
+    ),
 }
 
 # Functions whose native implementation EXISTS but is too inaccurate on galpy's
@@ -152,6 +155,66 @@ def ellipe(m):
     return _dispatch("ellipe", (m,), ellipe_fallback)
 
 
+# --- Tier 3: modified Bessel functions of the second kind (disk force paths) --
+def k0(x):
+    from ._fallback.bessel_k import k0_fallback
+
+    return _dispatch("k0", (x,), k0_fallback)
+
+
+def k1(x):
+    from ._fallback.bessel_k import k1_fallback
+
+    return _dispatch("k1", (x,), k1_fallback)
+
+
+def kn(n, x):
+    # Integer-order modified Bessel K_n; only the array arg x carries the namespace.
+    from ._fallback.bessel_k import kn_fallback
+
+    return _dispatch("kn", (n, x), kn_fallback, ns_args=(x,))
+
+
+# --- Tier 4: associated Legendre P_l^m (SCF / MultipoleExpansion) -------------
+def _scipy_assoc_legendre(L, M, x, deriv):
+    """numpy path: scipy.special.assoc_legendre_p_all reshaped to (...,L,M),
+    byte-identical to scipy (the convention used by util.special.compute_legendre)."""
+    import scipy.special as sp
+
+    arr = numpy.asarray(
+        sp.assoc_legendre_p_all(
+            L - 1, M - 1, numpy.asarray(x, dtype=float), branch_cut=2, diff_n=deriv
+        )
+    )  # (deriv+1, L, 2M-1, *x.shape)  -- m=0..M-1 are the first M columns
+    out = numpy.moveaxis(arr[:, :, :M], (1, 2), (-2, -1))  # (deriv+1, *x.shape, L, M)
+    return out[0] if deriv == 0 else tuple(out[i] for i in range(deriv + 1))
+
+
+def assoc_legendre(L, M, x, deriv=0):
+    """P_l^m(x) for 0<=l<L, 0<=m<M (Condon-Shortley phase), shape x.shape+(L,M).
+
+    deriv: 0 -> P; 1 -> (P, dP/dx); 2 -> (P, dP/dx, d2P/dx2). numpy routes to
+    scipy (byte-identical); jax/torch use the pure-backend Bonnet recurrence.
+    """
+    name, _ = _backend_special(get_namespace(x))
+    if name == "numpy":
+        return _scipy_assoc_legendre(L, M, x, deriv)
+    from ._fallback.assoc_legendre import assoc_legendre as _fb
+
+    return _fb(get_namespace(x), L, M, x, deriv)
+
+
+def gegenbauer(N, alpha, x):
+    """Gegenbauer polynomials C_n^alpha(x) for 0<=n<N, shape x.shape+(N,).
+
+    N static int, alpha scalar, x a backend array. Uses the three-term
+    recurrence on every backend (galpy's SCF radial basis never used a scipy
+    Gegenbauer, so there is no native to prefer)."""
+    from ._fallback.gegenbauer import gegenbauer as _fb
+
+    return _fb(get_namespace(x), N, alpha, x)
+
+
 def xlogy(x, y):
     # x * log(y), with the scipy/native convention 0 * log(0) = 0.
     from ._fallback.xlogy import xlogy_fallback