https://github.com/f0uriest/quadax
Numerical quadrature with JAX
https://github.com/f0uriest/quadax
Last synced: 3 months ago
JSON representation
Numerical quadrature with JAX
- Host: GitHub
- URL: https://github.com/f0uriest/quadax
- Owner: f0uriest
- License: mit
- Created: 2023-10-24T04:44:34.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-04-08T06:36:09.000Z (3 months ago)
- Last Synced: 2025-04-13T20:18:30.279Z (3 months ago)
- Language: Python
- Size: 291 KB
- Stars: 56
- Watchers: 4
- Forks: 2
- Open Issues: 5
-
Metadata Files:
- Readme: README.rst
- Changelog: CHANGELOG.md
- Contributing: CONTRIBUTING.rst
- License: LICENSE
- Citation: CITATION.cff
Awesome Lists containing this project
- awesome-jax - quadax - Numerical quadrature with JAX. <img src="https://img.shields.io/github/stars/f0uriest/quadax?style=social" align="center"> (Libraries)
- awesome-jax - quadax - Numerical quadrature with JAX. <img src="https://img.shields.io/github/stars/f0uriest/quadax?style=social" align="center"> (Libraries)
README
########
quadax
########
|License| |DOI| |Issues| |Pypi|
|Docs| |UnitTests| |Codecov|
quadax is a library for numerical quadrature and integration using JAX.
- ``vmap``-able, ``jit``-able, differentiable.
- 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.
Coming soon:
- Custom JVP/VJP rules (currently AD works by differentiating the loop which isn't the most efficient.)
- N-D quadrature (cubature)
- QMC methods
- Integration with weight functions
- Sparse grids (maybe, need to play with data structures and JAX)
Installation
============
quadax is installable with `pip`:
.. code-block:: console
pip install quadax
Usage
=====
.. code-block:: python
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)
For full details of various options see the `API documentation `__
.. |License| image:: https://img.shields.io/github/license/f0uriest/quadax?color=blue&logo=open-source-initiative&logoColor=white
:target: https://github.com/f0uriest/quadax/blob/master/LICENSE
:alt: License
.. |DOI| image:: https://zenodo.org/badge/709132830.svg
:target: https://zenodo.org/doi/10.5281/zenodo.10035983
:alt: DOI
.. |Docs| image:: https://img.shields.io/readthedocs/quadax?logo=Read-the-Docs
:target: https://quadax.readthedocs.io/en/latest/?badge=latest
:alt: Documentation
.. |UnitTests| image:: https://github.com/f0uriest/quadax/actions/workflows/unittest.yml/badge.svg
:target: https://github.com/f0uriest/quadax/actions/workflows/unittest.yml
:alt: UnitTests
.. |Codecov| image:: https://codecov.io/github/f0uriest/quadax/graph/badge.svg?token=MB11I7WE3I
:target: https://codecov.io/github/f0uriest/quadax
:alt: Coverage
.. |Issues| image:: https://img.shields.io/github/issues/f0uriest/quadax
:target: https://github.com/f0uriest/quadax/issues
:alt: GitHub issues
.. |Pypi| image:: https://img.shields.io/pypi/v/quadax
:target: https://pypi.org/project/quadax/
:alt: Pypi