Diffrax v0.5.1

New research paper

One of the new features of this release is the simulation of space-time Lévy area, over arbitrary intervals, deterministically with respect to a PRNG key. This is a required component for adaptive step-size and higher-order SDE solvers in particular.

Well, this turned out to be an unsolved research question! And so with huge credit to @andyElking for diligently figuring out all the details -- we now have a new paper on arXiv discussing all the details! So far as we know this sets a new state-of-the-art for numerical Brownian simulation, and is what now powers all of the numerical SDE solving inside Diffrax.

If you're interested in numerical methods for SDEs, then check out the arxiv paper here.

New features for SDEs

• Added a suite of Stochastic Runge--Kutta methods! These are higher-order solvers for SDEs, in particular when the noise has a particular form (additive, commutative, ...). A huge thank-you to @andyElking for implementing all of these:

  • GeneralShARK: recommended when the drift is expensive to evaluate;
  • SEA: recommended when the noise is additive and wanting a solve that is cheap-and-low-accuracy;
  • ShARK: recommended default choice when the noise is additive;
  • SlowRK: recommended for commutative noise;
  • SPaRK: recommended when performing adaptive time stepping;
  • SRA1: alternative to ShARK (this is a now-classical SRK method).

• Added support for simulating space-time Lévy Area to VirtualBrownianTree and UnsafeBrownianPath. This is the bit discussed in the "new research paper" section above! The main thing here is the ability to sample from random variables like space-time Lévy area, which is a doubly-indexed integral of Brownian motion over time: $H_{s,t} = \frac{1}{t-s} \int_s^t ((W_r - W_s) - \frac{r-s}{t-s} (W_t - W_s)) dr$.

New features for all differential equations

• Added TextProgressMeter and TqdmProgressMeter, which can be used to track how far through a differential equation solve things have progressed. (Thanks @abocquet! #357, #398)
• Added support for using adaptive step size controllers on TPUs (Thanks @stefanocortinovis! #366, #369)
• All AbstractPaths are now typing.Generics parameterised by their return type; all AbstractTerms are now typing.Generics parameterised by their vector field and control. (Thanks @tttc3! #359, #364)

Other

• Improved documentation for PIDController (Thanks @ParticularlyPythonicBS! #371, #372)
• Now have a py.typed file to declare that we're static-type-checking compatible. (Thanks @lockwo! #408)
• Bugfix for CubicInterpolation when there is nothing to interpolate. (Thanks @allen-adastra! #360)
• Compatibility with future versions of JAX by removing the now-deprecated jax.config import. (Thanks @jakevdp! #377)

New Contributors

• @allen-adastra made their first contribution in #360
• @stefanocortinovis made their first contribution in #369
• @jakevdp made their first contribution in #377
• @ParticularlyPythonicBS made their first contribution in #372

Full Changelog: v0.5.0...v0.5.1