Feat: Add DeepMD MLFF support (#999)

Author: MoseyQAQ

.gitignore

@@ -78,3 +78,5 @@ docs/reference/atomate2.*
 
 .ipynb_checkpoints
 .aider*
+
+tests/test_data/forcefields/deepmd_graph.pb

docs/user/codes/forcefields.md

@@ -0,0 +1,22 @@
+(codes.forcefields)=
+
+# Machine Learning forcefields / interatomic potentials
+
+`atomate2` includes an interface to a few common machine learning interatomic potentials (MLIPs), also known variously as machine learning forcefields (MLFFs), or foundation potentials (FPs) for universal variants.
+Support is provided for the following models, which can be selected using `atomate2.forcefields.utils.MLFF`, as shown in the table below.
+**You need only install packages for the forcefields you wish to use.**
+
+| Forcefield Name | `MLFF` | Reference | Description |
+| ---- | ---- | ---- | ---- |
+| CHGNet | `CHGNet` | [10.1038/s42256-023-00716-3](https://doi.org/10.1038/s42256-023-00716-3) | Available via the `chgnet` and `matgl` packages |
+| DeepMD | `MLFF.DeepMD` | [10.1103/PhysRevB.108.L180104](https://doi.org/10.1103/PhysRevB.108.L180104) | The Deep Potential model used for this test is `UniPero`, a universal interatomic potential for perovskite oxides. It can be downloaded [here](https://github.com/sliutheorygroup/UniPero) |
+| Gaussian Approximation Potential (GAP) | `GAP` | [10.1103/PhysRevLett.104.136403](https://doi.org/10.1103/PhysRevLett.104.136403) |  Relies on `quippy-ase` package |
+| M3GNet | `M3GNet` | [10.1038/s43588-022-00349-3](https://doi.org/10.1038/s43588-022-00349-3) | Relies on `matgl` package |
+| MACE-MP-0 | `MACE` or `MACE_MP_0` (recommended) | [10.1063/5.0297006](https://doi.org/10.1063/5.0297006) | Relies on `mace_torch` and optionally `torch_dftd` packages |
+| MACE-MP-0b3 | `MACE_MP_0B3` | [10.1063/5.0297006](https://doi.org/10.1063/5.0297006) | Relies on `mace_torch` and optionally `torch_dftd` packages |
+| MACE-MPA-0 | `MACE_MPA_0` | [10.1063/5.0297006](https://doi.org/10.1063/5.0297006) | Relies on `mace_torch` and optionally `torch_dftd` packages |
+| MatPES-PBE | `MATPES_PBE` | [10.48550/arXiv.2503.04070](https://doi.org/10.48550/arXiv.2503.04070) | Relies on `matgl`. Defaults to TensorNet architecture, but can also use M3GNet or CHGNet architectures via kwargs. See `atomate2.forcefields.utils._DEFAULT_CALCULATOR_KWARGS` for more options. |
+| MatPES-r<sup>2</sup>SCAN | `MATPES_R2SCAN`| [10.48550/arXiv.2503.04070](https://doi.org/10.48550/arXiv.2503.04070) | Relies on `matgl`. Defaults to TensorNet architecture, but can also use M3GNet or CHGNet architectures via kwargs. See `atomate2.forcefields.utils._DEFAULT_CALCULATOR_KWARGS` for more options. |
+| Neuroevolution Potential (NEP) | `NEP` | [10.1103/PhysRevB.104.104309](https://doi.org/10.1103/PhysRevB.104.104309) | Relies on `calorine` package |
+| Neural Equivariant Interatomic Potentials (Nequip) | `Nequip` | [10.1038/s41467-022-29939-5](https://doi.org/10.1038/s41467-022-29939-5) | Relies on the `nequip` package |
+| SevenNet | `SevenNet` | [10.1021/acs.jctc.4c00190](https://doi.org/10.1021/acs.jctc.4c00190) | Relies on the `sevenn` package |

docs/user/codes/index.md

@@ -7,4 +7,5 @@ The section gives the instructions for codes supported by atomate2.
 ```{toctree}
 vasp
 openmm
+forcefields
 ```

pyproject.toml

@@ -61,6 +61,7 @@ forcefields = [
     # quippy-ase support for py3.12 tracked in https://github.com/libAtoms/QUIP/issues/645
     "quippy-ase>=0.9.14; python_version < '3.12'",
     "sevenn>=0.9.3",
+    "deepmd-kit>=2.1.4",
 ]
 approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"]
 ase = ["ase>=3.26.0"]
@@ -111,6 +112,8 @@ strict-forcefields = [
     "sevenn==0.10.4",
     "torch==2.2.0",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
+    "deepmd-kit==2.2.11",
+    "tensorflow-cpu==2.16.2",
 ]
 strict = [
     "atomate2[strict-forcefields, docs, cclib, phonons, lobster, openmm, mp, defects, ase, ase-ext]",

src/atomate2/forcefields/schemas.py

@@ -119,6 +119,7 @@ def from_ase_compatible_result(
             MLFF.MACE_MP_0B3: "mace-torch",
             MLFF.GAP: "quippy-ase",
             MLFF.Nequip: "nequip",
+            MLFF.DeepMD: "deepmd-kit",
             MLFF.MATPES_PBE: "matgl",
             MLFF.MATPES_R2SCAN: "matgl",
         }

src/atomate2/forcefields/utils.py

@@ -40,6 +40,7 @@ class MLFF(Enum):  # TODO inherit from StrEnum when 3.11+
     SevenNet = "SevenNet"
     MATPES_R2SCAN = "MatPES-r2SCAN"
     MATPES_PBE = "MatPES-PBE"
+    DeepMD = "DeepMD"
 
     @classmethod
     def _missing_(cls, value: Any) -> Any:
@@ -273,6 +274,11 @@ def ase_calculator(
 
             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,15 @@
 from __future__ import annotations
 
+import hashlib
+import urllib.request
 from typing import TYPE_CHECKING
 
+import pytest
 import torch
+from emmet.core.utils import get_hash_blocked
 
 if TYPE_CHECKING:
+    from pathlib import Path
     from typing import Any
 
 
@@ -13,3 +18,18 @@ 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
+    if (file_md5 := get_hash_blocked(local_path, hasher=hashlib.md5())) != ref_md5:
+        raise RuntimeError(f"MD5 mismatch: {file_md5} != {ref_md5}")

tests/forcefields/test_jobs.py

@@ -522,6 +522,69 @@ def test_nequip_relax_maker(
     assert final_spg_num == 99
 
 
+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
+
+
 @pytest.mark.parametrize("ref_func", ["PBE", "r2SCAN"])
 def test_matpes_relax_makers(
     sr_ti_o3_structure: Structure,

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,
         MLFF.MATPES_PBE: -5.230762481689453,
         MLFF.MATPES_R2SCAN: -8.561729431152344,
     }
@@ -101,6 +102,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 == MLFF.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)
 

tests/forcefields/test_phonon.py

@@ -28,7 +28,7 @@ def test_supercell_orthorhombic(clean_dir, si_structure: Structure):
         min_length=5,
         max_length=10,
         prefer_90_degrees=False,
-        allow_orhtorhombic=True,
+        allow_orthorhombic=True,
     )
 
     # run the flow or job and ensure that it finished running successfully
@@ -43,14 +43,14 @@ def test_supercell_orthorhombic(clean_dir, si_structure: Structure):
         min_length=5,
         max_length=10,
         prefer_90_degrees=True,
-        allow_orhtorhombic=True,
+        allow_orthorhombic=True,
     )
 
     # run the flow or job and ensure that it finished running successfully
     responses = run_locally(job2, create_folders=True, ensure_success=True)
 
     assert_allclose(
-        responses[job2.uuid][1].output, [[2, -1, 0], [0, 3, 0], [-1, -1, 2]]
+        responses[job2.uuid][1].output, [[2, -1, 0], [0, 2, 0], [-1, -1, 2]]
     )
 
 
@@ -74,6 +74,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",