Special Functions
Special functions implemented by copulAX.
Currently the following are implemented: - kv - stdtr - igammainv - igammacinv - digamma - trigamma
Modified Bessel function of the second kind and Student’s t CDF.
References
Bessel function integral representation: https://dlmf.nist.gov/10.32#E10
Asymptotic forms: https://dlmf.nist.gov/10.30
- copulax._src.special.log_kv(v, x)[source]
Log of the modified Bessel function of the second kind, \(\log K_v(x)\).
Computes \(\log K_v(x)\) directly in log space, remaining finite and accurate for arbitrarily large \(x\), avoiding the underflow that occurs in \(K_v(x)\).
Pure JAX, JIT-compatible, and differentiable w.r.t. both v and x. Gradients use a hand-written
custom_jvp()rule based on the classical Bessel recurrence for∂/∂xand the integral representation for∂/∂v. Underjax.gradthe backward trace is therefore small (nolax.condunrolling, no differentiation through the 64-node quadrature).Four forward evaluation regimes, selected automatically:
v ≥ 15: Debye uniform asymptotic expansion (DLMF 10.41.3), 6-term series with Olver’s polynomials.
x < 10⁻⁸: small-x leading asymptotics (DLMF 10.30/10.31).
x > max(40, 2v²+20): large-x Hankel expansion (DLMF 10.40.2), 4-term series.
Otherwise: Gauss-Legendre quadrature (64 points) on the integral \(K_v(x) = \int_0^\infty \cosh(vt)\,e^{-x\cosh t}\,dt\) (DLMF 10.32.9), with saddle-point-centred integration interval.
- copulax._src.special.kv(v, x)[source]
Modified Bessel function of the second kind, \(K_v(x)\).
Convenience wrapper:
kv(v, x) = exp(log_kv(v, x)).Computes \(\log K_v(x)\) directly in log space, remaining finite and accurate for arbitrarily large \(x\), avoiding the underflow that occurs in \(K_v(x)\).
Pure JAX, JIT-compatible, and differentiable w.r.t. both v and x. Gradients use a hand-written
custom_jvp()rule based on the classical Bessel recurrence for∂/∂xand the integral representation for∂/∂v. Underjax.gradthe backward trace is therefore small (nolax.condunrolling, no differentiation through the 64-node quadrature).Four forward evaluation regimes, selected automatically:
v ≥ 15: Debye uniform asymptotic expansion (DLMF 10.41.3), 6-term series with Olver’s polynomials.
x < 10⁻⁸: small-x leading asymptotics (DLMF 10.30/10.31).
x > max(40, 2v²+20): large-x Hankel expansion (DLMF 10.40.2), 4-term series.
Otherwise: Gauss-Legendre quadrature (64 points) on the integral \(K_v(x) = \int_0^\infty \cosh(vt)\,e^{-x\cosh t}\,dt\) (DLMF 10.32.9), with saddle-point-centred integration interval.
- copulax._src.special.log_kv_plus_s_log_r(s, r)[source]
Numerically stable
log K_s(r) + s · log(r)fors > 0.Motivation
The skewed-t (and multivariate skewed-t) log-PDF contains the sum
\[\log K_s(r) + \tfrac{s}{2}\log\bigl((\nu+Q)\,R\bigr) = \log K_s(r) + s \log r\]where
r = sqrt((ν+Q)·R)andR = (γ/σ)². Asγ → 0both terms diverge individually (log K_s(r) → +∞,s log r → -∞) but the sum stays finite. Computing them separately as two log-space objects and then subtracting is lossy; this helper computes the sum as one object so the divergents log rparts cancel arithmetically in float64.Closed-form limit
From DLMF 10.30.2 (
K_ν(z) ~ ½ Γ(ν) (2/z)^νasz → 0⁺forν > 0):\[\log K_s(r) + s \log r = \log\Gamma(s) + (s-1)\log 2 + O(r^{2s}) \to \log\Gamma(s) + (s-1)\log 2 \quad \text{as } r \to 0.\]The
O(r^{2s})tail is the next-order correction from the subtractedI_sterm inK_s = (π/2)(I_{-s} - I_s)/sin(sπ); the helper returns the full combination (limit + tail), not the truncated limit, so the output is correct for everyr ≥ 0to the precision at whichlog_kv_small_xitself is accurate.Implementation
Computes
log_kv(s, r_safe) + s · log(r_safe)directly, where\[r_{\text{safe}} = \max(r, 10^{-12}).\]At
r = 0the floor picks up and the helper returnslgamma(s) + (s-1) log 2 + O(10^{-24s})— indistinguishable from the true limit at float64 precision for everys(includings = 0.5, where the tail isO(r^{2s}) = O(r)and soO(10^{-12})well below any fitting tolerance).For
rstrictly positive and above the floor the helper is exactly the direct sumlog K_s(r) + s log r.Catastrophic cancellation inside
log_kv_small_xfors ≥ 1(wherelog_kv(s, r) = lgamma(s) + (s-1) log 2 − s log rand the−s log r + s log rterms cancel arithmetically) costs at mosteps · s · log(r_safe)relative —≲ 1e-13even at the floor.
Gradient safety
The
jnp.maximumfloor ensureslog(0)andlog_kv’sx == 0special case never fire, sojax.gradthrough the helper does not encounternanorinf. Atr = 0the gradient∂output/∂ris0(the floor pinsr_safeto a constant), matching the analytic derivative of theO(r^{2s})tail in the limit.- type s:
Union[float,int,Array,ndarray,bool,number,bool,complex]- param s:
Bessel order (scalar,
s > 0).- type r:
Union[Array,ndarray,bool,number,bool,int,float,complex]- param r:
Non-negative argument (array-like).
- rtype:
- returns:
Array of
log K_s(r) + s · log(r)values, same shape asr, finite everywhere onr ≥ 0.
Notes
Not a drop-in replacement for
log_kv— it returns the sum, notlog K_salone. Use this only where the caller needs the cancellation-preserved combination (skewed-t log-PDF, and any future GH / GIG log-density rewrite).Requires
s > 0. The degenerate cases = 0(whereK_0is logarithmic, not power) is not handled here —lgamma(0) = +infwould leak into the small-x formula.
References
DLMF 10.30.2 (leading asymptotic), 10.30/10.31 (next-order correction, giving the
O(r^{2s})tail).
- copulax._src.special.digamma(x)[source]
Digamma function, \(\psi(x) = \frac{d}{dx} \ln \Gamma(x)\).
Computed via automatic differentiation of
gammaln, which is exact (not a finite-difference approximation).
- copulax._src.special.trigamma(x)[source]
Trigamma function, \(\psi'(x) = \frac{d^2}{dx^2} \ln \Gamma(x)\).
Computed via automatic differentiation of
gammaln, which is exact (not a finite-difference approximation).
- copulax._src.special.igammainv(a, p)[source]
Inverse of the regularized lower incomplete gamma function.
Finds \(x\) such that \(\mathrm{gammainc}(a, x) = p\).