Skip to content

Diffusion

Diffusion objects define the stochastic term in a continuous-time state evolution

\[ dx_t = f(x_t, u_t, t)\,dt + L(x_t, u_t, t)\,dW_t, \]

where:

  • \(x_t \in \mathbb{R}^{d_x}\) is the latent state,
  • \(u_t\) is an optional control,
  • \(W_t \in \mathbb{R}^{d_w}\) is Brownian motion,
  • bm_dim = d_w,
  • and \(L(x_t, u_t, t)\) is the diffusion coefficient.

Dynestyx exposes three structured diffusion classes:

  • If your diffusion has the form \(L = \sigma I_{d_w}\) with \(d_w = d_x\), use ScalarDiffusion(sigma, bm_dim=state_dim).
  • If your diffusion has the form \(L = \sigma \mathbf{1}_{d_x}\), use ScalarDiffusion(sigma, bm_dim=1).
  • If your diffusion has the form \(L = \mathrm{diag}(v)\), use DiagonalDiffusion(v, bm_dim=state_dim).
  • If your diffusion has the form \(L = v \in \mathbb{R}^{d_x \times 1}\), use DiagonalDiffusion(v, bm_dim=1).
  • If your diffusion is a general matrix \(L \in \mathbb{R}^{d_x \times d_w}\), use FullDiffusion(L).

The same constructors also accept callables (x, u, t) -> value instead of constants. For example, if \(L(x_t, u_t, t) = \sigma(x_t) I_{d_w}\), use ScalarDiffusion(lambda x, u, t: sigma(x), bm_dim=state_dim). The same pattern applies across ScalarDiffusion, DiagonalDiffusion, and FullDiffusion.

In practice, many models use a constant diffusion coefficient. In those cases, pass the matrix/vector/scalar value directly. Reserve callable diffusion coefficients for cases where the coefficient genuinely depends on state, control, or time.

Diffusion

Diffusion is the common base class for structured diffusion coefficients. Most users should instantiate one of its public subclasses instead of using Diffusion directly:

  • Use ScalarDiffusion when one scalar scale is enough. This is the most common choice for isotropic diffusion.
  • Use DiagonalDiffusion when each state coordinate should have its own scale but the loading remains axis-aligned.
  • Use FullDiffusion when you need to specify a genuinely matrix-valued loading.

Each class accepts either:

  • a constant coefficient, or
  • a callable (x, u, t) -> value.

FullDiffusion

Mathematically, FullDiffusion represents \(L(x_t, u_t, t) \in \mathbb{R}^{d_x \times d_w}\).

Use FullDiffusion(coefficient, bm_dim=None) when you want to specify this matrix-valued diffusion coefficient directly.

Accepted coefficient forms:

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

If coefficient is constant, bm_dim is inferred automatically when omitted.

Example:

import jax.numpy as jnp
from dynestyx import ContinuousTimeStateEvolution, FullDiffusion

state_evolution = ContinuousTimeStateEvolution(
    drift=lambda x, u, t: -x,
    diffusion=FullDiffusion(jnp.eye(2)),
)

DiagonalDiffusion

Mathematically, DiagonalDiffusion represents a vector-valued coefficient \(v(x_t, u_t, t) \in \mathbb{R}^{d_x}\). If bm_dim = d_x, this means \(L = \mathrm{diag}(v(x_t, u_t, t))\). If bm_dim = 1, this means \(L = v(x_t, u_t, t) \in \mathbb{R}^{d_x \times 1}\).

Use DiagonalDiffusion(coefficient, bm_dim) when the diffusion is naturally parameterized by a vector of state-wise loadings.

Accepted coefficient forms:

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

bm_dim is required and must be either 1 or state_dim.

Example:

import jax.numpy as jnp
from dynestyx import ContinuousTimeStateEvolution, DiagonalDiffusion

state_evolution = ContinuousTimeStateEvolution(
    drift=lambda x, u, t: -x,
    diffusion=DiagonalDiffusion(jnp.array([0.1, 0.2]), bm_dim=2),
)

ScalarDiffusion

Mathematically, ScalarDiffusion represents a scalar-valued coefficient \(\sigma(x_t, u_t, t) \in \mathbb{R}\). If bm_dim = d_x, this means \(L = \sigma(x_t, u_t, t)\,I_{d_w}\) with \(d_w = d_x\). If bm_dim = 1, this means \(L = \sigma(x_t, u_t, t)\,\mathbf{1}_{d_x}\), viewed as a column vector in \(\mathbb{R}^{d_x \times 1}\).

Use ScalarDiffusion(coefficient, bm_dim) when one scalar scale is enough. This is usually the simplest choice for isotropic diffusion.

Accepted coefficient forms:

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

bm_dim is required and must be either 1 or state_dim.

Constant example:

from dynestyx import ContinuousTimeStateEvolution, ScalarDiffusion

state_evolution = ContinuousTimeStateEvolution(
    drift=lambda x, u, t: -x,
    diffusion=ScalarDiffusion(0.1, bm_dim=2),
)

Callable example:

import jax.numpy as jnp
from dynestyx import ContinuousTimeStateEvolution, ScalarDiffusion

state_evolution = ContinuousTimeStateEvolution(
    drift=lambda x, u, t: -x,
    diffusion=ScalarDiffusion(
        lambda x, u, t: 0.1 + 0.05 * jnp.tanh(x[0]),
        bm_dim=2,
    ),
)

API

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.

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.

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.

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.