diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b9bb05469e..ebb0946e50 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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] diff --git a/pyproject.toml b/pyproject.toml index d2384e1689..17b256be0c 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/atomate2/forcefields/__init__.py b/src/atomate2/forcefields/__init__.py index e16c088654..9e578569d9 100644 --- a/src/atomate2/forcefields/__init__.py +++ b/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: diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index f2307e3a50..9be2874cf5 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/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"}, } diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index a66785b8b8..a986f26de5 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/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( diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 043b7b037a..28516e1f33 100644 --- a/src/atomate2/forcefields/utils.py +++ b/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) diff --git a/tests/forcefields/conftest.py b/tests/forcefields/conftest.py index 8d011ff423..2599dac6a8 100644 --- a/tests/forcefields/conftest.py +++ b/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}") diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 0dbb765311..ca7e3a492f 100644 --- a/tests/forcefields/test_jobs.py +++ b/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 diff --git a/tests/forcefields/test_md.py b/tests/forcefields/test_md.py index 3686b54a5e..683d6338ef 100644 --- a/tests/forcefields/test_md.py +++ b/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() diff --git a/tests/forcefields/test_phonon.py b/tests/forcefields/test_phonon.py index 06de77b866..8092a7b880 100644 --- a/tests/forcefields/test_phonon.py +++ b/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,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 @@ -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", diff --git a/tests/forcefields/test_utils.py b/tests/forcefields/test_utils.py index 023482ea2d..5999bdc4b1 100644 --- a/tests/forcefields/test_utils.py +++ b/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 == {} diff --git a/tests/test_data/forcefields/deepmd/README.md b/tests/test_data/forcefields/deepmd/README.md new file mode 100644 index 0000000000..f6363adc1d --- /dev/null +++ b/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