probdiffeq¤
Probabilistic solvers for differential equations in JAX¤
ProbDiffEq implements adaptive probabilistic numerical solvers for initial value problems.
It inherits automatic differentiation, vectorisation, and GPU capability from JAX.
Features include:
- Stable implementation
- Calibration, step-size adaptation, and checkpointing
- State-space model factorisations
- Dense output and posterior sampling
- Filtering, smoothing, and many other backends
- Custom information operators
- Physics-enhanced regression
- Compatibility with other JAX-based libraries such as Optax or BlackJAX.
and many more.
- AN EASY EXAMPLE: LINK
- EXAMPLES: LINK
- CHOOSING A SOLVER: LINK
- API DOCUMENTATION: LINK
- ISSUE TRACKER: LINK
- BENCHMARKS: LINK
- CHANGELOG: LINK
Installation¤
Get the most recent stable version from PyPi:
pip install probdiffeq
To install ProbDiffEq with jax[cpu]
, run
pip install probdiffeq[cpu]
WARNING: This is a research project. Expect rough edges and sudden API changes.
VERSIONING: As long as Probdiffeq is in its initial development phase (version 0.MINOR.PATCH), version numbers are increased as follows:
- Bugfixes and new features increase the PATCH version.
- Breaking changes increase the MINOR version.
See also: semantic versioning.
What's next?¤
Start with the quickstart, continue with the Solvers & Solutions
examples and only then move to the Parameter estimation
examples and the API documentation.
The examples show how to interact with the API, and explain some useful facts about probabilistic numerical solvers. While the API is not stable yet, the examples may be more instructive than the API docs.
The advanced examples show applications of probabilistic numerical solvers, often in conjunction with external libraries. For example, this notebook shows how to combine ProbDiffEq with Optax, and this notebook does the same with BlackJAX.
Contributing¤
Contributions are absolutely welcome! Start with checking the existing issues for a "good first issue" and have a look at the developer documentation.
If you have a feature that you would like to see implemented, create an issue!
Benchmarks¤
ProbDiffEq curates a range of benchmarks that includes various library-internal configurations but also other packages like SciPy, JAX, or Diffrax. To run the benchmark locally, install all dependencies via
pip install .[example,test]
docs/benchmarks
or execute all benchmarks via
make benchmarks-run
mkdocs serve
Similar projects¤
Check out how to transition from those packages: link.
Anything missing in this list? Please open an issue or make a pull request.
You might also like¤
- diffeqzoo: A library for example implementations of differential equations in NumPy and JAX
- probfindiff: Probabilistic numerical finite differences, in JAX.