jaxellip

Sina Atalay

2025

My PhD research at Princeton University focused on optimizing the geometry of high-temperature superconducting magnets to achieve higher magnetic fields and improved quench protection by leveraging the anisotropic properties of superconducting materials using computational, simulation-based methods.

Our work involved computing elliptic integrals. However, we were using JAX so that the objective functions could be automatically differentiated and optimized efficiently on GPUs.

JAX did not provide built-in implementations of these elliptic integrals, so I created a package implementing them.

The package provides JAX implementations of ellipk $K(m)$ , ellipkm1 $K(1 - p)$ , and ellipe $E(m)$ , matching the API of SciPy:

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

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

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

ellipk and ellipkm1 are the same function $K$ with different parameterizations: ellipk takes $m$ and returns $K(m)$ , while ellipkm1 takes $p = 1 - m$ , the distance from the singular point $m = 1$ , and returns $K(1 - p)$ . The split exists because $K(m)$ diverges logarithmically as $m \to 1$ , and near this point the result hinges on the exact value of $1 - m$ . Taking $p$ as the argument make ellipkm1 remain accurate arbitrarily close to the singularity, at the cost of a small amount of additional computation.

Under the hood, ellipk and ellipe use the same Cephes minimax polynomial approximations as SciPy: polynomials in the complementary parameter $1 - m$ plus logarithmic terms, with a Landen transformation for negative $m$ . ellipkm1 evaluates $K(1 - p)$ directly from $p$ , using a small- $p$ logarithmic expansion below $p = 10^{-8}$ .

Automatic differentiation is implemented with custom JVP rules. The derivative formulas reduce to Carlson’s symmetric integral $R_D$ , which is evaluated using a fixed-length duplication iteration with the zero argument handled analytically. The tests compare values against high-precision mpmath references and scipy.special, and check forward-, reverse-, and second-order JAX derivatives.

Check out the source code on GitHub: github.com/sinaatalay/jaxellip

Check out the package on PyPI: pypi.org/project/jaxellip