jaxellip 0.0.1

Some elliptic integrals in JAX

jaxellip: JAX implementations of some elliptic integrals

Elliptic integrals are available in scipy.special, but are not implemented in jax.scipy.special, making them neither differentiable nor compatible with JAX’s JIT compilation. This package provides JAX-compatible implementations of several elliptic integrals from scipy.special: ellipk, ellipkm1, and ellipe. The results have been tested against scipy.special to ensure accuracy and performance (see tests).

Usage

1. Install jaxellip.

pip install jaxellip

2. Import jaxellip in your Python code and use it.

import jaxellip

jaxellip.ellipk(0.5) # Complete elliptic integral of the first kind
jaxellip.ellipkm1(0.5) # Complete elliptic integral of the first kind around m = 1
jaxellip.ellipe(0.5) # Complete elliptic integral of the second kind

Elliptic integrals

The elliptic integrals in jaxellip follow the same parameter conventions as scipy.special. In particular, the input parameter $m$ is the elliptic parameter, defined as $m = k^2$, where $k$ is the modulus. This differs from some literature or libraries (e.g., Boost or Abramowitz & Stegun), where the integrals are sometimes expressed in terms of $k$ directly.

The function ellipkm1(m) computes the complete elliptic integral of the first kind with argument $1 - m$, i.e., $K(1 - m)$. This is particularly useful for improved numerical stability and precision when $m$ is close to 1, since $K(m)$ diverges logarithmically as $m \to 1$.

Complete elliptic integral of the first kind

$$ K(m) = \int_0^{\pi/2} \frac{1}{\sqrt{1 - m \sin^2 \theta}} d\theta $$

Complete elliptic integral of the second kind

$$ E(m) = \int_0^{\pi/2} \sqrt{1 - m \sin^2 \theta} d\theta $$

Complete elliptic integral of the third kind

$$ \Pi(n, m) = \int_0^{\pi/2} \frac{1}{(1 - n \sin^2 \theta) \sqrt{1 - m \sin^2 \theta}} d\theta $$