Feat: Add DeepMD MLFF support

pyproject.toml

@@ -60,6 +60,7 @@ forcefields = [
     "quippy-ase>=0.9.14; python_version < '3.12'",
     "sevenn>=0.9.3",
     "torchdata<=0.7.1",                            # TODO: remove when issue fixed
+    "deepmd-kit>=2.1.4",
 ]
 ase = ["ase>=3.23.0"]
 # tblite py3.12 support tracked in https://github.com/tblite/tblite/issues/198
@@ -127,6 +128,8 @@ strict-forcefields = [
     "sevenn==0.10.1",
     "torch==2.5.1",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
+    "deepmd-kit==2.2.11",
+    "tensorflow-cpu==2.16.2",
 ]
 
 [project.scripts]

src/atomate2/forcefields/__init__.py

@@ -16,6 +16,7 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     NEP = "NEP"
     Nequip = "Nequip"
     SevenNet = "SevenNet"
+    DeepMD = "DeepMD"
 
 
 def _get_formatted_ff_name(force_field_name: str | MLFF) -> str:

src/atomate2/forcefields/jobs.py

@@ -34,6 +34,7 @@
     MLFF.M3GNet: {"stress_weight": _GPa_to_eV_per_A3},
     MLFF.NEP: {"model_filename": "nep.txt"},
     MLFF.GAP: {"args_str": "IP GAP", "param_filename": "gap.xml"},
+    MLFF.DeepMD: {"model": "graph.pb"},
 }
 
 

src/atomate2/forcefields/schemas.py

@@ -142,6 +142,7 @@ def from_ase_compatible_result(
             MLFF.MACE: "mace-torch",
             MLFF.GAP: "quippy-ase",
             MLFF.Nequip: "nequip",
+            MLFF.DeepMD: "deepmd-kit",
         }
 
         if pkg_name := {str(k): v for k, v in model_to_pkg_map.items()}.get(

src/atomate2/forcefields/utils.py

@@ -95,6 +95,11 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
 
             calculator = SevenNetCalculator(**{"model": "7net-0"} | kwargs)
 
+        elif calculator_name == MLFF.DeepMD:
+            from deepmd.calculator import DP
+
+            calculator = DP(**kwargs)
+
     elif isinstance(calculator_meta, dict):
         calc_cls = MontyDecoder().process_decoded(calculator_meta)
         calculator = calc_cls(**kwargs)

tests/forcefields/conftest.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+import hashlib
+import urllib.request
 from typing import TYPE_CHECKING
 
+import pytest
 import torch
 
 if TYPE_CHECKING:
+    from pathlib import Path
     from typing import Any
 
 
@@ -13,3 +17,23 @@ def pytest_runtest_setup(item: Any) -> None:
     torch.set_default_dtype(torch.float32)
     # For consistent performance across hardware, explicitly set device to CPU
     torch.set_default_device("cpu")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def download_deepmd_pretrained_model(test_dir: Path) -> None:
+    # Download DeepMD pretrained model from GitHub
+    file_url = "https://raw.github.com/sliutheorygroup/UniPero/main/model/graph.pb"
+    local_path = test_dir / "forcefields" / "deepmd" / "graph.pb"
+    ref_md5 = "2814ae7f2eb1c605dd78f2964187de40"
+    _, http_message = urllib.request.urlretrieve(file_url, local_path)  # noqa: S310
+    if "Content-Type: text/html" in http_message:
+        raise RuntimeError(f"Failed to download from: {file_url}")
+
+    # Check MD5 to ensure file integrity
+    md5_hash = hashlib.md5()
+    with open(local_path, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            md5_hash.update(chunk)
+    file_md5 = md5_hash.hexdigest()
+    if file_md5 != ref_md5:
+        raise RuntimeError(f"MD5 mismatch: {file_md5} != {ref_md5}")

tests/forcefields/test_jobs.py

@@ -537,3 +537,66 @@ def test_nequip_relax_maker(
 
     with pytest.warns(FutureWarning):
         NequipRelaxMaker()
+
+
+def test_deepmd_static_maker(sr_ti_o3_structure: Structure, test_dir: Path):
+    importorskip("deepmd")
+
+    # generate job
+    job = ForceFieldStaticMaker(
+        force_field_name="DeepMD",
+        ionic_step_data=("structure", "energy"),
+        calculator_kwargs={"model": test_dir / "forcefields" / "deepmd" / "graph.pb"},
+    ).make(sr_ti_o3_structure)
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validate the outputs of the job
+    output1 = responses[job.uuid][1].output
+    assert isinstance(output1, ForceFieldTaskDocument)
+    assert output1.output.energy == approx(-3723.09868, rel=1e-4)
+    assert output1.output.n_steps == 1
+    assert output1.forcefield_version == get_imported_version("deepmd-kit")
+
+
+@pytest.mark.parametrize(
+    ("relax_cell", "fix_symmetry"),
+    [(True, False), (False, True)],
+)
+def test_deepmd_relax_maker(
+    sr_ti_o3_structure: Structure,
+    test_dir: Path,
+    relax_cell: bool,
+    fix_symmetry: bool,
+):
+    importorskip("deepmd")
+    # translate one atom to ensure a small number of relaxation steps are taken
+    sr_ti_o3_structure.translate_sites(0, [0, 0, 0.01])
+    # generate job
+    job = ForceFieldRelaxMaker(
+        force_field_name="DeepMD",
+        steps=25,
+        optimizer_kwargs={"optimizer": "BFGSLineSearch"},
+        relax_cell=relax_cell,
+        fix_symmetry=fix_symmetry,
+        calculator_kwargs={"model": test_dir / "forcefields" / "deepmd" / "graph.pb"},
+    ).make(sr_ti_o3_structure)
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(job, ensure_success=True)
+
+    # validate the outputs of the job
+    output1 = responses[job.uuid][1].output
+    assert isinstance(output1, ForceFieldTaskDocument)
+    if relax_cell:
+        assert output1.output.energy == approx(-3723.099519623731, rel=1e-3)
+        assert output1.output.n_steps == 3
+    else:
+        assert output1.output.energy == approx(-3723.0981880334643, rel=1e-4)
+        assert output1.output.n_steps == 3
+
+    # fix_symmetry makes no difference for this structure relaxer combo
+    # just testing that passing fix_symmetry doesn't break
+    final_spg_num = output1.output.structure.get_space_group_info()[1]
+    assert final_spg_num == 99

tests/forcefields/test_md.py

@@ -71,6 +71,7 @@ def test_ml_ff_md_maker(
         MLFF.NEP: -3.966232215741286,
         MLFF.Nequip: -8.84670181274414,
         MLFF.SevenNet: -5.394115447998047,
+        MLFF.DeepMD: -744.6197365326168,
     }
 
     # ASE can slightly change tolerances on structure positions
@@ -99,6 +100,14 @@ def test_ml_ff_md_maker(
             "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth"
         }
         unit_cell_structure = sr_ti_o3_structure.copy()
+    elif ff_name == MLFF.DeepMD:
+        calculator_kwargs = {"model": test_dir / "forcefields" / "deepmd" / "graph.pb"}
+        unit_cell_structure = sr_ti_o3_structure.copy()
+
+    elif ff_name == MLFF.MACE:
+        calculator_kwargs = {
+            "model": "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0/2023-12-10-mace-128-L0_epoch-199.model"
+        }
 
     structure = unit_cell_structure.to_conventional() * (2, 2, 2)
 
@@ -140,6 +149,7 @@ def test_ml_ff_md_maker(
         for key in ("energy", "forces", "stress", "velocities", "temperature")
         for step in task_doc.objects["trajectory"].frame_properties
     )
+
     if ff_maker := name_to_maker.get(ff_name):
         with pytest.warns(FutureWarning):
             ff_maker()

tests/forcefields/test_phonon.py

@@ -72,6 +72,7 @@ def test_phonon_maker_initialization_with_all_mlff(
         calc_kwargs = {
             MLFF.Nequip: {"model_path": f"{chk_pt_dir}/nequip/nequip_ff_sr_ti_o3.pth"},
             MLFF.NEP: {"model_filename": f"{test_dir}/forcefields/nep/nep.txt"},
+            MLFF.DeepMD: {"model": test_dir / "forcefields" / "deepmd" / "graph.pb"},
         }.get(mlff, {})
         static_maker = ForceFieldStaticMaker(
             name=f"{mlff} static",

tests/forcefields/test_utils.py

@@ -5,13 +5,17 @@
 
 
 @pytest.mark.parametrize(("force_field"), ["CHGNet", "MACE"])
-def test_ext_load(force_field: str):
+def test_ext_load(force_field: str, test_dir):
     decode_dict = {
         "CHGNet": {"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"},
         "MACE": {"@module": "mace.calculators", "@callable": "mace_mp"},
     }[force_field]
-    calc_from_decode = ase_calculator(decode_dict)
-    calc_from_preset = ase_calculator(str(MLFF(force_field)))
+    kwargs_calc = {
+        "CHGNet": {},
+        "MACE": {"model": test_dir / "forcefields" / "mace" / "MACE.model"},
+    }[force_field]
+    calc_from_decode = ase_calculator(decode_dict, **kwargs_calc)
+    calc_from_preset = ase_calculator(str(MLFF(force_field)), **kwargs_calc)
     assert type(calc_from_decode) is type(calc_from_preset)
     assert calc_from_decode.name == calc_from_preset.name
     assert calc_from_decode.parameters == calc_from_preset.parameters == {}

tests/test_data/forcefields/deepmd/README.md

@@ -0,0 +1,7 @@
+# About this model
+
+The Deep Potential model used for this test is `UniPero`, a universal interatomic potential for perovskite oxides.
+
+It can be downloaded from: https://github.com/sliutheorygroup/UniPero,
+
+For more details, refer to the original article: https://doi.org/10.1103/PhysRevB.108.L180104