[WIP] Add NEB, ApproxNEB jobs / workflows

.github/workflows/testing.yml

@@ -69,6 +69,7 @@ jobs:
           uv pip install .[strict,strict-forcefields,tests,abinit]
           uv pip install torch-runstats
           uv pip install --no-deps nequip==0.5.6
+          uv pip install --upgrade 'git+https://github.com/esoteric-ephemera/emmet.git@neb#subdirectory=emmet-core' # to-do: remove on emmet-core release
 
       - name: Install pymatgen from master if triggered by pymatgen repo dispatch
         if: github.event_name == 'repository_dispatch' && github.event.action == 'pymatgen-ci-trigger'
@@ -142,6 +143,7 @@ jobs:
           micromamba activate a2
           python -m pip install --upgrade pip
           uv pip install .[strict,tests]
+          uv pip install --upgrade 'git+https://github.com/esoteric-ephemera/emmet.git@neb#subdirectory=emmet-core' # to-do: remove on emmet-core release
 
       - name: Install pymatgen from master if triggered by pymatgen repo dispatch
         if: github.event_name == 'repository_dispatch' && github.event.action == 'pymatgen-ci-trigger'
@@ -181,6 +183,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install .[strict,strict-forcefields,docs]
+          pip install --upgrade 'git+https://github.com/esoteric-ephemera/emmet.git@neb#subdirectory=emmet-core' # to-do: remove on emmet-core release
 
       - name: Build
         run: sphinx-build docs docs_build

src/atomate2/ase/md.py

@@ -170,7 +170,7 @@ class AseMDMaker(AseMaker, metaclass=ABCMeta):
     ionic_step_data: tuple[str, ...] | None = None
     store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.PARTIAL
     traj_file: str | Path | None = None
-    traj_file_fmt: Literal["pmg", "ase"] = "ase"
+    traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase"
     traj_interval: int = 1
     mb_velocity_seed: int | None = None
     zero_linear_momentum: bool = False

src/atomate2/ase/neb.py

@@ -0,0 +1,102 @@
+"""Create NEB jobs with ASE."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import job
+
+from atomate2.ase.jobs import _ASE_DATA_OBJECTS, AseMaker
+from atomate2.ase.utils import AseNebInterface
+from atomate2.common.schemas.neb import NebResult
+
+if TYPE_CHECKING:
+    from pathlib import Path
+    from typing import Literal
+
+    from ase.calculators.calculator import Calculator
+    from pymatgen.core import Molecule, Structure
+
+
+@dataclass
+class AseNebMaker(AseMaker):
+    """Define scheme for performing ASE NEB calculations."""
+
+    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
+    relax_kwargs: dict = field(default_factory=dict)
+    optimizer_kwargs: dict = field(default_factory=dict)
+    traj_file: str | None = None
+    traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase"
+    traj_interval: int = 1
+    neb_doc_kwargs: dict = field(default_factory=dict)
+
+    def run_ase(
+        self,
+        images: list[Structure | Molecule],
+        prev_dir: str | Path | None = None,
+    ) -> NebResult:
+        """
+        Run an ASE NEB job from a list of images.
+
+        Parameters
+        ----------
+        images: list of pymatgen .Molecule or .Structure
+            pymatgen molecule or structure images
+        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 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,
+            **self.relax_kwargs,
+        )
+
+    @job(data=_ASE_DATA_OBJECTS, schema=NebResult)
+    def make(
+        self,
+        images: list[Structure | Molecule],
+        prev_dir: str | Path | None = None,
+    ) -> NebResult:
+        """
+        Run an ASE NEB job from a list of images.
+
+        Parameters
+        ----------
+        images: list of pymatgen .Molecule or .Structure
+            pymatgen molecule or structure images
+        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 self.run_ase(images, prev_dir=prev_dir)
+
+
+class LennardJonesNebMaker(AseNebMaker):
+    """Lennard-Jones NEB maker, primarily for testing/debugging."""
+
+    name: str = "Lennard-Jones 6-12 NEB"
+
+    @property
+    def calculator(self) -> Calculator:
+        """Lennard-Jones calculator."""
+        from ase.calculators.lj import LennardJones
+
+        return LennardJones(**self.calculator_kwargs)

src/atomate2/ase/utils.py

@@ -7,6 +7,7 @@
 import os
 import sys
 import time
+from copy import deepcopy
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -16,14 +17,17 @@
 from ase.constraints import FixSymmetry
 from ase.filters import FrechetCellFilter
 from ase.io import Trajectory as AseTrajectory
+from ase.mep.neb import NEB
 from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG
+from emmet.core.neb import NebMethod
 from monty.serialization import dumpfn
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory as PmgTrajectory
 from pymatgen.io.ase import AseAtomsAdaptor
 
 from atomate2.ase.schemas import AseResult
+from atomate2.common.schemas.neb import NebResult
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -45,6 +49,15 @@
     "BFGSLineSearch": BFGSLineSearch,
 }
 
+FORCE_BASED_OPTIMIZERS = {
+    "FIRE": FIRE,
+    "BFGS": BFGS,
+    "MDMin": MDMin,
+}
+
+# Parameters chosen for consistency with atomate2.vasp.sets.core.NebSetGenerator
+DEFAULT_NEB_KWARGS = {"k": 5.0, "climb": True, "method": "improvedtangent"}
+
 
 class TrajectoryObserver:
     """Trajectory observer.
@@ -328,6 +341,7 @@ def relax(
         fmax: float = 0.1,
         steps: int = 500,
         traj_file: str = None,
+        traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase",
         interval: int = 1,
         verbose: bool = False,
         cell_filter: Filter = FrechetCellFilter,
@@ -377,7 +391,7 @@ def relax(
             t_f = time.perf_counter()
             obs()
         if traj_file is not None:
-            obs.save(traj_file)
+            obs.save(traj_file, fmt=traj_file_fmt)
         if isinstance(atoms, cell_filter):
             atoms = atoms.atoms
 
@@ -398,3 +412,137 @@ def relax(
             dir_name=os.getcwd(),
             elapsed_time=t_f - t_i,
         )
+
+
+class AseNebInterface:
+    """Perform NEB using the Atomic Simulation Environment."""
+
+    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.
+        """
+        self.calculator = calculator
+
+        if isinstance(optimizer, str):
+            optimizer_obj = FORCE_BASED_OPTIMIZERS.get(optimizer)
+        elif optimizer is None:
+            raise ValueError("Optimizer cannot be None")
+        else:
+            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,
+        images: list[Atoms | Structure | Molecule],
+        fmax: float = 0.1,
+        steps: int = 500,
+        traj_file: str = None,
+        traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase",
+        interval: int = 1,
+        verbose: bool = False,
+        neb_doc_kwargs: dict | None = None,
+        **kwargs,
+    ) -> NebResult:
+        """
+        Perform NEB on a list of molecules or structures.
+
+        Parameters
+        ----------
+        images : list of ASE Atoms, pymatgen Structure, or pymatgen Molecule
+            The ordered list of atoms to perform NEB on.
+        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
+        -------
+            dict including optimized structure and the trajectory
+        """
+        is_mol = isinstance(images[0], Molecule) or (
+            isinstance(images[0], Atoms) and all(not pbc for pbc in images[0].pbc)
+        )
+        num_images = len(images)
+
+        for idx, image in enumerate(images):
+            if isinstance(image, Structure | Molecule):
+                images[idx] = self.ase_adaptor.get_atoms(image)
+
+            if self.fix_symmetry:
+                images[idx].set_constraint(FixSymmetry(image, symprec=self.symprec))
+            images[idx].calc = deepcopy(self.calculator)
+
+        neb_calc = NEB(images, **self.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)
+            for idx in range(num_images):
+                optimizer.attach(observers[idx], interval=interval, atoms=images[idx])
+            t_i = time.perf_counter()
+            optimizer.run(fmax=fmax, steps=steps)
+            t_f = time.perf_counter()
+            [observers[idx]() for idx in range(num_images)]
+
+        if traj_file is not None:
+            for idx in range(num_images):
+                traj_file_split = traj_file.split(".")
+                traj_file_prefix = ".".join(traj_file_split[:-1])
+                traj_file_ext = traj_file[-1]
+                observers[idx].save(
+                    f"{traj_file_prefix}-image-{idx+1}.{traj_file_ext}",
+                    fmt=traj_file_fmt,
+                )
+
+        images = [
+            self.ase_adaptor.get_structure(image, cls=Molecule if is_mol else Structure)
+            for image in images
+        ]
+        num_sites = len(images[0])
+        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)
+        )
+        return NebResult(
+            images=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)
+            else NebMethod.STANDARD,
+            is_force_converged=is_force_conv,
+            dir_name=os.getcwd(),
+            elapsed_time=t_f - t_i,
+            **neb_doc_kwargs,
+        )