Extending Diffraxยค
It's completely possible to extend Diffrax with your own custom solvers, step size controllers, and so on.
The main points of extension are as follows:
-
Custom solvers should inherit from
diffrax.AbstractSolver.- If you are making a new Runge--Kutta method then this is particularly easy; you can use the existing base classes Diffrax already uses for its own Runge--Kutta methods, and supply them with an appropriate
diffrax.ButcherTableau.- For explicit-Runge--Kutta methods (ERK) then inherit from
diffrax.AbstractERK. - For general diagonal-implicit-Runge--Kutta methods (DIRK) then inherit from
diffrax.AbstractDIRK. - For singly-diagonal-implicit-Runge--Kutta methods (SDIRK) then inherit from
diffrax.AbstractSDIRK. - For explicit-singly-diagonal-implicit-Runge--Kutta methods (ESDIRK) then inherit from
diffrax.AbstractESDIRK.
- For explicit-Runge--Kutta methods (ERK) then inherit from
- Several abstract base classes are available to specify the behaviour of the solver:
diffrax.AbstractImplicitSolverare those solvers that solve implicit problems (and therefore take aroot_finderargument).diffrax.AbstractAdaptiveSolverare those solvers capable of providing error estimates (and thus can be used with adaptive step size controllers).diffrax.AbstractItoSolveranddiffrax.AbstractStratonovichSolverare used to specify which SDE solution a particular solver is known to converge to.diffrax.AbstractWrappedSolverindicates that the solver is used to wrap another solver, and so e.g. it will be treated as an implicit solver/etc. if the wrapped solver is also an implicit solver/etc.
- If you are making a new Runge--Kutta method then this is particularly easy; you can use the existing base classes Diffrax already uses for its own Runge--Kutta methods, and supply them with an appropriate
-
Custom step size controllers should inherit from
diffrax.AbstractStepSizeController.- The abstract base class
diffrax.AbstractAdaptiveStepSizeControllercan be used to specify that this controller uses error estimates to adapt step sizes.
- The abstract base class
-
Custom Brownian motion simulations should inherit from
diffrax.AbstractBrownianPath. -
Custom controls (e.g. custom interpolation schemes analogous to
diffrax.CubicInterpolation) should inherit fromdiffrax.AbstractPath. -
Custom terms should inherit from
diffrax.AbstractTerm.- For example, if the vector field - control interaction is a matrix-vector product, but the matrix is known to have special structure, then you may wish to create a custom term that can calculate this interaction more efficiently than is given by a full matrix-vector product. Given the large suite of linear operators lineax implements (which are fully supported by
diffrax.ControlTerm), this is likely rarely necessary.
- For example, if the vector field - control interaction is a matrix-vector product, but the matrix is known to have special structure, then you may wish to create a custom term that can calculate this interaction more efficiently than is given by a full matrix-vector product. Given the large suite of linear operators lineax implements (which are fully supported by
In each case we recommend looking up existing solvers/etc. in Diffrax to understand how to implement them.
Contributions
If you implement a technique that you'd like to see merged into the main Diffrax library then open a pull request on GitHub. We're very happy to take contributions.