Simulation API¶

For a narrative introduction and workflow examples, see Simulation.

class SFI.langevin.OverdampedProcess(F, D, theta_F=None, theta_D=None, extras_global=None, extras_local=None, _structural_extras_prepared=False, _prepared_structural=None, _Dinv_const=None, _D_sf=None)[source]

Bases: LangevinBase

Overdamped Langevin simulator (Euler–Maruyama or stochastic Heun).

Parameters:

F (PSF | SF) – Force model with rank=vector, needs_v=False, and pdepth∈{0,1}. If a PSF is provided, bind parameters via set_params() prior to simulation.
D (float | Array | PSF | SF) – Diffusion model: scalar σ (interpreted as σ·I), constant (d×d) matrix, or a PSF/SF with rank=matrix, pdepth∈{0,1} compatible with F.
theta_F (Optional[Array], optional) – Parameter vectors for binding PSF → SF.
theta_D (Optional[Array], optional) – Parameter vectors for binding PSF → SF.
extras_global (Optional[dict], optional) –
Frozen, time-independent extras passed to both F and D at every call. Users should classify extras explicitly as:
- extras_global: system-wide objects (geometry, external field …)
- extras_local: per-particle objects (species labels, radii, …)
At runtime these are merged into a single extras mapping, with local keys overriding global ones, and passed identically to both models.
extras_local (Optional[dict], optional) –
Frozen, time-independent extras passed to both F and D at every call. Users should classify extras explicitly as:
- extras_global: system-wide objects (geometry, external field …)
- extras_local: per-particle objects (species labels, radii, …)
At runtime these are merged into a single extras mapping, with local keys overriding global ones, and passed identically to both models.
_structural_extras_prepared (bool)
_prepared_structural (Dict[str, Any] | None)
_Dinv_const (Array | None)
_D_sf (SF | None)

Notes

This class does not insert particle axes. The shapes must match the model contract:

If F.pdepth == 0, x0.shape == (d,).
If F.pdepth == 1, x0.shape == (P, d).

Notes

After the run (on recorded steps only), we compute:

information I approx 0.25 * sum_t <dx_t, D_inv(x_t) . F(x_t)>
entropy S approx sum_t <dx_t, D_inv(x_mid) . (F(x_t)+F(x_{t+dt}))/2>

where dx_t = x_{t+dt} - x_t and x_mid = (x_{t+dt}+x_t)/2. We evaluate F(x) exactly once per recorded step and reuse it for both terms.

D: float | Array | Basis | PSF | SF

F: Basis | PSF | SF

property diffusion_sf: SF | None

Bound diffusion state function (read-only), or None.

Returns the diffusion matrix as an SF when available. For constant-scalar or constant-matrix diffusion that was not built from a Basis/PSF, this returns None (since there is no callable SF).

Available after initialize() has been called.

Returns:: diffusion_sf(X) evaluates the diffusion matrix at X, or None if diffusion is not representable as an SF.
Return type:: SF or None

extras_global: Dict[str, Any] | None = None

extras_local: Dict[str, Any] | None = None

property force_sf: SF

Bound force state function (read-only).

Available after initialize() has been called. This is the same callable stored internally as _F_sf; exposing it publicly avoids callers reaching into private attributes.

Returns:: force_sf(X) evaluates the (vector) force at positions X.
Return type:: SF

initialize(x0)[source]

Initialize the process state.

Parameters:

x0 (Array) –

Initial position. Must satisfy:

If F.pdepth == 0: shape (d,)
If F.pdepth == 1: shape (P, d)

Return type:

None

Notes

Binds PSF parameters (if any), validates model contracts, and prepares diffusion shortcuts (constant vs state-dependent).

metadata: dict

set_extras(*, extras_global=None, extras_local=None)

Freeze or update extras dictionaries used when calling F and D.

Parameters:

extras_global (Dict[str, Any] | None) – System-wide extras (geometry, neighbor lists, drive protocols, …). Time-dependent entries are supported: a TimeSeriesExtra with one value per recorded frame of the next simulate call, or a plain callable f(t) of physical time (materialized at the frame times before the scan).
extras_local (Dict[str, Any] | None) – Per-particle extras (species labels, radii, …), with the same time-dependence options.

Return type:

None

Notes

Both dictionaries are merged into a single model-facing extras mapping that is passed as extras=… to both F and D. Local keys override global keys on conflicts. Time-dependent values are held constant across the oversampling substeps of each frame (zeroth-order hold); the prerun uses the frame-0 value.

set_params(*, theta_F=None, theta_D=None)

Bind PSF parameters to specialize models (PSF → SF).

If F or D are PSF, these will be consumed during initialize() when the subclass calls _bind_force() and _setup_diffusion().

Notes

We do not overwrite the user-provided F / D objects. Instead, we keep them unmodified and store specialized callables separately (e.g., _F_sf), derived from the pair (object, theta, extras).

Parameters:

theta_F (Array | None)
theta_D (Array | None)

Return type:

None

simulate(dt, Nsteps, key, *, oversampling=4, prerun=0, jit_compile=True, compute_observables=True, method='heun')[source]

Integrate overdamped Langevin dynamics and return a TrajectoryCollection of positions.

Parameters:

dt (float) – Time step between recorded frames.
Nsteps (int) – Number of recorded time steps.
key (Array) – PRNG key for the simulation.
oversampling (int) – Number of integration substeps between recorded frames. The effective substep size is dt / oversampling. Although all integrators have a consistent continuous limit, they introduce short-range, algorithm-specific temporal correlations at the scale of a single step. Downsampling by recording only every oversampling-th substep ensures these artefacts never reach the inference layer. The default of 4 is a safe minimum for typical use; increase it when dt is large or the process varies rapidly.
prerun (int) – Number of recorded steps to discard as burn-in, using the same dt and oversampling.
jit_compile (bool) – If True, JIT-compile the single-step integrator before scanning.
method (str) – Integration scheme. "heun" (default) selects the stochastic Heun predictor-corrector scheme, which achieves weak order 2 for constant (additive) diffusion — the dominant use case — at the cost of two force evaluations per substep. For state-dependent diffusion the Heun scheme still uses the Itô-correct left-point noise evaluation, giving weak order 1 but with better error constants than Euler–Maruyama. "euler" selects the classical Euler–Maruyama integrator (weak order 1).
compute_observables (bool) – If True, compute post-run information and entropy production estimates on the recorded trajectory and store them in the dataset metadata under the "observables" key.
physics: (..) –
Information functional & entropy production (overdamped): :label: info-entropy-overdamped :category: Observable

\[I \approx \tfrac{1}{4}\sum_t \mathrm{d}X_t^\top\, D^{-1}(x_t)\, F(x_t)\]

\[S \approx \sum_t \mathrm{d}X_t^\top\, D^{-1}(x_{\text{mid}})\, \tfrac{1}{2}\bigl[F(x_t)+F(x_{t+1})\bigr]\]

\(I\) estimates the information content; \(S\) the entropy production (time-reversal asymmetry).

Returns:

A collection with a single dataset containing the positions. The underlying dataset has:

X of shape (Nsteps, d) or (Nsteps, P, d),
metadata combining model info (kind, dimension, pdepth, etc.), run info (dt, Nsteps, oversampling, prerun), and optional observables.

Return type:

TrajectoryCollection

theta_D: Array | None = None

theta_F: Array | None = None

class SFI.langevin.UnderdampedProcess(F, D, theta_F=None, theta_D=None, extras_global=None, extras_local=None, _structural_extras_prepared=False, _prepared_structural=None, _D_needs_v=False, _D_sf=None)[source]

Bases: LangevinBase

Underdamped Langevin simulator.

Parameters:

F (PSF | SF) – Force model with rank=vector, needs_v=True, and pdepth∈{0,1}. If a PSF is provided, bind parameters via set_params() prior to simulation.
D (float | Array | PSF | SF) – Diffusion model acting on velocities: scalar σ (interpreted as σ·I), constant (d×d) matrix, or a PSF/SF with rank=matrix. If provided as PSF/SF, it may depend on (x) or (x, v), controlled by its needs_v flag.
theta_F (Array | None)
theta_D (Array | None)
extras_global (Dict[str, Any] | None)
extras_local (Dict[str, Any] | None)
_structural_extras_prepared (bool)
_prepared_structural (Dict[str, Any] | None)
_D_needs_v (bool)
_D_sf (SF | None)

Notes

This class does not insert particle axes; it follows the pdepth convention of the statefunc objects, similarly to OverdampedProcess.

D: float | Array | Basis | PSF | SF

F: Basis | PSF | SF

property diffusion_sf: SF | None

Bound diffusion state function (read-only), or None.

Available after initialize() has been called.

Returns:: diffusion_sf(X) evaluates the diffusion matrix at X, or None if diffusion is not representable as an SF.
Return type:: SF or None

extras_global: Dict[str, Any] | None = None

extras_local: Dict[str, Any] | None = None

property force_sf: SF

Bound force state function (read-only).

Available after initialize() has been called. This is the same callable stored internally as _F_sf; exposing it publicly avoids callers reaching into private attributes.

Returns:: force_sf(X) evaluates the (vector) force at positions X.
Return type:: SF

initialize(x0, v0=None)[source]

Initialize the process state.

Parameters:

x0 (Array) –
Initial position. Must satisfy:
- If F.pdepth == 0: shape (d,)
- If F.pdepth == 1: shape (P, d)
v0 (Array, optional) – Initial velocity. Must have the same shape as x0. Defaults to 0.

Return type:

None

metadata: dict

set_extras(*, extras_global=None, extras_local=None)

Freeze or update extras dictionaries used when calling F and D.

Parameters:

extras_global (Dict[str, Any] | None) – System-wide extras (geometry, neighbor lists, drive protocols, …). Time-dependent entries are supported: a TimeSeriesExtra with one value per recorded frame of the next simulate call, or a plain callable f(t) of physical time (materialized at the frame times before the scan).
extras_local (Dict[str, Any] | None) – Per-particle extras (species labels, radii, …), with the same time-dependence options.

Return type:

None

Notes

set_params(*, theta_F=None, theta_D=None)

Bind PSF parameters to specialize models (PSF → SF).

If F or D are PSF, these will be consumed during initialize() when the subclass calls _bind_force() and _setup_diffusion().

Notes

We do not overwrite the user-provided F / D objects. Instead, we keep them unmodified and store specialized callables separately (e.g., _F_sf), derived from the pair (object, theta, extras).

Parameters:

theta_F (Array | None)
theta_D (Array | None)

Return type:

None

simulate(dt, Nsteps, key, *, oversampling=4, prerun=0, jit_compile=True, compute_observables=False)[source]

Run the integrator and return a TrajectoryCollection of positions.

Parameters:

dt (float) – Time step between recorded frames.
Nsteps (int) – Number of recorded time steps.
key (Array) – PRNG key for the simulation.
oversampling (int) – Number of velocity-Verlet substeps between recorded frames. The effective substep size is dt / oversampling. Although all integrators have a consistent continuous limit, they introduce short-range, algorithm-specific temporal correlations at the scale of a single step. Downsampling by recording only every oversampling-th substep ensures these artefacts never reach the inference layer. The default of 4 is a safe minimum for typical use; increase it when dt is large or the process varies rapidly.
prerun (int) – Number of recorded steps to discard as burn-in.
jit_compile (bool) – If True, JIT-compile the single-step integrator before scanning.
compute_observables (bool) – If True, accumulate the information and entropy-production observables along the run using the true internal velocities (which live in the integrator carry even though they are not recorded), and store them in the dataset metadata under the "observables" key. Costs three extra force evaluations per substep.
physics: (..) –
Information functional & entropy production (underdamped): :label: info-entropy-underdamped :category: Observable

With \(F^{\pm}(x,v) = \tfrac12[F(x,v) \pm F(x,-v)]\) the time-reversal even/odd parts of the acceleration field (\(x\) even, \(v\) odd),

\[I \approx \tfrac{1}{4}\sum_t \mathrm{d}v_t^\top\, D^{-1}\, F(x_t, v_t)\]

\[S \approx \sum_t \bigl[\mathrm{d}v_t - \bar F^{+}_t\,\mathrm{d}t\bigr]^\top D^{-1}\, \bar F^{-}_t , \qquad \bar F^{\pm}_t = \tfrac12\bigl[F^{\pm}(x_t,v_t) + F^{\pm}(x_{t+1},v_{t+1})\bigr]\]

\(S\) is the Stratonovich log path-probability ratio (trajectory vs its time reverse): the entropy produced along the run, in units of \(k_B\). The trapezoid (Stratonovich) evaluation of \(\bar F^{\pm}\) is essential: it realises the exact cancellation \(\langle F^-\!\cdot D^{-1} F^-\rangle + \langle \nabla_v\!\cdot F^-\rangle = 0\) at equilibrium.

Returns:

A collection with a single dataset containing the positions only (velocities are not stored by design). The underlying dataset has:

X of shape (Nsteps, d) or (Nsteps, P, d),
metadata combining model info (kind, dimension, pdepth, etc.) and run info (dt, Nsteps, oversampling, prerun).

Return type:

TrajectoryCollection

theta_D: Array | None = None

theta_F: Array | None = None

Noise models¶

class SFI.langevin.noise.WhiteNoise(sigma, *, n_fields=1)[source]

Bases: NoiseModel

Spatially uncorrelated Gaussian noise: B = sqrt(2σ) I.

Each grid site receives i.i.d. N(0, 2σ dt) noise per component. This recovers the current D = σ (scalar constant) behaviour.

Parameters:

sigma (float) – Scalar diffusion coefficient (the D in dx = F dt + sqrt(2D) dW).
n_fields (int) – Number of field components per site.

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

property sigma: float

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

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