update momenta before integration in npt NH

[alphalm4]

Hello, thanks for offering such a great project!

I've got catastrophic failure in MD when using NPT Nose-Hoover integration. I believe the momenta update seems accidently dropped before position update. Adding momenta update gives no such failure.

Summary

• Add momenta update before position update
• Add stress attribute in NPT nose hoover state

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.

[orionarcher]

What model are you using?

[alphalm4]

What model are you using?

@orionarcher I'm using SevenNet-Omni with the inference channel omol25_low. I'm currently porting my MD jobs from LAMMPS to torch-sim to take advantage of its batch performance. For reference, I’ve attached trajectories from before/after the commit, along with the reference MD trajectory from LAMMPS.

for_ts_pr.zip

[CompRhys]

Looking at the data you share does suggest this is resulting in a trajectory closer to the lammps reference

[orionarcher]

@abhijeetgangan @thomasloux what are your thoughts?

[thomasloux]

I think he's right. I'll have to get a second read in the source equations. Because it's really hard to read.

Note: adding stress in state is potentially a good idea. Although you probably want it for all NPT MD. So potentially a NptMdState (terrible name) could be created to save the stress, all NPT integrator state would inherit from it. It would be easier to save it in trajectory.

[thomasloux]

We don't exactly have the exact same convention as the source. I added the source, I'm almost sure the source is the exact same equation for NVT nose hoover. But there is some slight differences with torch sim NPT nose hoover.
I think there is a swap of name between L1 and L2.

The equation should be checked carefully. I'm not even sure that this fix correspond exactly to the equations.

Thanks for finding this issue!

[thomasloux]

As a note, most of integrators except C/V-rescale come from JAX-MD. In JAX-MD the code actually indeed perform the same implementation as you suggested: https://github.com/jax-md/jax-md/blob/a2c11c18ce65606983824c5cb0d17f5569abb3a1/jax_md/simulate.py#L910

They don't use state.velocities but get the results of L2 (so momenta) and provide directly momenta/mass, but it is equivalent to your proposal @alphalm4

[thomasloux]

In the meantime, could you give isotropic NPT C-Rescale a try to check that it has a behaviour coherent with what you would expect. Of course it's going to be different than NH, but I would really like to know whether my implementation is indeed correct :)