Add updated MP workflows + cleanup (#1139)

* Update pyproject.toml

* Update pyproject.toml

* test a new pymatgen version

* Update pyproject.toml

* Update test_utils.py

* redraft mp flows

* split off openff tests

* remove deprecated flows (completely forgot at the start of the year

* fix pin in strict openff wf

* correctly partition tests?

* revert openff_md test ref change

* openff test over py310-312

---------

Co-authored-by: J. George <[email protected]>

Author: esoteric-ephemera

.github/workflows/testing.yml

@@ -58,7 +58,7 @@ jobs:
 
       - 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
+          micromamba install -n a2 -c conda-forge enumlib packmol bader --yes
 
       - name: Install dependencies
         run: |
@@ -87,7 +87,7 @@ jobs:
         # However this `splitting-algorithm` means that tests cannot depend sensitively on the order they're executed in.
         run: |
           micromamba activate a2
-          pytest --splits 3 --group ${{ matrix.split }} --durations-path tests/.pytest-split-durations --splitting-algorithm least_duration --ignore=tests/ase --cov=atomate2 --cov-report=xml
+          pytest --splits 3 --group ${{ matrix.split }} --durations-path tests/.pytest-split-durations --splitting-algorithm least_duration --ignore=tests/ase --ignore=tests/openff_md --ignore=tests/openmm_md --cov=atomate2 --cov-report=xml
 
 
       - uses: codecov/codecov-action@v1
@@ -97,6 +97,62 @@ jobs:
           name: coverage${{ matrix.split }}
           file: ./coverage.xml
 
+  test-openff:
+    # 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} # enables conda/mamba env activation by reading bash profile
+    strategy:
+      matrix:
+        python-version: ["3.10","3.11","3.12"]
+
+    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-openff,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: |
+          micromamba activate a2
+          uv pip install --upgrade 'git+https://github.com/materialsproject/pymatgen@${{ github.event.client_payload.pymatgen_ref }}'
+
+      - name: Test split ${{ matrix.split }}
+        env:
+          MP_API_KEY: ${{ secrets.MP_API_KEY }}
+
+        run: |
+          micromamba activate a2
+          pytest tests/{openff_md,openmm_md}
+
   test-notebooks-and-ase:
     # prevent this action from running on forks
     if: github.repository == 'materialsproject/atomate2'

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
     "custodian>=2024.4.18",
     "emmet-core>=0.84.3rc3",
     "jobflow>=0.1.11",
-    "monty>=2024.7.30",
+    "monty>=2024.12.10",
     "numpy",
     "pydantic-settings>=2.0.3",
     "pydantic>=2.0.1",
@@ -104,23 +104,27 @@ strict = [
     "ijson==3.3.0",
     "jobflow==0.1.19",
     "lobsterpy==0.4.9",
-    "mdanalysis==2.7.0",
-    "monty==2024.10.21",
-    "mp-api==0.43.0",
+    "monty==2025.1.9",
+    "mp-api==0.45.3",
     "numpy",
-    "openmm-mdanalysis-reporter==0.1.0",
-    "openmm==8.1.1",
     "phonopy==2.30.1",
     "pydantic-settings==2.7.0",
     "pydantic==2.9.2",
     "pymatgen-analysis-defects==2024.10.22",
-    "pymatgen==2024.11.13",
+    "pymatgen==2025.2.18",
     "pymongo==4.10.1",
     "python-ulid==3.0.0",
     "seekpath==2.1.0",
     "tblite==0.3.0; python_version < '3.12'",
     "typing-extensions==4.12.2",
 ]
+strict-openff = [
+    "mdanalysis==2.7.0",
+    "monty==2024.12.10",
+    "openmm-mdanalysis-reporter==0.1.0",
+    "openmm==8.1.1",
+    "pymatgen==2024.11.13",
+]
 strict-forcefields = [
     "calorine==3.0",
     "chgnet==0.4.0",

src/atomate2/common/jobs/utils.py

@@ -2,14 +2,20 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 from jobflow import Response, job
+from monty.os.path import zpath
+from pymatgen.command_line.bader_caller import bader_analysis_from_path
 from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
 
 from atomate2 import SETTINGS
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from pathlib import Path
+
     from pymatgen.core import Structure
 
 
@@ -137,3 +143,62 @@ def retrieve_structure_from_materials_project(
         output=structure,
         stored_data={"task_id": task_id, "database_version": database_version},
     )
+
+
+@job
+def remove_workflow_files(
+    directories: Sequence[str], file_names: Sequence[str], allow_zpath: bool = True
+) -> list[str]:
+    """
+    Remove files from previous jobs.
+
+    For example, at the end of an MP flow, WAVECAR files are generated
+    that take up a lot of disk space.
+    This utility can automatically remove them following a workflow.
+
+    Parameters
+    ----------
+        makers : Sequence[Maker, Flow, Job]
+            Jobs in the flow on which to remove files.
+        file_names : Sequence[str]
+            The list of file names to remove, ex. ["WAVECAR"] rather than a full path
+        allow_zpath : bool = True
+            Whether to allow checking for gzipped output using `monty.os.zpath`
+
+    Returns
+    -------
+        list[str] : list of removed files
+    """
+    abs_paths = [os.path.abspath(dir_name.split(":")[-1]) for dir_name in directories]
+
+    removed_files = []
+    for job_dir in abs_paths:
+        for file in file_names:
+            file_name = os.path.join(job_dir, file)
+            if allow_zpath:
+                file_name = zpath(file_name)
+
+            if os.path.isfile(file_name):
+                removed_files.append(file_name)
+                os.remove(file_name)
+
+    return removed_files
+
+
+@job
+def bader_analysis(dir_name: str | Path, suffix: str | None = None) -> dict:
+    """Run Bader charge analysis as a job.
+
+    Parameters
+    ----------
+    dir_name : str or Path
+        The name of the directory to run Bader in.
+    suffix : str or None (default)
+        Suffixes of the files to filter by.
+
+    Returns
+    -------
+    dict of bader charge analysis which is JSONable.
+    """
+    dir_name = os.path.abspath(str(dir_name).split(":")[-1])
+    return bader_analysis_from_path(dir_name, suffix=suffix or "")

src/atomate2/common/utils.py

@@ -194,3 +194,18 @@ def parse_additional_json(dir_name: Path) -> dict[str, Any]:
         if key not in ("custodian", "transformations", "FW"):
             additional_json[key] = loadfn(filename, cls=None)
     return additional_json
+
+
+def _recursive_get_dir_names(jobs: list, dir_names: list) -> None:
+    """Recursively get all `output.dir_name` from a list of jobs.
+
+    Parameters
+    ----------
+    jobs : list of jobs, Makers, Flows, etc.
+    dir_names : a list to add the `dir_name`'s to.
+    """
+    for a_job in jobs:
+        if (sub_jobs := getattr(a_job, "jobs", None)) is not None:
+            _recursive_get_dir_names(sub_jobs, dir_names)
+        else:
+            dir_names.append(a_job.output.dir_name)

src/atomate2/vasp/flows/mp.py

@@ -15,10 +15,15 @@
 from jobflow import Flow, Maker
 from pymatgen.io.vasp.sets import LobsterSet
 
+from atomate2.common.jobs.utils import remove_workflow_files
+from atomate2.common.utils import _recursive_get_dir_names
 from atomate2.lobster.jobs import LobsterMaker
 from atomate2.vasp.flows.core import DoubleRelaxMaker
 from atomate2.vasp.flows.lobster import VaspLobsterMaker
 from atomate2.vasp.jobs.mp import (
+    MP24PreRelaxMaker,
+    MP24RelaxMaker,
+    MP24StaticMaker,
     MPGGARelaxMaker,
     MPGGAStaticMaker,
     MPMetaGGARelaxMaker,
@@ -29,6 +34,7 @@
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from pathlib import Path
 
     from pymatgen.core.structure import Structure
@@ -194,6 +200,95 @@ def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow
         return Flow(jobs=jobs, output=output, name=self.name)
 
 
+@dataclass
+class MP24DoubleRelaxMaker(DoubleRelaxMaker):
+    """MP24 PBEsol + r2SCAN double relaxation workflow.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    relax_maker1 : .BaseVaspMaker
+        Maker to generate the first relaxation.
+    relax_maker2 : .BaseVaspMaker
+        Maker to generate the second relaxation.
+    """
+
+    name: str = "MP24 double relax"
+    relax_maker1: Maker | None = field(default_factory=MP24PreRelaxMaker)
+    relax_maker2: Maker = field(
+        default_factory=lambda: MP24RelaxMaker(
+            copy_vasp_kwargs={"additional_vasp_files": ("WAVECAR", "CHGCAR")}
+        )
+    )
+
+
+@dataclass
+class MP24DoubleRelaxStaticMaker(Maker):
+    """MP24 workflow to relax a structure with r2SCAN.
+
+    Optionally, files can be automatically cleaned following completion
+    of the workflow. By default, WAVECAR files are removed.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    relax_maker : .BaseVaspMaker
+        Maker to generate the relaxation.
+    static_maker : .BaseVaspMaker
+        Maker to generate the static calculation before the relaxation.
+    clean_files : Sequence of str or None
+        If a list of strings, names of files to remove following the workflow.
+        By default, this removes the WAVECAR files (gzipped or not).
+    """
+
+    name: str = "MP24 r2SCAN workflow"
+    relax_maker: Maker = field(default_factory=MP24DoubleRelaxMaker)
+    static_maker: Maker = field(
+        default_factory=lambda: MP24StaticMaker(
+            copy_vasp_kwargs={"additional_vasp_files": ("WAVECAR", "CHGCAR")}
+        )
+    )
+    clean_files: Sequence[str] | None = ("WAVECAR",)
+
+    def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow:
+        """Relax a structure with r2SCAN.
+
+        Parameters
+        ----------
+        structure : .Structure
+            A pymatgen structure object.
+        prev_dir : str or Path or None
+            A previous VASP calculation directory to copy output files from.
+
+        Returns
+        -------
+        Flow
+            A flow containing the MP relaxation workflow.
+        """
+        relax_flow = self.relax_maker.make(structure=structure, prev_dir=prev_dir)
+
+        static_job = self.static_maker.make(
+            structure=relax_flow.output.structure, prev_dir=relax_flow.output.dir_name
+        )
+
+        jobs = [relax_flow, static_job]
+
+        self.clean_files = self.clean_files or []
+        if len(self.clean_files) > 0:
+            directories: list[str] = []
+            _recursive_get_dir_names(jobs, directories)
+            cleanup = remove_workflow_files(
+                directories=directories,
+                file_names=self.clean_files,
+                allow_zpath=True,
+            )
+            jobs += [cleanup]
+
+        return Flow(jobs=jobs, output=static_job.output, name=self.name)
+
+
 # update potcars to 54, use correct W potcar
 # use staticmaker for compatibility
 @dataclass