Replace ase ExpCellFilter with FrechetCellFilter in Relaxer (#636)

* replace ExpCellFilter with FrechetCellFilter in Relaxer

* update test_(mace|chgnet)_relax_maker

* install ase from git since no release with FrechetCellFilter available

* fix typo and doc string format to numpy

* bump chgnet to v0.3.3 in strict deps

* fix failing test_elastic_wf was using m3gnet which predicts negative bulk modulus causing pydantic validation error, same with chgnet, MACE works fine

pyproject.toml

@@ -45,6 +45,7 @@ phonons = ["phonopy>=1.10.8", "seekpath"]
 lobster = ["ijson>=3.2.2", "lobsterpy>=0.3.2"]
 defects = ["dscribe>=1.2.0", "pymatgen-analysis-defects>=2022.11.30"]
 forcefields = [
+    "ase@git+https://gitlab.com/ase/ase",
     "chgnet>=0.2.2",
     "mace@git+https://github.com/ACEsuit/mace@develop",
     "matgl>=0.9.0",
@@ -66,8 +67,9 @@ dev = ["pre-commit>=2.12.1"]
 tests = ["FireWorks==2.0.3", "pytest-cov==4.1.0", "pytest==7.4.3"]
 strict = [
     "PyYAML==6.0.1",
+    "ase@git+https://gitlab.com/ase/ase@2bab58f4e",
     "cclib==1.8",
-    "chgnet==0.3.2",
+    "chgnet==0.3.3",
     "click==8.1.7",
     "custodian==2023.10.9",
     "dscribe==2.1.0",

src/atomate2/forcefields/utils.py

@@ -15,7 +15,7 @@
 import sys
 from typing import TYPE_CHECKING
 
-from ase.constraints import ExpCellFilter
+from ase.filters import FrechetCellFilter
 from ase.optimize.bfgs import BFGS
 from ase.optimize.bfgslinesearch import BFGSLineSearch
 from ase.optimize.fire import FIRE
@@ -163,13 +163,20 @@ def relax(
 
         Parameters
         ----------
-        atoms (Atoms): the atoms for relaxation
-        fmax (float): total force tolerance for relaxation convergence.
-        steps (int): max number of steps for relaxation
-        traj_file (str): the trajectory file for saving
-        interval (int): the step interval for saving the trajectories
-        verbose (bool): if True, screenoutput will be shown.
-        kwargs: further kwargs.
+        atoms : Atoms
+            The atoms for relaxation.
+        fmax : float
+            Total force tolerance for relaxation convergence.
+        steps : int
+            Max number of steps for relaxation.
+        traj_file : str
+            The trajectory file for saving.
+        interval : int
+            The step interval for saving the trajectories.
+        verbose : bool
+            If True, screen output will be shown.
+        **kwargs
+            Further kwargs.
 
         Returns
         -------
@@ -182,17 +189,15 @@ def relax(
         with contextlib.redirect_stdout(stream):
             obs = TrajectoryObserver(atoms)
             if self.relax_cell:
-                atoms = ExpCellFilter(atoms)
+                atoms = FrechetCellFilter(atoms)
             optimizer = self.opt_class(atoms, **kwargs)
             optimizer.attach(obs, interval=interval)
             optimizer.run(fmax=fmax, steps=steps)
             obs()
         if traj_file is not None:
             obs.save(traj_file)
-        if isinstance(atoms, ExpCellFilter):
+        if isinstance(atoms, FrechetCellFilter):
             atoms = atoms.atoms
 
-        return {
-            "final_structure": self.ase_adaptor.get_structure(atoms),
-            "trajectory": obs,
-        }
+        struct = self.ase_adaptor.get_structure(atoms)
+        return {"final_structure": struct, "trajectory": obs}

tests/forcefields/flows/test_elastic.py

@@ -3,17 +3,17 @@
 
 from atomate2.common.schemas.elastic import ElasticDocument
 from atomate2.forcefields.flows.elastic import ElasticMaker
-from atomate2.forcefields.jobs import M3GNetRelaxMaker
+from atomate2.forcefields.jobs import MACERelaxMaker
 
 
-def test_elastic_wf_with_m3gnet(clean_dir, si_structure):
+def test_elastic_wf_with_mace(clean_dir, si_structure):
     si_prim = SpacegroupAnalyzer(si_structure).get_primitive_standard_structure()
 
     job = ElasticMaker(
-        bulk_relax_maker=M3GNetRelaxMaker(
+        bulk_relax_maker=MACERelaxMaker(
             relax_cell=True, relax_kwargs={"fmax": 0.00001}
         ),
-        elastic_relax_maker=M3GNetRelaxMaker(
+        elastic_relax_maker=MACERelaxMaker(
             relax_cell=False, relax_kwargs={"fmax": 0.00001}
         ),
     ).make(si_prim)

tests/forcefields/test_jobs.py

@@ -51,7 +51,7 @@ def test_chgnet_relax_maker(si_structure, relax_cell: bool):
     assert output1.output.n_steps >= 12
     if relax_cell:
         assert output1.output.energy == approx(-10.62461, abs=1e-2)
-        assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.002964, rel=1e-1)
+        assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.00251674, rel=1e-1)
     else:
         assert output1.output.energy == approx(-10.6274, rel=1e-2)
         assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.00303572, rel=1e-2)
@@ -142,8 +142,8 @@ def test_mace_relax_maker(si_structure, test_dir, model, relax_cell: bool):
     output1 = responses[job.uuid][1].output
     assert isinstance(output1, ForceFieldTaskDocument)
     if relax_cell:
-        assert output1.output.energy == approx(-0.071052, rel=1e-1)
-        assert output1.output.n_steps >= 5
+        assert output1.output.energy == approx(-0.0526856, rel=1e-1)
+        assert output1.output.n_steps >= 4
     else:
         assert output1.output.energy == approx(-0.051912, rel=1e-4)
         assert output1.output.n_steps == 4