diff --git a/README.md b/README.md index 50ac3b0c..caa84574 100644 --- a/README.md +++ b/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() diff --git a/docs/reference/index.rst b/docs/reference/index.rst index a21d7418..9c6d172a 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -16,6 +16,7 @@ Overview of the TorchSim API. properties.correlations elastic io + integrators math models monte_carlo diff --git a/torch_sim/integrators/__init__.py b/torch_sim/integrators/__init__.py index 19767b6b..d9eb27a6 100644 --- a/torch_sim/integrators/__init__.py +++ b/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]]] ] = { diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py index 4d8f209a..b968e18f 100644 --- a/torch_sim/integrators/md.py +++ b/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]), diff --git a/torch_sim/integrators/nvt.py b/torch_sim/integrators/nvt.py index b4cb41ce..4021ce2b 100644 --- a/torch_sim/integrators/nvt.py +++ b/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: