Add CSVR / V-Rescale thermostat and anisotropic C rescale barostat

tests/test_integrators.py

@@ -485,6 +485,206 @@ def test_nvt_nose_hoover_multi_kt(
         assert invariant_std / invariant_traj.mean() < 0.1
 
 
+def test_nvt_vrescale(ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel):
+    n_steps = 100
+    dt = torch.tensor(0.001, dtype=DTYPE)
+    kT = torch.tensor(300, dtype=DTYPE) * MetalUnits.temperature
+
+    # Initialize integrator
+    state = ts.nvt_vrescale_init(
+        state=ar_double_sim_state, model=lj_model, kT=kT, seed=42
+    )
+    energies = []
+    temperatures = []
+    for _step in range(n_steps):
+        state = ts.nvt_vrescale_step(model=lj_model, state=state, dt=dt, kT=kT)
+
+        # Calculate instantaneous temperature from kinetic energy
+        temp = ts.calc_kT(
+            masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+        )
+        energies.append(state.energy)
+        temperatures.append(temp / MetalUnits.temperature)
+
+    # Convert temperatures list to tensor
+    temperatures_tensor = torch.stack(temperatures)
+    temperatures_list = [t.tolist() for t in temperatures_tensor.T]
+
+    energies_tensor = torch.stack(energies)
+    energies_list = [t.tolist() for t in energies_tensor.T]
+
+    # Basic sanity checks
+    assert len(energies_list[0]) == n_steps
+    assert len(temperatures_list[0]) == n_steps
+
+    # Check temperature is roughly maintained for each trajectory
+    mean_temps = torch.mean(temperatures_tensor, dim=0)  # Mean temp for each trajectory
+    for mean_temp in mean_temps:
+        assert (
+            abs(mean_temp - kT.item() / MetalUnits.temperature) < 100.0
+        )  # Allow for thermal fluctuations
+
+    # Check energy is stable for each trajectory
+    for traj in energies_list:
+        energy_std = torch.tensor(traj).std()
+        assert energy_std < 1.0  # Adjust threshold as needed
+
+    # Check positions and momenta have correct shapes
+    n_atoms = 8
+
+    # Verify the two systems remain distinct
+    pos_diff = torch.norm(
+        state.positions[:n_atoms].mean(0) - state.positions[n_atoms:].mean(0)
+    )
+    assert pos_diff > 0.0001  # Systems should remain separated
+
+
+def test_npt_anisotropic_crescale(
+    ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel
+) -> None:
+    n_steps = 200
+    dt = torch.tensor(0.001, dtype=DTYPE)
+    kT = torch.tensor(100.0, dtype=DTYPE) * MetalUnits.temperature
+    external_pressure = torch.tensor(0.0, dtype=DTYPE) * MetalUnits.pressure
+    tau_p = torch.tensor(0.1, dtype=DTYPE)
+    isothermal_compressibility = torch.tensor(1e-4, dtype=DTYPE)
+
+    # Initialize integrator using new direct API
+    state = ts.npt_crescale_init(
+        state=ar_double_sim_state,
+        model=lj_model,
+        dt=dt,
+        kT=kT,
+        tau_p=tau_p,
+        isothermal_compressibility=isothermal_compressibility,
+        seed=42,
+    )
+
+    # Run dynamics for several steps
+    energies = []
+    temperatures = []
+    for _step in range(n_steps):
+        state = ts.npt_crescale_anisotropic_step(
+            state=state,
+            model=lj_model,
+            dt=dt,
+            kT=kT,
+            external_pressure=external_pressure,
+        )
+
+        # Calculate instantaneous temperature from kinetic energy
+        temp = ts.calc_kT(
+            masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+        )
+        energies.append(state.energy)
+        temperatures.append(temp / MetalUnits.temperature)
+
+    # Convert temperatures list to tensor
+    temperatures_tensor = torch.stack(temperatures)
+    temperatures_list = [t.tolist() for t in temperatures_tensor.T]
+
+    energies_tensor = torch.stack(energies)
+    energies_list = [t.tolist() for t in energies_tensor.T]
+
+    # Basic sanity checks
+    assert len(energies_list[0]) == n_steps
+    assert len(temperatures_list[0]) == n_steps
+
+    # Check temperature is roughly maintained for each trajectory
+    mean_temps = torch.mean(temperatures_tensor, dim=0)  # Mean temp for each trajectory
+    for mean_temp in mean_temps:
+        assert (
+            abs(mean_temp - kT.item() / MetalUnits.temperature) < 150.0
+        )  # Allow for thermal fluctuations
+
+    # Check energy is stable for each trajectory
+    for traj in energies_list:
+        energy_std = torch.tensor(traj).std()
+        assert energy_std < 1.0  # Adjust threshold as needed
+
+    # Check positions and momenta have correct shapes
+    n_atoms = 8
+
+    # Verify the two systems remain distinct
+    pos_diff = torch.norm(
+        state.positions[:n_atoms].mean(0) - state.positions[n_atoms:].mean(0)
+    )
+    assert pos_diff > 0.0001  # Systems should remain separated
+
+
+def test_npt_isotropic_crescale(
+    ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel
+) -> None:
+    n_steps = 200
+    dt = torch.tensor(0.001, dtype=DTYPE)
+    kT = torch.tensor(100.0, dtype=DTYPE) * MetalUnits.temperature
+    external_pressure = torch.tensor(0.0, dtype=DTYPE) * MetalUnits.pressure
+    tau_p = torch.tensor(0.1, dtype=DTYPE)
+    isothermal_compressibility = torch.tensor(1e-4, dtype=DTYPE)
+
+    # Initialize integrator using new direct API
+    state = ts.npt_crescale_init(
+        state=ar_double_sim_state,
+        model=lj_model,
+        dt=dt,
+        kT=kT,
+        tau_p=tau_p,
+        isothermal_compressibility=isothermal_compressibility,
+        seed=42,
+    )
+
+    # Run dynamics for several steps
+    energies = []
+    temperatures = []
+    for _step in range(n_steps):
+        state = ts.npt_crescale_isotropic_step(
+            state=state,
+            model=lj_model,
+            dt=dt,
+            kT=kT,
+            external_pressure=external_pressure,
+        )
+
+        # Calculate instantaneous temperature from kinetic energy
+        temp = ts.calc_kT(
+            masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
+        )
+        energies.append(state.energy)
+        temperatures.append(temp / MetalUnits.temperature)
+
+    # Convert temperatures list to tensor
+    temperatures_tensor = torch.stack(temperatures)
+    temperatures_list = [t.tolist() for t in temperatures_tensor.T]
+
+    energies_tensor = torch.stack(energies)
+    energies_list = [t.tolist() for t in energies_tensor.T]
+
+    # Basic sanity checks
+    assert len(energies_list[0]) == n_steps
+    assert len(temperatures_list[0]) == n_steps
+
+    # Check temperature is roughly maintained for each trajectory
+    mean_temps = torch.mean(temperatures_tensor, dim=0)  # Mean temp for each trajectory
+    for mean_temp in mean_temps:
+        assert (
+            abs(mean_temp - kT.item() / MetalUnits.temperature) < 150.0
+        )  # Allow for thermal fluctuations
+
+    # Check energy is stable for each trajectory
+    for traj in energies_list:
+        energy_std = torch.tensor(traj).std()
+        assert energy_std < 1.0  # Adjust threshold as needed
+
+    # Check positions and momenta have correct shapes
+    n_atoms = 8
+
+    # Verify the two systems remain distinct
+    pos_diff = torch.norm(
+        state.positions[:n_atoms].mean(0) - state.positions[n_atoms:].mean(0)
+    )
+    assert pos_diff > 0.0001  # Systems should remain separated
+
+
 def test_npt_nose_hoover(ar_double_sim_state: ts.SimState, lj_model: LennardJonesModel):
     dtype = torch.float64
     n_steps = 100

torch_sim/__init__.py

@@ -34,10 +34,15 @@
     nvt_nose_hoover_init,
     nvt_nose_hoover_invariant,
     nvt_nose_hoover_step,
+    nvt_vrescale_init,
+    nvt_vrescale_step,
 )
 from torch_sim.integrators.npt import (
     NPTLangevinState,
     NPTNoseHooverState,
+    npt_crescale_anisotropic_step,
+    npt_crescale_init,
+    npt_crescale_isotropic_step,
     npt_langevin_init,
     npt_langevin_step,
     npt_nose_hoover_init,

torch_sim/integrators/__init__.py

@@ -8,23 +8,42 @@
 NVE:
     - Velocity Verlet integrator for constant energy simulations :func:`nve.nve_step`
 NVT:
+    - Velocity Rescaling thermostat integrator
+        :func:`nvt.nvt_vrescale_step` [1]
     - Langevin thermostat integrator :func:`nvt.nvt_langevin_step`
-        using BAOAB scheme [1]
-    - Nosé-Hoover thermostat integrator :func:`nvt.nvt_nose_hoover_step` from [2]
+        using BAOAB scheme [2]
+    - Nosé-Hoover thermostat integrator :func:`nvt.nvt_nose_hoover_step` from [3]
 NPT:
-    - Langevin barostat integrator :func:`npt.npt_langevin_step` [3, 4]
-    - Nosé-Hoover barostat integrator :func:`npt.npt_nose_hoover_step` from [2]
+    - Langevin barostat integrator :func:`npt.npt_langevin_step` [4, 5]
+    - Nosé-Hoover barostat integrator :func:`npt.npt_nose_hoover_step` from [3]
+    - Isotropic C-Rescale barostat integrator :func:`npt.npt_crescale_isotropic_step`
+    from [6, 8, 9]
+    - C-Rescale barostat integrator :func:`npt.npt_crescale_anisotropic_step`
+    from [7, 8, 9]. Available implementations include isotropic and
+    anisotropic cell rescaling, allowing to change cell lengths, and potentially angles
+    as well.
 
 References:
-    [1] Leimkuhler B, Matthews C.2016 Efficient molecular dynamics using geodesic
+    [1] Bussi G, Donadio D, Parrinello M. "Canonical sampling through velocity rescaling."
+        The Journal of chemical physics, 126(1), 014101 (2007).
+    [2] 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).
+    [3] 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.
-    [3] Grønbech-Jensen, N., & Farago, O. (2014).
+    [4] Grønbech-Jensen, N., & Farago, O. (2014).
         Constant pressure and temperature discrete-time Langevin molecular dynamics.
         The Journal of chemical physics, 141(19).
-    [4] LAMMPS: https://docs.lammps.org/fix_press_langevin.html
+    [5] LAMMPS: https://docs.lammps.org/fix_press_langevin.html
+    [6] Bernetti, Mattia, and Giovanni Bussi.
+        "Pressure control using stochastic cell rescaling."
+        The Journal of Chemical Physics 153.11 (2020).
+    [7] Del Tatto, Vittorio, et al. "Molecular dynamics of solids at
+        constant pressure and stress using anisotropic stochastic cell rescaling."
+        Applied Sciences 12.3 (2022): 1139.
+    [8] Bussi Anisotropic C-Rescale SimpleMD implementation:
+        https://github.com/bussilab/crescale/blob/master/simplemd_anisotropic/simplemd.cpp
+    [9] Supplementary Information for [6].
 
 
 Examples:
@@ -53,6 +72,9 @@
 from .npt import (
     NPTLangevinState,
     NPTNoseHooverState,
+    npt_crescale_anisotropic_step,
+    npt_crescale_init,
+    npt_crescale_isotropic_step,
     npt_langevin_init,
     npt_langevin_step,
     npt_nose_hoover_init,
@@ -67,6 +89,8 @@
     nvt_nose_hoover_init,
     nvt_nose_hoover_invariant,
     nvt_nose_hoover_step,
+    nvt_vrescale_init,
+    nvt_vrescale_step,
 )
 
 
@@ -79,11 +103,16 @@ class Integrator(StrEnum):
 
     Available options:
         - ``nve``: Constant energy (microcanonical) ensemble.
+        - ``nvt_vrescale``: Velocity rescaling thermostat for constant temperature.
         - ``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.
+        - ``npt_isotropic_crescale``: Isotropic C-Rescale barostat for constant
+                temperature and pressure with fixed cell shape.
+        - ``npt_anisotropic_crescale``: Anisotropic C-Rescale barostat for constant
+                temperature and pressure with variable cell shape.
 
     Example:
         >>> integrator = Integrator.nvt_langevin
@@ -93,10 +122,13 @@ class Integrator(StrEnum):
     """
 
     nve = "nve"
+    nvt_vrescale = "nvt_vrescale"
     nvt_langevin = "nvt_langevin"
     nvt_nose_hoover = "nvt_nose_hoover"
     npt_langevin = "npt_langevin"
     npt_nose_hoover = "npt_nose_hoover"
+    npt_isotropic_crescale = "npt_isotropic_crescale"
+    npt_anisotropic_crescale = "npt_anisotropic_crescale"
 
 
 #: Integrator registry - maps integrator names to (init_fn, step_fn) pairs.
@@ -116,18 +148,27 @@ class Integrator(StrEnum):
 #: The available integrators are:
 #:
 #: - ``Integrator.nve``: Velocity Verlet (microcanonical)
+#: - ``Integrator.nvt_vrescale``: V-Rescale thermostat
 #: - ``Integrator.nvt_langevin``: Langevin thermostat
 #: - ``Integrator.nvt_nose_hoover``: Nosé-Hoover thermostat
 #: - ``Integrator.npt_langevin``: Langevin barostat
 #: - ``Integrator.npt_nose_hoover``: Nosé-Hoover barostat
+#: - ``Integrator.npt_isotropic_crescale``: Isotropic NPT C-Rescale barostat
+#: - ``Integrator.npt_anisotropic_crescale``: Anisotropic NPT C-Rescale barostat
 #:
 #: :type: dict[Integrator, tuple[Callable[..., Any], Callable[..., Any]]]
 INTEGRATOR_REGISTRY: Final[
     dict[Integrator, tuple[Callable[..., Any], Callable[..., Any]]]
 ] = {
     Integrator.nve: (nve_init, nve_step),
+    Integrator.nvt_vrescale: (nvt_vrescale_init, nvt_vrescale_step),
     Integrator.nvt_langevin: (nvt_langevin_init, nvt_langevin_step),
     Integrator.nvt_nose_hoover: (nvt_nose_hoover_init, nvt_nose_hoover_step),
     Integrator.npt_langevin: (npt_langevin_init, npt_langevin_step),
     Integrator.npt_nose_hoover: (npt_nose_hoover_init, npt_nose_hoover_step),
+    Integrator.npt_isotropic_crescale: (npt_crescale_init, npt_crescale_isotropic_step),
+    Integrator.npt_anisotropic_crescale: (
+        npt_crescale_init,
+        npt_crescale_anisotropic_step,
+    ),
 }

torch_sim/integrators/md.py

@@ -7,7 +7,7 @@
 
 from torch_sim import transforms
 from torch_sim.models.interface import ModelInterface
-from torch_sim.quantities import calc_temperature
+from torch_sim.quantities import calc_kT, calc_temperature
 from torch_sim.state import SimState
 from torch_sim.units import MetalUnits
 
@@ -76,6 +76,19 @@ def calc_temperature(
             units=units,
         )
 
+    def calc_kT(self) -> torch.Tensor:  # noqa: N802
+        """Calculate kT from momenta, masses, and system indices.
+
+        Returns:
+            torch.Tensor: Calculated kT
+        """
+        return calc_kT(
+            masses=self.masses,
+            momenta=self.momenta,
+            system_idx=self.system_idx,
+            dof_per_system=self.get_number_of_degrees_of_freedom(),
+        )
+
 
 def calculate_momenta(
     positions: torch.Tensor,