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

[thomasloux]

Summary

Add Canonical sampling through velocity rescaling (CSVR) or v-rescale as called on Gromacs.

Note: removed the optimization to sample from Gamma distribution which I doubt is useful on a high level framework like pytorch.

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

We highly recommended installing the prek hooks running in CI locally to speedup the development process. Simply run pip install prek && prek install to install the hooks which will check your code before each commit.

[thomasloux]

And is working well. Temperature set is indeed 200K, script modified from https://github.com/TorchSim/torch-sim/blob/main/examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py

[thomasloux]

Note: change the name to follow GROMACS, I prefer this name

[thomasloux]

I need to check very carefully for the number of degrees of freedom used. I think the V-rescale implementation in SimpleMD removes the 3 dof linked to the COM.

[thomasloux]

Observation on some simple implementations:

• V-rescale seems to work great
• C-Rescale seems interestingly more difficult to use. Note that I haven't tried to use an estimation of the isothermal compressibility. In my tests it was rather unstable depending on the used parameters.

If you want to experiment with C-rescale:

"""Lennard-Jones simulation in NPT ensemble using Langevin thermostat."""

# /// script
# dependencies = ["scipy>=1.15"]
# ///
import itertools
import os

import torch

import torch_sim as ts
from torch_sim.models.lennard_jones import LennardJonesModel
from torch_sim.units import MetalUnits as Units
from torch_sim.units import UnitConversion


# Set up the device and data type
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32

a_len = 5.26  # Lattice constant

# Generate base FCC unit cell positions (scaled by lattice constant)
base_positions = torch.tensor(
    [
        [0.0, 0.0, 0.0],  # Corner
        [0.0, 0.5, 0.5],  # Face centers
        [0.5, 0.0, 0.5],
        [0.5, 0.5, 0.0],
    ],
    device=device,
    dtype=dtype,
)

# Create 4x4x4 supercell of FCC Argon manually
positions = []
for i, j, k in itertools.product(range(4), range(4), range(4)):
    for base_pos in base_positions:
        # Add unit cell position + offset for supercell
        pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype)
        positions.append(pos)

# Stack the positions into a tensor
positions = torch.stack(positions)

# Scale by lattice constant
positions = positions * a_len

# Create the cell tensor
cell = torch.tensor(
    [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype
)

# Create the atomic numbers tensor (Argon = 18)
atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int)
# Create the masses tensor (Argon = 39.948 amu)
masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype)

# Initialize the Lennard-Jones model
# Parameters:
#  - sigma: distance at which potential is zero (3.405 Å for Ar)
#  - epsilon: depth of potential well (0.0104 eV for Ar)
#  - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma)
model = LennardJonesModel(
    use_neighbor_list=False,
    sigma=3.405,
    epsilon=0.0104,
    cutoff=2.5 * 3.405,
    device=device,
    dtype=dtype,
    compute_forces=True,
    compute_stress=True,
)
model.compile()
state = ts.SimState(
    positions=positions,
    masses=masses,
    cell=cell.unsqueeze(0),
    atomic_numbers=atomic_numbers,
    pbc=True,
)
# state = ts.concatenate_states([state.clone() for _ in range(2)])  # Duplicate for 2 systems

n_steps = 2_500
dt = torch.tensor(0.001, dtype=dtype) * ts.units.MetalUnits.time
kT = torch.tensor(100.0, dtype=dtype) * ts.units.MetalUnits.temperature
external_pressure = torch.tensor(10_0.0, dtype=dtype) * ts.units.MetalUnits.pressure
# tau_p = torch.tensor(0.1, dtype=dtype)
tau_p = 100 * dt
tau = 100 * dt
isothermal_compressibility = torch.tensor(5e-6, dtype=dtype) / ts.units.MetalUnits.pressure

# Initialize integrator using new direct API
state = ts.npt_crescale_init(
    state=state,
    model=model,
    dt=dt,
    kT=kT,
    tau_p=tau_p,
    isothermal_compressibility=isothermal_compressibility,
    seed=42,
)

# Run dynamics for several steps
energies = []
temperatures = []
pressures = []
volumes = []
for _step in range(n_steps):
    state = ts.npt_crescale_step(
        state=state,
        model=model,
        dt=dt,
        kT=kT,
        external_pressure=external_pressure,
        tau=tau
    )
    if _step % 100 == 0:
        pressure = ts.get_pressure(
            model(state)["stress"],
            ts.calc_kinetic_energy(
                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
            ),
            torch.linalg.det(state.cell),
        )

        # 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 / ts.units.MetalUnits.temperature)
        pressures.append(pressure / ts.units.MetalUnits.pressure)
        volumes.append(torch.linalg.det(state.cell))

# 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]

pressures_tensor = torch.stack(pressures)
pressures_list = [t.tolist() for t in pressures_tensor.T]

[thomasloux]

Current implementation with an exterior isotropic pressure only.
Could be easily extended, the equations are in paper and are rather simple.

The isotropic barostat should probably be implemented as well

[abhijeetgangan]

Looks great. Left a few comments. Choosing the correct value if isothermal compressibility was also an issue in previous version of NPT in ase so we could leave that to the user for now.