Add GRACE

docs/source/user_guide/models.rst

@@ -402,18 +402,31 @@ mattersim-1M
 GRACE
 =====
 
-GRACE-2L-OAM
-------------
+GRACE-2L-OMAT
+-------------
 
 .. code-block:: yaml
 
-   GRACE-2L-OAM:
-     module: tensorpotential.calculator
-     class_name: TPCalculator
-     device: "cpu"
-     trained_on_dispersion: false
-     kwargs:
-       model: "<local path to the GRACE-2L-OAM model>"
+  grace-2l-omat:
+    module: tensorpotential.calculator
+    class_name: grace_fm
+    device: "cpu"
+    trained_on_dispersion: false
+    kwargs:
+      model: GRACE-2L-OMAT
+
+
+GRACE-2L-SMAX-OMAT-medium
+-------------------------
+.. code-block:: yaml
+
+  grace-2l-smax-omat-medium:
+    module: tensorpotential.calculator
+    class_name: grace_fm
+    device: "cpu"
+    trained_on_dispersion: false
+    kwargs:
+      model: GRACE-2L-SMAX-OMAT-medium
 
 
 SevenNet

ml_peg/models/get_models.py

@@ -138,6 +138,7 @@ def load_models(
     from ml_peg.models.models import (
         FairChemCalc,
         GenericASECalc,
+        GraceCalc,
         MatterSimCalc,
         MockCalc,
         OrbCalc,
@@ -228,6 +229,16 @@ def load_models(
                     trained_on_dispersion=cfg.get("trained_on_dispersion", False),
                     dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
                 )
+            case "grace_fm":
+                loaded_models[name] = GraceCalc(
+                    module=cfg["module"],
+                    class_name=cfg["class_name"],
+                    device=cfg.get("device", "cpu"),
+                    kwargs=cfg.get("kwargs", {}),
+                    trained_on_dispersion=cfg.get("trained_on_dispersion", False),
+                    dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
+                    default_dtype=cfg.get("overwrite_dtype", None),
+                )
             case _:
                 loaded_models[name] = GenericASECalc(
                     module=cfg["module"],

ml_peg/models/models.py

@@ -7,6 +7,7 @@
 import dataclasses
 from functools import wraps
 from typing import TYPE_CHECKING, Any
+from warnings import warn
 
 from mlipx import GenericASECalculator as MlipxGenericASECalc
 from mlipx.nodes.generic_ase import Device
@@ -474,3 +475,46 @@ def get_calculator(self, **kwargs) -> Calculator:
         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(device=device_str, **self.kwargs)
+
+
+@dataclasses.dataclass(kw_only=True)
+class GraceCalc(GenericASECalc):
+    """Dataclass for GRACE calculator."""
+
+    device: Device | None = None
+    kwargs: dict = dataclasses.field(default_factory=dict)
+
+    def get_calculator(self, precision="high", **kwargs) -> Calculator:
+        """
+        Prepare and load the calculator.
+
+        Parameters
+        ----------
+        precision
+            Level of precision to evaluate the model.
+        **kwargs
+            Additional keyword arguments (ignored).
+
+        Returns
+        -------
+        Calculator
+            Loaded GRACE ASE calculator.
+        """
+        from tensorpotential.calculator.foundation_models import MODELS_NAME_LIST
+
+        precision_map = {"low": "", "high": "-fp64"}
+        suffix = precision_map[precision]
+
+        if self.default_dtype is not None:
+            suffix = self.default_dtype
+
+        model_name = f"{self.kwargs['model']}{suffix}"
+        if model_name in MODELS_NAME_LIST:
+            self.kwargs["model"] = model_name
+        else:
+            warn(
+                "Unable to find model with requested precision, using default",
+                stacklevel=2,
+            )
+
+        return MlipxGenericASECalc.get_calculator(self, **kwargs)

ml_peg/models/models.yml

@@ -122,13 +122,21 @@ orb-v3-consv-omol:
 #     model_name: "uma-m-1p1"
 #     task_name: "omat"
 
-# GRACE-2L-OAM:
-#   module: tensorpotential.calculator
-#   class_name: TPCalculator
-#   device: "cpu"
-#   trained_on_dispersion: false
-#   kwargs:
-#     model: "/Users/joehart/Desktop/0_Cambridge/0_MPhil_Scientific_Computing/MPhil_project/mlipx_testing/models/GRACE-2L-OAM_28Jan25"
+grace-2l-omat:
+  module: tensorpotential.calculator
+  class_name: grace_fm
+  device: "cpu"
+  trained_on_dispersion: false
+  kwargs:
+    model: GRACE-2L-OMAT
+
+grace-2l-smax-omat-medium:
+  module: tensorpotential.calculator
+  class_name: grace_fm
+  device: "cpu"
+  trained_on_dispersion: false
+  kwargs:
+    model: GRACE-2L-SMAX-OMAT-medium
 
 # MACE-OFF23(L):
 #   module: mace.calculators

pyproject.toml

@@ -55,7 +55,7 @@ dpa3 = [
     "deepmd-kit == 3.1.0",
 ]
 grace = [
-    "tensorpotential == 0.5.1; python_version < '3.13'",
+    "tensorpotential == 0.6.0",
 ]
 mace = [
     "mace-torch==0.3.15",
@@ -211,22 +211,10 @@ conflicts = [
         { extra = "uma" },
         { extra = "mace" },
     ],
-    [
-        { extra = "uma" },
-        { extra = "grace" },
-    ],
     [
         { extra = "mattersim" },
         { extra = "mace" },
     ],
-    [
-        { extra = "mattersim" },
-        { extra = "grace" },
-    ],
-    [
-        { extra = "mlipaudit" },
-        { extra = "grace" },
-    ],
     [
         { extra = "mace" },
         { extra = "sevenn" },