SDE solvers¤
See also How to choose a solver.
Term structure
The type of solver chosen determines how the terms
argument of diffeqsolve
should be laid out. Most of them operate in the same way whether they are solving an ODE or an SDE, and as such expected that it should be a single AbstractTerm
. For SDEs that typically means a diffrax.MultiTerm
wrapping together a drift (diffrax.ODETerm
) and diffusion (diffrax.ControlTerm
). (Although you could also include any other term, e.g. an exogenous forcing term, if you wished.) For example:
drift = lambda t, y, args: -y
diffusion = lambda t, y, args: y[..., None]
bm = UnsafeBrownianPath(shape=(1,), key=...)
terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, bm))
diffeqsolve(terms, solver=Euler(), ...)
Some solvers are SDE-specific. For these, such as for example diffrax.StratonovichMilstein
, then terms
should be a 2-tuple (AbstractTerm, AbstractTerm)
, representing the drift and diffusion separately.
For those SDE-specific solvers then this is documented below, and the term structure is available programmatically under <solver>.term_structure
.
Explicit Runge--Kutta (ERK) methods¤
Each of these takes a scan_stages
argument at initialisation, which behaves the same as as the explicit Runge--Kutta methods for ODEs.
diffrax.Euler (AbstractItoSolver)
¤
Euler's method.
1st order explicit Runge--Kutta method. Does not support adaptive step sizing.
When used to solve SDEs, converges to the Itô solution.
diffrax.Heun (AbstractERK, AbstractStratonovichSolver)
¤
Heun's method.
2nd order explicit Runge--Kutta method. Has an embedded Euler method for adaptive step sizing.
Also sometimes known as either the "improved Euler method", "modified Euler method" or "explicit trapezoidal rule".
Should not be confused with Heun's third order method, which is a different (higher order) method occasionally also just referred to as "Heun's method".
When used to solve SDEs, converges to the Stratonovich solution.
diffrax.Midpoint (AbstractERK, AbstractStratonovichSolver)
¤
Midpoint method.
2nd order explicit Runge--Kutta method. Has an embedded Euler method for adaptive step sizing.
Also sometimes known as the "modified Euler method".
When used to solve SDEs, converges to the Stratonovich solution.
diffrax.Ralston (AbstractERK, AbstractStratonovichSolver)
¤
Ralston's method.
2nd order explicit Runge--Kutta method. Has an embedded Euler method for adaptive step sizing.
When used to solve SDEs, converges to the Stratonovich solution.
Info
In addition to the solvers above, then most higher-order ODE solvers can actually also be used as SDE solvers. They will typically converge to the Stratonovich solution. In practice this is computationally wasteful as they will not obtain more accurate solutions when applied to SDEs.
Reversible methods¤
These are reversible in the same way as when applied to ODEs. See here.
diffrax.ReversibleHeun (AbstractAdaptiveSolver, AbstractStratonovichSolver)
¤
Reversible Heun method.
Algebraically reversible 2nd order method. Has an embedded 1st order method for adaptive step sizing.
When used to solve SDEs, converges to the Stratonovich solution.
Reference
@article{kidger2021efficient,
author={Kidger, Patrick and Foster, James and Li, Xuechen and Lyons, Terry},
title={{E}fficient and {A}ccurate {G}radients for {N}eural {SDE}s},
year={2021},
journal={Advances in Neural Information Processing Systems}
}
SDE-only solvers¤
Term structure
For these SDE-specific solvers, the terms (given by the value of terms
to diffrax.diffeqsolve
) must be a 2-tuple (AbstractTerm, AbstractTerm)
, representing the drift and diffusion respectively. Typically that means (ODETerm(...), ControlTerm(..., ...))
.
diffrax.EulerHeun (AbstractStratonovichSolver)
¤
Euler-Heun method.
Used to solve SDEs, and converges to the Stratonovich solution.
diffrax.ItoMilstein (AbstractItoSolver)
¤
Milstein's method; Itô version.
Used to solve SDEs, and converges to the Itô solution.
Warning
Requires commutative noise. Note that this commutativity condition is not checked.
diffrax.StratonovichMilstein (AbstractStratonovichSolver)
¤
Milstein's method; Stratonovich version.
Used to solve SDEs, and converges to the Stratonovich solution.
Warning
Requires commutative noise. Note that this commutativity condition is not checked.
Wrapper solvers¤
diffrax.HalfSolver (AbstractAdaptiveSolver, AbstractWrappedSolver)
¤
Wraps another solver, trading cost in order to provide error estimates. (That is, it means the solver can be used with an adaptive step size controller, regardless of whether the underlying solver supports adaptive step sizing.)
For every step of the wrapped solver, it does this by also making two half-steps, and comparing the results between the full step and the two half steps. Hence the name "HalfSolver".
As such each step costs 3 times the computational cost of the wrapped solver, whilst producing results that are roughly twice as accurate, in addition to producing error estimates.
Tip
Many solvers already provide error estimates, making HalfSolver
primarily
useful when using a solver that doesn't provide error estimates -- e.g.
diffrax.Euler
. Such solvers are most common when solving SDEs.
__init__(self, solver: AbstractSolver)
¤
Arguments:
solver
: The solver to wrap.