NVT Nosé Hoover does not work for chain_length=1

[thomasloux]

The thermostat fails when running half_step_chain_fn:

scale = torch.exp(-delta_8 * p_xi[1] / Q[1]) # Around line 403 in md.py

The patch is easy probably, for instance by comparing with ASE implementation (mostly the same but slightly better for edge cases like this one): https://ase-lib.org/_modules/ase/md/nose_hoover_chain.html in NoseHooverChainThermostat._integrate_p_eta_j

Code to reproduce:

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

import os

import torch
from ase.build import bulk

import torch_sim as ts
from torch_sim.integrators.nvt_batched import nvt_nose_hoover as nvt_nose_hoover_batched
from torch_sim.integrators.nvt_batched import nvt_nose_hoover_invariant as nvt_nose_hoover_invariant_batched
# from torch_sim.integrators.nvt import nvt_nose_hoover, nvt_nose_hoover_invariant
from orb_models.forcefield import pretrained

from torch_sim.models.orb import OrbModel
from torch_sim.quantities import calc_kT
from torch_sim.units import MetalUnits as Units
from torch_sim.state import concatenate_states
from torch_sim.state import initialize_state
from torch_sim.integrators.md import calculate_momenta
from torch_sim.quantities import calc_kinetic_energy


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


# Number of steps to run
N_steps = 100

# Create diamond cubic Silicon
si_dc = bulk("Si", "diamond", a=5.43, cubic=True)

orbff = pretrained.orb_v3_conservative_inf_omat(
        device=device,
        # precision="float32-highest",
        precision="float64",
        compile=False
    )
dtype = torch.get_default_dtype()
orb_model = OrbModel(model=orbff, dtype=torch.get_default_dtype(), device=device)

dt = 0.001 * Units.time  # Timestep (ps)
kT = (
    torch.tensor(1000, device=device, dtype=dtype) * Units.temperature
)  # Initial temperature (K)


# state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype)
state = initialize_state(si_dc, device=device, dtype=dtype)

# Run initial inference
results = orb_model(state)



nvt_init, nvt_update = nvt_nose_hoover_batched(model=orb_model, kT=kT, dt=dt, chain_length=1)
state = nvt_init(state=state, kT=kT, seed=1)

for step in range(N_steps):
    if step % 10 == 0:
        temp = (
            calc_kT(
                masses=state.masses, momenta=state.momenta, system_idx=state.system_idx
            )
            / Units.temperature
        )
        invariant = nvt_nose_hoover_invariant_batched(state, kT=kT)
        print(f"{step}: Temperature: {temp}: {invariant}")
    state = nvt_update(state=state, kT=kT)

final_temp = (
    calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx)
    / Units.temperature
)
print(f"Final temperature: {final_temp}")

[abhijeetgangan]

@thomasloux I implemented this following jax-md and it compares well with LAMMPS so would follow that instead of ASE.

[thomasloux]

I didn't know that but JAX out of bound for array standard behaviour is to output the last item of the array. So at least jax_md can run without throwing an error.

Example:

from jax_md import simulate
from jax_md import energy, quantity
from jax import random
from jax_md import space
box_size = 25.0
displacement_fn, shift_fn = space.periodic(box_size)

N = 2
spatial_dimension = 2
key = random.PRNGKey(0)
R = random.uniform(key, (N, spatial_dimension), minval=0.0, maxval=1.0)
energy_fn = energy.lennard_jones_pair(displacement_fn)

temperature = 1.0
dt = 1e-3
init, update = simulate.nvt_nose_hoover(energy_fn, shift_fn, dt, temperature, chain_length=1)
state = init(key, R)
for _ in range(100):
  state = update(state)
R = state.position
print(state.chain.momentum, state.chain.momentum[0], state.chain.momentum[1])
# [1.5175745] 1.5175745 1.5175745

[thomasloux]

After thinking about it, I tend to think that jax-md implementation for chain_length=1 is false. The algorithm is modified in a undesirable way with this out of bound mechanism

[abhijeetgangan]

I need to check my notes but from my old tests with the same settings for the nose-hoover chains I get comparable results with LAMMPS with the same settings. Don't remember if I ran with chain_length=1 with both LAMMPS and jax-md. I can take a look at the ASE case you mentioned. Did ASE compare with LAMMPS?

[thomasloux]

At first sight, it seems to be similar and actually I think that the half update function is strictly the same. But actually at initialization, the code adds a dummy thermostat that probably handles the edge case.

https://github.com/lammps/lammps/blob/0d479cae298abb73ffd8b1460197a8a74c63bc5e/src/fix_nh.cpp#L529

[abhijeetgangan]

But actually at initialization, the code adds a dummy thermostat that probably handles the edge case.

torch-sim/torch_sim/integrators/nvt.py

Line 359 in d7551b6

chain_fns = construct_nose_hoover_chain(

This would be a right place to replicate the behavior. Either ways we need to be consistent with other implementations so will revisit the batched version of NVT again.