SPDE API reference¶

This page documents the public API for grid-based SPDE operators and helpers. For narrative examples, see Spatial field inference (SPDE).

Grid extras¶

SFI.bases.spde.square_grid_extras(grid_shape, dx=1.0, *, prefix='box')[source]

Return minimal box extras needed by all square-grid SPDE operators.

Creates two arrays in the returned dict:

{prefix}/grid_shape – integer grid dimensions (Nx, Ny, ...)
{prefix}/dx – float grid spacing per axis

Parameters:

grid_shape (sequence of int) – Number of grid points along each spatial axis.
dx (float or sequence of float) – Grid spacing (uniform or per-axis).
prefix (str) – Namespace prefix for the extras keys (default "box").

Return type:

dict[str, Array]

Composable operators¶

class SFI.bases.spde.StencilOp(*, ndim, bc='noflux', order='C', key_grid_shape='box/grid_shape', key_dx='box/dx')[source]

Bases: ABC

Composable finite-difference operator on a regular Cartesian grid.

A StencilOp wraps a stencil-based finite-difference formula and can be applied to any pointwise StateExpr to produce a new Basis:

Lap = Laplacian(ndim=2, bc="pbc")
phi = field_component(0, n_fields=1)
lap_phi  = Lap(phi)       # standard  ∇²φ
lap_phi3 = Lap(phi**3)    # ∇²(φ³) — exactly conservative on PBC

The operator evaluates the inner expression on each stencil neighbor, then applies the finite-difference formula to the transformed values. This guarantees exact discrete conservation properties (e.g., zero spatial sum for the Laplacian on periodic grids) regardless of the inner nonlinearity.

Operator composition is fully supported via automatic stencil fusion. Nesting operators like Lap(Lap(phi)) computes the Minkowski sum of the two stencils at construction time so that only a single dispatcher pass is needed at evaluation time:

bih = Lap(Lap(phi))      # ≡ ∇⁴φ, fused 13-point stencil
div_grad = Div(Grad(phi))  # ≡ ∇²φ via composition

Parameters:

ndim (int) – Number of spatial dimensions.
bc (str) – Boundary condition: "pbc" (periodic) or "noflux" (reflecting).
order (str) – Flattening order for the spatial grid ("C" row-major or "F" column-major). Must match the layout used to build the extras["box/grid_shape"] array passed at evaluation time.
key_grid_shape (str) – Extras key for the grid shape array.
key_dx (str) – Extras key for the grid spacing array.

visualize_stencil()[source]

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

Scalar operators¶

class SFI.bases.spde.Laplacian(*, ndim, bc='noflux', order='C', key_grid_shape='box/grid_shape', key_dx='box/dx')[source]

Bases: StencilOp

Composable Laplacian operator \(\nabla^2\) on a Cartesian grid.

\[[\nabla^2 f(\phi)]_i = \sum_{\alpha=1}^{n_\mathrm{dim}} \frac{f(\phi_{i+e_\alpha}) + f(\phi_{i-e_\alpha}) - 2\,f(\phi_i)} {\Delta x_\alpha^2}\]

where \(f\) is the inner expression evaluated at each grid site.

On periodic grids, the Laplacian matrix has zero column sums, so \(\sum_i [\nabla^2 g]_i = 0\) for any function \(g\). This guarantees exact discrete conservation for expressions like \(\nabla^2(\phi^3)\).

Examples

>>> from SFI.bases import field_component
>>> from SFI.bases.spde import Laplacian
>>> Lap = Laplacian(ndim=2, bc="pbc")
>>> phi = field_component(0, n_fields=1)
>>> lap_phi  = Lap(phi)       # ∇²φ
>>> lap_phi3 = Lap(phi**3)    # ∇²(φ³) — conservative

Parameters:

ndim (int)
bc (str)
order (str)
key_grid_shape (str)
key_dx (str)

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.Biharmonic(**kwargs)[source]

Bases: StencilOp

Composable biharmonic operator \(\nabla^4\) on a Cartesian grid.

Uses the standard 13-point stencil (2D) or 25-point stencil (3D). Requires ndim >= 2.

On periodic grids, \(\sum_i [\nabla^4 g]_i = 0\) for any \(g\), so compositions like \(\nabla^4(\phi^2)\) are exactly conservative.

Examples

>>> Bih = Biharmonic(ndim=2, bc="pbc")
>>> phi = field_component(0, n_fields=1)
>>> bih_phi = Bih(phi)         # ∇⁴φ

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.LaplacianOfGradientSquared(**kwargs)[source]

Bases: StencilOp

Composable operator \(\nabla^2(|\nabla f(\phi)|^2)\).

This is the genuine Active Model B+ term that breaks time-reversal symmetry. It cannot be written as \(\nabla^2(g(\phi))\) for any pointwise function \(g\), because \(|\nabla\phi|^2\) depends on spatial derivatives, not just the local field value.

The computation proceeds in two stages on the biharmonic (radius-2) stencil:

Evaluate the inner expression \(f(\phi)\) at all stencil neighbors (13 points in 2D, 25 in 3D).
Compute \(|\nabla f|^2\) at each of the Laplacian’s 5 stencil points (center + nearest neighbours) using central differences among the gathered values.
Apply the standard Laplacian formula to the resulting \(|\nabla f|^2\) field.

On periodic grids, the outer Laplacian guarantees \(\sum_i [\nabla^2 G]_i = 0\) exactly, so the operator is conservative even though \(G = |\nabla f|^2\) is nonlinear in the spatial derivatives.

\[[\nabla^2(|\nabla f|^2)]_i = \sum_\alpha \frac{G_{i+e_\alpha} + G_{i-e_\alpha} - 2\,G_i} {\Delta x_\alpha^2}\]

where

\[G_j = |\nabla f(\phi)|^2_j = \sum_\beta \left(\frac{f(\phi_{j+e_\beta}) - f(\phi_{j-e_\beta})} {2\,\Delta x_\beta}\right)^{\!2}\]

Requires ndim >= 2.

Examples

>>> LGS = LaplacianOfGradientSquared(ndim=2, bc="pbc")
>>> phi = field_component(0, n_fields=1)
>>> lgs_phi  = LGS(phi)       # ∇²(|∇φ|²)  (AMB+ active term)
>>> lgs_phi2 = LGS(phi**2)    # ∇²(|∇(φ²)|²)

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

Vector operators¶

class SFI.bases.spde.Gradient(*, ndim, bc='noflux', order='C', key_grid_shape='box/grid_shape', key_dx='box/dx')[source]

Bases: StencilOp

Composable gradient operator \(\nabla\) on a Cartesian grid.

\[[\nabla f(\phi)]_{i,\alpha} = \frac{f(\phi_{i+e_\alpha}) - f(\phi_{i-e_\alpha})}{2\,\Delta x_\alpha}\]

Output features = n_features_inner × ndim, ordered as (feature_0/dx_0, feature_0/dx_1, ..., feature_1/dx_0, ...).

Examples

>>> Grad = Gradient(ndim=2, bc="pbc")
>>> phi  = field_component(0, n_fields=1)
>>> grad_phi  = Grad(phi)       # ∇φ  (2 features)
>>> grad_phi2 = Grad(phi**2)    # ∇(φ²) ≈ 2φ∇φ

Parameters:

ndim (int)
bc (str)
order (str)
key_grid_shape (str)
key_dx (str)

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.Divergence(*, ndim, bc='noflux', order='C', key_grid_shape='box/grid_shape', key_dx='box/dx')[source]

Bases: StencilOp

Composable divergence \(\nabla\!\cdot\) on a Cartesian grid.

\[[\nabla\!\cdot\mathbf{f}]_i = \sum_{\alpha} \frac{f_\alpha(\phi_{i+e_\alpha}) - f_\alpha(\phi_{i-e_\alpha})} {2\,\Delta x_\alpha}\]

The inner expression must produce exactly ndim features.

Examples

>>> Div = Divergence(ndim=2, bc="pbc")
>>> v   = vector_field(n_fields=2)
>>> div_v = Div(v)           # scalar divergence (1 feature)

Composition with Gradient recovers the Laplacian:

>>> Grad = Gradient(ndim=2, bc="pbc")
>>> phi  = field_component(0, n_fields=1)
>>> lap  = Div(Grad(phi))    # ∇·(∇φ) = ∇²φ

Parameters:

ndim (int)
bc (str)
order (str)
key_grid_shape (str)
key_dx (str)

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.Curl(**kwargs)[source]

Bases: StencilOp

Composable curl operator on a Cartesian grid.

2D (scalar curl / vorticity, 1 output feature):

\[[\nabla\!\times\!\mathbf{f}]_i = \frac{f_1(i\!+\!e_0) - f_1(i\!-\!e_0)}{2\Delta x_0} - \frac{f_0(i\!+\!e_1) - f_0(i\!-\!e_1)}{2\Delta x_1}\]

3D (vector curl, 3 output features):

\[(\nabla\!\times\!\mathbf{f})_\alpha = \varepsilon_{\alpha\beta\gamma}\, \frac{f_\gamma(i\!+\!e_\beta) - f_\gamma(i\!-\!e_\beta)} {2\,\Delta x_\beta}\]

The inner expression must produce exactly ndim features. Requires ndim >= 2.

Examples

>>> Curl2D = Curl(ndim=2, bc="pbc")
>>> v = vector_field(n_fields=2)
>>> omega = Curl2D(v)   # scalar vorticity

Identity check — curl of a gradient vanishes:

>>> Grad = Gradient(ndim=2, bc="pbc")
>>> phi  = field_component(0, n_fields=1)
>>> should_vanish = Curl2D(Grad(phi))

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.SymGrad(**kwargs)[source]

Bases: StencilOp

Symmetric gradient (strain-like) on a Cartesian grid.

Computes the symmetric part of the gradient tensor of a vector field:

\[S_{ij} = \tfrac{1}{2}\!\left( \frac{\partial f_i}{\partial x_j} + \frac{\partial f_j}{\partial x_i} \right)\]

For velocity fields this is the strain-rate tensor; for displacement fields the strain tensor.

Output features = ndim*(ndim+1)//2 (upper triangle), ordered: (0,0), (0,1), ..., (0,d-1), (1,1), (1,2), ..., (d-1,d-1).

The inner expression must produce exactly ndim features. Requires ndim >= 2.

Examples

>>> SG = SymGrad(ndim=2, bc="pbc")
>>> v  = vector_field(n_fields=2)
>>> strain = SG(v)   # 3 features: S_00, S_01, S_11

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

class SFI.bases.spde.SkewGrad(**kwargs)[source]

Bases: StencilOp

Antisymmetric gradient (rotation-like) on a Cartesian grid.

Computes the antisymmetric part of the gradient tensor:

\[W_{ij} = \tfrac{1}{2}\!\left( \frac{\partial f_i}{\partial x_j} - \frac{\partial f_j}{\partial x_i} \right)\]

For velocity fields this is the rotation-rate (vorticity) tensor.

Output features = ndim*(ndim-1)//2 (strict upper triangle), ordered: (0,1), (0,2), ..., (1,2), ....

In 2D the single component equals half the scalar curl.

The inner expression must produce exactly ndim features. Requires ndim >= 2.

Examples

>>> W  = SkewGrad(ndim=2, bc="pbc")
>>> v  = vector_field(n_fields=2)
>>> omega = W(v)   # 1 feature: W_01

visualize_stencil()

Return an ASCII picture of the stencil pattern (2D only).

For ndim != 2 a simple offset listing is returned.

Returns:: Multi-line string.
Return type:: str

Advection¶

class SFI.bases.spde.AdvectionBy(v_expr, *, ndim, bc='pbc', **stencil_kwargs)[source]

Bases:

Return a stencil operator for advection \((\mathbf{v}\!\cdot\!\nabla)f\).

Parameters:

v_expr (StateExpr) – Pointwise expression with n_features == ndim giving the advecting vector field.
ndim (int) – Number of spatial dimensions.
bc (str) – Boundary condition ("pbc" or "noflux").
**stencil_kwargs – Forwarded to StencilOp.__init__.

Returns:

Operator computing \(\sum_\alpha v_\alpha\,\partial f / \partial x_\alpha\).

Return type:

StencilOp

Notes

Results of this operator are not further composable because the advection FD formula references a second expression (the velocity).

Examples

>>> v   = vector_field(n_fields=2)
>>> phi = field_component(0, n_fields=2)
>>> Adv = AdvectionBy(v, ndim=2, bc="pbc")
>>> adv_phi = Adv(phi)   # (v·∇)φ

Convenience functions¶

SFI.bases.spde.vector_field(n_fields, *, labels=None)[source]

All field components as a rank-0 expression with n_features == n_fields.

Shorthand for concatenating field_component() over all indices. Useful as input to same-space operators (Divergence, Curl, SymGrad, …).

Parameters:

n_fields (int) – Number of field components per grid site.
labels (list of str, optional) – Human-readable labels; defaults to field[0], field[1], ….

Returns:

n_features == n_fields, rank == 0.

Return type:

Basis

Stencil composition utilities¶

SFI.bases.spde.minkowski_sum_offsets(offsets_a, offsets_b)[source]

Compute the Minkowski sum of two stencil offset arrays.

Given outer offsets O_a and inner offsets O_b, computes the deduplicated set {a + b : a in O_a, b in O_b} and an index mapping suitable for composing the two FD formulas.

Parameters:

offsets_a (array, shape (Ka, ndim)) – Outer operator stencil offsets.
offsets_b (array, shape (Kb, ndim)) – Inner operator stencil offsets.

Returns:

fused (jnp.ndarray, shape (K_fused, ndim)) – Deduplicated sorted offset vectors.
index_maps (jnp.ndarray, shape (Ka, Kb)) – index_maps[i, j] is the position of offsets_a[i] + offsets_b[j] in fused.

Noise models¶

class SFI.langevin.noise.ConservedNoise(sigma, *, grid_shape, dx=1.0, n_fields=1)[source]

Bases: NoiseModel

Conserved (divergence-form) noise on a periodic square grid.

Implements noise of the form

\[\eta(x, t) = \nabla \cdot \bigl(\sigma\, \vec{\Lambda}(x,t)\bigr)\]

where \(\vec{\Lambda}\) is spatiotemporal white vector noise. In Fourier space this is equivalent to multiplying each mode by \(|k|\):

\[\hat{\eta}_k = \sigma\,|k|\,\hat{\xi}_k\]

This noise conserves the spatial average of the field (\(\sum_i \phi_i\) is a constant of the noise process), as required by Model B / Active Model B+ dynamics.

Parameters:

sigma (float) – Noise amplitude (the \(\sigma\) in the equations above). This is the continuum amplitude; the grid discretisation is handled internally.
grid_shape (sequence of int) – Grid dimensions (Nx, Ny, ...) — must match the simulation grid.
dx (float or sequence of float) – Grid spacing (uniform or per-axis).
n_fields (int) – Number of field components per site.

Notes

The sample method uses real FFT (rfftn / irfftn) for efficiency. It draws white noise in real space, transforms to Fourier space, multiplies by \(\sigma\,|k|\,\sqrt{2/\Delta V}\) (where \(\Delta V = \prod \Delta x_\alpha\) is the cell volume), and transforms back.

The factor \(1/\sqrt{\Delta V}\) provides the correct continuum limit: the noise covariance \(\langle\eta_i\,\eta_j\rangle = -\sigma^2 \nabla^2 \delta_{ij} / \Delta V\) is independent of grid resolution when sigma is held fixed.

property dim: int

effective_D_per_site(extras)[source]

Return an approximate per-site diffusion matrix (d, d).

This is used by the inference pipeline as a pragmatic approximation when the noise is not white-in-space. For WhiteNoise(σ) this returns exactly σ·I. For ConservedNoise it returns the spatially-averaged effective variance per site.

Return type:: Array, shape (d, d)
Parameters:: extras (dict)

property grid_shape: Tuple[int, ...]

property n_fields: int: Number of field components per site.

property noise_kind: str: Short string tag for the noise type.

sample(key, x, extras)[source]

Draw one conserved-noise increment.

Steps: 1. Draw white noise ξ ~ N(0,1) on the grid, shape (Nx, Ny, …, d) 2. rFFT along spatial axes 3. Multiply by \(\sigma\,|k|\,\sqrt{2/\Delta V}\) (the _multiplier) 4. iFFT back to real space 5. Flatten back to (P, d)

The k=0 mode of _multiplier is zero, so sum_i η_i = 0 exactly.

Parameters:

key (Array)
x (Array)
extras (dict)

Return type:

Array

property sigma: float

class SFI.langevin.noise.CompositeNoise(*, components, n_fields)[source]

Bases: NoiseModel

Apply different noise models to different field components.

Useful when some fields have conserved dynamics (e.g. concentration) and others have non-conserved dynamics (e.g. velocity).

Parameters:

components (list of (NoiseModel, field_indices) pairs) – Each element specifies a noise model and the field indices it applies to. field_indices is a list of ints. Together the indices must cover range(n_fields) exactly once.
n_fields (int)

Example

>>> conserved = ConservedNoise(sigma=0.3, grid_shape=(64, 64), n_fields=1)
>>> white = WhiteNoise(sigma=0.1, n_fields=2)
>>> composite = CompositeNoise(
...     components=[(conserved, [0]), (white, [1, 2])],
...     n_fields=3,
... )

property dim: int

effective_D_per_site(extras)[source]

Return an approximate per-site diffusion matrix (d, d).

Return type:: Array, shape (d, d)
Parameters:: extras (dict)

property n_fields: int: Number of field components per site.

property noise_kind: str: Short string tag for the noise type.

sample(key, x, extras)[source]

Draw one noise increment.

Parameters:

key (PRNG key)
x (Array, shape (P, d) or (d,)) – Current state (used only for shape; not accessed by white/conserved noise, but may be needed by multiplicative noise subclasses).
extras (dict) – Process extras (contains grid geometry, etc.).

Returns:

Noise increment already multiplied by sqrt(2) so the integrator applies x += dt*F + sqrt(dt) * sample(...).

Return type:

Array, same shape as x

Factory¶

SFI.bases.spde.conserved_noise_pbc(sigma, *, grid_shape, dx=1.0, n_fields=1)[source]

Build a conserved-noise model for a periodic square grid.

Returns a ConservedNoise instance configured for the given grid geometry. The noise implements

\[\eta(\mathbf{x}, t) = \nabla\!\cdot\!(\sigma\,\vec{\Lambda})\]

where \(\vec{\Lambda}\) is spatiotemporal white vector noise with periodic boundary conditions. This conserves the spatial average: \(\sum_i \eta_i = 0\) exactly at every time step.

Parameters:

sigma (float) – Continuum noise amplitude.
grid_shape (sequence of int) – Grid dimensions (Nx, Ny, ...).
dx (float or sequence of float) – Grid spacing (uniform or per-axis).
n_fields (int) – Number of field components per site.

Returns:

Ready to pass as D= to OverdampedProcess.

Return type:

ConservedNoise

Example

>>> from SFI.bases.spde import conserved_noise_pbc, square_grid_extras
>>> noise = conserved_noise_pbc(sigma=0.3, grid_shape=(64, 64), dx=1.0)
>>> proc  = OverdampedProcess(basis, D=noise)