add missing metadata in ForceFieldTaskDoc, MACE-MPA-0, test for it

Author: esoteric-ephemera

src/atomate2/ase/schemas.py

@@ -247,10 +247,10 @@ def from_ase_task_doc(
             Additional keyword args passed to :obj:`.AseStructureTaskDoc()`.
         """
         task_document_kwargs.update(
-            {k: getattr(ase_task_doc, k) for k in _task_doc_translation_keys}
+            {k: getattr(ase_task_doc, k) for k in _task_doc_translation_keys},
+            structure = ase_task_doc.mol_or_struct,
         )
-        task_document_kwargs["structure"] = ase_task_doc.mol_or_struct
-        return cls(**task_document_kwargs)
+        return cls.from_structure(meta_structure = ase_task_doc.mol_or_struct, **task_document_kwargs)
 
 
 class AseMoleculeTaskDoc(MoleculeMetadata):

src/atomate2/forcefields/__init__.py

@@ -12,7 +12,8 @@
 class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     """Names of ML force fields."""
 
-    MACE = "MACE"
+    MACE = "MACE" # This is MACE-MP-0-medium
+    MACE_MPA_0 = "MACE_MPA_0"
     GAP = "GAP"
     M3GNet = "M3GNet"
     CHGNet = "CHGNet"

src/atomate2/forcefields/jobs.py

@@ -35,6 +35,7 @@
     MLFF.NEP: {"model_filename": "nep.txt"},
     MLFF.GAP: {"args_str": "IP GAP", "param_filename": "gap.xml"},
     MLFF.MACE: {"model": "medium"},
+    MLFF.MACE_MPA_0 : {"model": "medium-mpa-0",},
 }
 
 

src/atomate2/forcefields/utils.py

@@ -43,7 +43,8 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
     calculator = None
 
     if isinstance(calculator_meta, str | MLFF) and calculator_meta in map(str, MLFF):
-        calculator_name = MLFF(calculator_meta.split("MLFF.")[-1])
+
+        calculator_name = MLFF[calculator_meta.split("MLFF.")[-1]]
 
         if calculator_name == MLFF.CHGNet:
             from chgnet.model.dynamics import CHGNetCalculator
@@ -58,7 +59,7 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
             potential = matgl.load_model(path)
             calculator = PESCalculator(potential, **kwargs)
 
-        elif calculator_name == MLFF.MACE:
+        elif calculator_name in {MLFF.MACE, MLFF.MACE_MPA_0}:
             from mace.calculators import MACECalculator, mace_mp
 
             model = kwargs.get("model")

tests/forcefields/test_jobs.py

@@ -318,6 +318,25 @@ def test_mace_relax_maker(
         assert output1.output.energy == approx(-0.06772976, rel=1e-4)
         assert output1.output.n_steps == 7
 
+def test_mace_mpa_0_relax_maker(
+    si_structure: Structure,
+):
+    job = ForceFieldRelaxMaker(
+        force_field_name="MACE_MPA_0",
+        steps=25,
+        relax_kwargs={"fmax": 0.005},
+    ).make(si_structure)
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validating the outputs of the job
+    output = responses[job.uuid][1].output
+
+    assert output.ase_calculator_name == "MLFF.MACE_MPA_0"
+    assert output.output.energy == pytest.approx(-10.829493522644043)
+    assert output.output.structure.volume == pytest.approx(40.87471552602735)
+    assert len(output.output.ionic_steps) == 4
+    assert output.structure.volume == output.output.structure.volume
 
 def test_gap_static_maker(si_structure: Structure, test_dir):
     importorskip("quippy")