Probabilistic solvers in JAX

Probdiffeq implements adaptive probabilistic numerical solvers for differential equations (ODEs). It builds on JAX, thus inheriting automatic differentiation, vectorisation, and GPU acceleration.

⚠️ Probdiffeq is an active research project. Expect rough edges and sudden API changes.

Features:

• ⚡ Automatic calibration and step-size adaptation
• ⚡ Stable implementations of filtering, smoothing, and other estimation strategies
• ⚡ Custom information operators, dense output, posterior sampling, and prior distributions.
• ⚡ Efficient handling of high-dimensional problems through state-space model factorisations
• ⚡ Parameter estimation
• ⚡ Taylor-series estimation with and without automatic differentiation
• ⚡ Seamless interoperability with Optax, BlackJAX, and other JAX-based libraries
• ⚡ Numerous examples (basic and advanced) -- see the documentation

Quickstart: See here for a minimal example to get you started.

Contributing: Contributions are very welcome!

• Browse open issues (look for “good first issue”).
• Check the developer documentation.
• Open an issue for feature requests or ideas.

Acknowledgements:

• Tornadox: One of Probdiffeq's precursors.
• ProbNumDiffEq.jl: A similar library in Julia
• ProbNum: Probabilistic numerics in Numpy.

The docs include guidance on migrating from these packages. Missing something? Open an issue or pull request!

You might also like:

• diffeqzoo: reference implementations of differential equations in NumPy and JAX
• probfindiff: probabilistic finite-difference methods in JAX