Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat: Add DeepMD MLFF support #999

Open
wants to merge 18 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ default_language_version:
exclude: ^(.github/|tests/test_data/abinit/)
repos:
- repo: https://github.com/charliermarsh/ruff-pre-commit
rev: v0.7.3
rev: v0.7.4
hooks:
- id: ruff
args: [--fix]
Expand Down
3 changes: 3 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down
1 change: 1 addition & 0 deletions src/atomate2/forcefields/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
1 change: 1 addition & 0 deletions src/atomate2/forcefields/jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"},
}


Expand Down
1 change: 1 addition & 0 deletions src/atomate2/forcefields/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 5 additions & 0 deletions src/atomate2/forcefields/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
24 changes: 24 additions & 0 deletions tests/forcefields/conftest.py
Original file line number Diff line number Diff line change
@@ -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


Expand All @@ -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}")
63 changes: 63 additions & 0 deletions tests/forcefields/test_jobs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
10 changes: 10 additions & 0 deletions tests/forcefields/test_md.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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()
Expand Down
5 changes: 3 additions & 2 deletions tests/forcefields/test_phonon.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -43,7 +43,7 @@ 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
Expand Down Expand Up @@ -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",
Expand Down
10 changes: 7 additions & 3 deletions tests/forcefields/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 == {}
Expand Down
7 changes: 7 additions & 0 deletions tests/test_data/forcefields/deepmd/README.md
Original file line number Diff line number Diff line change
@@ -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
Loading