ML forcefields bug fixes (#1220)

* ff bug fixes

* make filter_kwargs an optional kwarg for AseRelaxer

* temporarily disable ASE energy checks - runs are inconsistent even with same initial input

* skip optimizer instantiation when steps <= 1 in ase/mlff flows

* ensure lj test runs one step

* bump ase bc of inconsistent behavior between 3.24 and 3.25

* bump emmet-core

Author: esoteric-ephemera

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "PyYAML",
     "click",
     "custodian>=2024.4.18",
-    "emmet-core>=0.84.3rc3",
+    "emmet-core>=0.84.8",
     "jobflow>=0.1.11",
     "monty>=2024.12.10",
     "numpy",
@@ -51,7 +51,7 @@ defects = [
     "python-ulid>=2.7",
 ]
 forcefields = [
-    "ase>=3.23.0",
+    "ase>=3.25.0",
     "calorine>=3.0",
     "chgnet>=0.2.2",
     "mace-torch>=0.3.3",
@@ -62,7 +62,7 @@ forcefields = [
     "sevenn>=0.9.3",
     "torchdata<=0.7.1",                            # TODO: remove when issue fixed
 ]
-ase = ["ase>=3.23.0"]
+ase = ["ase>=3.25.0"]
 ase-ext = ["tblite>=0.3.0; platform_system=='Linux'"]
 openmm = [
     "mdanalysis>=2.8.0",
@@ -94,12 +94,12 @@ tests = [
 ]
 strict = [
     "PyYAML==6.0.2",
-    "ase==3.24.0",
+    "ase==3.25.0",
     "cclib==1.8.1",
     "click==8.2.0",
     "custodian==2025.4.14",
     "dscribe==2.1.1",
-    "emmet-core==0.84.5",
+    "emmet-core==0.84.8",
     "ijson==3.3.0",
     "jobflow==0.1.19",
     "lobsterpy==0.4.9",

src/atomate2/ase/jobs.py

@@ -254,7 +254,7 @@ def run_ase(
         if self.steps < 0:
             logger.warning(
                 "WARNING: A negative number of steps is not possible. "
-                "Behavior may vary..."
+                "Defaulting to a static calculation."
             )
 
         relaxer = AseRelaxer(

src/atomate2/ase/utils.py

@@ -16,7 +16,7 @@
 from ase.constraints import FixSymmetry
 from ase.filters import FrechetCellFilter
 from ase.io import Trajectory as AseTrajectory
-from ase.io import write
+from ase.io import write as ase_write
 from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
 from monty.serialization import dumpfn
@@ -329,15 +329,18 @@ def relax(
         fmax: float = 0.1,
         steps: int = 500,
         traj_file: str = None,
-        final_atoms_object_file: str = "final_atoms_object.xyz",
+        final_atoms_object_file: str | os.PathLike[str] = "final_atoms_object.xyz",
         interval: int = 1,
         verbose: bool = False,
         cell_filter: Filter = FrechetCellFilter,
+        filter_kwargs: dict | None = None,
         **kwargs,
     ) -> AseResult:
         """
         Relax the molecule or structure.
 
+        If steps <= 1, this will perform a single-point calculation.
+
         Parameters
         ----------
         atoms : ASE Atoms, pymatgen Structure, or pymatgen Molecule
@@ -348,7 +351,7 @@ def relax(
             Max number of steps for relaxation.
         traj_file : str
             The trajectory file for saving.
-        final_atoms_object_file: str
+        final_atoms_object_file: str | os.PathLike
             The final atoms object file for saving.
         interval : int
             The step interval for saving the trajectories.
@@ -374,14 +377,16 @@ def relax(
         atoms.calc = self.calculator
         with contextlib.redirect_stdout(sys.stdout if verbose else io.StringIO()):
             obs = TrajectoryObserver(atoms)
-            if self.relax_cell and (not is_mol):
-                atoms = cell_filter(atoms)
-            optimizer = self.opt_class(atoms, **kwargs)
-            optimizer.attach(obs, interval=interval)
             t_i = time.perf_counter()
-            optimizer.run(fmax=fmax, steps=steps)
-            t_f = time.perf_counter()
+            if steps > 1:
+                if self.relax_cell and (not is_mol):
+                    atoms = cell_filter(atoms, **(filter_kwargs or {}))
+                optimizer = self.opt_class(atoms, **kwargs)
+                optimizer.attach(obs, interval=interval)
+                optimizer.run(fmax=fmax, steps=steps)
             obs()
+            t_f = time.perf_counter()
+
         if traj_file is not None:
             obs.save(traj_file)
         if isinstance(atoms, cell_filter):
@@ -402,7 +407,9 @@ def relax(
                 write_atoms.calc = self.calculator
             else:
                 write_atoms = atoms
-            write(final_atoms_object_file, write_atoms, format="extxyz", append=True)
+            ase_write(
+                final_atoms_object_file, write_atoms, format="extxyz", append=True
+            )
 
         return AseResult(
             final_mol_or_struct=struct,

src/atomate2/forcefields/jobs.py

@@ -38,8 +38,16 @@
     MLFF.MACE_MP_0: {"model": "medium"},
     MLFF.MACE_MPA_0: {"model": "medium-mpa-0"},
     MLFF.MACE_MP_0B3: {"model": "medium-0b3"},
-    MLFF.MATPES_PBE: {"architecture": "TensorNet", "version": "2025.1"},
-    MLFF.MATPES_R2SCAN: {"architecture": "TensorNet", "version": "2025.1"},
+    MLFF.MATPES_PBE: {
+        "architecture": "TensorNet",
+        "version": "2025.1",
+        "stress_unit": "eV/A3",
+    },
+    MLFF.MATPES_R2SCAN: {
+        "architecture": "TensorNet",
+        "version": "2025.1",
+        "stress_unit": "eV/A3",
+    },
 }
 
 

src/atomate2/forcefields/utils.py

@@ -17,7 +17,9 @@
     from ase.calculators.calculator import Calculator
 
 
-def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | None:
+def ase_calculator(
+    calculator_meta: str | MLFF | dict, **kwargs: Any
+) -> Calculator | None:
     """
     Create an ASE calculator from a given set of metadata.
 
@@ -42,8 +44,10 @@ def ase_calculator(calculator_meta: str | dict, **kwargs: Any) -> Calculator | N
     """
     calculator = None
 
-    if isinstance(calculator_meta, str | MLFF) and calculator_meta in map(str, MLFF):
-        calculator_name = MLFF[calculator_meta.split("MLFF.")[-1]]
+    if (
+        isinstance(calculator_meta, str) and calculator_meta in map(str, MLFF)
+    ) or isinstance(calculator_meta, MLFF):
+        calculator_name = MLFF(calculator_meta)
 
         if calculator_name == MLFF.CHGNet:
             from chgnet.model.dynamics import CHGNetCalculator

tests/ase/test_jobs.py

@@ -12,6 +12,7 @@
 
 from atomate2.ase.jobs import (
     AseMaker,
+    AseRelaxMaker,
     GFNxTBRelaxMaker,
     GFNxTBStaticMaker,
     LennardJonesRelaxMaker,
@@ -34,6 +35,15 @@ def calculator(self):
         return EMT()
 
 
+@dataclass
+class EMTRelaxMaker(AseRelaxMaker):
+    name: str = "EMT relax maker"
+
+    @property
+    def calculator(self):
+        return EMT()
+
+
 def test_base_maker(test_dir):
     structure = Structure.from_file(test_dir / "structures" / "Al2Au.cif")
     ase_res = EMTStaticMaker().run_ase(structure)
@@ -47,6 +57,24 @@ def test_base_maker(test_dir):
     assert isinstance(output, AseStructureTaskDoc)
 
 
+@pytest.mark.parametrize("constant_vol", [True, False])
+def test_filters_and_kwargs(test_dir, constant_vol):
+    structure = Structure.from_file(test_dir / "structures" / "Al2Au.cif")
+    structure = structure.scale_lattice(1.1 * structure.volume)
+
+    job = EMTRelaxMaker(
+        relax_kwargs={"filter_kwargs": {"constant_volume": constant_vol}}
+    ).make(structure)
+    resp = run_locally(job)
+    output = resp[job.uuid][1].output
+
+    assert len(output.output.ionic_steps) > 1
+    if constant_vol:
+        assert output.structure.volume == pytest.approx(structure.volume)
+    else:
+        assert abs(output.structure.volume - structure.volume) > 1e-2
+
+
 def test_lennard_jones_relax_maker(lj_fcc_ne_pars, fcc_ne_structure):
     job = LennardJonesRelaxMaker(
         calculator_kwargs=lj_fcc_ne_pars, relax_kwargs={"fmax": 0.001}
@@ -70,6 +98,7 @@ def test_lennard_jones_static_maker(lj_fcc_ne_pars, fcc_ne_structure):
     response = run_locally(job)
     output = response[job.uuid][1].output
 
+    assert len(output.output.ionic_steps) == 1
     assert output.output.energy == pytest.approx(-0.0179726955438795)
     assert output.structure.volume == pytest.approx(24.334)
     assert isinstance(output, AseStructureTaskDoc)

tests/ase/test_md.py

@@ -27,10 +27,12 @@
 
 @pytest.mark.parametrize("calculator_name", list(name_to_maker))
 def test_ase_nvt_maker(calculator_name, lj_fcc_ne_pars, fcc_ne_structure, clean_dir):
-    reference_energies = {
-        "LJ": -0.0179726955438795,
-        "GFN-xTB": -160.93692979071128,
-    }
+    # Langevin thermostat no longer works with single atom structures in ase>3.24.x
+    structure = fcc_ne_structure * (2, 2, 2)
+    # reference_energies_per_atom = {
+    #     "LJ": -0.0179726955438795,
+    #     "GFN-xTB": -160.93692979071128,
+    # }
 
     md_job = name_to_maker[calculator_name](
         calculator_kwargs=lj_fcc_ne_pars if calculator_name == "LJ" else {},
@@ -40,24 +42,27 @@ def test_ase_nvt_maker(calculator_name, lj_fcc_ne_pars, fcc_ne_structure, clean_
         n_steps=100,
         tags=["test"],
         store_trajectory="partial",
-    ).make(fcc_ne_structure)
+    ).make(structure)
 
     response = run_locally(md_job)
     output = response[md_job.uuid][1].output
 
     assert isinstance(output, AseStructureTaskDoc)
     assert output.tags == ["test"]
-    assert output.output.energy_per_atom == pytest.approx(
-        reference_energies[calculator_name]
-    )
-    assert output.structure.volume == pytest.approx(fcc_ne_structure.volume)
+
+    # TODO: ASE MD runs very inconsistent
+    # assert output.output.energy_per_atom == pytest.approx(
+    #     reference_energies_per_atom[calculator_name],
+    #     abs=1e-3,
+    # )
+    assert output.structure.volume == pytest.approx(structure.volume)
 
 
 @pytest.mark.parametrize("calculator_name", ["LJ"])
 def test_ase_npt_maker(calculator_name, lj_fcc_ne_pars, fcc_ne_structure, tmp_dir):
     os.environ["OMP_NUM_THREADS"] = "1"
 
-    reference_energies = {
+    reference_energies_per_atom = {
         "LJ": 0.01705592581943574,
     }
 
@@ -82,10 +87,10 @@ def test_ase_npt_maker(calculator_name, lj_fcc_ne_pars, fcc_ne_structure, tmp_di
 
     assert isinstance(output, AseStructureTaskDoc)
     assert output.output.energy_per_atom == pytest.approx(
-        reference_energies[calculator_name]
+        reference_energies_per_atom[calculator_name]
     )
 
     # TODO: improve XDATCAR parsing test when class is fixed in pmg
     assert os.path.isfile("XDATCAR")
 
-    assert len(output.objects["trajectory"]) == n_steps
+    assert len(output.objects["trajectory"]) == n_steps + 1

tests/ase/test_utils.py

@@ -122,7 +122,10 @@ def test_fix_symmetry(fix_symmetry):
     atoms_al = atoms_al * (2, 2, 2)
     atoms_al.positions[0, 0] += 1e-7
     symmetry_init = check_symmetry(atoms_al, 1e-6)
-    final_struct: Structure = relaxer.relax(atoms=atoms_al, steps=1).final_mol_or_struct
+    final_struct: Structure = relaxer.relax(
+        atoms=atoms_al,
+        steps=2,
+    ).final_mol_or_struct
     symmetry_final = check_symmetry(final_struct.to_ase_atoms(), 1e-6)
     if fix_symmetry:
         assert symmetry_init["number"] == symmetry_final["number"] == 229

tests/conftest.py

@@ -50,7 +50,7 @@ def clean_dir(debug_mode):
         shutil.rmtree(new_path)
 
 
-@pytest.fixture
+@pytest.fixture(autouse=True)
 def tmp_dir():
     """Same as clean_dir but is fresh for every test"""