Add calculator test (#6)

* add files

* resolve dgl issue

* fix typos; remove poetry

* miss dgl in test

* fix farichem version

* spearate mace from other dependencies

* use scripts to install pyg and dgl

* fix typo

* fix another typo

* add pynanoflann dependency

* install pynanoflann separately in test workflow

* change alignn downlaod method; update all diatomic curve stats

.github/workflows/test.yaml

@@ -0,0 +1,34 @@
+name: Python Test
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v2
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch==2.2.0
+        bash scripts/install-pyg.sh
+        bash scripts/install-dgl.sh
+        pip install .[test]
+        pip install mace-torch
+        pip install "pynanoflann@git+https://github.com/dwastberg/pynanoflann#egg=af434039ae14bedcbb838a7808924d6689274168"
+
+    - name: Run tests
+      run: |
+        pytest tests

mlip_arena/models/externals.py

@@ -7,7 +7,7 @@ from typing import Literal
 import matgl
 import requests
 import torch
-from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff
+from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff, default_path
 from ase import Atoms
 from chgnet.model.dynamics import CHGNetCalculator
 from chgnet.model.model import CHGNet as CHGNetModel
@@ -232,11 +232,13 @@ class eSCN(OCPCalculator):
 
 
 class ALIGNN(AlignnAtomwiseCalculator):
-    def __init__(self, device=None, dir_path: str = "/tmp/alignn/", **kwargs) -> None:
+    def __init__(self, device=None, **kwargs) -> None:
         # TODO: cannot control version
-        _ = get_figshare_model_ff(dir_path=dir_path)
+        # _ = get_figshare_model_ff(dir_path=dir_path)
+        model_path = default_path()
+
         device = device or get_freer_device()
-        super().__init__(path=dir_path, device=device, **kwargs)
+        super().__init__(path=model_path, device=device, **kwargs)
 
 
 class SevenNet(SevenNetCalculator):

pyproject.toml

@@ -41,10 +41,10 @@ run = [
     "torch==2.2.0",
     "e3nn==0.5.1",
     "matgl==1.1.2",
-    "dgl==2.4.0+cu121",
+    "dgl==2.4.0",
     "mace-torch==0.3.4",
     "chgnet==0.3.8",
-    "fairchem-core==0.1.0",
+    "fairchem-core==1.1.0",
     "sevenn==0.9.3.post1",
     "orb-models==0.3.1",
     "alignn==2024.5.27"
@@ -59,15 +59,17 @@ test = [
     "torch==2.2.0",
     "e3nn==0.5.1",
     "matgl==1.1.2",
-    "dgl==2.4.0+cu121",
-    "mace-torch==0.3.4",
+    "dgl==2.4.0",
     "chgnet==0.3.8",
-    "fairchem-core==0.1.0",
+    "fairchem-core==1.1.0",
     "sevenn==0.9.3.post1",
     "orb-models==0.3.1",
     "alignn==2024.5.27",
     "pytest"
 ]
+mace = [
+    "mace-torch==0.3.4",
+]
 
 [project.urls]
 Homepage = "https://github.com/atomind-ai/mlip-arena"

scripts/install-dgl.sh

@@ -0,0 +1,6 @@
+# DGL (M3GNet, ALIGNN)
+
+TORCH=2.2
+CUDA=cu121
+
+pip install dgl -U -f https://data.dgl.ai/wheels/torch-${TORCH}/${CUDA}/repo.html

scripts/install-pyg.sh

@@ -1,11 +1,6 @@
-
-
 # PyTorch Geometric (OCP)
-TORCH=2.3.1
+TORCH=2.2.0
 CUDA=cu121
 
 pip install --verbose --no-cache torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 pip install --verbose --no-cache torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
-
-# DGL (M3GNet, ALIGNN)
-pip install dgl -U -f https://data.dgl.ai/wheels/torch-2.3/cu121/repo.html

tests/test_external_calculators.py

@@ -0,0 +1,21 @@
+import pytest
+from ase import Atoms
+
+from mlip_arena.models.utils import MLIPEnum
+
+
+@pytest.mark.parametrize("model", MLIPEnum)
+def test_calculate(model):
+    calc = MLIPEnum[model.name].value()
+
+    atoms = Atoms(
+        "OO",
+        positions=[[0, 0, 0], [1.5, 0, 0]],
+        cell=[10, 10 + 0.001, 10 + 0.002],
+        pbc=True,
+    )
+
+    atoms.calc = calc
+
+    print(atoms.get_potential_energy())
+    assert isinstance(atoms.get_potential_energy(), float)