Add PET-MAD model

ml_peg/models/get_models.py

@@ -97,7 +97,7 @@ def load_models(models: None | str | Iterable = None) -> dict[str, Any]:
     dict[str, Any]
         Loaded models from models.yml.
     """
-    from ml_peg.models.models import FairChemCalc, GenericASECalc, OrbCalc
+    from ml_peg.models.models import FairChemCalc, GenericASECalc, OrbCalc, PetMadCalc
 
     loaded_models = {}
 
@@ -137,6 +137,16 @@ def load_models(models: None | str | Iterable = None) -> dict[str, Any]:
                 trained_on_d3=cfg.get("trained_on_d3", False),
                 d3_kwargs=cfg.get("d3_kwargs", {}),
             )
+        elif cfg["class_name"] == "PETMADCalculator":
+            loaded_models[name] = PetMadCalc(
+                module=cfg["module"],
+                class_name=cfg["class_name"],
+                device=cfg.get("device", "cpu"),
+                default_dtype=cfg.get("default_dtype", "float32"),
+                kwargs=cfg.get("kwargs", {}),
+                trained_on_d3=cfg.get("trained_on_d3", False),
+                d3_kwargs=cfg.get("d3_kwargs", {}),
+            )
         else:
             loaded_models[name] = GenericASECalc(
                 module=cfg["module"],

ml_peg/models/models.py

@@ -94,6 +94,32 @@ def get_calculator(self, **kwargs) -> Calculator:
         return MlipxGenericASECalc.get_calculator(self, **kwargs)
 
 
+@dataclasses.dataclass(kw_only=True)
+class PetMadCalc(GenericASECalc):
+    """Dataclass for PET-MAD calculator."""
+
+    def get_calculator(self, **kwargs) -> Calculator:
+        """
+        Prepare and load the calculator.
+
+        Parameters
+        ----------
+        **kwargs
+            Any keyword arguments to pass to `get_calculator`.
+
+        Returns
+        -------
+        Calculator
+            Loaded ASE Calculator.
+        """
+        if self.default_dtype is not None:
+            kwargs["dtype"] = self.default_dtype
+        else:
+            kwargs["dtype"] = self.default_dtype
+
+        return MlipxGenericASECalc.get_calculator(self, **kwargs)
+
+
 # https://github.com/orbital-materials/orb-models
 @dataclasses.dataclass(kw_only=True)
 class OrbCalc(SumCalc):

ml_peg/models/models.yml

@@ -156,3 +156,15 @@ orb-v3-consv-inf-omat:
 #   kwargs:
 #     model_name: "uma-s-1p1"
 #     task_name: "omol"
+
+pet-mad:
+  module: pet_mad.calculator
+  class_name: PETMADCalculator
+  device: "cpu"
+  default_dtype: float32
+  trained_on_d3: false
+  level_of_theory: PBEsol
+  kwargs:
+    version: "v1.0.2"
+  d3_kwargs:
+    xc: pbesol

pyproject.toml

@@ -58,7 +58,9 @@ mattersim = [
 orb = [
     "orb-models == 0.5.5; sys_platform != 'win32' and python_version < '3.13'",
 ]
-
+pet-mad = [
+    "pet-mad == 1.4.4; sys_platform != 'win32'"
+]
 uma = [
     "fairchem-core == 2.10.0",
 ]