fea: make test match other models using foundation nequip

Author: CompRhys

tests/models/test_nequip_framework.py

@@ -1,16 +1,19 @@
 import traceback
-import urllib.request
-from enum import StrEnum
 from pathlib import Path
 
 import pytest
 
-from tests.conftest import DEVICE
-from tests.models.conftest import make_model_calculator_consistency_test
+from tests.conftest import DEVICE, DTYPE
+from tests.models.conftest import (
+    consistency_test_simstate_fixtures,
+    make_model_calculator_consistency_test,
+    make_validate_model_outputs_test,
+)
 
 
 try:
     from nequip.ase import NequIPCalculator
+    from nequip.scripts.compile import main
 
     from torch_sim.models.nequip_framework import (
         NequIPFrameworkModel,
@@ -22,29 +25,34 @@
     )
 
 
-class NequIPUrls(StrEnum):
-    """Checkpoint download URLs for NequIP models."""
-
-    Si = "https://github.com/abhijeetgangan/pt_model_checkpoints/raw/refs/heads/main/nequip/Si.nequip.pth"
-
-
 @pytest.fixture(scope="session")
-def model_path_nequip(tmp_path_factory: pytest.TempPathFactory) -> Path:
-    tmp_path = tmp_path_factory.mktemp("nequip_checkpoints")
-    model_name = "Si.nequip.pth"
-    model_path = Path(tmp_path) / model_name
-
-    if not model_path.is_file():
-        urllib.request.urlretrieve(NequIPUrls.Si, model_path)  # noqa: S310
+def compiled_nequip_model_path(tmp_path_factory: pytest.TempPathFactory) -> Path:
+    """Compile NequIP OAM-L model from nequip.net."""
+    tmp_path = tmp_path_factory.mktemp("nequip_compiled")
+    output_model_name = "mir-group__NequIP-OAM-L__0.1.nequip.pt2"
+    output_path = Path(tmp_path) / output_model_name
+
+    main(
+        args=[
+            "nequip.net:mir-group/NequIP-OAM-L:0.1",
+            str(output_path),
+            "--mode",
+            "aotinductor",
+            "--device",
+            "cuda",
+            "--target",
+            "ase",
+        ]
+    )
 
-    return model_path
+    return output_path
 
 
 @pytest.fixture
-def nequip_model(model_path_nequip: Path) -> NequIPFrameworkModel:
+def nequip_model(compiled_nequip_model_path: Path) -> NequIPFrameworkModel:
     """Create an NequIPModel wrapper for the pretrained model."""
     compiled_model, (r_max, type_names) = from_compiled_model(
-        model_path_nequip, device=DEVICE
+        compiled_nequip_model_path, device=DEVICE
     )
     return NequIPFrameworkModel(
         model=compiled_model,
@@ -74,11 +82,18 @@ def test_nequip_initialization(model_path_nequip: Path) -> None:
     assert model._device == DEVICE  # noqa: SLF001
 
 
-test_nequip_consistency = make_model_calculator_consistency_test(
+test_metatomic_consistency = make_model_calculator_consistency_test(
     test_name="nequip",
     model_fixture_name="nequip_model",
     calculator_fixture_name="nequip_calculator",
-    sim_state_names=("si_sim_state", "rattled_si_sim_state"),
+    sim_state_names=consistency_test_simstate_fixtures,
+    energy_atol=5e-5,
+    dtype=DTYPE,
+    device=DEVICE,
 )
 
-# TODO (AG): Test multi element models
+test_metatomic_model_outputs = make_validate_model_outputs_test(
+    model_fixture_name="nequip_model",
+    dtype=DTYPE,
+    device=DEVICE,
+)