restrict mace version, update eos test

atomind-ai · Oct 8, 2024 · 79e5bd9 · 79e5bd9
1 parent 18d7cd4
commit 79e5bd9
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 77 deletions.
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
@@ -26,7 +26,7 @@ jobs:
         bash scripts/install-pyg.sh
         bash scripts/install-dgl.sh
         pip install .[test]
-        pip install mace-torch
+        pip install .[mace]
         pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
 
     - name: Run tests

diff --git a/mlip_arena/models/utils.py b/mlip_arena/models/utils.py
@@ -47,78 +47,4 @@ def get_freer_device() -> torch.device:
         print("No GPU or MPS available. Using CPU.")
         device = torch.device("cpu")
 
-    return device
-
-
-# class EXTMLIPEnum(Enum):
-#     """Enumeration class for EXTMLIP models.
-
-#     Attributes:
-#         M3GNet (str): M3GNet model.
-#         CHGNet (str): CHGNet model.
-#         MACE (str): MACE model.
-#     """
-
-#     M3GNet = "M3GNet"
-#     CHGNet = "CHGNet"
-#     MACE = "MACE"
-#     Equiformer = "Equiformer"
-
-
-# def get_freer_device() -> torch.device:
-#     """Get the GPU with the most free memory.
-
-#     Returns:
-#         torch.device: The selected GPU device.
-
-#     Raises:
-#         ValueError: If no GPU is available.
-#     """
-#     device_count = torch.cuda.device_count()
-#     if device_count == 0:
-#         print("No GPU available. Using CPU.")
-#         return torch.device("cpu")
-
-#     mem_free = [
-#         torch.cuda.get_device_properties(i).total_memory
-#         - torch.cuda.memory_allocated(i)
-#         for i in range(device_count)
-#     ]
-
-#     free_gpu_index = mem_free.index(max(mem_free))
-
-#     print(
-#         f"Selected GPU {free_gpu_index} with {mem_free[free_gpu_index] / 1024**2:.2f} MB free memory from {device_count} GPUs",
-#     )
-
-#     return torch.device(f"cuda:{free_gpu_index}")
-
-
-
-# def external_ase_calculator(name: EXTMLIPEnum, **kwargs: Any) -> Calculator:
-#     """Construct an ASE calculator from an external third-party MLIP packages"""
-#     calculator = None
-#     device = get_freer_device()
-
-#     if name == EXTMLIPEnum.MACE:
-#         from mace.calculators import mace_mp
-
-#         calculator = mace_mp(device=str(device), **kwargs)
-
-#     elif name == EXTMLIPEnum.CHGNet:
-#         from chgnet.model.dynamics import CHGNetCalculator
-
-#         calculator = CHGNetCalculator(use_device=str(device), **kwargs)
-
-#     elif name == EXTMLIPEnum.M3GNet:
-#         import matgl
-#         from matgl.ext.ase import PESCalculator
-
-#         potential = matgl.load_model("M3GNet-MP-2021.2.8-PES")
-#         calculator = PESCalculator(potential, **kwargs)
-
-
-
-#     calculator.__setattr__("name", name.value)
-
-#     return calculator
+    return device
diff --git a/tests/test_eos.py b/tests/test_eos.py
@@ -2,7 +2,6 @@
 from ase.build import bulk
 
 from mlip_arena.tasks.eos.run import fit as EOS
-from mlip_arena.models.utils import MLIPEnum
 
 atoms = bulk("Cu", "fcc", a=3.6)