nvt_langevin_step

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

Perform one complete Langevin dynamics integration step.

This function implements the BAOAB splitting scheme for Langevin dynamics, which provides accurate sampling of the canonical ensemble. The integration sequence is: 1. Half momentum update using forces (B step) 2. Half position update using updated momenta (A step) 3. Full stochastic update with noise and friction (O step) 4. Half position update using updated momenta (A step) 5. Half momentum update using new forces (B step)

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

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

  • 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

Notes

  • Uses BAOAB splitting scheme for Langevin dynamics

  • Preserves detailed balance for correct NVT sampling

  • Handles periodic boundary conditions if enabled in state

  • Friction coefficient gamma controls the thermostat coupling strength

  • Weak coupling (small gamma) preserves dynamics but with slower thermalization

  • Strong coupling (large gamma) faster thermalization but may distort dynamics