Restarts with MACE_MP+D3

[trdurrant]

Summary

I am seeing a problem restarting MACE_MP calculations, that include dispersion:

from janus_core.calculations.md import NVT

from ase.build import molecule
mol = molecule('H2O')

opt = NVT(
    struct=mol,
    arch="mace_mp",
    model="medium",
    device="cpu",
    calc_kwargs={"default_dtype":"float64", "dispersion":True},
    temp=300.0,
    timestep=1.0,
    steps=100,
    restart_every=100,
    stats_every=100,
    traj_every=100,
    restart=False,
    )

opt.run()

opt = NVT(
    struct=mol,
    arch="mace_mp",
    model="medium",
    device="cpu",
    calc_kwargs={"default_dtype":"float64", "dispersion":True},
    temp=300.0,
    timestep=1.0,
    steps=200,
    restart_every=100,
    stats_every=100,
    traj_every=100,
    restart=True,
    )

opt.run()

which produces the following error:

Traceback (most recent call last):
  File "/home/tom/janus_bug/md.py", line 23, in <module>
    opt = NVT(
          ^^^^
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/calculations/md.py", line 1706, in __init__
    super().__init__(*args, ensemble=ensemble, **kwargs)
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/calculations/md.py", line 545, in __init__
    self._prepare_restart()
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/calculations/md.py", line 822, in _prepare_restart
    self.struct, self.struct_path = input_structs(
                                    ^^^^^^^^^^^^^^
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/helpers/struct_io.py", line 221, in input_structs
    attach_calculator(
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/helpers/struct_io.py", line 98, in attach_calculator
    calculator = choose_calculator(
                 ^^^^^^^^^^^^^^^^^^
  File "/home/tom/janus_bug/local/lib/python3.12/site-packages/janus_core/helpers/mlip_calculators.py", line 487, in choose_calculator
    raise ValueError(
ValueError: Unrecognized arch='mace_mp_d3'. Suported architectures are mace, mace_mp, mace_off, m3gnet, chgnet, alignn, sevennet, nequip, dpa3, orb, mattersim, grace, esen, equiformer, pet_mad, uma, mace_omol

I haven't seen this problem using other MLIPs.

Platform

Linux 6.6.87.1-microsoft-standard-WSL2 x86_64 GNU/Linux

Version

0.8.6

Python version

3.12.3

[ElliottKasoar]

Interesting, thanks for spotting this!

What's happening is that when we initialise NVT (super().__init__) we set self.arch to be [arch]_d3 to match the value stored under Atoms.calc.parameters["arch"].

Normally, this isn't a problem, but we then re-use self.arch to set up the structure to continue running MD for, but then the value doesn't match expected architectures.

We could of course set it such that any [arch]_d3 is also valid, but I think the only place self.arch is actually used currently is to set the default_columns that we write structures during MD, so a neater solution may be to keep self.arch as mace_mp (for example), and use Atoms.calc.parameters["arch"] instead to set default_columns.