Probdiffeq implements adaptive probabilistic numerical solvers for ordinary differential equations (ODEs). It builds on JAX, thus inheriting automatic differentiation, vectorisation, and GPU acceleration.
- ⚡ Calibration and step-size adaptation
- ⚡ Stable implementations of filtering, smoothing, and other estimation strategies
- ⚡ Custom information operators, dense output, and posterior sampling
- ⚡ State-space model factorisations
- ⚡ Parameter estimation
- ⚡ Taylor-series estimation with and without Jets
- ⚡ Seamless interoperability with Optax, BlackJAX, and other JAX-based libraries
- ⚡ Numerous tutorials (basic and advanced) -- see the documentation
Install the latest release from PyPI:
pip install probdiffeqThis assumes JAX is already installed.
To install with JAX (CPU backend):
pip install probdiffeq[cpu]We maintain benchmarks comparing Probdiffeq against other solvers and libraries, including SciPy, JAX, and Diffrax.
Run benchmarks locally:
pip install .[example,test]
make benchmarks-runContributions are very welcome!
- Browse open issues (look for “good first issue”).
- Check the developer documentation.
- Open an issue for feature requests or ideas.
If you use Probdiffeq in your research, please cite:
@phdthesis{kramer2024implementing,
title={Implementing probabilistic numerical solvers for differential equations},
author={Kr{"a}mer, Peter Nicholas},
year={2024},
school={Universit{"a}t T{"u}bingen}
}The PDF explains the mathematics and algorithms behind this library.
For the solve-and-save-at functionality, cite:
@InProceedings{kramer2024adaptive,
title = {Adaptive Probabilistic ODE Solvers Without Adaptive Memory Requirements},
author = {Kr"{a}mer, Nicholas},
booktitle = {Proceedings of the First International Conference on Probabilistic Numerics},
pages = {12--24},
year = {2025},
editor = {Kanagawa, Motonobu and Cockayne, Jon and Gessner, Alexandra and Hennig, Philipp},
volume = {271},
series = {Proceedings of Machine Learning Research},
publisher = {PMLR},
url = {https://proceedings.mlr.press/v271/kramer25a.html}
}Link to the paper: PDF.
Link to the experiments: Code for experiments.
📌 Algorithms in Probdiffeq are based on multiple research papers. If you’re unsure which to cite, feel free to reach out.
Probdiffeq follows 0.MINOR.PATCH until its first stable release:
- PATCH → bugfixes & new features
- MINOR → breaking changes
See semantic versioning.
The docs include guidance on migrating from these packages. Missing something? Open an issue or pull request!
- diffeqzoo — reference implementations of differential equations in NumPy and JAX
- probfindiff — probabilistic finite-difference methods in JAX