Add integrator tests to check correct sampling of thermodynamic ensemble

[thomasloux]

It would be great to test if NVT integrators correctly sample the canonical ensemble, same for NPT.

I suggest using physical_validation, used to check consistency of Gromacs (https://gitlab.com/gromacs/gromacs/-/tree/main/tests?ref_type=heads).

I started doing this but did not have time to fully write and run the tests.
It's a great opportunity to actually have a good understanding of the time convergence of the integrators and understand better the behavior of these integrators.

"""NVT simulation with MACE and Nose-Hoover thermostat."""


import torch
from ase.build import bulk

import torch_sim as ts
from torch_sim.quantities import calc_kT

# import partial
from torch_sim.units import MetalUnits
from torch_sim.units import MetalUnits as Units
from functools import partial
import matplotlib.pyplot as plt
import numpy as np
import torch_sim as ts
from torch_sim.models.lennard_jones import LennardJonesModel

DEVICE = torch.device("cpu")
DTYPE = torch.float32


def get_energies(atoms, model, init_fn, step_fn, dt, temperature, n_steps=100_000, seed=42, device="cpu"):
    dt = torch.tensor(dt, device=device, dtype=torch.float32) * Units.time  # Timestep (ps)
    kT = (
        torch.tensor(temperature, device=device, dtype=torch.float32) * Units.temperature
    )  # Initial temperature (K)
    # Initialize integrator

    # step_fn = torch.compile(step_fn)
    state =  init_fn(state=atoms, seed=seed, model=model, kT=kT, dt=dt)
    energies = []
    kinetic_energies = []
    volumes = []
    # temperatures = []
    for _step in range(n_steps):
        state = step_fn(model=model, state=state, dt=dt, kT=kT)

        # Calculate instantaneous temperature from kinetic energy
        # temp = calc_kT(
        #     masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
        # )
        kinetic_energy = ts.calc_kinetic_energy(
            momenta=state.momenta, masses=state.masses, system_idx=state.system_idx
        )
        kinetic_energies.append(kinetic_energy)
        energies.append(state.energy)
        volumes.append(torch.det(state.cell))
        # temperatures.append(temp / MetalUnits.temperature)

    energies = torch.stack(energies) / state.n_atoms
    kinetic_energies = torch.stack(kinetic_energies) / state.n_atoms
    volumes = torch.stack(volumes)
    return energies, kinetic_energies, energies + kinetic_energies, volumes


atoms = ts.io.atoms_to_state(bulk("Cu", "fcc", a=3.6, cubic=True)*4, device=DEVICE, dtype=DTYPE)
atoms.pbc = True
model = LennardJonesModel(
    use_neighbor_list=True,
    sigma=3.405,
    epsilon=0.0104,
    device=DEVICE,
    dtype=DTYPE,
    compute_forces=True,
    compute_stress=True,
    cutoff=2.5 * 3.405,
)

n_steps = 10_000
get_energies_langevin_temperature = partial(
    get_energies,
    atoms=atoms,
    model=model,
    init_fn=ts.integrators.nvt_nose_hoover_init,
    step_fn=ts.integrators.nvt_nose_hoover_step,
    n_steps=n_steps,
    dt=0.001,
    seed=42,
    device=DEVICE
)
# dt = [0.0002, 0.0005, 0.001, 0.002, 0.003, 0.004, 0.005]
# dt = np.arange(0.001, 0.010, 0.001)
temperatures = np.array([300, 330])
beta = 1 / (temperatures * ts.units.MetalUnits.temperature)
energies_langevin = [get_energies_langevin_temperature(temperature=temperature)[0] for temperature in temperatures]
energies = torch.stack(energies_langevin)
energy_low, energy_high = energies_langevin
energy_low = energy_low.flatten().numpy()
energy_high = energy_high.flatten().numpy()
1/(ts.units.UnitConversion.eV_to_J / 1000 * ts.units.bc.n_av)
import physical_validation

units_torchsim = physical_validation.data.UnitData(
    # energy_conversion=(ts.units.UnitConversion.eV_to_J / 1000 * ts.units.bc.n_av),
    energy_conversion=1,
    length_conversion=0.1,
    time_conversion=1,
    temperature_conversion=1,
    volume_conversion= 0.1**3,
    kb=ts.units.MetalUnits.temperature,
    pressure_conversion=1,
    energy_str="eV/atom",
    length_str="Angstrom",
    time_str="ps",
    temperature_str="K",
    volume_str="Angstrom^3",
    pressure_str="bar",
)

simulation_data_nvt_nh_low = physical_validation.data.SimulationData(
    # Example simulations were performed using GROMACS
    units=units_torchsim,
    ensemble=physical_validation.data.EnsembleData(
        ensemble="NVT",
        natoms=atoms.n_atoms,
        volume=torch.det(atoms.cell).item(),
        temperature=temperatures[0],
    ),
    observables=physical_validation.data.ObservableData(
        # This test requires only the potential energy
        potential_energy=energy_low * (ts.units.UnitConversion.eV_to_J / 1000 * ts.units.bc.n_av)
    ),
)

simulation_data_nvt_nh_high = physical_validation.data.SimulationData(
    # Example simulations were performed using GROMACS
    units=units_torchsim,
    ensemble=physical_validation.data.EnsembleData(
        ensemble="NVT",
        natoms=atoms.n_atoms,
        volume=torch.det(atoms.cell).item(),
        temperature=temperatures[1],
    ),
    observables=physical_validation.data.ObservableData(
        # This test requires only the potential energy
        potential_energy=energy_high * (ts.units.UnitConversion.eV_to_J / 1000 * ts.units.bc.n_av)
    ),
)
physical_validation.ensemble.estimate_interval(
    data=simulation_data_nvt_nh_low,
)
physical_validation.ensemble.check(
    data_sim_one=simulation_data_nvt_nh_low,
    data_sim_two=simulation_data_nvt_nh_high,
    screen=True,
    filename="test2.png"
)
```

[RishikeshMagar]

I can take point on this