API reference¤

Compute the modified Bessel function of the second kind of order 0.

Equivalent to scipy.special.kn(0, z). See Zhang and Jin, SPECIAL_FUNCTIONS in FORTRAN77, for the algorithm: an ascending series below z = 9 and an asymptotic expansion above it.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k0(1.0)), 8)
0.42102444

Array input is evaluated elementwise, spanning both branches:

>>> [round(float(k), 8) for k in sp.k0(jnp.asarray([0.5, 5.0, 20.0]))]
[0.92441907, 0.0036911, 0.0]

Compute the exponentially scaled :math:e^z K_0(z).

Equivalent to scipy.special.k0e(z), which has no JAX counterpart. This is the form to reach for beyond z = 705, where :math:K_0(z) itself is smaller than any normal double and unrepresentable; :math:e^z K_0(z) decays only as :math:1/\sqrt{z} and stays accurate at any z.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k0e(1.0)), 8)
1.14446308

Where k0 has underflowed to zero, the scaled form is still exact:

>>> float(sp.k0(800.0))
0.0
>>> round(float(sp.k0e(800.0)), 10)
0.0443044275

Compute the modified Bessel function of the second kind of order 1.

Obtained from k0 through the Wronskian :math:I_0(z) K_1(z) + I_1(z) K_0(z) = 1/z.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k1(1.0)), 8)
0.60190723

>>> [round(float(k), 8) for k in sp.k1(jnp.asarray([0.5, 5.0, 20.0]))]
[1.65644112, 0.00404461, 0.0]

Compute the exponentially scaled :math:e^z K_1(z).

Equivalent to scipy.special.k1e(z), which has no JAX counterpart. Obtained from k0e through the Wronskian :math:I_0(z) K_1(z) + I_1(z) K_0(z) = 1/z, in the scaled variables.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k1e(1.0)), 8)
1.63615349
>>> round(float(sp.k1e(800.0)), 10)
0.0443321091

Compute the modified Bessel function of the second kind of order 2.

Obtained from the recurrence :math:K_2(z) = K_0(z) + (2/z) K_1(z).

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k2(1.0)), 8)
1.6248389

>>> [round(float(k), 8) for k in sp.k2(jnp.asarray([0.5, 5.0, 20.0]))]
[7.55018355, 0.00530894, 0.0]

Compute the exponentially scaled :math:e^z K_2(z).

Equivalent to scipy.special.kve(2, z), which has no JAX counterpart. Obtained from the recurrence :math:K_2(z) = K_0(z) + (2/z) K_1(z), which the scaling leaves unchanged.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k2e(1.0)), 8)
4.41677005
>>> round(float(sp.k2e(800.0)), 10)
0.0444152578

Compute the polylogarithm :math:\mathrm{Li}_n(z).

There is no scipy.special counterpart; mpmath.polylog is the reference used by the test suite.

Three series are stitched together: the defining sum for :math:|z| \le 1/2, the Hurwitz-zeta expansion in :math:\log z for :math:1/2 < |z| < 2, and the inversion formula for :math:|z| \ge 2.

Parameters:

Returns:

Examples:

>>> import spexial as sp

Li_1(z) == -log(1 - z):

>>> round(float(sp.polylog(1, 0.5)), 12)
0.69314718056

Li_2(1) == zeta(2):

>>> round(float(sp.polylog(2, 1.0)), 10)
1.6449340668

>>> round(float(sp.polylog(3, -1.0)), 8)
-0.90154268

Compute the number of combinations of N things taken k at a time.

This is the "N choose k" factor :math:N! / (k! (N-k)!). It is the inexact variant of scipy.special.comb -- the exact=False one -- and is evaluated through jax.scipy.special.gammaln, so N and k need not be integers and the result is a float, not a Python int.

Reference: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.comb.html

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.comb(5, 2)), 5)
10.0

Rounded to five places, not twelve, because these examples have to be true at the library's default precision as well as under x64. Without x64 this is 9.999995: jax.scipy.special.gammaln returns 2**-21 rather than 0 at a float32 argument of exactly 1 (spexial#29), so comb(N, 0) -- which reduces to exp(-gammaln(1)) -- is up to 8 ulps below the exact 1 it should be.

Up to 8, and not a single value: measured over N from 0 to FLT_MAX, the shortfall is 0, 4, 6 or 8 ulps depending on N, as the two branch cross-overs and gammaln's own rounding move under it. It is 0 from about N = 8.9e6, where the asymptotic form takes over, and 0 again at a scattering of smaller N where the subtraction happens to round away. Naming individual N here would be describing the rounding rather than the library. float16 and bfloat16 round it away entirely, at every N. float64 is exact throughout.

Out-of-range k gives 0, not nan:

>>> [
...     round(float(c), 5)
...     for c in sp.comb(jnp.asarray([5, 5, 5]), jnp.asarray([2, 6, -1]))
... ]
[10.0, 0.0, 0.0]

Non-integer arguments are the generalized binomial coefficient:

>>> round(float(sp.comb(5.5, 2)), 5)
12.375

Evaluate the Gegenbauer polynomial :math:C_n^{(\alpha)}(x).

The Gegenbauer polynomials can be defined via the Gauss hypergeometric function :math:{}_2F_1 as

.. math::

C_n^{(\alpha)} = \frac{(2\alpha)_n}{\Gamma(n + 1)}
  {}_2F_1(-n, 2\alpha + n; \alpha + 1/2; (1 - z)/2).

When :math:n is an integer the result is a polynomial of degree :math:n. See 22.5.46 in [AS]_ for details.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> float(sp.eval_gegenbauer(3, 1.0, 0.5))
-1.0

It broadcasts over x:

>>> sp.eval_gegenbauer(2, 1.0, jnp.asarray([-1.0, 0.0, 1.0])).tolist()
[3.0, -1.0, 3.0]

Return the Gegenbauer polynomial of degree n and all lower ones.

There is no scipy.special counterpart; it is the by-product of the three-term recurrence that eval_gegenbauer runs anyway, and is much cheaper than n + 1 separate calls.

The Gegenbauer polynomials can be defined via the Gauss hypergeometric function :math:{}_2F_1 as

.. math::

C_n^{(\alpha)} = \frac{(2\alpha)_n}{\Gamma(n + 1)}
  {}_2F_1(-n, 2\alpha + n; \alpha + 1/2; (1 - z)/2).

When :math:n is an integer the result is a polynomial of degree :math:n. See 22.5.46 in [AS]_ for details.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> sp.eval_gegenbauers(3, 1.0, 0.5).tolist()
[1.0, 1.0, 0.0, -1.0]

The degree-0 case is a single value, not two:

>>> sp.eval_gegenbauers(0, 1.0, 0.5).tolist()
[1.0]

alpha and x broadcast, so a whole table of orders against parameters comes out of the one recurrence:

>>> import jax.numpy as jnp
>>> alpha = jnp.asarray([1.0, 2.0])[:, None]
>>> x = jnp.asarray([0.0, 0.5, 1.0])
>>> sp.eval_gegenbauers(3, alpha, x).shape
(4, 2, 3)

Compute the gamma function :math:\Gamma(x).

The value is jax.scipy.special.gamma, called directly, so it cannot drift from upstream. What this adds is the derivative: JAX differentiates its own implementation term by term, while :math:\Gamma'(x) = \Gamma(x)\,\psi(x) is one extra call. The saving is in memory, not time: measured over 10,000 points, jax.grad keeps 240 kB of residuals through the backward pass against this rule's 80 kB, while wall-clock is a wash (171 µs against 166 µs, i.e. 1.03x -- neutral within noise, reproduced across grad, vmap(grad) and jvp harnesses). An earlier revision of this docstring claimed 4.3x faster; that was a measurement error, and it is not plausible either -- JAX's gamma is sign * exp(gammaln(x)), whose autodiff already is the same product, so there is no arithmetic to save.

Reference: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.gamma.html

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.gamma(5.0)), 10)
24.0

It broadcasts, and handles negative arguments:

>>> [round(float(g), 10) for g in sp.gamma(jnp.asarray([0.5, 1.0, -0.5]))]
[1.7724538509, 1.0, -3.5449077018]

Complex input works on jax >= 0.10.2 and returns the complex gamma function; it is not shown here because this example is executed against every supported JAX, including ones that predate it.

Unregularized incomplete beta function :math:B(a, b, z).

.. math::

B(a, b, z) = \int_0^z t^{a-1}(1-t)^{b-1}\,\mathrm{d}t

the DLMF 8.17.1 form. There is no scipy.special counterpart: SciPy's betainc is the regularized :math:I_z(a, b), and the two differ by the complete beta function :math:B(a, b).

Reconstructing this as beta(a, b) * betainc(a, b, z) works only for :math:b > 0. At :math:b \le 0 the complete beta function has a pole while the product does not, so that route returns nan over a domain this function handles without difficulty.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import jax.scipy.special as jsp
>>> import spexial as sp

It agrees with the regularized form wherever that is defined, including close to the :math:z \to 1 endpoint:

>>> a, b = 2.0, 1.5
>>> z = jnp.asarray([0.3, 0.999])
>>> bool(
...     jnp.allclose(
...         sp.incomplete_beta(a, b, z), jsp.beta(a, b) * jsp.betainc(a, b, z)
...     )
... )
True

But unlike that product it stays finite for b <= 0:

>>> round(float(sp.incomplete_beta(2.0, 0.0, jnp.asarray(0.5))), 8)
0.19314718

>>> bool(jnp.isnan(jsp.beta(2.0, 0.0) * jsp.betainc(2.0, 0.0, 0.5)))
True

Compute the modified Bessel function of the second kind of order 0.

Equivalent to scipy.special.kn(0, z). See Zhang and Jin, SPECIAL_FUNCTIONS in FORTRAN77, for the algorithm: an ascending series below z = 9 and an asymptotic expansion above it.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k0(1.0)), 8)
0.42102444

Array input is evaluated elementwise, spanning both branches:

>>> [round(float(k), 8) for k in sp.k0(jnp.asarray([0.5, 5.0, 20.0]))]
[0.92441907, 0.0036911, 0.0]

Compute the exponentially scaled :math:e^z K_0(z).

Equivalent to scipy.special.k0e(z), which has no JAX counterpart. This is the form to reach for beyond z = 705, where :math:K_0(z) itself is smaller than any normal double and unrepresentable; :math:e^z K_0(z) decays only as :math:1/\sqrt{z} and stays accurate at any z.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k0e(1.0)), 8)
1.14446308

Where k0 has underflowed to zero, the scaled form is still exact:

>>> float(sp.k0(800.0))
0.0
>>> round(float(sp.k0e(800.0)), 10)
0.0443044275

Compute the modified Bessel function of the second kind of order 1.

Obtained from k0 through the Wronskian :math:I_0(z) K_1(z) + I_1(z) K_0(z) = 1/z.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k1(1.0)), 8)
0.60190723

>>> [round(float(k), 8) for k in sp.k1(jnp.asarray([0.5, 5.0, 20.0]))]
[1.65644112, 0.00404461, 0.0]

Compute the exponentially scaled :math:e^z K_1(z).

Equivalent to scipy.special.k1e(z), which has no JAX counterpart. Obtained from k0e through the Wronskian :math:I_0(z) K_1(z) + I_1(z) K_0(z) = 1/z, in the scaled variables.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k1e(1.0)), 8)
1.63615349
>>> round(float(sp.k1e(800.0)), 10)
0.0443321091

Compute the modified Bessel function of the second kind of order 2.

Obtained from the recurrence :math:K_2(z) = K_0(z) + (2/z) K_1(z).

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.k2(1.0)), 8)
1.6248389

>>> [round(float(k), 8) for k in sp.k2(jnp.asarray([0.5, 5.0, 20.0]))]
[7.55018355, 0.00530894, 0.0]

Compute the exponentially scaled :math:e^z K_2(z).

Equivalent to scipy.special.kve(2, z), which has no JAX counterpart. Obtained from the recurrence :math:K_2(z) = K_0(z) + (2/z) K_1(z), which the scaling leaves unchanged.

Parameters:

Returns:

Examples:

>>> import spexial as sp

>>> round(float(sp.k2e(1.0)), 8)
4.41677005
>>> round(float(sp.k2e(800.0)), 10)
0.0444152578

Compute the polylogarithm :math:\mathrm{Li}_n(z).

There is no scipy.special counterpart; mpmath.polylog is the reference used by the test suite.

Three series are stitched together: the defining sum for :math:|z| \le 1/2, the Hurwitz-zeta expansion in :math:\log z for :math:1/2 < |z| < 2, and the inversion formula for :math:|z| \ge 2.

Parameters:

Returns:

Examples:

>>> import spexial as sp

Li_1(z) == -log(1 - z):

>>> round(float(sp.polylog(1, 0.5)), 12)
0.69314718056

Li_2(1) == zeta(2):

>>> round(float(sp.polylog(2, 1.0)), 10)
1.6449340668

>>> round(float(sp.polylog(3, -1.0)), 8)
-0.90154268

Compute Spence's function -- the dilogarithm -- for real or complex input.

jax.scipy.special.spence covers the real case at every JAX version spexial supports, and raises on complex input. This accepts both, and carries an analytic derivative: differentiating the series instead costs 444x the time and 1652x the residual memory (26.4 MB against 16 kB over 2000 points; the near_root fallback added ~2 MB to the differentiated series, which is exactly the cost a custom rule avoids).

.. math::

\int_{1}^{z} dt \frac{\log(t)}{1 - t}

There is no upper limit on a real argument. Above z = 1/tiny the value comes from the inversion formula rather than from upstream, whose opening 1/z is subnormal there and is flushed into the wrong branch; see the comment on that branch. spence(inf) is nan, as in SciPy.

Spherical harmonic :math:Y_n^m(\theta, \phi).

.. math::

Y_n^m(\theta, \phi) = \bar{P}_n^m(\cos\theta)\, e^{im\phi}

with sph_legendre_p supplying :math:\bar{P}_n^m. Same convention as scipy.special.sph_harm_y and jax.scipy.special.sph_harm_y: theta is the polar angle and phi the azimuth, and the Condon-Shortley phase is included.

This differs from jax.scipy.special.sph_harm_y in two ways that are the reason it exists, both measured on JAX 0.11.1. Upstream pairs n[i] with theta[i] positionally instead of broadcasting, so a scalar degree against a batch of angles is right only at index 0 -- up to 1.18 absolute error for :math:n \le 3 -- and it rejects 0-d input. And upstream's derivatives are nan at :math:\theta = 0, \pi for every :math:n \ge 1. Here n and m are static, so there is nothing to mispair, and theta/phi broadcast against each other at any rank.

.. note::

The **Cartesian** gradient of an :math:`m \ge 1` harmonic is still
exactly zero on the z-axis if it is obtained by differentiating through
:math:`\theta` and :math:`\phi`, which have no directional derivative
there. That is a property of the coordinates, not of any
implementation. Use `sph_harm_y_cart` when the gradient on the axis
matters.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> complex(sp.sph_harm_y(0, 0, 0.4, 1.2)).real
0.28209479177387814

It broadcasts, which upstream does not:

>>> theta = jnp.asarray([0.3, 1.1, 2.0])
>>> sp.sph_harm_y(2, 1, theta, 0.0).real.round(6).tolist()
[-0.218107, -0.312301, 0.292333]

Spherical harmonic :math:Y_n^m from a Cartesian unit direction.

.. math::

Y_n^m = q_n^m(\hat{z}) \left(\hat{x} + i\hat{y}\right)^m,
\qquad q_n^m(u) = N_{nm} \frac{P_n^m(u)}{(1-u^2)^{m/2}}

Mathematically identical to sph_harm_y(n, m, arccos(uz), arctan2(uy, ux)), and numerically better behaved in the one place it matters: the right-hand side is a polynomial in :math:\hat{x} and :math:\hat{y}, so it is smooth on the z-axis, where :math:\theta and :math:\phi are singular.

That singularity is not cosmetic. Differentiating the :math:(\theta, \phi) form through the chain rule gives a Cartesian gradient of exactly zero for every :math:m \ge 1 term on the axis -- neither angle has a directional derivative there -- against a non-zero true limit. There is no counterpart to this function in SciPy or JAX; the name is scipy.special.sph_harm_y plus the argument convention that distinguishes it.

Parameters:

Returns:

Examples:

>>> import jax, jax.numpy as jnp
>>> import spexial as sp

>>> uvec = jnp.asarray([0.0, 0.0, 1.0])
>>> round(complex(sp.sph_harm_y_cart(1, 0, uvec)).real, 10)
0.4886025119

On the z-axis the gradient of an m = 1 term is finite and non-zero, where the spherical form gives exactly zero:

>>> g = jax.grad(lambda v: sp.sph_harm_y_cart(1, 1, v).real)(uvec)
>>> round(float(g[0]), 6)
-0.345494

Every :math:Y_l^k with :math:l \le n and :math:\lvert k \rvert \le m.

The same values sph_harm_y_cart returns pair by pair, but with both recurrences carried across the table instead of restarted for each entry: one pass per order advances :math:((x+iy)/r)^k by a single complex multiply and walks the Legendre recurrence up in :math:l from its seed at :math:l = k. That makes the table :math:O(n^2) rather than cubic.

Since the degree and order are static the saving is in traced operations, so it shows up as a smaller HLO -- and hence faster tracing and compilation -- rather than as faster execution, which XLA's fusion had already largely recovered. Measured on the equivalent code in galax, tracing was 3-4.6x faster for :math:n \le 20 with run time flat.

Layout and argument order follow scipy.special.sph_harm_y_all, which has no JAX counterpart; the difference is only that the direction is Cartesian.

.. warning::

Reach for this when you want the **table**. If instead you are about
to index it and reduce -- summing :math:`\\sum_{lm} c_{lm} Y_l^m`,
say -- call `sph_harm_y_cart` per pair and fold each term into the sum
as it is produced. Indexing a stacked table defeats XLA's fusion, so
the whole thing is materialized: measured on a multipole expansion at
:math:`n = 12` over a million directions, the table form ran in 17.7 s
against 10 ms for per-pair calls, for identical values. The saving
here is in *traced* operations -- roughly 2.3x less tracing and
compiling at :math:`n = 20` -- which is worth having only when the
table itself is the thing you need.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> uz = jnp.asarray([0.0, 0.0, 1.0])
>>> Y = sp.sph_harm_y_cart_all(2, 2, uz)
>>> Y.shape
(3, 5)

Entries with |j| > i are zero, since no such harmonic exists:

>>> bool((Y[0, 1] == 0) & (Y[1, 2] == 0))
True

Each entry matches the one-pair function, negative orders included:

>>> bool(jnp.isclose(Y[2, -1], sp.sph_harm_y_cart(2, -1, uz)))
True

Every :math:Y_l^k as a separate array, sharing the recurrences.

Same values, same indexing and the same layout as sph_harm_y_cart_all -- terms[i][j] is :math:Y_i^j, with negative orders reachable at the end of the inner tuple by ordinary negative indexing, as in SciPy. The difference is entirely the container: a nested tuple of arrays rather than one stacked array.

That difference is a performance decision, and a large one. Indexing a stacked table stops XLA folding each term into a caller's reduction as it is produced, so the whole table is materialized instead. Measured on a multipole expansion at :math:n = 12 over a million directions, summing :math:\sum_{lm} c_{lm} Y_l^m from sph_harm_y_cart_all ran in 17.7 s against 10 ms from these terms, for identical values.

So: use this when you are going to reduce over the table, and sph_harm_y_cart_all when you want the table itself.

Deliberately not jax.jit-decorated, unlike everything else here. A jitted function returning a pytree materializes each leaf as its own output buffer at the call boundary, which is exactly the fusion this exists to preserve. It is a pure trace-time helper: it runs inside the caller's trace, and the caller is free to jit around it.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> uz = jnp.asarray([0.0, 0.0, 1.0])
>>> terms = sp.sph_harm_y_cart_all_terms(2, 2, uz)
>>> len(terms), len(terms[0])
(3, 5)

Indexed exactly as the stacked table is, negative orders included:

>>> bool(jnp.isclose(terms[2][-1], sp.sph_harm_y_cart(2, -1, uz)))
True

Spherical Legendre function :math:\bar{P}_n^m(\cos\theta).

.. math::

\bar{P}_n^m(\cos\theta) = \sqrt{\frac{2n+1}{4\pi}
    \frac{(n-m)!}{(n+m)!}}\; P_n^m(\cos\theta)

the normalization for which :math:Y_n^m(\theta,\phi) = \bar{P}_n^m(\cos\theta)\,e^{im\phi}. Includes the Condon-Shortley phase, matching scipy.special.sph_legendre_p, scipy.special.lpmv and GSL's gsl_sf_legendre_sphPlm.

jax.scipy.special has no counterpart at any version -- the normalized Legendre function is reachable only through jax.scipy.special.sph_harm_y, and then only as part of a complex harmonic.

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

:math:\bar{P}_0^0 = 1/\sqrt{4\pi} everywhere:

>>> round(float(sp.sph_legendre_p(0, 0, 0.7)), 10)
0.2820947918

Negative orders follow SciPy, :math:\bar{P}_n^{-m} = (-1)^m \bar{P}_n^{m}:

>>> bool(jnp.isclose(sp.sph_legendre_p(2, -1, 0.3), -sp.sph_legendre_p(2, 1, 0.3)))
True

Unlike upstream's, the derivative is finite at the pole:

>>> import jax
>>> float(jax.grad(lambda t: sp.sph_legendre_p(1, 0, t))(0.0))
-0.0

Compute the Riemann zeta function :math:\zeta(n).

Differs from jax.scipy.special.zeta in that negative arguments are supported, through the functional equation :math:\zeta(-k) = (-1)^k B_{k+1} / (k+1).

Reference: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.zeta.html

Parameters:

Returns:

Examples:

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> round(float(sp.zeta(2.0)), 10)
1.6449340668

Negative integers use the functional equation:

>>> [round(float(z), 12) for z in sp.zeta(jnp.asarray([0.0, -1.0, -2.0, -3.0]))]
[-0.5, -0.083333333333, 0.0, 0.008333333333]

zeta(-3) == 1 / 120:

>>> float(1 / 120)
0.008333333333333333

Large arguments are exactly 1, where jax.scipy.special.zeta gives nan:

>>> float(sp.zeta(1e16))
1.0