Generalize forcefields for generic ASE calculator support (#940)

- Add ASE base calculator classes from forcefield library
- Add tests for ASE calculators, refactor forcefields
- Add static makers
- Add tblite dependence to ASE in toml
- Migrate forcefield MD base class to ASE
- Add barebones tests for ASE MD jobs
- Generalize to allow molecules in ASE calculations
- Make tblite dependence optional
- Move forcefield strict version pins to separate toml block
- Refactor naming, enumify and lazy load MD dynamics modules
- Make Ase{Structure,Molecule}TaskDoc construction a single function
- Change ASE and forcefield result from dict to BaseModel
- Ensure ASE and forcefield result docs are subscriptable
- Add "ionic_steps" to DATA_OBJECTS
- Add 'elapsed_time' and 'tags' to ASE schemas
- Add timing to ASE runs
- Define basic ASE job schema; remove task_doc_kwargs from ASE in favor of explicit kwargs
- Revise ionic_step storage (only for MD)
- Add option to pass forcefield name as str without MLFF prefix or as enum
- Add deprecation warnings to all predefined MLFF makers (static, relax, MD)
- Correct typing of atoms in AseRelaxer
- Add tblite to optional ase-ext and strict requirements
- Remove dependence on deprecated forcefield makers
- Make default calc kwargs and ensure these are loaded when not specified by user
- Remove Frechet cell filter safety checks
- Make MD ensembles an enum
- Move ASE tests to separate test run
- Add ASE to phonon supported codes, enforce string literal in BasePhononMaker
- Fix potential Pydantic validation error if Ase(MD|Relax)Maker.calculator has no name

---------

Co-authored-by: orionarcher <ocohen@radical-ai.com>
Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

.github/workflows/testing.yml

@@ -62,7 +62,7 @@ jobs:
           python -m pip install --upgrade pip
           mkdir -p ~/.abinit/pseudos
           cp -r tests/test_data/abinit/pseudos/ONCVPSP-PBE-SR-PDv0.4 ~/.abinit/pseudos
-          uv pip install .[strict,tests,abinit]
+          uv pip install .[strict,strict-forcefields,tests,abinit]
           uv pip install torch-runstats
           uv pip install --no-deps nequip==0.5.6
 
@@ -72,17 +72,78 @@ jobs:
           micromamba activate a2
           uv pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
 
+      - name: Test
+        env:
+          MP_API_KEY: ${{ secrets.MP_API_KEY }}
+        run: |
+          micromamba activate a2
+          pytest --ignore=tests/ase --cov=atomate2 --cov-report=xml
+
+      - uses: codecov/codecov-action@v1
+        if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          file: ./coverage.xml
+
+  test-notebooks-and-ase:
+    # It seems like anything torch-dependent and tblite can't be installed in the same environment
+    # without the tblite tests failing in CI, see, e.g.:
+    # https://github.com/tblite/tblite/issues/116
+    # Outside of CI, having torch installed but not loaded seems not to affect tblite
+    # Set off ASE tests here, where tblite-dependent tests live
+    services:
+      local_mongodb:
+        image: mongo:4.0
+        ports:
+          - 27017:27017
+
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0} # enables conda/mamba env activation by reading bash profile
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - name: Check out repo
+        uses: actions/checkout@v4
+
+      - name: Set up micromamba
+        uses: mamba-org/setup-micromamba@main
+
+      - name: Create mamba environment
+        run: |
+          micromamba create -n a2 python=${{ matrix.python-version }} --yes
+
+      - name: Install uv
+        run: micromamba run -n a2 pip install uv
+
+      - name: Install conda dependencies
+        run: |
+          micromamba install -n a2 -c conda-forge enumlib packmol bader openbabel openff-toolkit==0.16.2 openff-interchange==0.3.22 --yes
+
+      - name: Install dependencies
+        run: |
+          micromamba activate a2
+          python -m pip install --upgrade pip
+          uv pip install .[strict,tests]
+
+      - 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 }}'
+
       - name: Test Notebooks
-        run: | 
+        run: |
           micromamba activate a2
           pytest --nbmake ./tutorials --ignore=./tutorials/openmm_tutorial.ipynb
 
-      - name: Test
+      - name: Test ASE
         env:
           MP_API_KEY: ${{ secrets.MP_API_KEY }}
         run: |
           micromamba activate a2
-          pytest --cov=atomate2 --cov-report=xml
+          pytest --cov=atomate2 --cov-report=xml tests/ase
 
       - uses: codecov/codecov-action@v1
         if: matrix.python-version == '3.10' && github.repository == 'materialsproject/atomate2'
@@ -105,7 +166,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install .[strict,docs]
+          pip install .[strict,strict-forcefields,docs]
 
       - name: Build
         run: sphinx-build docs docs_build

pyproject.toml

@@ -59,6 +59,12 @@ forcefields = [
     "quippy-ase>=0.9.14",
     "sevenn>=0.9.3",
 ]
+ase = [
+    "ase>=3.23.0",
+]
+ase-ext = [
+    "tblite>=0.3.0",
+]
 openmm = [
     "mdanalysis>=2.7.0",
     "openmm>=8.1.0",
@@ -87,19 +93,14 @@ tests = [
 strict = [
     "PyYAML==6.0.2",
     "ase==3.23.0",
-    "calorine==2.2.1",
     "cclib==1.8.1",
-    "chgnet==0.3.8",
     "click==8.1.7",
     "custodian==2024.6.24",
     "dscribe==2.1.1",
     "emmet-core==0.84.2",
     "ijson==3.3.0",
     "jobflow==0.1.18",
     "lobsterpy==0.4.5",
-    "mace-torch>=0.3.3",
-    "torchdata<=0.7.1", # TODO: remove when issue fixed
-    "matgl==1.1.3",
     "monty==2024.7.30",
     "mp-api==0.42.2",
     "numpy",
@@ -109,15 +110,23 @@ strict = [
     "pymatgen-analysis-defects==2024.7.19",
     "pymatgen==2024.6.10",
     "python-ulid==2.7.0",
-    "quippy-ase==0.9.14",
     "seekpath==2.1.0",
-    "sevenn==0.9.3.post1",
-    "torch==2.2.1",
+    "tblite==0.3.0",
     "typing-extensions==4.12.2",
     "mdanalysis==2.7.0",
     "openmm==8.1.1",
     "openmm-mdanalysis-reporter==0.1.0",
 ]
+strict-forcefields = [
+    "calorine==2.2.1",
+    "chgnet==0.3.8",
+    "mace-torch>=0.3.3",
+    "torchdata==0.7.1", # TODO: remove when issue fixed
+    "matgl==1.1.3",
+    "quippy-ase==0.9.14",
+    "sevenn==0.9.3.post1",
+    "torch==2.2.1",
+]
 
 [project.scripts]
 atm = "atomate2.cli:cli"

src/atomate2/aims/flows/phonons.py

@@ -14,6 +14,8 @@
 from atomate2.common.flows.phonons import BasePhononMaker
 
 if TYPE_CHECKING:
+    from typing import Literal
+
     from atomate2.aims.jobs.base import BaseAimsMaker
 
 
@@ -118,7 +120,7 @@ class PhononMaker(BasePhononMaker):
     min_length: float | None = 20.0
     prefer_90_degrees: bool = True
     get_supercell_size_kwargs: dict = field(default_factory=dict)
-    use_symmetrized_structure: str | None = None
+    use_symmetrized_structure: Literal["primitive", "conventional"] | None = None
     create_thermal_displacements: bool = True
     generate_frequencies_eigenvectors_kwargs: dict = field(default_factory=dict)
     kpath_scheme: str = "seekpath"

src/atomate2/ase/__init__.py

@@ -0,0 +1 @@
+"""Module for Atomic Simulation Environment workflows."""