Fairchem v2 support

.github/workflows/test.yml

@@ -82,14 +82,6 @@ jobs:
       - name: Check out repo
         uses: actions/checkout@v4
 
-      - name: Check out fairchem repository
-        if: ${{ matrix.model.name == 'fairchem' }}
-        uses: actions/checkout@v4
-        with:
-          repository: FAIR-Chem/fairchem
-          path: fairchem-repo
-          ref: fairchem_core-1.10.0
-
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
@@ -98,24 +90,14 @@ jobs:
       - name: Set up uv
         uses: astral-sh/setup-uv@v6
 
-      - name: Install fairchem repository and dependencies
+      - name: Install fairchem and dependencies
         if: ${{ matrix.model.name == 'fairchem' }}
         env:
           HF_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
         run: |
-          uv pip install huggingface_hub --system
-          if [ -n "$HF_TOKEN" ]; then
-            huggingface-cli login --token "$HF_TOKEN"
-          else
-            echo "HF_TOKEN is not set. Skipping login."
-          fi
-          if [ -f fairchem-repo/packages/requirements.txt ]; then
-            uv pip install -r fairchem-repo/packages/requirements.txt --system
-          fi
-          if [ -f fairchem-repo/packages/requirements-optional.txt ]; then
-            uv pip install -r fairchem-repo/packages/requirements-optional.txt --system
-          fi
-          uv pip install -e fairchem-repo/packages/fairchem-core[dev] --system
+          uv pip install "torch>=2.6" --index-url https://download.pytorch.org/whl/cpu --system
+          uv pip install "fairchem-core>=2.2.0" --system
+          uv pip install "huggingface_hub[cli]" --system
           uv pip install -e .[test] --resolution=${{ matrix.version.resolution }} --system
 
       - name: Install torch_sim with model dependencies
@@ -124,7 +106,12 @@ jobs:
           uv pip install -e .[test,${{ matrix.model.name }}] --resolution=${{ matrix.version.resolution }} --system
 
       - name: Run ${{ matrix.model.test_path }} tests
+        env:
+          HF_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
         run: |
+          if [ "${{ matrix.model.name }}" == "fairchem" ]; then
+            huggingface-cli login --token "$HF_TOKEN"
+          fi
           pytest --cov=torch_sim --cov-report=xml ${{ matrix.model.test_path }}
 
       - name: Upload coverage to Codecov
@@ -168,4 +155,11 @@ jobs:
         uses: astral-sh/setup-uv@v6
 
       - name: Run example
-        run: uv run --with . ${{ matrix.example }}
+        env:
+          HF_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
+        run: |
+          if [[ "${{ matrix.example }}" == *"fairchem"* ]]; then
+            uv pip install "huggingface_hub[cli]" --system
+            huggingface-cli login --token "$HF_TOKEN"
+          fi
+          uv run --with . ${{ matrix.example }}

examples/scripts/1_Introduction/1.3_Fairchem.py → examples/scripts/1_Introduction/1.3_fairchem.py

@@ -3,12 +3,10 @@
 
 # /// script
 # dependencies = [
-#     "fairchem-core==1.10.0",
+#     "fairchem-core>=2.2.0",
 # ]
 # ///
 
-import sys
-
 import torch
 from ase.build import bulk
 
@@ -19,38 +17,35 @@
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float32
 
-try:
-    from fairchem.core.models.model_registry import model_name_to_local_file
-except ImportError:
-    print("Skipping example due to missing fairchem dependency")
-    sys.exit(0)
-
-MODEL_PATH = model_name_to_local_file(
-    "EquiformerV2-31M-S2EF-OC20-All+MD", local_cache="."
-)
+# UMA = Unified Machine Learning for Atomistic simulations
+MODEL_NAME = "uma-s-1"
 
 # Create diamond cubic Silicon
 si_dc = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2))
 atomic_numbers = si_dc.get_atomic_numbers()
 model = FairChemModel(
-    model=MODEL_PATH,
+    model=None,
+    model_name=MODEL_NAME,
+    task_name="omat",  # Open Materials task for crystalline systems
     cpu=False,
-    seed=0,
 )
 atoms_list = [si_dc, si_dc]
-state = ts.io.atoms_to_state(atoms_list)
+state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype)
 
 results = model(state)
 
 print(results["energy"].shape)
 print(results["forces"].shape)
-print(results["stress"].shape)
+if stress := results.get("stress"):
+    print(stress.shape)
 
 print(f"Energy: {results['energy']}")
 print(f"Forces: {results['forces']}")
-print(f"Stress: {results['stress']}")
+if stress := results.get("stress"):
+    print(f"{stress=}")
 
 # Check if the energy, forces, and stress are the same for the Si system across the batch
 print(torch.max(torch.abs(results["energy"][0] - results["energy"][1])))
 print(torch.max(torch.abs(results["forces"][0] - results["forces"][1])))
-print(torch.max(torch.abs(results["stress"][0] - results["stress"][1])))
+if stress := results.get("stress"):
+    print(torch.max(torch.abs(stress[0] - stress[1])))