probdiffeq 0.8.0

Probabilistic numerical solvers for differential equations

probdiffeq

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.

---

Related projects

• Tornadox
• ProbNumDiffEq.jl
• ProbNum

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