SFI.bases.spde module

SPDE-oriented basis helpers.

Grid operators here are designed to be box-polymorphic: geometry is provided via per-dataset extras (typically collection.extras_global), so the same basis/PSF object can run on different grid sizes without re-creation.

Important conventions

We represent a regular grid as a “particle system”:

  • particle axis corresponds to flattened grid index p = 0...P-1;

  • n_fields (historically called dim) is the number of field components per grid site (e.g. 2 for Gray-Scott U/V, 1 for a scalar concentration).

Grid-site positions are purely implicit: they are reconstructed from box/grid_shape and box/dx stored in extras. The state vector X[p, :] contains only field values, never spatial coordinates.

The Laplacian helper returns a scalar-rank basis with n_features = n_fields: feature a is the Laplacian of field component X[..., a].

This is deliberate: it composes cleanly with generic StateExpr machinery (slice a component and re-embed it into a vector drift with unit vectors).

Box extras

Every operator reads two global extras (namespaced box by default):

  • {box}/grid_shape : int array, shape (ndim,)

  • {box}/dx : float array, shape (ndim,)

Use square_grid_extras() to create them.

Tip

For GPU acceleration with JAX, ensure a CUDA-enabled jaxlib is installed. All operators are JIT-compiled and run transparently on any JAX backend (CPU / GPU / TPU).

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

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·∇)φ
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.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.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.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.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.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:

  1. Evaluate the inner expression \(f(\phi)\) at all stencil neighbors (13 points in 2D, 25 in 3D).

  2. Compute \(|\nabla f|^2\) at each of the Laplacian’s 5 stencil points (center + nearest neighbours) using central differences among the gathered values.

  3. 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

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

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

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

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

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]

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