diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 3ff79b68c..901f0e129 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -60,6 +60,7 @@ jobs: - { python: '3.12', resolution: lowest-direct } - { python: '3.14', resolution: highest } model: + - { name: deepmd, test_path: "tests/models/test_deepmd.py" } - { name: fairchem, test_path: "tests/models/test_fairchem.py" } - { name: mace, test_path: "tests/models/test_mace.py" } - { name: mace, test_path: "tests/test_elastic.py" } diff --git a/docs/conf.py b/docs/conf.py index b162d0c8c..05243c402 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -62,6 +62,7 @@ ] autodoc_mock_imports = [ + "deepmd", "fairchem", "mace", "mattersim", diff --git a/pyproject.toml b/pyproject.toml index 4e49380d9..b0bfbd507 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ mace = ["mace-torch>=0.3.16"] # phono3py>=4.0.0: some older 3.x sdists (e.g. 3.2.0) fail to build and break uv resolution. mattersim = ["mattersim>=1.2.5", "phono3py>=4.0.0"] metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0", "upet>=0.2.0"] +deepmd = ["deepmd_torchsim[deepmd]>=0.1.1"] orb = ["orb-models>=0.6.2"] sevenn = ["sevenn[torchsim]>=0.12.1"] nequip = ["nequip>=0.17.1"] @@ -177,6 +178,10 @@ conflicts = [ { extra = "mace" }, { extra = "sevenn" }, ], + [ + { extra = "deepmd" }, + { extra = "fairchem" }, + ], ] [dependency-groups] diff --git a/tests/models/test_deepmd.py b/tests/models/test_deepmd.py new file mode 100644 index 000000000..8f67a0e7c --- /dev/null +++ b/tests/models/test_deepmd.py @@ -0,0 +1,128 @@ +"""Tests for the bundled DeePMD-kit integration (torch_sim.models.deepmd). + +Uses the DPA-3.1-3M universal foundation checkpoint (full periodic-table +type_map, "Omat24" head). +""" + +from __future__ import annotations + +import time +import traceback +import urllib.error +import urllib.request + +import pytest +import torch + +from tests.conftest import DEVICE +from tests.models.conftest import ( + make_model_calculator_consistency_test, + make_validate_model_outputs_test, +) +from torch_sim.testing import SIMSTATE_BULK_GENERATORS, SIMSTATE_MOLECULE_GENERATORS + + +try: + from deepmd.calculator import DP + + from torch_sim.models.deepmd import DeepmdModel + + _IMPORT_ERROR: str | None = None +except ImportError: + _IMPORT_ERROR = traceback.format_exc() + +pytestmark = pytest.mark.skipif( + _IMPORT_ERROR is not None, reason=f"deepmd not installed: {_IMPORT_ERROR}" +) + +DTYPE = torch.float64 +MAX_RETRIES = 3 +RETRY_DELAY = 30 + +_MODEL_URL = ( + "https://store.aissquare.com/models/35b4ce45-4f59-4868-9fd7-a0c0f5ad9464/" + "DPA-3.1-3M.pt" +) +_MODEL_HEAD = "Omat24" + + +@pytest.fixture(scope="session") +def model_path(tmp_path_factory: pytest.TempPathFactory) -> str: + """Download the DPA-3.1-3M checkpoint once per session, with retries.""" + dest = tmp_path_factory.mktemp("deepmd") / "DPA-3.1-3M.pt" + for attempt in range(MAX_RETRIES): + try: + urllib.request.urlretrieve(_MODEL_URL, dest) # noqa: S310 + except (urllib.error.URLError, TimeoutError) as exc: + if attempt == MAX_RETRIES - 1: + pytest.skip(f"could not download DPA-3.1-3M from {_MODEL_URL}: {exc}") + time.sleep(RETRY_DELAY * (attempt + 1)) + else: + break + return str(dest) + + +@pytest.fixture +def deepmd_model(model_path: str) -> DeepmdModel: + return DeepmdModel( + model_path=model_path, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=True, + head=_MODEL_HEAD, + ) + + +@pytest.fixture +def deepmd_calculator(model_path: str) -> DP: + return DP(model=model_path, head=_MODEL_HEAD) + + +def test_deepmd_initialization(deepmd_model: DeepmdModel) -> None: + assert deepmd_model.device == DEVICE + assert deepmd_model.dtype == DTYPE + assert deepmd_model.compute_forces is True + assert deepmd_model.compute_stress is True + assert "Cu" in deepmd_model.type_map # universal periodic-table type_map + + +test_deepmd_consistency = make_model_calculator_consistency_test( + test_name="deepmd", + model_fixture_name="deepmd_model", + calculator_fixture_name="deepmd_calculator", + sim_state_names=tuple(SIMSTATE_BULK_GENERATORS.keys()), + device=DEVICE, + dtype=DTYPE, +) + + +@pytest.fixture +def deepmd_molecule_model(model_path: str) -> DeepmdModel: + """Stress disabled (mirroring the mace_off molecule test): ASE's DP + calculator raises PropertyNotImplementedError for stress on non-periodic + systems, so molecules are checked on energy/forces only.""" + return DeepmdModel( + model_path=model_path, + device=DEVICE, + dtype=DTYPE, + compute_forces=True, + compute_stress=False, + head=_MODEL_HEAD, + ) + + +test_deepmd_molecule_consistency = make_model_calculator_consistency_test( + test_name="deepmd_molecule", + model_fixture_name="deepmd_molecule_model", + calculator_fixture_name="deepmd_calculator", + sim_state_names=tuple(SIMSTATE_MOLECULE_GENERATORS.keys()), + device=DEVICE, + dtype=DTYPE, +) + +test_deepmd_model_outputs = make_validate_model_outputs_test( + model_fixture_name="deepmd_model", + device=DEVICE, + dtype=DTYPE, +) diff --git a/torch_sim/models/deepmd.py b/torch_sim/models/deepmd.py new file mode 100644 index 000000000..e1ccf56d8 --- /dev/null +++ b/torch_sim/models/deepmd.py @@ -0,0 +1,53 @@ +"""Wrapper for DeePMD-kit models in TorchSim. + +This module provides :class:`DeepmdModel`, the TorchSim +`ModelInterface` implementation for the PyTorch backend of DeePMD-kit. +The underlying implementation is maintained in the standalone +`deepmd_torchsim` package, available from +`GitHub `_ and +`PyPI `_. + +`DeepmdModel` evaluates DeePMD-kit interatomic potential models and +provides energies, atomic forces, and stress tensors derived from the +virial. It supports custom-trained `se_e2_a` models as well as +multitask and multidomain foundation-model checkpoints, including DPA-3 +models through the `head=` argument. See the `deepmd_torchsim` +documentation for usage examples, installation instructions, and +requirements for a compatible `deepmd-kit` backend. + +If `deepmd_torchsim` is not installed, this module will throw a warning and +provides a placeholder :class:`DeepmdModel` that raises the underlying +`ImportError` when instantiated. + +References: + - DeePMD-kit: https://github.com/deepmodeling/deepmd-kit + - deepmd_torchsim: https://github.com/rahulumrao/deepmd_torchsim +""" + +import traceback +import warnings +from typing import Any + + +try: + from deepmd_torchsim import DeepmdModel +except ImportError as exc: + warnings.warn( + f"deepmd_torchsim import failed: {traceback.format_exc()}", stacklevel=2 + ) + + from torch_sim.models.interface import ModelInterface + + class DeepmdModel(ModelInterface): + """Placeholder when deepmd_torchsim is not installed.""" + + def __init__(self, err: ImportError = exc, *_args: Any, **_kwargs: Any) -> None: + """Raise the original ImportError.""" + raise err + + def forward(self, *_args: Any, **_kwargs: Any) -> Any: + """Unreachable — __init__ always raises.""" + raise NotImplementedError + + +__all__ = ["DeepmdModel"]