Added DeePMD-kit PyTorch backend integration

.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" }

docs/conf.py

@@ -62,6 +62,7 @@
 ]
 
 autodoc_mock_imports = [
+    "deepmd",
     "fairchem",
     "mace",
     "mattersim",

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]

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,
+)

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 <https://github.com/rahulumrao/deepmd_torchsim>`_ and
+`PyPI <https://pypi.org/project/deepmd-torchsim/>`_.
+
+`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"]