Add support for Torchsim

.github/workflows/testing.yml

@@ -206,14 +206,16 @@ jobs:
 
       - name: Install pymatgen from master if triggered by pymatgen repo dispatch
         if: github.event_name == 'repository_dispatch' && github.event.action == 'pymatgen-ci-trigger'
-        run: uv pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
+        run: |
+          micromamba activate a2
+          uv pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
 
       - name: Test Notebooks
         env:
           MP_API_KEY: ${{ secrets.MP_API_KEY }}
         run: |
           micromamba activate a2
-          pytest -n auto --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields
+          pytest -n auto --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb --ignore=./tutorials/force_fields --ignore=./tutorials/torchsim_tutorial.ipynb
 
       - name: Test ASE
         env:
@@ -304,6 +306,64 @@ jobs:
 
 
 
+  test-torchsim:
+    # prevent this action from running on forks
+    if: github.repository == 'materialsproject/atomate2'
+
+    services:
+      local_mongodb:
+        image: mongo:4.0
+        ports:
+          - 27017:27017
+
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+
+    steps:
+      - name: Check out repo
+        uses: actions/checkout@v4
+
+      - name: Set up micromamba
+        uses: mamba-org/setup-micromamba@main
+        with:
+          environment-name: a2
+          cache-environment: false
+          create-args: >-
+            python=3.12
+
+      - name: Install uv
+        run: micromamba run -n a2 pip install uv
+
+      - name: Install dependencies
+        run: |
+          micromamba activate a2
+          python -m pip install --upgrade pip
+          uv pip install torch --index-url https://download.pytorch.org/whl/cpu
+          uv pip install .[torchsim,ase,phonons] --group tests
+          uv pip install mace-torch>=0.3.3
+
+      - name: Test TorchSim
+        env:
+          MP_API_KEY: ${{ secrets.MP_API_KEY }}
+        run: |
+          micromamba activate a2
+          pytest -n auto tests/torchsim --cov=atomate2 --cov-report=xml
+
+      - name: Test TorchSim notebook
+        env:
+          MP_API_KEY: ${{ secrets.MP_API_KEY }}
+        run: |
+          micromamba activate a2
+          pytest -n auto --nbmake ./tutorials/torchsim_tutorial.ipynb
+
+      - uses: codecov/codecov-action@v1
+        if: github.repository == 'materialsproject/atomate2'
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: ./coverage.xml
+
   docs:
     runs-on: ubuntu-latest
 
@@ -325,7 +385,7 @@ jobs:
         run: sphinx-build docs docs_build
 
   automerge:
-    needs: [docs, lint, test-force-field-notebook, test-non-ase, test-notebooks-and-ase, test-openff]
+    needs: [docs, lint, test-force-field-notebook, test-non-ase, test-notebooks-and-ase, test-openff, test-torchsim]
     runs-on: ubuntu-latest
 
     permissions:

.gitignore

@@ -29,6 +29,7 @@ develop-eggs
 .installed.cfg
 lib
 lib64
+uv.lock
 
 # Installer logs
 pip-log.txt

pyproject.toml

@@ -65,6 +65,9 @@ forcefields = [
     "sevenn>=0.9.3",
     "deepmd-kit>=2.1.4",
 ]
+torchsim = [
+    "torch-sim-atomistic==0.5.0; python_version >= '3.12'"
+]
 approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"]
 ase = ["ase>=3.26.0"]
 ase-ext = ["tblite>=0.3.0; platform_system=='Linux'"]

src/atomate2/common/flows/phonons.py

@@ -29,7 +29,7 @@
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
     from atomate2.vasp.jobs.base import BaseVaspMaker
 
-SUPPORTED_CODES = frozenset(("vasp", "aims", "forcefields", "ase"))
+SUPPORTED_CODES = frozenset(("vasp", "aims", "forcefields", "ase", "torchsim"))
 
 
 @dataclass

src/atomate2/common/schemas/phonons.py

@@ -53,7 +53,7 @@ def get_factor(code: str) -> float:
     ValueError
         If code is not defined
     """
-    if code in ["ase", "forcefields", "vasp"]:
+    if code in ["ase", "forcefields", "vasp", "torchsim"]:
         return VaspToTHz
     if code == "aims":
         return omegaToTHz  # Based on CODATA 2002

src/atomate2/torchsim/__init__.py

@@ -0,0 +1,9 @@
+"""TorchSim module for atomate2."""
+
+from atomate2.torchsim.core import (
+    TorchSimIntegrateMaker,
+    TorchSimOptimizeMaker,
+    TorchSimStaticMaker,
+)
+
+__all__ = ["TorchSimIntegrateMaker", "TorchSimOptimizeMaker", "TorchSimStaticMaker"]