SFI.langevin.underdamped module

Underdamped Langevin simulator (velocity-Verlet-like, generic F(x,v) and D(x[,v])).

This mirrors overdamped as closely as possible, but simulates the phase-space SDE

dx = v dt dv = F(x, v) dt + sqrt(2 D(x, v)) dW

where diffusion acts on velocity increments. The returned TrajectoryCollection stores positions only by design.

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

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

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