Skip to content

probdiffeq¤

CI PyPI version License Python versions

Probabilistic ODE solvers in JAX¤

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

Features¤

  • ⚡ 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

Installation¤

Install the latest release from PyPI:

pip install probdiffeq

This assumes JAX is already installed.

To install with JAX (CPU backend): 

pip install probdiffeq[cpu]

⚠️ Note: This is an active research project. Expect rough edges and breaking API changes.

Benchmarks¤

We maintain benchmarks comparing Probdiffeq against other solvers and libraries, including SciPy, JAX, and Diffrax.

Run benchmarks locally:

pip install .[example,test]
make benchmarks-run

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.

Citing¤

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.

Versioning¤

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!

You might also like¤

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