SFI.langevin package¶

SFI.langevin¶

Langevin simulators (overdamped, underdamped, …).

This submodule provides simulation classes built on top of the statefunc API. They accept PSF/SF force and diffusion models, enforce shape contracts, and integrate trajectories using stochastic schemes.

Public classes¶

OverdampedProcessoverdamped Langevin simulator (Euler–Maruyama or
stochastic Heun) with entropy/information observables.
UnderdampedProcess : underdamped Langevin simulator (velocity-Verlet).

class SFI.langevin.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).

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

class SFI.langevin.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.NoiseModel(*, n_fields)[source]¶

Bases: ABC

Abstract base for noise models used by Langevin simulators.

Subclasses must implement sample() and effective_D_per_site(). The simulator calls sample(key, x, extras) once per Euler–Maruyama substep to obtain the noise increment (already scaled by sqrt(2) so that the step becomes x += dt*F + sqrt(dt)*sample(key, x, extras)).

Parameters:: n_fields (int) – Number of field components per grid site (= dim in the force contract). E.g. 1 for a scalar field, 2 for a 2-component system.

property dim: int¶

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

abstractmethod 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

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) – Not yet implemented for the underdamped case.

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¶

class SFI.langevin.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).

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¶

SFI.langevin package¶

SFI.langevin¶

Public classes¶

Submodules¶