Add MatPES forcefields (#1158)

* add matpes models / bump matgl

* ensure MD ensemble is correctly rehydrated during fireworks style runs

* pin mdanalysis, add matpes tests

* modify schemas slightly to accommodate single ase calculator addition; pin openmdanalysis version

* precommit

* resolve torch conflict

* pin chgnet version to slightly older release for torch interop

* fix ase utils test and pin typing to see if that fixes docs

* pin upper bound on setuptools from major breaking changes

* pin upper bound on setuptools from major breaking changes

* pin upper bound on setuptools from major breaking changes

* keep experimenting to see how to mitigate setuptools breakage

* bump ase version for newer md calc

* fix num threads in forcefields tests

* modify test tolerance

* change ref data for matpes pbe stress

Author: esoteric-ephemera

.github/workflows/docs.yml

@@ -33,6 +33,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
+          pip install "setuptools==77.0.3"
           pip install .[strict,docs]
 
       - name: Copy tutorials

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools >= 42", "versioningit >= 1,< 4", "wheel"]
+requires = ["setuptools >= 42, <78", "versioningit >= 1,< 4", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -55,7 +55,7 @@ forcefields = [
     "calorine>=3.0",
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
-    "matgl>=1.1.3",
+    "matgl>=1.2.1",
     "torchdata<=0.7.1",
     # quippy-ase support for py3.12 tracked in https://github.com/libAtoms/QUIP/issues/645
     "quippy-ase>=0.9.14; python_version < '3.12'",
@@ -95,7 +95,7 @@ tests = [
 ]
 strict = [
     "PyYAML==6.0.2",
-    "ase==3.23.0",
+    "ase==3.24.0",
     "cclib==1.8.1",
     "click==8.1.7",
     "custodian==2024.10.16",
@@ -119,6 +119,7 @@ strict = [
     "typing-extensions==4.12.2",
 ]
 strict-openff = [
+    "mdanalysis==2.7.0",
     "monty==2025.3.3",
     "openmm-mdanalysis-reporter==0.1.0",
     "openmm==8.1.1",
@@ -127,12 +128,12 @@ strict-openff = [
 ]
 strict-forcefields = [
     "calorine==3.0",
-    "chgnet==0.4.0",
+    "chgnet==0.3.8",
     "mace-torch==0.3.10",
-    "matgl==1.1.3",
+    "matgl==1.2.1",
     "quippy-ase==0.9.14; python_version < '3.12'",
     "sevenn==0.10.3",
-    "torch==2.6.0",
+    "torch==2.2.0",
     "torchdata==0.7.1",                            # TODO: remove when issue fixed
 ]
 

src/atomate2/ase/jobs.py

@@ -3,14 +3,17 @@
 from __future__ import annotations
 
 import logging
+import time
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from ase.io import Trajectory as AseTrajectory
 from emmet.core.vasp.calculation import StoreTrajectoryOption
 from jobflow import Maker, job
+from pymatgen.core import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory as PmgTrajectory
+from pymatgen.io.ase import AseAtomsAdaptor
 
 from atomate2.ase.schemas import AseResult, AseTaskDoc
 from atomate2.ase.utils import AseRelaxer
@@ -21,7 +24,6 @@
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core import Molecule, Structure
 
     from atomate2.ase.schemas import AseMoleculeTaskDoc, AseStructureTaskDoc
 
@@ -33,11 +35,33 @@ class AseMaker(Maker, metaclass=ABCMeta):
     """
     Define basic template of ASE-based jobs.
 
-    This class defines two functions relevant attributes
-    for the ASE TaskDoc schemas, as well as two methods
-    that must be implemented in subclasses:
-        1. `calculator`: the ASE .Calculator object
-        2. `run_ase`: which actually makes the call to ASE.
+    This class defines relevant attributes for the ASE TaskDoc
+    schemas, and one method that must be implemented in subclasses:
+    `calculator`: the ASE .Calculator object
+
+    The intent of this class is twofold: if users wish to have a
+    high-throughput way to access a calculator, they need only
+    subclass this class with a calculator defined, e.g., the following
+    is sufficient to define an EMT static calculator with basic I/O:
+
+    ```python
+    from ase.calculators.emt import EMT
+
+
+    @dataclass
+    class EMTStaticMaker(AseMaker):
+        name: str = "EMT static maker"
+
+        @property
+        def calculator(self):
+            return EMT()
+    ```
+
+    Note that the user should adapt `run_ase`, which is not a job
+    and makes a call to ASE, and `make`, which is a job, to their uses.
+
+    `run_ase` should return an `AseResult` which has basic calculation info.
+    `make` should return a pydantic-based document model with more details.
 
     Parameters
     ----------
@@ -69,17 +93,35 @@ class AseMaker(Maker, metaclass=ABCMeta):
     store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO
     tags: list[str] | None = None
 
-    @abstractmethod
+    @job(data=_ASE_DATA_OBJECTS)
+    def make(
+        self,
+        mol_or_struct: Molecule | Structure,
+        prev_dir: str | Path | None = None,
+    ) -> AseStructureTaskDoc | AseMoleculeTaskDoc:
+        """
+        Run ASE as job, can be re-implemented in subclasses.
+
+        Parameters
+        ----------
+        mol_or_struct: .Molecule or .Structure
+            pymatgen molecule or structure
+        prev_dir : str or Path or None
+            A previous calculation directory to copy output files from. Unused, just
+                added to match the method signature of other makers.
+        """
+        return AseTaskDoc.to_mol_or_struct_metadata_doc(
+            getattr(self.calculator, "name", type(self.calculator).__name__),
+            self.run_ase(mol_or_struct, prev_dir=prev_dir),
+        )
+
     def run_ase(
         self,
         mol_or_struct: Structure | Molecule,
         prev_dir: str | Path | None = None,
     ) -> AseResult:
         """
-        Run ASE, method to be implemented in subclasses.
-
-        This method exists to permit subclasses to redefine `make`
-        for different output schemas.
+        Run ASE, can be re-implemented in subclasses.
 
         Parameters
         ----------
@@ -89,7 +131,20 @@ def run_ase(
             A previous calculation directory to copy output files from. Unused, just
                 added to match the method signature of other makers.
         """
-        raise NotImplementedError
+        is_mol = isinstance(mol_or_struct, Molecule)
+        adaptor = AseAtomsAdaptor()
+        atoms = adaptor.get_atoms(mol_or_struct)
+        atoms.calc = self.calculator
+        t_i = time.perf_counter()
+        final_energy = atoms.get_potential_energy()
+        t_f = time.perf_counter()
+        return AseResult(
+            final_mol_or_struct=getattr(
+                adaptor, f"get_{'molecule' if is_mol else 'structure'}"
+            )(atoms),
+            final_energy=final_energy,
+            elapsed_time=t_f - t_i,
+        )
 
     @property
     @abstractmethod

src/atomate2/ase/md.py

@@ -180,7 +180,7 @@ class AseMDMaker(AseMaker, metaclass=ABCMeta):
     def __post_init__(self) -> None:
         """Ensure that ensemble is an enum."""
         if isinstance(self.ensemble, str):
-            self.ensemble = MDEnsemble(self.ensemble)
+            self.ensemble = MDEnsemble(self.ensemble.split("MDEnsemble.")[-1])
 
     @staticmethod
     def _interpolate_quantity(values: Sequence | np.ndarray, n_pts: int) -> np.ndarray: