Add NEB, ApproxNEB jobs / workflows

.github/workflows/testing.yml

@@ -66,7 +66,7 @@ jobs:
           python -m pip install --upgrade pip
           mkdir -p ~/.abinit/pseudos
           cp -r tests/test_data/abinit/pseudos/ONCVPSP-PBE-SR-PDv0.4 ~/.abinit/pseudos
-          uv pip install .[strict,strict-forcefields,tests,abinit]
+          uv pip install .[strict,strict-forcefields,tests,abinit,approxneb]
           uv pip install torch-runstats
           uv pip install --no-deps nequip==0.5.6
 
@@ -321,7 +321,7 @@ jobs:
         run: sphinx-build docs docs_build
 
   automerge:
-    needs: [lint, test-non-ase, test-notebooks-and-ase, test-force-field-notebook, docs]
+    needs: [docs, lint, test-force-field-notebook, test-non-ase, test-notebooks-and-ase, test-openff]
     runs-on: ubuntu-latest
 
     permissions:

pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "PyYAML",
     "click",
     "custodian>=2024.4.18",
-    "emmet-core>=0.84.8",
+    "emmet-core>=0.84.9rc",
     "jobflow>=0.1.11",
     "monty>=2024.12.10",
     "numpy",
@@ -62,6 +62,7 @@ forcefields = [
     "sevenn>=0.9.3",
     "torchdata<=0.7.1",                            # TODO: remove when issue fixed
 ]
+approxneb = ["pymatgen-analysis-diffusion>=2024.7.15"]
 ase = ["ase>=3.25.0"]
 ase-ext = ["tblite>=0.3.0; platform_system=='Linux'"]
 openmm = [
@@ -99,7 +100,7 @@ strict = [
     "click==8.2.1",
     "custodian==2025.4.14",
     "dscribe==2.1.1",
-    "emmet-core==0.84.8",
+    "emmet-core==0.84.9rc",
     "ijson==3.4.0",
     "jobflow==0.2.0",
     "lobsterpy==0.4.9",
@@ -121,7 +122,7 @@ strict-openff = [
     "monty==2025.3.3",
     "openmm-mdanalysis-reporter==0.1.0",
     "openmm==8.1.1",
-    "pymatgen==2025.6.14", # TODO: open ff is extremely sensitive to pymatgen version
+    "pymatgen==2024.11.13", # TODO: open ff is extremely sensitive to pymatgen version
 ]
 strict-forcefields = [
     "calorine==3.1",

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,100 @@
+"""Create NEB jobs with ASE."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from emmet.core.neb import NebResult
+from jobflow import job
+
+from atomate2.ase.jobs import _ASE_DATA_OBJECTS, AseMaker
+from atomate2.ase.utils import AseNebInterface
+
+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)
+    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,
+            symprec=self.symprec,
+        ).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,
+        )
+
+    @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,8 @@
 import os
 import sys
 import time
+from copy import deepcopy
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -17,8 +19,10 @@
 from ase.filters import FrechetCellFilter
 from ase.io import Trajectory as AseTrajectory
 from ase.io import write as ase_write
+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, NebResult
 from monty.serialization import dumpfn
 from pymatgen.core.structure import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory as PmgTrajectory
@@ -46,6 +50,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.
@@ -149,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"
@@ -329,6 +344,7 @@ def relax(
         fmax: float = 0.1,
         steps: int = 500,
         traj_file: str = None,
+        traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase",
         final_atoms_object_file: str | os.PathLike[str] = "final_atoms_object.xyz",
         interval: int = 1,
         verbose: bool = False,
@@ -388,7 +404,7 @@ def relax(
             t_f = time.perf_counter()
 
         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
 
@@ -420,3 +436,157 @@ 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",
+        fix_symmetry: bool = False,
+        symprec: float = 1e-2,
+    ) -> None:
+        """Initialize the interface.
+
+        Parameters
+        ----------
+        calculator (ase Calculator): an ase calculator
+        optimizer (str or ase Optimizer): the optimization algorithm.
+        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.ase_adaptor = AseAtomsAdaptor()
+        self.fix_symmetry = fix_symmetry
+        self.symprec = symprec
+
+    def run_neb(
+        self,
+        images: list[Atoms | Structure | Molecule],
+        fmax: float = 0.1,
+        steps: int = 500,
+        traj_file: str | Path | list[str | Path] = None,
+        traj_file_fmt: Literal["pmg", "ase", "xdatcar"] = "ase",
+        interval: int = 1,
+        verbose: bool = False,
+        neb_doc_kwargs: dict | None = None,
+        neb_kwargs: dict = DEFAULT_NEB_KWARGS,
+        optimizer_kwargs: dict | None = None,
+    ) -> 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, Path, or a list of str / Path
+            The trajectory file for saving. If a single str or Path,
+            this specifies the file name prefix. For example,
+                `traj_file = "traj_mp-149.json.gz"`
+            will yield individual trajectory file names:
+                traj_mp-149-image-1.json.gz
+                traj_mp-149-image-2.json.gz
+                ...
+            Alternately, if this is a list of str / Path, this specifies the
+            file name for each image.
+        interval : int
+            The step interval for saving the trajectories.
+        verbose : bool
+            If True, screen output will be shown.
+        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
+        -------
+            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)
+        initial_images = [img.copy() for img in 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, **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, **(optimizer_kwargs or {}))
+            for idx in range(num_images):
+                optimizer.attach(observers[idx], interval=interval)
+            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:
+            if isinstance(traj_file, str | Path):
+                traj_file = Path(traj_file)
+                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)
+                ]
+            elif isinstance(traj_file, list | tuple):
+                traj_files = [str(f) for f in traj_file]
+
+            for idx, f in enumerate(traj_files):
+                observers[idx].save(f, 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])
+
+        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 neb_kwargs.get("climb", False)
+            else NebMethod.STANDARD,
+            # 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,
+        )