add mattersim calculator fixes stfc#425

alinelena · Feb 14, 2025 · a8da9a9 · a8da9a9
1 parent f3d7f66
commit a8da9a9
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 1 deletion.
diff --git a/janus_core/helpers/janus_types.py b/janus_core/helpers/janus_types.py
@@ -115,7 +115,7 @@ class CorrelationKwargs(TypedDict, total=True):
 
 # Janus specific
 Architectures = Literal[
-    "mace", "mace_mp", "mace_off", "m3gnet", "chgnet", "alignn", "sevennet"
+    "mace", "mace_mp", "mace_off", "m3gnet", "chgnet", "alignn", "sevennet", "mattersim"
 ]
 Devices = Literal["cpu", "cuda", "mps", "xpu"]
 Ensembles = Literal["nph", "npt", "nve", "nvt", "nvt-nh", "nvt-csvr", "npt-mtk"]

diff --git a/janus_core/helpers/mlip_calculators.py b/janus_core/helpers/mlip_calculators.py
@@ -224,6 +224,17 @@ def choose_calculator(
         kwargs.setdefault("sevennet_config", None)
         calculator = SevenNetCalculator(model=model_path, device=device, **kwargs)
 
+    elif arch == "mattersim":
+        from mattersim.forcefield import MatterSimCalculator
+        from mattersim import __version__
+
+        if isinstance(model_path, Path):
+            model_path = str(model_path)
+        elif not isinstance(model_path, str):
+            model_path = "mattersim-v1.0.0-5M"
+
+        calculator = MatterSimCalculator(load_path=model_path, device=device, **kwargs)
+
     else:
         raise ValueError(
             f"Unrecognized {arch=}. Suported architectures "

diff --git a/pyproject.toml b/pyproject.toml
@@ -55,11 +55,15 @@ m3gnet = [
 sevennet = [
     "sevenn == 0.10.3",
 ]
+mattersim = [
+    "mattersim == 1.1.1",
+]
 all = [
     "janus-core[alignn]",
     "janus-core[chgnet]",
     "janus-core[m3gnet]",
     "janus-core[sevennet]",
+      "janus-core[mattersim]",
 ]
 
 [project.scripts]

diff --git a/tests/test_mlip_calculators.py b/tests/test_mlip_calculators.py
@@ -119,6 +119,7 @@ def test_invalid_device(arch):
         ("sevennet", "cpu", {"model_path": SEVENNET_PATH}),
         ("sevennet", "cpu", {}),
         ("sevennet", "cpu", {"model": "sevennet-0"}),
+        ("mattersim", "cpu", {}),
     ],
 )
 def test_extra_mlips(arch, device, kwargs):

diff --git a/tests/test_single_point.py b/tests/test_single_point.py
@@ -284,6 +284,7 @@ def test_mlips(arch, device, expected_energy):
     ("sevennet", "cpu", -27.061979293823242, {"model_path": SEVENNET_PATH}),
     ("sevennet", "cpu", -27.061979293823242, {}),
     ("sevennet", "cpu", -27.061979293823242, {"model_path": "SevenNet-0_11July2024"}),
+    ("mattersim", "cpu", -27.06208038330078, {}),
 ]