DOCSEARCHbyAlgolia
Submit Docs

diffrax

repository·main·Indexed 24 days ago

https://github.com/patrick-kidger/diffrax

A JAX-based library for numerical differential equation solvers, including ODE (Ordinary), SDE (Stochastic), and CDE (Controlled) solvers. It is designed to be autodifferentiable, GPU-capable, and compatible with JAX transformations like vmap. Diffrax provides a variety of explicit and implicit solvers (e.g., Dopri5, Tsit5), step size controllers, adjoint methods for backpropagation, and support for dense solutions via interpolation.

Tokens
35.1K
Snippets
49
Records
170
Agent score
85%

What's inside diffrax

  1. Overview of Diffrax features

    main

    Diffrax is a JAX-based library for numerical differential equation solvers. Key capabilities include:

    • Solver Types: Supports ODE (Ordinary), SDE (Stochastic), and CDE (Controlled) solvers.
    • Solver Variety: Includes Tsit5, Dopri8, symplectic solvers, and implicit solvers.
    • JAX Integration: Everything is vmap-able (including the integration region), supports PyTrees as state, and provides multiple adjoint methods for backpropagation.
    • Advanced Use Cases: Supports dense solutions and neural differential equations.

  2. Use SubSaveAt for nested state saving

    main
    The diffrax.SubSaveAt class is used when you need to specify saving behavior for sub-components of a complex state. It allows you to apply SaveAt logic to specific parts of a structured state (like a PyTree), enabling fine-grained control over what is recorded during the integration process.

  3. How to structure terms for SDE solvers

    main

    The choice of solver determines the required layout of the terms argument in diffeqsolve.

    General-purpose solvers

    Most solvers handle both ODEs and SDEs using a single term structure. For an SDE, you should use MultiTerm to combine an ODETerm (for drift) and a ControlTerm (for diffusion and Brownian motion).

    SDE-only solvers

    SDE-specific solvers require terms to be specifically of the form MultiTerm(ODETerm(...), SomeOtherTerm(...)). Typically, SomeOtherTerm is a ControlTerm representing the diffusion.

    Advanced users can implement custom AbstractTerm objects for diffusion to achieve more efficient behavior for specific diffusion matrix structures.

    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(), ...)

  4. How to create a custom path by subclassing AbstractPath

    main

    In Diffrax, a path is a piecewise continuous function $f : [t_0, t_1] o ext{R}^d$. This abstraction is used for objects like Brownian motion or interpolations for neural controlled differential equations (CDEs).

    To implement your own custom path, you must subclass diffrax.AbstractPath and provide implementations for the following members:

    • t0: The start time of the path.
    • t1: The end time of the path.
    • evaluate: A method to evaluate the path at a given time $t$.
    • derivative: A method to compute the derivative of the path at a given time $t$.

  5. How terms work in Diffrax

    main

    In Diffrax, "terms" are the building blocks of differential equations. A term consists of a vector field (the function describing the rate of change) and a control (the differential, e.g., $\mathrm{d}t$ or $\mathrm{d}w(t)$), along with a rule for how they interact (the product).

    Depending on the type of equation you are solving, you will use different term combinations:

    • ODEs: Use a single diffrax.ODETerm.
    • SDEs: Use diffrax.MultiTerm to group an ODETerm (drift) and a diffrax.ControlTerm (diffusion).
    • Independent terms (e.g., Hamiltonian systems): If different equations affect different parts of the state, group them in a PyTree structure, such as a tuple of terms.

  6. How ODEs, SDEs, and CDEs relate in Diffrax

    main

    In Diffrax, ODEs and SDEs are conceptually unified. Under the hood, both are solved by lowering them to Controlled Differential Equations (CDEs).

    • ODEs are the simplest case where the control signal is just $x(t) = t$.
    • SDEs are implemented by providing a drift term (ODETerm) and a diffusion term (ControlTerm) bundled via MultiTerm.
    • CDEs allow for more complex control signals by implementing the AbstractPath interface, which is useful for time-series data where the control signal is an external input.

  7. How solver term structures affect `diffeqsolve`

    main

    The choice of solver determines the required layout of the terms argument in diffrax.diffeqsolve. While most solvers expect a single AbstractTerm, certain solver types require specific structures:

    • Standard Solvers: Expect a single AbstractTerm.
    • IMEX Methods: Require terms=MultiTerm(explicit_term, implicit_term).
    • Symplectic Methods (e.g., SemiImplicitEuler): Require terms to be a 2-tuple of AbstractTerms.
    • Specialized Solvers: Some solvers like SemiImplicitEuler may also require the state y0 to be a 2-tuple.

    You can programmatically check the required structure for any solver instance using the <solver>.term_structure attribute.

  10. Understand the `diffrax.Solution` object

    main
    The diffrax.Solution object is the primary container returned by Diffrax solvers. It encapsulates the entire state of a differential equation integration, including the time steps, the solution values, and metadata about the solver's performance and convergence.

  11. Implement custom Brownian motion, paths, and terms

    main

    Diffrax provides several abstraction points for extending core functionality:

    • Brownian Motion: Inherit from diffrax.AbstractBrownianPath to implement custom Brownian motion simulations.
    • Paths and Interpolation: Inherit from diffrax.AbstractPath to implement custom controls, such as custom interpolation schemes (similar to diffrax.CubicInterpolation).
    • Terms: Inherit from diffrax.AbstractTerm to implement custom terms. This is useful if you have a specialized interaction (like a matrix-vector product with specific structure) that can be calculated more efficiently than the default implementation. Note that diffrax.ControlTerm already supports many linear operators via the lineax library.

  12. Extend Diffrax with custom solver subclasses

    main

    When defining custom solvers, you can inherit from specific subclasses of diffrax.AbstractSolver to signal particular behaviors to the library. This allows your custom solver to be recognized by relevant Diffrax components.

    Available behavioral subclasses include:

    • diffrax.AbstractImplicitSolver: For implicit solvers.
    • diffrax.AbstractAdaptiveSolver: For solvers that adapt their step size.
    • diffrax.AbstractItoSolver: For solvers using the Ito interpretation.
    • diffrax.AbstractStratonovichSolver: For solvers using the Stratonovich interpretation.
    • diffrax.AbstractWrappedSolver: For solvers that wrap other solvers.