Skip to content

Diffusions

This page documents the diffusion internals used by continuous-time solvers, discretizers, and backend integrations.

Overview

Public model code should usually work with FullDiffusion, DiagonalDiffusion, and ScalarDiffusion through ContinuousTimeStateEvolution(diffusion=...).

Developer-facing code sometimes needs the lower-level evaluation layer:

  • Diffusion.evaluate(...) evaluates a structured diffusion object at a concrete (x, u, t).
  • EvaluatedDiffusion.as_matrix(...) returns the corresponding matrix L.
  • EvaluatedDiffusion.apply(...) applies that matrix to a Brownian increment.
  • resolve_metadata(...) resolves bm_dim and validates the coefficient shape against state_dim.

In other words, the public classes describe the structure of the diffusion, while EvaluatedDiffusion is the solver-facing object used after a concrete state, control, and time have been chosen.

API

EvaluatedDiffusion

Bases: NamedTuple

A diffusion coefficient evaluated at a specific (x, u, t).

This is primarily a developer-facing helper used by solvers and backend integrations. It pairs a structured Diffusion object with the concrete value of its coefficient at one state, control, and time.

apply(dw: Float[Array, ' bm_dim'], *, state_dim: int) -> Real[Array, ' state_dim'] | Real[Array, '']

Return L @ dw using the structured diffusion representation.

as_matrix(*, state_dim: int) -> Real[Array, '*plate state_dim bm_dim']

Return the evaluated diffusion coefficient as a matrix L.

gram_matrix(*, state_dim: int) -> Float[Array, '*plate state_dim state_dim']

L L^T.

Diffusion

Bases: Module

Base class for diffusion coefficients in SDEs.

A diffusion encapsulates both the numeric coefficient L(x, u, t) and the structural interpretation of that coefficient inside the SDE

dx_t = f(x_t, u_t, t) dt + L(x_t, u_t, t) dW_t.

Most users should instantiate one of the concrete subclasses:

  • FullDiffusion for an arbitrary matrix-valued coefficient.
  • DiagonalDiffusion for axis-aligned loadings.
  • ScalarDiffusion for isotropic or shared-driver noise.

The coefficient may be either:

  • a constant array or scalar, for time-homogeneous diffusion, or
  • a callable (x, u, t) -> value for state-, control-, or time-dependent diffusion.

bm_dim is the Brownian dimension d_w. For FullDiffusion it can be inferred from the matrix shape. For DiagonalDiffusion and ScalarDiffusion it must be specified explicitly and must be either 1 or state_dim.

coefficient_event_rank: int | None property

The per-member event rank of a constant coefficient (None if callable).

The event rank is the number of trailing axes that make up one member's coefficient, independent of any leading plate batch axes (FullDiffusion is matrix-valued so rank 2, DiagonalDiffusion vector-valued so rank 1, ScalarDiffusion scalar so rank 0). It lets the plate machinery tell a shared coefficient from a genuinely per-member one even when its shape coincides with the plate sizes. None for callable coefficients, whose value is not a sliceable pytree leaf.

This is a derived property rather than a stored field beacuse of interactions with effectful and equinox.

apply(dw: Float[Array, ' bm_dim'], *, x: Real[Array, ' state_dim'] | Real[Array, ''] | None, u: Real[Array, ' control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, ''], state_dim: int) -> Real[Array, ' state_dim'] | Real[Array, '']

Apply the diffusion coefficient to a Brownian increment dw.

as_matrix(*, x: Real[Array, ' state_dim'] | Real[Array, ''] | None, u: Real[Array, ' control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, ''], state_dim: int) -> Real[Array, '*plate state_dim bm_dim']

Return the matrix-valued diffusion coefficient L(x, u, t).

evaluate(*, x: Real[Array, ' state_dim'] | Real[Array, ''] | None, u: Real[Array, ' control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, '']) -> EvaluatedDiffusion

Evaluate the diffusion at (x, u, t).

evaluate_value(*, x: Real[Array, ' state_dim'] | Real[Array, ''] | None, u: Real[Array, ' control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, '']) -> Array

Return the raw coefficient value at (x, u, t).

gram_matrix(*, x: Real[Array, ' state_dim'] | Real[Array, ''] | None, u: Real[Array, ' control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, ''], state_dim: int) -> Float[Array, '*plate state_dim state_dim']

Return the diffusion Gram matrix L(x,u,t) L(x,u,t)^T.

resolve_metadata(*, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''] | None, u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: float | int | Real[Array, '']) -> Self

Check coefficient shape and resolve bm_dim if needed.

FullDiffusion

Bases: Diffusion

General matrix-valued diffusion coefficient.

Use FullDiffusion when you want to specify the diffusion matrix L(x, u, t) directly.

Parameters:

Name Type Description Default
coefficient DiffusionSpec

Either a constant array with trailing shape (..., state_dim, bm_dim) or a callable (x, u, t) -> array with that trailing shape.

 required
bm_dim int | None

Optional Brownian dimension. If omitted for a constant coefficient, it is inferred from the trailing matrix dimension.

 None

This is the most general public diffusion class. Prefer it when your model genuinely needs a dense or otherwise unstructured loading matrix.

resolve_metadata(*, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''] | None, u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: float | int | Real[Array, '']) -> Self

Check matrix shape and infer bm_dim from the trailing dimension if needed.

DiagonalDiffusion

Bases: Diffusion

Vector-valued diffusion with diagonal or shared-driver interpretation.

Use DiagonalDiffusion(v, bm_dim=...) when the diffusion is naturally parameterized by a vector v of length state_dim.

Parameters:

Name Type Description Default
coefficient DiffusionSpec

Either a constant vector with trailing shape (..., state_dim) or a callable (x, u, t) -> array with that trailing shape.

 required
bm_dim int | None

Brownian dimension. Must be either 1 or state_dim.

 None
  • If bm_dim = state_dim, the vector is interpreted as the diagonal of L = diag(v).
  • If bm_dim = 1, the vector is interpreted as a column loading vector, so all state coordinates share one Brownian driver.

This is often the right public API choice when each state coordinate has its own scale but you do not want to write out a full matrix.

resolve_metadata(*, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''] | None, u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: float | int | Real[Array, '']) -> Self

Check vector shape and verify that bm_dim is either 1 or state_dim.

ScalarDiffusion

Bases: Diffusion

Scalar-valued diffusion with isotropic or shared-driver interpretation.

Use ScalarDiffusion(sigma, bm_dim=...) when the diffusion is controlled by a single scalar scale sigma.

Parameters:

Name Type Description Default
coefficient DiffusionSpec

Either a scalar, a constant array with trailing shape (..., 1), or a callable (x, u, t) -> scalar_or_length_1_array.

 required
bm_dim int | None

Brownian dimension. Must be either 1 or state_dim.

 None
  • If bm_dim = state_dim, this represents isotropic diffusion L = sigma I.
  • If bm_dim = 1, this represents a shared scalar driver applied equally to every state coordinate.

This is typically the simplest public API choice for constant isotropic diffusion, and is usually preferable to writing sigma * eye(state_dim) by hand.

resolve_metadata(*, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''] | None, u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: float | int | Real[Array, '']) -> Self

Check scalar shape and verify that bm_dim is either 1 or state_dim.