SFI.langevin.overdamped module

Overdamped Langevin simulator (Euler–Maruyama / Heun) with post-run observables.

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

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)

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

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