Update NequIP (#604)

Author: ElliottKasoar

.github/workflows/ci.yml

@@ -42,7 +42,7 @@ jobs:
 
       - name: Install updated e3nn dependencies
         run: |
-          uv sync --extra mattersim --extra fairchem --extra sevennet --extra d3
+          uv sync --extra mattersim --extra fairchem --extra sevennet --extra nequip --extra d3
           uv pip install --reinstall pynvml
           uv pip install fairchem-core[torch-extras] --no-build-isolation
 

.github/workflows/mac.yml

@@ -42,7 +42,7 @@ jobs:
 
       - name: Install updated e3nn dependencies
         run: |
-          uv sync --extra mattersim --extra fairchem --extra sevennet --extra d3
+          uv sync --extra mattersim --extra fairchem --extra sevennet --extra nequip --extra d3
           uv pip install --reinstall pynvml
           uv pip install "fairchem-core[torch-extras]" --no-build-isolation
 

.github/workflows/windows.yml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Install updated e3nn dependencies
         run: |
-          uv sync --extra fairchem --extra sevennet
+          uv sync --extra fairchem --extra sevennet --extra nequip
           uv pip install --reinstall pynvml
           uv pip install "fairchem-core[torch-extras]" --no-build-isolation
 

janus_core/helpers/mlip_calculators.py

@@ -340,8 +340,8 @@ def choose_calculator(
 
             model = str(model)
 
-            calculator = NequIPCalculator.from_deployed_model(
-                model_path=model, device=device, **kwargs
+            calculator = NequIPCalculator.from_compiled_model(
+                compile_path=model, device=device, **kwargs
             )
 
         case "dpa3":

pyproject.toml

@@ -57,9 +57,6 @@ mace = [
     "mace-torch==0.3.14",
     "janus-core[d3]",
 ]
-nequip = [
-    "nequip == 0.6.1",
-]
 orb = [
     "orb-models == 0.5.5; sys_platform != 'win32'",
 ]
@@ -77,7 +74,6 @@ all = [
     "janus-core[grace]",
     "janus-core[d3]",
     "janus-core[mace]",
-    "janus-core[nequip]",
     "janus-core[orb]",
     "janus-core[pet-mad]",
     "janus-core[plumed]",
@@ -91,10 +87,12 @@ fairchem = [
 mattersim = [
     "mattersim == 1.1.2; sys_platform != 'win32'",
 ]
+nequip = [
+    "nequip == 0.14.0",
+]
 sevennet = [
     "sevenn == 0.11.2.post1",
 ]
-
 uma = [
     "fairchem-core == 2.3.0",
 ]
@@ -326,4 +324,12 @@ conflicts = [
       { extra = "sevennet" },
       { extra = "all" },
     ],
+    [
+      { extra = "nequip" },
+      { extra = "mace" },
+    ],
+    [
+      { extra = "nequip" },
+      { extra = "all" },
+    ],
 ]

tests/models/toluene.pth renamed to tests/models/toluene.nequip.pth

@@ -47,7 +47,7 @@
     EQUIFORMER_PATH = None
     ESEN_PATH = None
 
-NEQUIP_PATH = MODEL_PATH / "toluene.pth"
+NEQUIP_PATH = MODEL_PATH / "toluene.nequip.pth"
 
 ORB_WEIGHTS_PATH = MODEL_PATH / "orb-d3-xs-v2-20241011.ckpt"
 

tests/test_mlip_calculators.py

@@ -23,7 +23,7 @@
 EQUIFORMER_LABEL = "EquiformerV2-83M-S2EF-OC20-2M"
 ESEN_LABEL = "eSEN-30M-MP"
 MACE_PATH = MODEL_PATH / "mace_mp_small.model"
-NEQUIP_PATH = MODEL_PATH / "toluene.pth"
+NEQUIP_PATH = MODEL_PATH / "toluene.nequip.pth"
 PET_MAD_CHECKPOINT = (
     "https://huggingface.co/lab-cosmo/pet-mad/resolve/v1.1.0/models/pet-mad-v1.1.0.ckpt"
 )