define MaceUrls StrEnum to avoid breaking tests when "small" checkpoints get redirected in mace-torch

Author: janosh

tests/conftest.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Any, Final
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any
 
 import pytest
 import torch
@@ -19,6 +20,13 @@
     from mace.calculators import MACECalculator
 
 
+class MaceUrls(StrEnum):
+    mace_small = "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b/mace_agnesi_small.model"
+    mace_off_small = (
+        "https://github.com/ACEsuit/mace-off/blob/main/mace_off23/MACE-OFF23_small.model"
+    )
+
+
 @pytest.fixture
 def device() -> torch.device:
     return torch.device("cpu")
@@ -324,18 +332,13 @@ def lj_model(device: torch.device, dtype: torch.dtype) -> LennardJonesModel:
     )
 
 
-MACE_CHECKPOINT_URL: Final[str] = (
-    "https://github.com/ACEsuit/mace-mp/releases/download/mace_mp_0b/mace_agnesi_small.model"
-)
-
-
 @pytest.fixture
 def ase_mace_mpa() -> "MACECalculator":
     """Provides an ASE MACECalculator instance using mace_mp."""
     from mace.calculators.foundations_models import mace_mp
 
     # Ensure dtype matches the one used in the torchsim fixture (float64)
-    return mace_mp(model=MACE_CHECKPOINT_URL, default_dtype="float64")
+    return mace_mp(model=MaceUrls.mace_small, default_dtype="float64")
 
 
 @pytest.fixture
@@ -346,7 +349,7 @@ def torchsim_mace_mpa() -> MaceModel:
     # Use float64 for potentially higher precision needed in optimization
     dtype = getattr(torch, dtype_str := "float64")
     raw_mace = mace_mp(
-        model=MACE_CHECKPOINT_URL, return_raw_model=True, default_dtype=dtype_str
+        model=MaceUrls.mace_small, return_raw_model=True, default_dtype=dtype_str
     )
     return MaceModel(
         model=raw_mace,

tests/models/test_mace.py

@@ -3,6 +3,7 @@
 from ase.atoms import Atoms
 
 import torch_sim as ts
+from tests.conftest import MaceUrls
 from tests.models.conftest import (
     consistency_test_simstate_fixtures,
     make_model_calculator_consistency_test,
@@ -19,8 +20,8 @@
     pytest.skip("MACE not installed", allow_module_level=True)
 
 
-mace_model = mace_mp(model="small", return_raw_model=True)
-mace_off_model = mace_off(model="small", return_raw_model=True)
+mace_model = mace_mp(model=MaceUrls.mace_small, return_raw_model=True)
+mace_off_model = mace_off(model=MaceUrls.mace_off_small, return_raw_model=True)
 
 
 @pytest.fixture
@@ -32,7 +33,7 @@ def dtype() -> torch.dtype:
 @pytest.fixture
 def ase_mace_calculator() -> MACECalculator:
     return mace_mp(
-        model="small",
+        model=MaceUrls.mace_small,
         device="cpu",
         default_dtype="float32",
         dispersion=False,
@@ -96,7 +97,7 @@ def benzene_system(
 @pytest.fixture
 def ase_mace_off_calculator() -> MACECalculator:
     return mace_off(
-        model="small",
+        model=MaceUrls.mace_off_small,
         device="cpu",
         default_dtype="float32",
         dispersion=False,

tests/unbatched/test_unbatched_mace.py

@@ -3,6 +3,7 @@
 from ase.atoms import Atoms
 
 import torch_sim as ts
+from tests.conftest import MaceUrls
 from tests.unbatched.conftest import make_unbatched_model_calculator_consistency_test
 
 
@@ -15,8 +16,8 @@
     pytest.skip("MACE not installed", allow_module_level=True)
 
 
-mace_model = mace_mp(model="small", return_raw_model=True)
-mace_off_model = mace_off(model="small", return_raw_model=True)
+mace_model = mace_mp(model=MaceUrls.mace_small, return_raw_model=True)
+mace_off_model = mace_off(model=MaceUrls.mace_off_small, return_raw_model=True)
 
 
 @pytest.fixture
@@ -28,7 +29,7 @@ def dtype() -> torch.dtype:
 @pytest.fixture
 def ase_mace_calculator() -> MACECalculator:
     return mace_mp(
-        model="small",
+        model=MaceUrls.mace_small,
         device="cpu",
         default_dtype="float32",
         dispersion=False,