Add integrators in docs (#290)

Author: thomasloux

README.md

@@ -58,7 +58,7 @@ final_state = ts.integrate(
     n_steps=50,
     timestep=0.002,
     temperature=1000,
-    integrator=ts.integrators.nvt_langevin,
+    integrator=ts.Integrator.nvt_langevin,
     trajectory_reporter=dict(filenames=trajectory_files, state_frequency=10),
 )
 final_atoms_list = final_state.to_atoms()

docs/reference/index.rst

@@ -16,6 +16,7 @@ Overview of the TorchSim API.
     properties.correlations
     elastic
     io
+    integrators
     math
     models
     monte_carlo

torch_sim/integrators/__init__.py

@@ -5,6 +5,23 @@
 Each integrator handles batched simulations efficiently using PyTorch tensors and
 supports periodic boundary conditions.
 
+NVE:
+    - Velocity Verlet integrator for constant energy simulations :func:`nve.nve_step`
+NVT:
+    - Langevin thermostat integrator :func:`nvt.nvt_langevin_step`
+        using BAOAB scheme [1]
+    - Nosé-Hoover thermostat integrator :func:`nvt.nvt_nose_hoover_step` from [2]
+NPT:
+    - Langevin barostat integrator :func:`npt.npt_langevin_step`
+    - Nosé-Hoover barostat integrator :func:`npt.npt_nose_hoover_step` from [2]
+
+References:
+    [1] Leimkuhler B, Matthews C.2016 Efficient molecular dynamics using geodesic
+        integration and solvent-solute splitting. Proc. R. Soc. A 472: 20160138
+    [2] Martyna, G. J., Tuckerman, M. E., Tobias, D. J., & Klein, M. L. (1996).
+        Explicit reversible integrators for extended systems dynamics.
+        Molecular Physics, 87(5), 1117-1157.
+
 Examples:
     >>> import torch_sim as ts
     >>> state = ts.nvt_langevin_init(model, initial_state, kT=300.0 * units.temperature)
@@ -16,6 +33,8 @@
 Notes:
     All integrators support batched operations for efficient parallel simulation
     of multiple systems.
+
+
 """
 
 # ruff: noqa: F401
@@ -47,7 +66,26 @@
 
 
 class Integrator(StrEnum):
-    """Flavor of molecular dynamics simulation."""
+    """Enumeration of available molecular dynamics (MD) integrators.
+
+    Each member represents a different simulation ensemble or thermostat/barostat
+    scheme. These values are used as keys in :data:`INTEGRATOR_REGISTRY`
+    to select the corresponding initialization and stepping functions.
+
+    Available options:
+        - ``nve``: Constant energy (microcanonical) ensemble.
+        - ``nvt_langevin``: Langevin thermostat for constant temperature.
+        - ``nvt_nose_hoover``: Nosé-Hoover thermostat for constant temperature.
+        - ``npt_langevin``: Langevin barostat for constant temperature and pressure.
+        - ``npt_nose_hoover``: Nosé-Hoover barostat for constant temperature
+                and constant pressure.
+
+    Example:
+        >>> integrator = Integrator.nvt_langevin
+        >>> print(integrator.value)
+        'nvt_langevin'
+
+    """
 
     nve = "nve"
     nvt_langevin = "nvt_langevin"
@@ -56,7 +94,29 @@ class Integrator(StrEnum):
     npt_nose_hoover = "npt_nose_hoover"
 
 
-# Integrator registry - maps integrator names to (init_fn, step_fn) pairs
+#: Integrator registry - maps integrator names to (init_fn, step_fn) pairs.
+#:
+#: This dictionary associates each :class:`Integrator` enum value with a pair
+#: of callables:
+#:
+#: - **init_fn**: A function used to initialize the integrator state.
+#: - **step_fn**: A function that advances the state by one simulation step.
+#:
+#: Example:
+#:
+#:     >>> init_fn, step_fn = INTEGRATOR_REGISTRY[Integrator.nvt_langevin]
+#:     >>> state = init_fn(...)
+#:     >>> new_state = step_fn(state, ...)
+#:
+#: The available integrators are:
+#:
+#: - ``Integrator.nve``: Velocity Verlet (microcanonical)
+#: - ``Integrator.nvt_langevin``: Langevin thermostat
+#: - ``Integrator.nvt_nose_hoover``: Nosé-Hoover thermostat
+#: - ``Integrator.npt_langevin``: Langevin barostat
+#: - ``Integrator.npt_nose_hoover``: Nosé-Hoover barostat
+#:
+#: :type: dict[Integrator, tuple[Callable[..., Any], Callable[..., Any]]]
 INTEGRATOR_REGISTRY: Final[
     dict[Integrator, tuple[Callable[..., Any], Callable[..., Any]]]
 ] = {

torch_sim/integrators/md.py

@@ -21,14 +21,14 @@ class MDState(SimState):
 
     Attributes:
         positions (torch.Tensor): Particle positions [n_particles, n_dim]
-        momenta (torch.Tensor): Particle momenta [n_particles, n_dim]
-        energy (torch.Tensor): Total energy of the system [n_systems]
-        forces (torch.Tensor): Forces on particles [n_particles, n_dim]
         masses (torch.Tensor): Particle masses [n_particles]
         cell (torch.Tensor): Simulation cell matrix [n_systems, n_dim, n_dim]
         pbc (bool): Whether to use periodic boundary conditions
         system_idx (torch.Tensor): System indices [n_particles]
         atomic_numbers (torch.Tensor): Atomic numbers [n_particles]
+        momenta (torch.Tensor): Particle momenta [n_particles, n_dim]
+        energy (torch.Tensor): Total energy of the system [n_systems]
+        forces (torch.Tensor): Forces on particles [n_particles, n_dim]
 
     Properties:
         velocities (torch.Tensor): Particle velocities [n_particles, n_dim]
@@ -184,7 +184,7 @@ def velocity_verlet[T: MDState](state: T, dt: torch.Tensor, model: ModelInterfac
         Updated state after one complete velocity Verlet step
 
     Notes:
-        - Time-reversible and symplectic integrator
+        - Time-reversible and symplectic integrator of second order accuracy
         - Conserves energy in the absence of numerical errors
         - Handles periodic boundary conditions if enabled in state
     """
@@ -240,7 +240,25 @@ class NoseHooverChainFns:
     update_mass: Callable
 
 
-# Suzuki-Yoshida weights for multi-timestep integration
+#: Suzuki-Yoshida composition weights for higher-order symplectic integrators.
+#:
+#: These coefficients are used to construct high-order operator-splitting
+#: schemes (Suzuki-Yoshida compositions) in molecular dynamics and Hamiltonian
+#: simulations.
+#:
+#: The coefficients define how lower-order symplectic integrators (e.g., leapfrog)
+#: can be recursively composed to achieve higher-order accuracy while preserving
+#: symplectic structure.
+#:
+#: References:
+#:     - M. Suzuki, *General Decomposition Theory of Ordered Exponentials*,
+#:       Proc. Japan Acad. 69, 161 (1993).
+#:     - H. Yoshida, *Construction of higher order symplectic integrators*,
+#:       Phys. Lett. A 150, 262-268 (1990).
+#:     - M. Tuckerman, *Statistical Mechanics: Theory and Molecular Simulation*,
+#:       Oxford University Press (2010). Section 4.11
+#:
+#: :type: dict[int, torch.Tensor]
 SUZUKI_YOSHIDA_WEIGHTS = {
     1: torch.tensor([1.0]),
     3: torch.tensor([0.828981543588751, -0.657963087177502, 0.828981543588751]),

torch_sim/integrators/nvt.py

@@ -207,12 +207,12 @@ class NVTNoseHooverState(MDState):
 
     Attributes:
         positions: Particle positions with shape [n_particles, n_dimensions]
-        momenta: Particle momenta with shape [n_particles, n_dimensions]
-        energy: Energy of the system
-        forces: Forces on particles with shape [n_particles, n_dimensions]
         masses: Particle masses with shape [n_particles]
         cell: Simulation cell matrix with shape [n_dimensions, n_dimensions]
         pbc: Whether to use periodic boundary conditions
+        momenta: Particle momenta with shape [n_particles, n_dimensions]
+        energy: Energy of the system
+        forces: Forces on particles with shape [n_particles, n_dimensions]
         chain: State variables for the Nose-Hoover chain thermostat
 
     Properties: