From a5d98a3fee77a011f32a2c2a30100f4846189395 Mon Sep 17 00:00:00 2001 From: rahulumrao Date: Tue, 1 Sep 2026 14:54:17 -0700 Subject: [PATCH 1/4] Added DeePMD-kit PyTorch backend integration --- .github/workflows/test.yml | 1 + docs/conf.py | 1 + docs/reference/deepmd.rst | 72 ++++++++++++++++++ docs/reference/index.rst | 8 ++ pyproject.toml | 1 + tests/models/test_deepmd.py | 142 ++++++++++++++++++++++++++++++++++++ torch_sim/models/deepmd.py | 42 +++++++++++ 7 files changed, 267 insertions(+) create mode 100644 docs/reference/deepmd.rst create mode 100644 tests/models/test_deepmd.py create mode 100644 torch_sim/models/deepmd.py 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/docs/reference/deepmd.rst b/docs/reference/deepmd.rst new file mode 100644 index 000000000..91fc9b3e5 --- /dev/null +++ b/docs/reference/deepmd.rst @@ -0,0 +1,72 @@ +.. _deepmd: + +DeePMD-kit +========== + +A `torch-sim `_ ``ModelInterface`` implementation +for `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: diff --git a/docs/reference/index.rst b/docs/reference/index.rst index 9c6d172a0..27524e5da 100644 --- a/docs/reference/index.rst +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 4e49380d9..fc0f08f22 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>=0.1.1"] orb = ["orb-models>=0.6.2"] sevenn = ["sevenn[torchsim]>=0.12.1"] nequip = ["nequip>=0.17.1"] diff --git a/tests/models/test_deepmd.py b/tests/models/test_deepmd.py new file mode 100644 index 000000000..9069132b6 --- /dev/null +++ b/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) diff --git a/torch_sim/models/deepmd.py b/torch_sim/models/deepmd.py new file mode 100644 index 000000000..3c978302d --- /dev/null +++ b/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"] From 822d41f7e0e85ab7e803345f5a4e3f791c0b6c99 Mon Sep 17 00:00:00 2001 From: rahulumrao Date: Wed, 2 Sep 2026 09:07:45 -0700 Subject: [PATCH 2/4] torchsim standard test --- .github/workflows/test.yml | 4 + docs/reference/deepmd.rst | 72 -------------- docs/reference/index.rst | 8 -- tests/models/test_deepmd.py | 190 +++++++++++++++++------------------- torch_sim/models/deepmd.py | 33 ++++--- 5 files changed, 114 insertions(+), 193 deletions(-) delete mode 100644 docs/reference/deepmd.rst diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 901f0e129..1048a3274 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -110,6 +110,10 @@ jobs: if [[ "${{ matrix.model.name }}" == *"fairchem"* ]]; then uv pip install huggingface_hub --system fi + if [[ "${{ matrix.model.name }}" == *"deepmd"* ]]; then + # deepmd-kit backend (exact torch pin, kept out of torch-sim's extras) + uv pip install "deepmd-torchsim[deepmd]" --system + fi pytest -vv -ra -rs --cov=torch_sim --cov-report=xml ${{ matrix.model.test_path }} - name: Upload coverage to Codecov diff --git a/docs/reference/deepmd.rst b/docs/reference/deepmd.rst deleted file mode 100644 index 91fc9b3e5..000000000 --- a/docs/reference/deepmd.rst +++ /dev/null @@ -1,72 +0,0 @@ -.. _deepmd: - -DeePMD-kit -========== - -A `torch-sim `_ ``ModelInterface`` implementation -for `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: diff --git a/docs/reference/index.rst b/docs/reference/index.rst index 27524e5da..9c6d172a0 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -29,14 +29,6 @@ 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 diff --git a/tests/models/test_deepmd.py b/tests/models/test_deepmd.py index 9069132b6..67353ead7 100644 --- a/tests/models/test_deepmd.py +++ b/tests/models/test_deepmd.py @@ -1,142 +1,128 @@ """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. +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 warnings -import numpy as np import pytest import torch -import torch_sim as ts +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 = "deepmd_torchsim not installed" + _IMPORT_ERROR = traceback.format_exc() pytestmark = pytest.mark.skipif( - _IMPORT_ERROR is not None, reason=f"deepmd_torchsim not installed: {_IMPORT_ERROR}" + _IMPORT_ERROR is not None, reason=f"deepmd not installed: {_IMPORT_ERROR}" ) -FLOAT64_DTYPE = torch.float64 -DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +DTYPE = torch.float64 +MAX_RETRIES = 3 +RETRY_DELAY = 30 _MODEL_URL = ( - "https://raw.githubusercontent.com/rahulumrao/deepmd_torchsim/main/" - "tests/model/frozen_model.pth" + "https://store.aissquare.com/models/35b4ce45-4f59-4868-9fd7-a0c0f5ad9464/" + "DPA-3.1-3M.pt" ) - -# 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], - ], -} +_MODEL_HEAD = "Omat24" @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}") + """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) + 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) -def build_system() -> ts.SimState: - """A tetrahedral CH4 molecule, centered in a 10 A cubic box.""" - from ase import Atoms +@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) + - 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 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 -def compute(model_path: str, device: torch.device) -> dict: - """Load the example CH4 model on ``device`` and return one single-point result.""" - model = DeepmdModel( +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=FLOAT64_DTYPE, + device=DEVICE, + dtype=DTYPE, compute_forces=True, compute_stress=False, + head=_MODEL_HEAD, ) - 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) + + +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 index 3c978302d..e1ccf56d8 100644 --- a/torch_sim/models/deepmd.py +++ b/torch_sim/models/deepmd.py @@ -1,16 +1,27 @@ -"""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. +"""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-kit documentation: - https://deepmd-kit.readthedocs.io/ + - DeePMD-kit: https://github.com/deepmodeling/deepmd-kit + - deepmd_torchsim: https://github.com/rahulumrao/deepmd_torchsim """ import traceback From 1b50a1d9574ef9ab6e392074105f10ff162ca8f7 Mon Sep 17 00:00:00 2001 From: rahulumrao Date: Wed, 2 Sep 2026 19:30:28 -0700 Subject: [PATCH 3/4] deepmd extra to a working backend --- .github/workflows/test.yml | 4 ---- pyproject.toml | 6 +++++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1048a3274..901f0e129 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -110,10 +110,6 @@ jobs: if [[ "${{ matrix.model.name }}" == *"fairchem"* ]]; then uv pip install huggingface_hub --system fi - if [[ "${{ matrix.model.name }}" == *"deepmd"* ]]; then - # deepmd-kit backend (exact torch pin, kept out of torch-sim's extras) - uv pip install "deepmd-torchsim[deepmd]" --system - fi pytest -vv -ra -rs --cov=torch_sim --cov-report=xml ${{ matrix.model.test_path }} - name: Upload coverage to Codecov diff --git a/pyproject.toml b/pyproject.toml index fc0f08f22..b0bfbd507 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +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"] +deepmd = ["deepmd_torchsim[deepmd]>=0.1.1"] orb = ["orb-models>=0.6.2"] sevenn = ["sevenn[torchsim]>=0.12.1"] nequip = ["nequip>=0.17.1"] @@ -178,6 +178,10 @@ conflicts = [ { extra = "mace" }, { extra = "sevenn" }, ], + [ + { extra = "deepmd" }, + { extra = "fairchem" }, + ], ] [dependency-groups] From 4d42bcd81ca224b8e3ad3ca1b00820a21f30e0b9 Mon Sep 17 00:00:00 2001 From: rahulumrao Date: Sun, 20 Sep 2026 11:15:28 -0700 Subject: [PATCH 4/4] deepmd extra to a working backend --- tests/models/test_deepmd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/test_deepmd.py b/tests/models/test_deepmd.py index 67353ead7..8f67a0e7c 100644 --- a/tests/models/test_deepmd.py +++ b/tests/models/test_deepmd.py @@ -52,7 +52,7 @@ def model_path(tmp_path_factory: pytest.TempPathFactory) -> str: dest = tmp_path_factory.mktemp("deepmd") / "DPA-3.1-3M.pt" for attempt in range(MAX_RETRIES): try: - urllib.request.urlretrieve(_MODEL_URL, dest) + 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}")