Add MatterSim (#25)

* add mattersim diatomic pecs

* change e3nn to 0.5.0 in test

* update model info

* clean notebook

mlip_arena/models/externals/chgnet.py

@@ -18,7 +18,7 @@ class CHGNet(CHGNetCalculator):
         on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
         **kwargs,
     ) -> None:
-        use_device = device or str(get_freer_device())
+        use_device = str(device or get_freer_device())
         super().__init__(
             model=checkpoint,
             use_device=use_device,

mlip_arena/models/externals/mattersim.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+from mattersim.forcefield import MatterSimCalculator
+
+from mlip_arena.models.utils import get_freer_device
+
+with open(Path(__file__).parents[1] / "registry.yaml", encoding="utf-8") as f:
+    REGISTRY = yaml.safe_load(f)
+
+class MatterSim(MatterSimCalculator):
+    def __init__(
+        self,
+        checkpoint=REGISTRY["MatterSim"]["checkpoint"],
+        device=None,
+        **kwargs,
+    ):
+        super().__init__(
+            load_path=checkpoint, 
+            device=str(device or get_freer_device()), 
+            **kwargs
+        )

mlip_arena/models/registry.yaml

@@ -67,6 +67,28 @@ M3GNet:
   npt: true
   license: BSD-3-Clause
 
+MatterSim:
+  module: externals
+  class: MatterSim
+  family: mattersim
+  package: mattersim==1.0.0rc9
+  checkpoint: MatterSim-v1.0.0-5M.pth
+  username: 
+  last-update: 2024-10-29T00:00:00
+  datetime: 2024-10-29T00:00:00 # TODO: Fake datetime
+  datasets:
+    - MPTrj
+    - Alexandria
+  gpu-tasks:
+    - homonuclear-diatomics
+  github: https://github.com/microsoft/mattersim
+  doi: https://arxiv.org/abs/2405.04967
+  date: 2024-12-05
+  prediction: EFS
+  nvt: true
+  npt: true
+  license: MIT
+
 ORB:
   module: externals
   class: ORB
@@ -133,7 +155,7 @@ eqV2(OMat):
     - homonuclear-diatomics
   prediction: EFS
   nvt: true
-  npt: true
+  npt: false # https://github.com/FAIR-Chem/fairchem/issues/888, https://github.com/atomind-ai/mlip-arena/issues/17
   date: 2024-10-18
   github: https://github.com/FAIR-Chem/fairchem
   doi: https://arxiv.org/abs/2410.12771

mlip_arena/tasks/diatomics/alignn/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3f880af442288d79b925e96723fa07c5e31c1d46e2888ed64382c3828c8b6893
-size 2383737
+oid sha256:325979c740279016c3652f570b742d0ed28065156e553ae6c878a775951c344b
+size 2383784

mlip_arena/tasks/diatomics/chgnet/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d3bec21876e97d0841afbebe3b9adc87e61c4b090917f46842ddfcc507b4870e
-size 2003358
+oid sha256:5f3595e1ae02dbc7c30a8ff39227b2a2e46ea940be17e482200f8b3d518cda01
+size 2003404

mlip_arena/tasks/diatomics/equiformer/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fa2354642db212a30b8779e1a9069c58834e690b5d3f8757177676a7236d9cff
-size 3764416
+oid sha256:d278cac9a2d7ce073735dcf6e2df6637510efcb8d89289706b54e8b540c942ec
+size 3764511

mlip_arena/tasks/diatomics/escn/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f24fc8f74dc3e9cc77c2099aefe81226764056db3bcb4a9077d42db68ca821d8
-size 1960014
+oid sha256:2bfb8f0a5424259c9e9686e3e8778adb36e9b08937296e2a127d6007a1adf3bf
+size 1960063

mlip_arena/tasks/diatomics/fairchem/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:65bb77b1dca8bddb5ae2ad13aa1e336a30dbefc95953074ccbe077009ab73e08
-size 2136843
+oid sha256:06fe63898d03644e491ef12e275884c110cf993b8182403408d65fdd23e51816
+size 2136895

mlip_arena/tasks/diatomics/mace-mp/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f36170ebca7c2c2c876cad7ff9fc4ed5d5043df3da4000089465b31ca2b5544a
-size 1915526
+oid sha256:1f6f4a2d4f36071625db988dde933674cdf5478951cf227e6eacc0d818c13a1f
+size 1915573

mlip_arena/tasks/diatomics/mace-off/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1202ad5dcf6434cefd25246deeb5785d58cb70ad83bc54577c9f2d711a6b3974
-size 174038
+oid sha256:010d9215f9c1765fd2d807c0ef57a95b14832b6bda841e25948c6ee342232579
+size 173998

mlip_arena/tasks/diatomics/matgl/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:217346241aaa0509e95e5f0ec6eb6dafbcc9ef954fd1f3077a778f0ac420f3f2
-size 1858366
+oid sha256:32a7001f8635327e166f7747942ab8fdf4bfb20e68fc632b24005d2017daf5f4
+size 1858414

mlip_arena/tasks/diatomics/mattersim/homonuclear-diatomics.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0181c4173a7ec6cd0ddb469817bc6ff1befec5b9a4a122b5876622ad5006fd49
+size 1929036

mlip_arena/tasks/diatomics/orb/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b6e846339ca128f4e2a0463f057917b51764e36278f43bf218f1717a01f79692
-size 5197369
+oid sha256:6799123dc798c713ccd100cec528a71b3e3ec4633de8f92c37c0c2d61b5c7801
+size 5197473

mlip_arena/tasks/diatomics/sevennet/homonuclear-diatomics.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:29542777bafc7425e0de89e9d6fe45341c414da269420658033d6072864ed6e0
-size 1859815
+oid sha256:1cb205f2c271f1799a03ff33ff23e9a1acd9a2d8162d232ba7d3aa59f0a1b30e
+size 1859865

pyproject.toml

@@ -52,7 +52,7 @@ app = [
 test = [
     "torch==2.2.0",
     "nvidia-ml-py==12.560.30",
-    "e3nn==0.5.1",
+    "e3nn==0.5.0",
     "matgl==1.1.2",
     "dgl==2.4.0",
     "chgnet==0.3.8",
@@ -61,6 +61,7 @@ test = [
     "orb-models==0.4.0",
     "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168",
     "alignn==2024.5.27",
+    "mattersim==1.0.0rc9",
     "pytest",
     "prefect>=3.0.4,<3.1.0",
 ]

serve/ranks/homonuclear-diatomics.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from ase.data import chemical_symbols
 
 import numpy as np
 import pandas as pd
@@ -66,11 +65,12 @@ table.sort_values(
 )
 table["Rank"] += np.argsort(table["Spearman's coeff. (F: descending)"].to_numpy())
 
-table.sort_values("PBE energy MAE [eV]", ascending=True, inplace=True)
-table["Rank"] += np.argsort(table["PBE energy MAE [eV]"].to_numpy())
+# NOTE: it's not fair to models trained on different level of theory
+# table.sort_values("PBE energy MAE [eV]", ascending=True, inplace=True)
+# table["Rank"] += np.argsort(table["PBE energy MAE [eV]"].to_numpy())
 
-table.sort_values("PBE force MAE [eV/Å]", ascending=True, inplace=True)
-table["Rank"] += np.argsort(table["PBE force MAE [eV/Å]"].to_numpy())
+# table.sort_values("PBE force MAE [eV/Å]", ascending=True, inplace=True)
+# table["Rank"] += np.argsort(table["PBE force MAE [eV/Å]"].to_numpy())
 
 table.sort_values("Tortuosity", ascending=True, inplace=True)
 table["Rank"] += np.argsort(table["Tortuosity"].to_numpy())
@@ -119,13 +119,13 @@ s = (
         ],
         # vmin=-1, vmax=-0.5
     )
-    .background_gradient(
-        cmap="BuPu",
-        subset=[
-            "PBE energy MAE [eV]",
-            "PBE force MAE [eV/Å]",
-        ],
-    )
+    # .background_gradient(
+    #     cmap="Greys",
+    #     subset=[
+    #         "PBE energy MAE [eV]",
+    #         "PBE force MAE [eV/Å]",
+    #     ],
+    # )
     .background_gradient(
         cmap="RdPu",
         subset=["Tortuosity", "Energy jump [eV]", "Force flips"],

serve/tasks/homonuclear-diatomics.py

@@ -30,7 +30,7 @@ valid_models = [
 mlip_methods = container.multiselect(
     "MLIPs",
     valid_models,
-    ["MACE-MP(M)", "CHGNet", "M3GNet", "SevenNet", "ORB", "ORBv2", "eqV2(OMat)"],
+    ["MACE-MP(M)", "CHGNet", "M3GNet", "MatterSim", "SevenNet", "ORBv2", "eqV2(OMat)"],
 )
 dft_methods = container.multiselect("DFT Methods", ["PBE"], ["PBE"])
 

tests/test_external_calculators.py

@@ -1,5 +1,7 @@
 import pytest
 from ase import Atoms
+from ase.calculators.calculator import PropertyNotImplementedError
+import numpy as np
 
 from mlip_arena.models import MLIPEnum
 
@@ -31,5 +33,22 @@ def test_calculate(model: MLIPEnum):
 
     atoms.calc = calc
 
-    print(atoms.get_potential_energy())
-    assert isinstance(atoms.get_potential_energy(), float)
+    energy = atoms.get_potential_energy()
+
+    assert isinstance(energy, (float, np.float64, np.float32))
+
+    forces = atoms.get_forces()
+    assert isinstance(forces, (np.ndarray, list))
+    assert len(forces) == len(atoms)
+
+    try:
+        stress = atoms.get_stress()
+    except PropertyNotImplementedError:
+        stress = None
+
+    if stress is None:
+        pytest.xfail("Stress calculation is not supported by the model")
+    else:
+        assert isinstance(stress, (np.ndarray, list))
+
+