Feat: Add DeepMD MLFF support

pyproject.toml

@@ -59,6 +59,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
@@ -125,6 +126,8 @@ strict-forcefields = [
     "sevenn==0.9.3.post1",
     "torch==2.2.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

@@ -83,6 +83,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().decode(json.dumps(calculator_meta))
         calculator = calc_cls(**kwargs)

tests/forcefields/conftest.py

@@ -2,9 +2,12 @@
 
 from typing import TYPE_CHECKING
 
+import pytest
+import requests
 import torch
 
 if TYPE_CHECKING:
+    from pathlib import Path
     from typing import Any
 
 
@@ -13,3 +16,16 @@ 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"
+    response = requests.get(file_url, timeout=10)
+    if response.status_code == 200:
+        with open(local_path, "wb") as file:
+            file.write(response.content)
+    else:
+        raise requests.RequestException(f"Failed to download: {file_url}")

tests/forcefields/test_jobs.py

@@ -535,3 +535,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

@@ -49,7 +49,7 @@ def test_maker_initialization():
 
 @pytest.mark.parametrize(
     "ff_name",
-    ["CHGNet", "M3GNet", "MACE", "GAP", "NEP", "Nequip"],
+    ["CHGNet", "M3GNet", "MACE", "GAP", "NEP", "Nequip", "DeepMD"],
 )
 def test_ml_ff_md_maker(
     ff_name, si_structure, sr_ti_o3_structure, al2_au_structure, test_dir, clean_dir
@@ -68,6 +68,7 @@ def test_ml_ff_md_maker(
         "GAP": -5.391255755606209,
         "NEP": -3.966232215741286,
         "Nequip": -8.84670181274414,
+        "DeepMD": -744.6197365326168,
     }
 
     # ASE can slightly change tolerances on structure positions
@@ -96,6 +97,9 @@ 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 == "DeepMD":
+        calculator_kwargs = {"model": test_dir / "forcefields" / "deepmd" / "graph.pb"}
+        unit_cell_structure = sr_ti_o3_structure.copy()
 
     structure = unit_cell_structure.to_conventional() * (2, 2, 2)
 
@@ -138,8 +142,10 @@ def test_ml_ff_md_maker(
         for step in task_doc.objects["trajectory"].frame_properties
     )
 
-    with pytest.warns(FutureWarning):
-        name_to_maker[ff_name]()
+    # Skip the following test for DeepMD, since it doesn't have concrete implementations
+    if ff_name != "DeepMD":
+        with pytest.warns(FutureWarning):
+            name_to_maker[ff_name]()
 
 
 @pytest.mark.parametrize("traj_file", ["trajectory.json.gz", "atoms.traj"])

tests/forcefields/test_phonon.py

@@ -67,6 +67,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/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