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?

[bishopfunc]

@gabor1 Thank you for the replay. However, I’m a CS student with very limited experience in protein engineering. Is it possible to automate the process mentioned above using libraries or code? For example, could something like the following command be used to achieve this? obabel -i pdb your_structure.pdb -o pdb -O clean_structure.pdb --gen3d

[jharrymoore]

Hi, I had a look at your structure and technically it is fine (no residues missing atoms etc). Is this an alphafold structure? Worth noting the extended strand, in case you want to simulate in water with periodic boundary conditions.

I did try to run dynamics with mace-md but find the structure explodes in the minimiser, which might be to do with how many charged residues your structure has (this model has not been trained on charged chemistry). I will check what happens when I try to run the system with gromacs and update you

[jharrymoore]

ok, so I ran through a gromacs simulation to see if it's a mace issue or a structure issue - turns out that this resulted in unstable simulation. On closer inspection of your pdbfile, there are duplicate atoms that are practically on top of each other, hence the unstable simulation

I'm not sure if your starting structure is from the PDB where an ensemble of side chain positions can be included, or if AlphaFold is doing something similar , but you need to remove these superimposed side chains before you feed the structure file to an MD package

[gabor1]

Use a visualiser to find this yourself (VMD or Ovito), and identify the offending atoms by their index. This is not a mace problem, no MD model would run this structure as it is.

[ilyes319]

Closing, please re post if this does not fix it.