Add upet and sevennet models

ml_peg/models/get_models.py

@@ -141,6 +141,8 @@ def load_models(
         MockCalc,
         OrbCalc,
         PetMadCalc,
+        SevenNetCalc,
+        UPETCalc,
     )
 
     if run_mock is None:
@@ -201,6 +203,25 @@ def load_models(
                     trained_on_dispersion=cfg.get("trained_on_dispersion", False),
                     dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
                 )
+            case "UPETCalculator":
+                loaded_models[name] = UPETCalc(
+                    module=cfg["module"],
+                    class_name=cfg["class_name"],
+                    device=cfg.get("device", "cpu"),
+                    default_dtype=cfg.get("overwrite_dtype", None),
+                    kwargs=cfg.get("kwargs", {}),
+                    trained_on_dispersion=cfg.get("trained_on_dispersion", False),
+                    dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
+                )
+            case "SevenNetCalculator":
+                kwargs = cfg.get("kwargs", {})
+                loaded_models[name] = SevenNetCalc(
+                    model=kwargs["model"],
+                    device=cfg.get("device", "cpu"),
+                    kwargs={k: v for k, v in kwargs.items() if k != "model"},
+                    trained_on_dispersion=cfg.get("trained_on_dispersion", False),
+                    dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
+                )
             case _:
                 loaded_models[name] = GenericASECalc(
                     module=cfg["module"],

ml_peg/models/models.py

@@ -7,9 +7,10 @@
 import dataclasses
 from typing import TYPE_CHECKING
 
-from mlipx import GenericASECalculator as MlipxGenericASECalc
 from mlipx.nodes.generic_ase import Device
 
+from mlipx import GenericASECalculator as MlipxGenericASECalc
+
 if TYPE_CHECKING:
     from ase.calculators.calculator import Calculator
     from ase.calculators.mixing import SumCalculator
@@ -283,3 +284,61 @@ def get_calculator(self, **kwargs) -> Calculator:
         from ml_peg.models.mock import MockCalculator
 
         return MockCalculator()
+
+
+@dataclasses.dataclass(kw_only=True)
+class UPETCalc(GenericASECalc):
+    """Dataclass for upet (PET-MAD / PET-OAM) calculator."""
+
+    def get_calculator(self, precision="high", **kwargs) -> Calculator:
+        """
+        Prepare and load the calculator.
+
+        Parameters
+        ----------
+        precision
+            Level of precision to evaluate the model.
+        **kwargs
+            Any keyword arguments to pass to `get_calculator`.
+
+        Returns
+        -------
+        Calculator
+            Loaded upet ASE calculator.
+        """
+        precision_map = {"low": "float32", "high": "float64"}
+        kwargs["dtype"] = precision_map[precision]
+
+        if self.default_dtype is not None:
+            kwargs["dtype"] = self.default_dtype
+
+        return MlipxGenericASECalc.get_calculator(self, **kwargs)
+
+
+@dataclasses.dataclass(kw_only=True)
+class SevenNetCalc(SumCalc):
+    """Dataclass for SevenNet calculator."""
+
+    model: str
+    device: Device | None = None
+    kwargs: dict = dataclasses.field(default_factory=dict)
+
+    def get_calculator(self, **kwargs) -> Calculator:
+        """
+        Prepare and load the calculator.
+
+        Parameters
+        ----------
+        **kwargs
+            Additional keyword arguments (ignored).
+
+        Returns
+        -------
+        Calculator
+            Loaded SevenNet ASE calculator.
+        """
+        from sevenn.sevennet_calculator import SevenNetCalculator
+
+        device = Device.resolve_auto() if self.device == Device.AUTO else self.device
+        device_str = device.value if isinstance(device, Device) else (device or "cpu")
+        return SevenNetCalculator(model=self.model, device=device_str, **self.kwargs)

ml_peg/models/models.yml

@@ -179,6 +179,45 @@ pet-mad:
   dispersion_kwargs:
     xc: pbesol
 
+pet-mad-1.5-s:
+  module: upet.calculator
+  class_name: UPETCalculator
+  device: "cpu"
+  trained_on_dispersion: false
+  level_of_theory: r2SCAN
+  kwargs:
+    model: "pet-mad-s"
+    version: "1.5.0"
+
+pet-oam-xl:
+  module: upet.calculator
+  class_name: UPETCalculator
+  device: "cpu"
+  trained_on_dispersion: False
+  level_of_theory: PBE
+  kwargs:
+    model: "pet-oam-xl"
+    version: "1.0.0"
+
+sevennet-l3i5:
+  module: sevenn.sevennet_calculator
+  class_name: SevenNetCalculator
+  device: "cpu"
+  trained_on_dispersion: false
+  level_of_theory: PBE
+  kwargs:
+    model: "7net-l3i5"
+
+sevennet-omni-i12-mpa:
+  module: sevenn.sevennet_calculator
+  class_name: SevenNetCalculator
+  device: "cpu"
+  trained_on_dispersion: false
+  level_of_theory: PBE
+  kwargs:
+    model: "7net-omni"
+    modal: "mpa"
+
 # mace-polar-1-s:
 #   module: mace.calculators
 #   class_name: mace_polar

pyproject.toml

@@ -73,6 +73,12 @@ pet-mad = [
 uma = [
     "fairchem-core == 2.19.0; python_version >= '3.11'",
 ]
+upet = [
+    "upet",
+]
+sevenn = [
+    "sevenn",
+]
 
 [project.scripts]
 ml_peg = "ml_peg.cli.cli:app"
@@ -206,6 +212,10 @@ conflicts = [
       { extra = "mattersim" },
       { extra = "grace" },
     ],
+    [
+      { extra = "pet-mad" },
+      { extra = "upet" },
+    ],
 ]
 
 constraint-dependencies = [