Add PET-MAD (#575)

* Add PET-MAD dependency

* Add PET-MAD architecture

* Test PET-MAD architecture

* Add PET-MAD to docs

* Save PET-MAD model version

* Remove Windows support for PET-MAD

* Update docs for PET-MAD

Author: ElliottKasoar

README.md

@@ -92,6 +92,7 @@ Current and planned features include:
   - GRACE
   - EquiformerV2
   - eSEN
+  - PET-MAD
 - [x] Single point calculations
 - [x] Geometry optimisation
 - [x] Molecular Dynamics

docs/source/user_guide/get_started.rst

@@ -60,10 +60,11 @@ Currently supported MLIP ``extras`` are:
 - ``mattersim``: `MatterSim <https://github.com/microsoft/mattersim>`_
 - ``grace``: `GRACE <https://github.com/ICAMS/grace-tensorpotential>`_
 - ``fairchem``: `eqV2 DeNS/eSEN <https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core>`_
+- ``pet-mad``: `PET-MAD <https://github.com/lab-cosmo/pet-mad>`_
 
 .. note::
 
-    ``orb`` and ``mattersim`` are not currently compatible with Windows natively,
+    ``orb``, ``mattersim``, and ``pet-mad`` are not currently compatible with Windows natively,
     but can be installed and run via Windows Subsystem for Linux.
 
 

docs/source/user_guide/installation.rst

@@ -158,5 +158,5 @@ MLIPs with limited OS support
 -----------------------------
 
 Several MLIP packages have limited support on Windows. We are currently unable to
-support ``orb``, ``mattersim``, ``alignn`` or ``matgl`` as ``extras`` on Windows, so they
+support ``orb``, ``mattersim``, ``pet-mad``, ``alignn`` or ``matgl`` as ``extras`` on Windows, so they
 must be installed manually.

janus_core/helpers/janus_types.py

@@ -131,6 +131,7 @@ class Correlation(TypedDict, total=True):
     "grace",
     "esen",
     "equiformer",
+    "pet_mad",
 ]
 Devices = Literal["cpu", "cuda", "mps", "xpu"]
 Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]

janus_core/helpers/mlip_calculators.py

@@ -370,6 +370,18 @@ def choose_calculator(
                 **kwargs,
             )
 
+        case "pet_mad":
+            from pet_mad import __version__
+            from pet_mad._version import LATEST_VERSION
+            from pet_mad.calculator import PETMADCalculator
+
+            calculator = PETMADCalculator(
+                checkpoint_path=model, device=device, **kwargs
+            )
+
+            if model is None:
+                model = LATEST_VERSION
+
         case _:
             raise ValueError(
                 f"Unrecognized {arch=}. Suported architectures "

pyproject.toml

@@ -59,6 +59,9 @@ nequip = [
 orb = [
     "orb-models == 0.5.4; sys_platform != 'win32'",
 ]
+pet-mad = [
+    "pet-mad == 1.3.1; sys_platform != 'win32'"
+]
 plumed = [
     "plumed<3.0.0,>=2.9.0; sys_platform != 'win32'",
 ]
@@ -75,6 +78,7 @@ all = [
     "janus-core[mace]",
     "janus-core[nequip]",
     "janus-core[orb]",
+    "janus-core[pet-mad]",
     "janus-core[plumed]",
     "janus-core[sevennet]",
     "janus-core[visualise]",

tests/models/pet-mad-1.1.0.pt

@@ -70,6 +70,10 @@
 except ImportError:
     M3GNET_POTENTIAL = None
 
+PET_MAD_CHECKPOINT = (
+    "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt"
+)
+
 
 @pytest.mark.parametrize(
     "arch, device, kwargs",
@@ -109,6 +113,9 @@
         ("nequip", "cpu", {"model": NEQUIP_PATH}),
         ("orb", "cpu", {}),
         ("orb", "cpu", {"model": ORB_MODEL}),
+        ("pet_mad", "cpu", {}),
+        ("pet_mad", "cpu", {"model": PET_MAD_CHECKPOINT}),
+        ("pet_mad", "cpu", {"checkpoint_path": PET_MAD_CHECKPOINT}),
         ("sevennet", "cpu", {"model": SEVENNET_PATH}),
         ("sevennet", "cpu", {"path": SEVENNET_PATH}),
         ("sevennet", "cpu", {"model_path": SEVENNET_PATH}),
@@ -167,6 +174,7 @@ def test_invalid_arch():
         ("mattersim", "/invalid/path"),
         ("nequip", "/invalid/path"),
         ("orb", "/invalid/path"),
+        ("pet_mad", "/invalid/path"),
         ("sevenn", "/invalid/path"),
         ("alignn", "invalid/path"),
         ("m3gnet", "/invalid/path"),
@@ -228,6 +236,11 @@ def test_invalid_model(arch, model):
         {"arch": "nequip", "model": NEQUIP_PATH, "path": NEQUIP_PATH},
         {"arch": "orb", "model_path": ORB_MODEL, "model": ORB_MODEL},
         {"arch": "orb", "model": ORB_MODEL, "path": ORB_MODEL},
+        {
+            "arch": "pet_mad",
+            "model": PET_MAD_CHECKPOINT,
+            "checkpoint_path": PET_MAD_CHECKPOINT,
+        },
         {"arch": "sevennet", "model_path": SEVENNET_PATH, "model": SEVENNET_PATH},
         {"arch": "sevennet", "model": SEVENNET_PATH, "path": SEVENNET_PATH},
     ],

tests/test_mlip_calculators.py

@@ -23,9 +23,11 @@
 ESEN_LABEL = "eSEN-30M-MP"
 MACE_PATH = MODEL_PATH / "mace_mp_small.model"
 NEQUIP_PATH = MODEL_PATH / "toluene.pth"
+PET_MAD_CHECKPOINT = (
+    "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt"
+)
 SEVENNET_PATH = MODEL_PATH / "sevennet_0.pth"
 
-
 test_data = [
     ("benzene.xyz", -76.0605725422795, "energy", "energy", {}, None),
     (
@@ -131,6 +133,14 @@ def test_potential_energy(struct, expected, properties, prop_key, calc_kwargs, i
         ),
         ("orb", "cpu", -27.08186149597168, "NaCl.cif", {}),
         ("orb", "cpu", -27.089094161987305, "NaCl.cif", {"model": "orb-v2"}),
+        ("pet_mad", "cpu", -27.47624969482422, "NaCl.cif", {}),
+        (
+            "pet_mad",
+            "cpu",
+            -27.47624969482422,
+            "NaCl.cif",
+            {"model": PET_MAD_CHECKPOINT},
+        ),
         (
             "sevennet",
             "cpu",

tests/test_single_point.py

@@ -149,6 +149,8 @@ def skip_extras(arch: str):
             pytest.importorskip("nequip")
         case "orb":
             pytest.importorskip("orb_models")
+        case "pet_mad":
+            pytest.importorskip("pet_mad")
         case "sevennet":
             pytest.importorskip("sevenn")
         case "alignn":