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",

docs/reference/deepmd.rst

@@ -0,0 +1,72 @@
+.. _deepmd:
+
+DeePMD-kit
+==========
+
+A `torch-sim <https://github.com/TorchSim/torch-sim>`_ ``ModelInterface`` implementation
+for `DeePMD-kit <https://github.com/deepmodeling/deepmd-kit>`_'s PyTorch backend.
+
+Install
+-------
+
+The package (``DeepmdModel``, torch-sim ``ModelInterface`` wrapper) can be installed
+with either ``pip`` or ``uv``:
+
+.. code-block:: bash
+
+    # from PyPI or a local checkout
+    pip install deepmd-torchsim
+    uv pip install deepmd-torchsim
+    uv add deepmd-torchsim
+
+    # editable, from a local checkout
+    pip install -e .
+    uv pip install -e .
+
+Getting a working ``deepmd-kit`` backend
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: bash
+
+    pip install "deepmd-torchsim[deepmd]"
+    uv pip install "deepmd-torchsim[deepmd]"
+
+The ``deepmd`` extra pins a working ``deepmd-kit==3.1.3``, with ``torch==2.10.0``.
+
+Usage
+-----
+
+.. code-block:: python
+
+    import torch
+    from deepmd_torchsim import DeepmdModel
+    import torch_sim as ts
+    from ase.build import molecule
+
+    model = DeepmdModel(
+        model_path="frozen_model.pth",
+        device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+        compute_forces=True,
+        compute_stress=True,
+    )
+
+    state = ts.io.atoms_to_state([molecule("H2O")], model.device, model.dtype)
+    results = model(state)
+    print(results["energy"])  # [n_systems]
+    print(results["forces"])  # [n_atoms, 3]
+    print(results["stress"])  # [n_systems, 3, 3]
+
+For multitask/multi-domain foundation checkpoints (e.g. DPA-3), pass ``head=`` to
+select which trained domain to evaluate with:
+
+.. code-block:: python
+
+    model = DeepmdModel(model_path="DPA-3.1-3M.pt", head="Omat24")
+
+API
+---
+
+.. autoclass:: torch_sim.models.deepmd.DeepmdModel
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/reference/index.rst

@@ -29,6 +29,14 @@ Overview of the TorchSim API.
     transforms
     units
 
+Model integrations
+-------------------
+
+.. toctree::
+    :titlesonly:
+
+    deepmd
+
 
 TorchSim module treemap. Each node represents a Python module. Arrows indicate
 imports between modules. Node color indicates connectedness: blue nodes have fewer

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>=0.1.1"]
 orb = ["orb-models>=0.6.2"]
 sevenn = ["sevenn[torchsim]>=0.12.1"]
 nequip = ["nequip>=0.17.1"]

tests/models/test_deepmd.py

@@ -0,0 +1,142 @@
+"""Tests for the bundled DeePMD-kit integration (torch_sim.models.deepmd).
+
+Fetches the small CH4 example model from deepmd_torchsim's own GitHub repo
+(``tests/model/frozen_model.pth``) at test time, once per session, rather
+than keeping a duplicate copy of the binary checked into this repo. Computes
+a single-point energy/forces evaluation on the default device (CUDA if
+available, else CPU; CPU is also checked when CUDA is the default) and
+compares against a hardcoded reference. A mismatch does NOT fail the test --
+the model having loaded and produced finite values is the actual pass/fail
+gate -- it warns instead, so numerical drift is visible without breaking CI.
+"""
+
+from __future__ import annotations
+
+import urllib.error
+import urllib.request
+import warnings
+
+import numpy as np
+import pytest
+import torch
+
+import torch_sim as ts
+
+
+try:
+    from torch_sim.models.deepmd import DeepmdModel
+
+    _IMPORT_ERROR: str | None = None
+except ImportError:
+    _IMPORT_ERROR = "deepmd_torchsim not installed"
+
+pytestmark = pytest.mark.skipif(
+    _IMPORT_ERROR is not None, reason=f"deepmd_torchsim not installed: {_IMPORT_ERROR}"
+)
+
+FLOAT64_DTYPE = torch.float64
+DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+_MODEL_URL = (
+    "https://raw.githubusercontent.com/rahulumrao/deepmd_torchsim/main/"
+    "tests/model/frozen_model.pth"
+)
+
+# Computed once on a machine with both CPU and GPU (deepmd_torchsim's own
+# tests/energy_forces_reference.json) -- CPU and GPU agreed to full float64
+# precision there, so one reference is used for both devices here.
+_REFERENCE = {
+    "energy_eV": -1099.4124512539038,
+    "forces_eV_per_A": [
+        [0.0, 0.0, 0.0],
+        [0.17491856887788382, 0.17491856887788393, 0.17491856887788393],
+        [0.17491856887788382, -0.17491856887788393, -0.17491856887788393],
+        [-0.17491856887788393, 0.17491856887788393, -0.17491856887788393],
+        [-0.17491856887788393, -0.17491856887788382, 0.17491856887788393],
+    ],
+}
+
+
+@pytest.fixture(scope="session")
+def model_path(tmp_path_factory: pytest.TempPathFactory) -> str:
+    """Download the example CH4 model once per test session; skip if unreachable."""
+    dest = tmp_path_factory.mktemp("deepmd") / "frozen_model.pth"
+    try:
+        urllib.request.urlretrieve(_MODEL_URL, dest)  # noqa: S310
+    except (urllib.error.URLError, TimeoutError) as exc:
+        pytest.skip(f"could not download example model from {_MODEL_URL}: {exc}")
+    return str(dest)
+
+
+def build_system() -> ts.SimState:
+    """A tetrahedral CH4 molecule, centered in a 10 A cubic box."""
+    from ase import Atoms
+
+    symbols = ["C", "H", "H", "H", "H"]
+    positions = [
+        [0.000000, 0.000000, 0.000000],
+        [0.627581, 0.627581, 0.627581],
+        [0.627581, -0.627581, -0.627581],
+        [-0.627581, 0.627581, -0.627581],
+        [-0.627581, -0.627581, 0.627581],
+    ]
+    box_size = 10.0
+    atoms = Atoms(symbols=symbols, positions=positions, cell=[box_size] * 3, pbc=True)
+    atoms.positions += box_size / 2  # center in the box
+    return atoms
+
+
+def compute(model_path: str, device: torch.device) -> dict:
+    """Load the example CH4 model on ``device`` and return one single-point result."""
+    model = DeepmdModel(
+        model_path=model_path,
+        device=device,
+        dtype=FLOAT64_DTYPE,
+        compute_forces=True,
+        compute_stress=False,
+    )
+    state = ts.io.atoms_to_state([build_system()], device, FLOAT64_DTYPE)
+    output = model.forward(state)
+    return {
+        "energy_eV": output["energy"][0].item(),
+        "forces_eV_per_A": output["forces"].detach().cpu().tolist(),
+    }
+
+
+def compare_to_reference(label: str, result: dict) -> None:
+    """Compare against the hardcoded reference; warn (don't fail) on mismatch."""
+    tolerance = 1e-5
+    energy_diff = abs(result["energy_eV"] - _REFERENCE["energy_eV"])
+    forces_diff = (
+        torch.tensor(result["forces_eV_per_A"])
+        - torch.tensor(_REFERENCE["forces_eV_per_A"])
+    ).abs().max().item()
+
+    if energy_diff >= tolerance or forces_diff >= tolerance:
+        with np.printoptions(precision=4, suppress=True, floatmode="fixed"):
+            warnings.warn(
+                f"{label} energy/forces do not match reference (tol {tolerance:.0e}):\n"
+                f"  energy: computed={result['energy_eV']:.4f} eV, "
+                f"reference={_REFERENCE['energy_eV']:.4f} eV, diff={energy_diff:.4f} eV\n"
+                f"  forces: computed=\n{np.array(result['forces_eV_per_A'])}\n"
+                f"  forces: reference=\n{np.array(_REFERENCE['forces_eV_per_A'])}\n"
+                f"  max abs force diff={forces_diff:.4f} eV/A",
+                stacklevel=2,
+            )
+
+
+def test_deepmd_energy_forces(model_path: str) -> None:
+    """Model loads and produces finite energy/forces on the default device
+    (CUDA if available, else CPU) and the results are compared against a
+    checked-in reference.
+    """
+    default_result = compute(model_path, DEFAULT_DEVICE)
+    assert torch.isfinite(torch.tensor(default_result["energy_eV"]))
+    assert torch.isfinite(torch.tensor(default_result["forces_eV_per_A"])).all()
+    compare_to_reference(DEFAULT_DEVICE.type, default_result)
+
+    if DEFAULT_DEVICE.type == "cuda":
+        cpu_result = compute(model_path, torch.device("cpu"))
+        assert torch.isfinite(torch.tensor(cpu_result["energy_eV"]))
+        assert torch.isfinite(torch.tensor(cpu_result["forces_eV_per_A"])).all()
+        compare_to_reference("cpu", cpu_result)

torch_sim/models/deepmd.py

@@ -0,0 +1,42 @@
+"""TorchSim interface for DeePMD-kit atomistic machine-learning models.
+
+This module exposes :class:`DeepmdModel` as the TorchSim interface to DeePMD-kit models. The model implementation is provided by the ``deepmd_torchsim`` package and is re-exported here to provide a consistent import path within TorchSim.
+
+The integration supports evaluation of DeePMD-kit potential-energy models within TorchSim simulations, including the computation of energies, atomic forces, and virial-derived stresses when requested.
+
+If ``deepmd_torchsim`` is unavailable, importing this module emits a warning and provides a placeholder :class:`DeepmdModel` that raises the original ``ImportError`` upon instantiation.
+
+References:
+    DeePMD-kit:
+        https://github.com/deepmodeling/deepmd-kit
+    DeePMD-kit documentation:
+        https://deepmd-kit.readthedocs.io/
+"""
+
+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"]