quadax

quadax is a library for numerical quadrature and integration using JAX.

  • vmap-able, jit-able, differentiable in both forward and reverse mode.

  • Pluggable adjoint methods controlling how derivatives are computed.

  • Scalar or vector valued integrands.

  • Finite or infinite domains with discontinuities or singularities within the domain of integration.

  • Globally adaptive Gauss-Kronrod and Clenshaw-Curtis quadrature for smooth integrands (similar to scipy.integrate.quad)

  • Adaptive tanh-sinh quadrature for singular or near singular integrands.

  • Quadrature from sampled values using trapezoidal and Simpsons methods.

  • Globally adaptive cubature over a region in several dimensions, whose limits may depend on the other coordinates (similar to scipy.integrate.nquad).

Coming soon:

  • QMC methods

  • Integration with weight functions

  • Sparse grids (maybe, need to play with data structures and JAX)

Installation

quadax is installable with pip:

pip install quadax

Usage

import jax.numpy as jnp
import numpy as np
from quadax import quadgk

fun = lambda t: t * jnp.log(1 + t)

epsabs = epsrel = 1e-5 # by default jax uses 32 bit, higher accuracy requires going to 64 bit
a, b = 0, 1
y, info = quadgk(fun, [a, b], epsabs=epsabs, epsrel=epsrel)
assert info.err < max(epsabs, epsrel*abs(y))
np.testing.assert_allclose(y, 1/4, rtol=epsrel, atol=epsabs)

User Guide

Reference

Indices and tables