add forcefield approx/neb tests

Author: esoteric-ephemera

src/atomate2/ase/neb.py

@@ -25,7 +25,6 @@ class AseNebMaker(AseMaker):
 
     name: str = "ASE NEB maker"
     neb_kwargs: dict = field(default_factory=dict)
-    relax_cell: bool = True
     fix_symmetry: bool = False
     symprec: float | None = 1e-2
     steps: int = 500
@@ -55,17 +54,16 @@ def run_ase(
         return AseNebInterface(
             calculator=self.calculator,
             fix_symmetry=self.fix_symmetry,
-            relax_cell=self.relax_cell,
             symprec=self.symprec,
-            neb_kwargs=self.neb_kwargs,
-            **self.optimizer_kwargs,
         ).run_neb(
             images,
             steps=self.steps,
             traj_file=self.traj_file,
             traj_file_fmt=self.traj_file_fmt,
             interval=self.traj_interval,
             neb_doc_kwargs=self.neb_doc_kwargs,
+            neb_kwargs=self.neb_kwargs,
+            optimizer_kwargs=self.optimizer_kwargs,
             **self.relax_kwargs,
         )
 

src/atomate2/ase/utils.py

@@ -162,6 +162,8 @@ def save(
             self.to_pymatgen_trajectory(filename=filename, file_format=fmt)  # type: ignore[arg-type]
         elif fmt == "ase":
             self.to_ase_trajectory(filename=filename)
+        else:
+            raise ValueError(f"Unknown trajectory format {fmt}.")
 
     def to_ase_trajectory(
         self, filename: str | None = "atoms.traj"
@@ -443,18 +445,15 @@ def __init__(
         self,
         calculator: Calculator,
         optimizer: Optimizer | str = "FIRE",
-        relax_cell: bool = True,
         fix_symmetry: bool = False,
         symprec: float = 1e-2,
-        neb_kwargs: dict | None = None,
     ) -> None:
         """Initialize the interface.
 
         Parameters
         ----------
         calculator (ase Calculator): an ase calculator
         optimizer (str or ase Optimizer): the optimization algorithm.
-        relax_cell (bool): if True, cell parameters will be optimized.
         fix_symmetry (bool): if True, symmetry will be fixed during relaxation.
         symprec (float): Tolerance for symmetry finding in case of fix_symmetry.
         """
@@ -468,11 +467,9 @@ def __init__(
             optimizer_obj = optimizer
 
         self.opt_class: Optimizer = optimizer_obj
-        self.relax_cell = relax_cell
         self.ase_adaptor = AseAtomsAdaptor()
         self.fix_symmetry = fix_symmetry
         self.symprec = symprec
-        self.neb_kwargs = neb_kwargs or DEFAULT_NEB_KWARGS.copy()
 
     def run_neb(
         self,
@@ -484,7 +481,8 @@ def run_neb(
         interval: int = 1,
         verbose: bool = False,
         neb_doc_kwargs: dict | None = None,
-        **kwargs,
+        neb_kwargs: dict = DEFAULT_NEB_KWARGS,
+        optimizer_kwargs: dict | None = None,
     ) -> NebResult:
         """
         Perform NEB on a list of molecules or structures.
@@ -511,8 +509,10 @@ def run_neb(
             The step interval for saving the trajectories.
         verbose : bool
             If True, screen output will be shown.
-        **kwargs
-            Further kwargs.
+        neb_kwargs : dict, defaults to DEFAULT_NEB_KWARGS
+            kwargs to pass to ASE's NEB.
+        optimizer_kwargs : dict or None (default)
+            kwargs to pass to the optimizer.
 
         Returns
         -------
@@ -522,6 +522,7 @@ def run_neb(
             isinstance(images[0], Atoms) and all(not pbc for pbc in images[0].pbc)
         )
         num_images = len(images)
+        initial_images = [img.copy() for img in images]
 
         for idx, image in enumerate(images):
             if isinstance(image, Structure | Molecule):
@@ -531,11 +532,11 @@ def run_neb(
                 images[idx].set_constraint(FixSymmetry(image, symprec=self.symprec))
             images[idx].calc = deepcopy(self.calculator)
 
-        neb_calc = NEB(images, **self.neb_kwargs)
+        neb_calc = NEB(images, **neb_kwargs)
 
         with contextlib.redirect_stdout(sys.stdout if verbose else io.StringIO()):
             observers = [TrajectoryObserver(image) for image in images]
-            optimizer = self.opt_class(neb_calc, **kwargs)
+            optimizer = self.opt_class(neb_calc, **(optimizer_kwargs or {}))
             for idx in range(num_images):
                 optimizer.attach(observers[idx], interval=interval)
             t_i = time.perf_counter()
@@ -546,8 +547,10 @@ def run_neb(
         if traj_file is not None:
             if isinstance(traj_file, str | Path):
                 traj_file = Path(traj_file)
-                traj_file_suffix = "".join(traj_file.suffixes)
-                traj_file_prefix = str(traj_file).split(traj_file_suffix)[0]
+                if traj_file_suffix := "".join(traj_file.suffixes):
+                    traj_file_prefix = str(traj_file).split(traj_file_suffix)[0]
+                else:
+                    traj_file_prefix = str(traj_file)
                 traj_files = [
                     f"{traj_file_prefix}-image-{idx + 1}{traj_file_suffix}"
                     for idx in range(num_images)
@@ -563,21 +566,27 @@ def run_neb(
             for image in images
         ]
         num_sites = len(images[0])
+
+        tags = [os.getcwd()]
         is_force_conv = all(
             np.linalg.norm(observers[image_idx].forces[-1][site_idx]) < abs(fmax)
             for site_idx in range(num_sites)
             for image_idx in range(num_images)
         )
+        tags += ["force converged" if is_force_conv else "forces not converged"]
+        tags += [f"elapsed time {t_f - t_i} seconds"]
+
         return NebResult(
             images=images,
+            initial_images=initial_images,
             energies=[
                 observers[image_idx].energies[-1] for image_idx in range(num_images)
             ],
             method=NebMethod.CLIMBING_IMAGE
-            if self.neb_kwargs.get("climb", False)
+            if neb_kwargs.get("climb", False)
             else NebMethod.STANDARD,
-            is_force_converged=is_force_conv,
-            dir_name=os.getcwd(),
-            elapsed_time=t_f - t_i,
+            # dir_name=os.getcwd(), # NB: this should be migrated in emmet-core
+            state="successful" if is_force_conv else "failed",
+            tags=tags,
             **neb_doc_kwargs,
         )

src/atomate2/forcefields/flows/approx_neb.py

@@ -14,7 +14,7 @@
 
 
 @dataclass
-class MLFFApproxNebFromEndpointsMaker(ApproxNebFromEndpointsMaker):
+class ForceFieldApproxNebFromEndpointsMaker(ApproxNebFromEndpointsMaker):
     """
     Perform ApproxNEB on a single hop using ML forcefields.
 

tests/forcefields/flows/test_approx_neb.py

@@ -0,0 +1,51 @@
+import pytest
+from emmet.core.neb import NebResult
+from jobflow import run_locally
+from pymatgen.core import Structure
+
+from atomate2.forcefields.flows.approx_neb import ForceFieldApproxNebFromEndpointsMaker
+from atomate2.forcefields.jobs import ForceFieldStaticMaker
+from atomate2.utils.testing.common import get_job_uuid_name_map
+
+
+def test_approx_neb_from_endpoints(test_dir, clean_dir):
+    vasp_aneb_dir = test_dir / "vasp" / "ApproxNEB"
+
+    endpoints = [
+        Structure.from_file(
+            vasp_aneb_dir / f"ApproxNEB_image_relax_endpoint_{idx}/inputs/POSCAR.gz"
+        )
+        for idx in (0, 3)
+    ]
+
+    flow = ForceFieldApproxNebFromEndpointsMaker(
+        image_relax_maker=ForceFieldStaticMaker(force_field_name="MATPES_R2SCAN")
+    ).make("Zn", endpoints, vasp_aneb_dir / "host_structure_relax_2/outputs/CHGCAR.bz2")
+
+    response = run_locally(flow)
+    output = {
+        job_name: response[uuid][1].output
+        for uuid, job_name in get_job_uuid_name_map(flow).items()
+    }
+
+    assert isinstance(output["collate_images_single_hop"], NebResult)
+    assert all(
+        output["collate_images_single_hop"].energies[i] == pytest.approx(energy)
+        for i, energy in enumerate(
+            [
+                -1559.5150146484375,
+                -1554.3154296875,
+                -1529.771484375,
+                -1525.8846435546875,
+                -1533.1453857421875,
+                -1554.561767578125,
+                -1559.5150146484375,
+            ]
+        )
+    )
+
+    assert len(output["collate_images_single_hop"].images) == 7
+    assert all(
+        image.volume == pytest.approx(endpoints[0].volume)
+        for image in output["collate_images_single_hop"].images
+    )

tests/forcefields/test_neb.py

@@ -0,0 +1,86 @@
+from pathlib import Path
+
+import pytest
+from jobflow import run_locally
+from monty.serialization import loadfn
+from pymatgen.core import Structure
+from pymatgen.io.vasp.outputs import Xdatcar
+
+from atomate2.forcefields.neb import ForceFieldNebMaker
+
+
+def test_neb_from_images(test_dir, clean_dir):
+    endpoints = [
+        Structure.from_file(
+            test_dir
+            / "vasp"
+            / "Si_NEB"
+            / f"relax_endpoint_{1 + i}"
+            / "inputs"
+            / "POSCAR.gz"
+        )
+        for i in range(2)
+    ]
+
+    images = endpoints[0].interpolate(endpoints[1], nimages=4, autosort_tol=0.5)
+
+    job = ForceFieldNebMaker(
+        force_field_name="MATPES_PBE",
+        traj_file="XDATCAR_si_self_diffusion",
+        traj_file_fmt="xdatcar",
+        relax_kwargs={"fmax": 0.5},
+    ).make(images)
+
+    response = run_locally(job)
+    output = response[job.uuid][1].output
+
+    cwd = next(Path(p) for p in output.tags if Path(p).exists())
+    xdatcars = [
+        Xdatcar(cwd / f"XDATCAR_si_self_diffusion-image-{i + 1}") for i in range(5)
+    ]
+
+    # Check that trajectory initial and final images are consistent with document
+    assert all(
+        xdatcars[i].structures[0] == image
+        for i, image in enumerate(output.initial_images)
+    )
+
+    all(xdatcars[i].structures[-1] == image for i, image in enumerate(output.images))
+
+    assert all(
+        output.energies[i] == pytest.approx(energy)
+        for i, energy in enumerate(
+            [
+                -339.3764953613281,
+                -339.2301025390625,
+                -338.865234375,
+                -339.23004150390625,
+                -339.37652587890625,
+            ]
+        )
+    )
+
+    if output.state.value == "successful":
+        assert "force converged" in output.tags
+    else:
+        assert "forces not converged" in output.tags
+
+    images = endpoints[0].interpolate(endpoints[1], nimages=2, autosort_tol=0.5)
+    job = ForceFieldNebMaker(
+        force_field_name="MACE",
+        traj_file="si_self_diffusion.json.gz",
+        traj_file_fmt="pmg",
+        relax_kwargs={"fmax": 0.5},
+    ).make(images)
+
+    response = run_locally(job)
+    output = response[job.uuid][1].output
+
+    trajectories = [
+        loadfn(cwd / f"si_self_diffusion-image-{idx + 1}.json.gz") for idx in range(3)
+    ]
+
+    assert all(
+        trajectories[i].frame_properties[-1]["energy"] == pytest.approx(energy)
+        for i, energy in enumerate(output.energies)
+    )

tests/vasp/flows/test_approx_neb.py

@@ -8,6 +8,7 @@
 from monty.serialization import loadfn
 from pymatgen.core import Structure
 
+from atomate2.utils.testing.common import get_job_uuid_name_map
 from atomate2.vasp.flows.approx_neb import ApproxNebMaker
 
 
@@ -66,9 +67,8 @@ def test_approx_neb_flow(mock_vasp, clean_dir, vasp_test_dir):
     # ApproxNEB image relax hop 3+1 image 2
     responses = run_locally(flow, create_folders=True, ensure_success=False)
     output = {
-        job.name: responses[job.uuid][1].output
-        for job in flow.jobs
-        if job.uuid in responses
+        job_name: responses[uuid][1].output
+        for uuid, job_name in get_job_uuid_name_map(flow).items()
     }
 
     assert len(output["collate_results"].hops) == 2
@@ -113,10 +113,6 @@ def test_approx_neb_flow(mock_vasp, clean_dir, vasp_test_dir):
         for idx, energy in enumerate(ref_hop.energies)
     )
 
-    from monty.serialization import dumpfn
-
-    dumpfn(output["collate_results"], "/Users/aaronkaplan/Desktop/temp_aneb.json.gz")
-
     assert all(
         getattr(output["collate_results"], f"{direction}_barriers")[k]
         == pytest.approx(getattr(ref_results, f"{direction}_barriers")[k])