add dpa3 calculator fixes #411 (#428)

* add dpa3 calculator fixes #411

* Update docs for DPA3

* Separate missing and invalid DPA3 model

* Add test for invalid DPA3 model

* Add reference for example model

---------

Co-authored-by: ElliottKasoar <[email protected]>

Author: alinelena

README.md

@@ -35,6 +35,7 @@ Tools for machine learnt interatomic potentials
 - sevenn = 0.10.3 (optional)
 - alignn = 2024.5.27 (optional)
 - nequip = 0.6.1 (optional)
+- deepmd-kit = dpa3-alpha (optional)
 
 All required and optional dependencies can be found in [pyproject.toml](pyproject.toml).
 
@@ -94,6 +95,7 @@ Current and planned features include:
   - ALIGNN (experimental)
   - SevenNet (experimental)
   - NequIP (experimental)
+  - DPA3 (experimental)
 - [x] Single point calculations
 - [x] Geometry optimisation
 - [x] Molecular Dynamics

docs/source/getting_started/getting_started.rst

@@ -15,6 +15,7 @@ Dependencies
 - sevenn = 0.10.3 (optional)
 - alignn = 2024.5.27 (optional)
 - nequip = 0.6.1 (optional)
+- deepmd-kit = dpa3-alpha (optional)
 
 All required and optional dependencies can be found in `pyproject.toml <https://github.com/stfc/janus-core/blob/main/pyproject.toml>`_.
 
@@ -58,5 +59,6 @@ Currently supported extras are:
 - ``m3gnet``: `M3GNet <https://github.com/materialsvirtuallab/matgl/>`_
 - ``sevenn``: `SevenNet <https://github.com/MDIL-SNU/SevenNet/>`_
 - ``nequip``: `NequIP <https://github.com/mir-group/nequip>`_
+- ``dpa3``: `DPA3 <https://github.com/deepmodeling/deepmd-kit/tree/dpa3-alpha>`_
 
 ``extras`` are also listed in `pyproject.toml <https://github.com/stfc/janus-core/blob/main/pyproject.toml>`_ under ``[project.optional-dependencies]``.

janus_core/helpers/janus_types.py

@@ -115,7 +115,15 @@ class CorrelationKwargs(TypedDict, total=True):
 
 # Janus specific
 Architectures = Literal[
-    "mace", "mace_mp", "mace_off", "m3gnet", "chgnet", "alignn", "sevennet", "nequip"
+    "mace",
+    "mace_mp",
+    "mace_off",
+    "m3gnet",
+    "chgnet",
+    "alignn",
+    "sevennet",
+    "nequip",
+    "dpa3",
 ]
 Devices = Literal["cpu", "cuda", "mps", "xpu"]
 Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]

janus_core/helpers/mlip_calculators.py

@@ -240,6 +240,23 @@ def choose_calculator(
             model_path=model_path, device=device, **kwargs
         )
 
+    elif arch == "dpa3":
+        from deepmd import __version__
+        from deepmd.calculator import DP
+
+        # No default `model_path`
+        if model_path is None:
+            # From https://matbench-discovery.materialsproject.org/models/dpa3-v1-mptrj
+            raise ValueError(
+                "Please specify `model_path`, as there is no "
+                f"default model for {arch} "
+                "e.g. https://bohrium-api.dp.tech/ds-dl/dpa3openlam-74ng-v3.zip"
+            )
+
+        model_path = str(model_path)
+
+        calculator = DP(model=model_path, **kwargs)
+
     else:
         raise ValueError(
             f"Unrecognized {arch=}. Suported architectures "

pyproject.toml

@@ -48,22 +48,26 @@ alignn = [
 chgnet = [
     "chgnet == 0.3.8",
 ]
+dpa3 = [
+    "deepmd-kit",
+]
 m3gnet = [
     "matgl == 1.1.3",
     "dgl == 2.1.0",
 ]
-sevennet = [
-    "sevenn == 0.10.3",
-]
 nequip = [
     "nequip == 0.6.1",
 ]
+sevennet = [
+    "sevenn == 0.10.3",
+]
 all = [
     "janus-core[alignn]",
     "janus-core[chgnet]",
+    "janus-core[dpa3]",
     "janus-core[m3gnet]",
-    "janus-core[sevennet]",
     "janus-core[nequip]",
+    "janus-core[sevennet]",
 ]
 
 [project.scripts]
@@ -179,3 +183,6 @@ default-groups = [
     "docs",
     "pre-commit",
 ]
+
+[tool.uv.sources]
+deepmd-kit = { git = "https://github.com/deepmodeling/deepmd-kit.git", rev = "dpa3-alpha" }

tests/models/2025-01-10-dpa3-mptrj.pth

@@ -30,6 +30,8 @@
 
 NEQUIP_PATH = MODEL_PATH / "toluene.pth"
 
+DPA3_PATH = MODEL_PATH / "2025-01-10-dpa3-mptrj.pth"
+
 
 @pytest.mark.parametrize(
     "arch, device, kwargs",
@@ -80,6 +82,7 @@ def test_invalid_arch():
         ("m3gnet", "/invalid/path"),
         ("chgnet", "/invalid/path"),
         ("nequip", "/invalid/path"),
+        ("dpa3", "/invalid/path"),
     ],
 )
 def test_invalid_model_path(arch, model_path):
@@ -126,6 +129,8 @@ def test_invalid_device(arch):
         ("sevennet", "cpu", {"model": "sevennet-0"}),
         ("nequip", "cpu", {"model_path": NEQUIP_PATH}),
         ("nequip", "cpu", {"model": NEQUIP_PATH}),
+        ("dpa3", "cpu", {"model_path": DPA3_PATH}),
+        ("dpa3", "cpu", {"model": DPA3_PATH}),
     ],
 )
 def test_extra_mlips(arch, device, kwargs):
@@ -176,6 +181,16 @@ def test_extra_mlips(arch, device, kwargs):
             "model_path": NEQUIP_PATH,
             "path": NEQUIP_PATH,
         },
+        {
+            "arch": "dpa3",
+            "model_path": DPA3_PATH,
+            "model": DPA3_PATH,
+        },
+        {
+            "arch": "dpa3",
+            "model_path": DPA3_PATH,
+            "path": DPA3_PATH,
+        },
     ],
 )
 def test_extra_mlips_invalid(kwargs):

tests/test_mlip_calculators.py

@@ -18,6 +18,7 @@
 SEVENNET_PATH = MODEL_PATH / "sevennet_0.pth"
 ALIGNN_PATH = MODEL_PATH / "v5.27.2024"
 NEQUIP_PATH = MODEL_PATH / "toluene.pth"
+DPA3_PATH = MODEL_PATH / "2025-01-10-dpa3-mptrj.pth"
 
 test_data = [
     (DATA_PATH / "benzene.xyz", -76.0605725422795, "energy", "energy", {}, None),
@@ -293,6 +294,7 @@ def test_mlips(arch, device, expected_energy):
         {"model_path": "SevenNet-0_11July2024"},
     ),
     ("nequip", "cpu", -169815.1282456301, "toluene.xyz", {"model_path": NEQUIP_PATH}),
+    ("dpa3", "cpu", -27.053507387638092, "NaCl.cif", {"model_path": DPA3_PATH}),
 ]