nvt_langevin_step

torch_sim.integrators.nvt.nvt_langevin_step(state, model, *, dt, kT, gamma=None)[source]

Perform one complete Langevin dynamics integration step using the BAOAB scheme.

Implements the BAOAB splitting of the Langevin equation from Leimkuhler & Matthews (2013, 2016) [2]. The Langevin SDE is:

\[\begin{split}d\mathbf{q} &= M^{-1}\mathbf{p}\,dt \\ d\mathbf{p} &= -\nabla U(\mathbf{q})\,dt - \gamma\,\mathbf{p}\,dt + \sigma M^{1/2}\,d\mathbf{W}\end{split}\]

where \(\sigma = \sqrt{2\gamma k_BT}\) (fluctuation-dissipation relation).

BAOAB splitting (B = kick, A = drift, O = Ornstein-Uhlenbeck):

\[\begin{split}\text{B:}\quad \mathbf{p} &\leftarrow \mathbf{p} + \tfrac{\Delta t}{2}\,\mathbf{F}(\mathbf{q}) \\ \text{A:}\quad \mathbf{q} &\leftarrow \mathbf{q} + \tfrac{\Delta t}{2}\,M^{-1}\mathbf{p} \\ \text{O:}\quad \mathbf{p} &\leftarrow c_1\,\mathbf{p} + c_2\,M^{1/2}\mathbf{R}, \quad \mathbf{R}\sim\mathcal{N}(0,I) \\ \text{A:}\quad \mathbf{q} &\leftarrow \mathbf{q} + \tfrac{\Delta t}{2}\,M^{-1}\mathbf{p} \\ \text{B:}\quad \mathbf{p} &\leftarrow \mathbf{p} + \tfrac{\Delta t}{2}\,\mathbf{F}(\mathbf{q})\end{split}\]

with \(c_1 = e^{-\gamma\Delta t}\) and \(c_2 = \sqrt{k_BT\,(1-c_1^2)}\).

Variable mapping (equation -> code):

Equation symbol

Code variable

\(\mathbf{q}\) (positions)

state.positions

\(\mathbf{p}\) (momenta)

state.momenta

\(M\) (mass matrix)

state.masses

\(\mathbf{F}\) (forces)

state.forces

\(\gamma\) (friction coefficient)

gamma

\(k_BT\) (thermal energy)

kT

\(\Delta t\) (timestep)

dt

\(c_1\)

c1 in _ou_step

\(c_2\)

c2 in _ou_step
Parameters:

  • state (MDState) – Current system state containing positions, momenta, forces

  • model (ModelInterface) – Neural network model that computes energies and forces. Must return a dict with ‘energy’ and ‘forces’ keys.

  • dt (float | Tensor) – Integration timestep, either scalar or shape [n_systems]

  • kT (float | Tensor) – Target temperature in energy units, either scalar or with shape [n_systems]

  • gamma (float | Tensor | None) – Friction coefficient for Langevin thermostat, either scalar or with shape [n_systems]. Defaults to 1/(100*dt).

Returns:

Updated state after one complete Langevin step with new positions,

momenta, forces, and energy

Return type:

MDState

References