Add Vivace model (#606)

Author: joehart2001

ml_peg/models/get_models.py

@@ -143,6 +143,7 @@ def load_models(
         OrbCalc,
         SevenNetCalc,
         UPETCalc,
+        VivaceCalc,
     )
 
     if run_mock is None:
@@ -220,6 +221,13 @@ def load_models(
                     trained_on_dispersion=cfg.get("trained_on_dispersion", False),
                     dispersion_kwargs=cfg.get("dispersion_kwargs", {}),
                 )
+            case "MLFFCalculator":
+                loaded_models[name] = VivaceCalc(
+                    device=cfg.get("device", "auto"),
+                    kwargs=cfg.get("kwargs", {}),
+                    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

@@ -169,6 +169,44 @@ def get_calculator(self, precision="high", **kwargs) -> Calculator:
         return MlipxGenericASECalc.get_calculator(self, **kwargs)
 
 
+@dataclasses.dataclass(kw_only=True)
+class VivaceCalc(SumCalc):
+    """Dataclass for Vivace 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
+            Unused precision argument, kept for the common model API.
+        **kwargs
+            Keyword arguments passed to the Vivace calculator.
+
+        Returns
+        -------
+        Calculator
+            Loaded ASE calculator.
+        """
+        from simpoly.vivace.calculator import MLFFCalculator
+
+        kwargs.update(self.kwargs)
+        calc = MLFFCalculator(**kwargs)
+
+        # Vivace sets dtype from checkpoint metadata inside MLFFCalculator.
+        # Leave precision/overwrite_dtype untouched unless SimPoly exposes it.
+        device = Device.resolve_auto() if self.device == Device.AUTO else self.device
+        if device is not None:
+            calc.device = device
+            calc.model = calc.model.to(device=device)
+
+        return calc
+
+
 # https://github.com/orbital-materials/orb-models
 @dataclasses.dataclass(kw_only=True)
 class OrbCalc(SumCalc):

ml_peg/models/models.yml

@@ -231,6 +231,15 @@ sevennet-omni-mpa:
     model: "7net-omni"
     modal: "mpa"
 
+# Vivace:
+#   module: simpoly.vivace.calculator
+#   class_name: MLFFCalculator
+#   device: "cuda"
+#   trained_on_dispersion: true
+#   level_of_theory: r2SCAN+D3
+#   kwargs:
+#     model_path: /path/to/model.pt
+
 # mace-polar-1-s:
 #   module: mace.calculators
 #   class_name: mace_polar

pyproject.toml

@@ -83,6 +83,12 @@ upet = [
 sevenn = [
     "sevenn==0.12.1",
 ]
+vivace = [
+    "simpoly[vivace,cuda13] @ git+https://github.com/microsoft/simpoly.git@da100ad32d971214cf9343d9d853e8ce61604dc0; sys_platform == 'linux'",
+    "simpoly[vivace] @ git+https://github.com/microsoft/simpoly.git@da100ad32d971214cf9343d9d853e8ce61604dc0; sys_platform != 'linux'",
+    "cuequivariance ==0.8.1; sys_platform != 'linux'",
+    "cuequivariance-torch ==0.8.1; sys_platform != 'linux'",
+]
 
 [project.scripts]
 ml_peg = "ml_peg.cli.cli:app"
@@ -225,6 +231,10 @@ conflicts = [
         { extra = "mace" },
         { extra = "sevenn" },
     ],
+    [
+      { extra = "vivace" },
+      { extra = "mace" },
+    ],
 ]
 
 constraint-dependencies = [