MD Simulation with MACE Results in NaN Temperature and Unchanging Energy

[bishopfunc]

Describe the bug

I am running a molecular dynamics (MD) simulation using mace_off for a PET and PETase docking result. After starting the simulation, the temperature becomes NaN, and the energy does not change over the course of the MD steps.

To Reproduce

Steps to reproduce the behavior:

Use a PDB file for PET-PETase docking results (e.g., 2pet_petase_wt_dockpose.pdb).
Set up the mace_off calculator with the small model (mace_off(model="small", device="cuda")).
Initialize MD simulation using the Andersen thermostat and a 1.0 fs timestep.
Run 10 steps of MD simulation and observe the output.

Expected behavior

I expected the temperature to remain stable and the energy to vary as the MD simulation progressed. Instead, the temperature quickly became NaN and the energy did not change after the first step.

Code

The code is almost identical to the 'Mace Tutorial - Molecular Dynamics with MACE,' which can be found here.

from ase.io import read, write
from ase import units
from mace.calculators import mace_off, mace_mp
from ase.md.langevin import Langevin
from ase.md.verlet import VelocityVerlet
from ase.md.andersen import Andersen
from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution

import random
import os
import time

def simpleMD(init_conf, T_init, calc, fname, s, T):
    init_conf.calc = calc

    random.seed(701)
    MaxwellBoltzmannDistribution(init_conf, temperature_K=T_init)
    Stationary(init_conf)
    ZeroRotation(init_conf)

    dyn = Andersen(atoms=init_conf, timestep=1.0*units.fs, temperature_K=T_init, andersen_prob=1e-2)

    os.system('rm -rfv '+fname)

    def write_frame():
        dyn.atoms.write(fname, append=True)
        time_fs = dyn.get_time()/units.fs
        temperature = dyn.atoms.get_temperature()
        energy = dyn.atoms.get_potential_energy()/len(dyn.atoms)
        print("Time: {0:.2f} fs, Temperature: {1:.2f} K, Energy: {2:.2f} eV".format(time_fs, temperature, energy))
        time.sleep(0.01)

    dyn.attach(write_frame, interval=s)
    dyn.run(T)

pdb_path = "2pet_petase_wt_dockpose.pdb"
molecule = read(pdb_path)
mace_calc = mace_off(model="small", device="cuda")

simpleMD(molecule, 300, mace_calc, "traj.pdb", s=1, T=10)

Output

Time: 0.00 fs, Temperature: 308.81 K, Energy: nan eV/atom 
Time: 1.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 2.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 3.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 4.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 5.00 fs, Temperature: nan K, Energy: -863.39 eV/atom
Time: 6.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 7.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 8.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 9.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
Time: 10.00 fs, Temperature: nan K, Energy: -863.39 eV/atom 
MD finished in 0.74 minutes!

MD finished in 0.74 minutes!
Desktop (please complete the following information):

OS: Ubuntu 24.04 LTS
MACE Version: 0.3.6
ACE Version: 3.23.0
Torch Version: 2.4.0+cu124
GPU Info: V100
CUDA Version: 12.5

Support Info

2pet_petase_wt_dockpose.pdb.zip

[gabor1]

You might want to geometry-optimise the structure before you start the MD. Also make sure your initial structure is "clean", no atoms on top of each other, (e.g. hydrogens) etc.

[ilyes319]

You are using a quite large system. I would suggest you to try to run in OpenMM instead of lammps. You can see the documentation here: https://mace-docs.readthedocs.io/en/latest/guide/openmm.html. Try to see if you can reproduce your errors there. It might be a memory handling probem.

What MACE version are you using?