WIP: Added ase Nudged Elastic Band in the NewtonNet recipe

pyproject.toml

@@ -49,7 +49,7 @@ defects = ["pymatgen-analysis-defects>=2023.8.22", "shakenbreak>=3.2.0"]
 jobflow = ["jobflow[fireworks]>=0.1.14", "jobflow-remote>=0.1.0"]
 mlp = ["matgl>=1.0.0", "chgnet>=0.3.3", "mace-torch>=0.3.3", "torch-dftd>=0.4.0"]
 mp = ["atomate2>=0.0.14"]
-newtonnet = ["newtonnet>=1.1"]
+newtonnet = ["newtonnet>=1.1", "geodesic-interpolate @ git+https://github.com/virtualzx-nad/geodesic-interpolate.git"]
 parsl = ["parsl[monitoring]>=2023.10.23; platform_system!='Windows'"]
 phonons = ["phonopy>=2.20.0", "seekpath>=2.1.0"]
 prefect = ["prefect[dask]>=2.19.0", "dask-jobqueue>=0.8.2"]

src/quacc/recipes/newtonnet/ts.py

@@ -2,8 +2,17 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
+from ase import Atoms
+from ase.atoms import Atoms
+from ase.io import Trajectory, read, write
+from ase.mep.neb import NEBOptimizer
+from ase.neb import NEB
+from geodesic_interpolate.fileio import write_xyz
+from geodesic_interpolate.geodesic import Geodesic
+from geodesic_interpolate.interpolation import redistribute
 from monty.dev import requires
 
 from quacc import SETTINGS, change_settings, job, strip_decorator
@@ -17,6 +26,11 @@
 except ImportError:
     Sella = None
 
+try:
+    from ase.mep import neb
+except ImportError:
+    neb = None
+
 try:
     from newtonnet.utils.ase_interface import MLAseCalculator as NewtonNet
 except ImportError:
@@ -25,7 +39,7 @@
 if TYPE_CHECKING:
     from typing import Any, Literal
 
-    from ase.atoms import Atoms
+    from ase.optimize.optimize import Optimizer
     from numpy.typing import NDArray
 
     from quacc.recipes.newtonnet.core import FreqSchema
@@ -287,3 +301,250 @@
     ml_calculator.calculate(atoms)
 
     return ml_calculator.results["hessian"].reshape((-1, 3 * len(atoms)))
+
+
+"""
+@job
+@requires(NewtonNet, "NewtonNet must be installed. Refer to the quacc documentation.")
+def neb_job(
+    atoms: Atoms,
+    relax_endpoints: bool = True,
+    run_geodesic: bool = True,
+    run_single_ended: bool = True,
+    run_freq: bool = True,
+    neb_job_kwargs: dict[str, Any] | None = None,
+    relax_job_kwargs: dict[str, Any] | None = None,
+    freq_job_kwargs: dict[str, Any] | None = None,
+):
+    relax_job_kwargs = relax_job_kwargs or {}
+    freq_job_kwargs = freq_job_kwargs or {}
+
+    irc_job_defaults = {"max_steps": 5}
+    irc_job_kwargs = recursive_dict_merge(irc_job_defaults, irc_job_kwargs)
+
+    # Run IRC
+    irc_summary = strip_decorator(irc_job)(
+        atoms, direction=direction, run_freq=False, **irc_job_kwargs
+    )
+
+    # Run opt
+    relax_summary = strip_decorator(relax_job)(irc_summary["atoms"], **relax_job_kwargs)
+
+    # Run frequency
+    freq_summary = (
+        strip_decorator(freq_job)(relax_summary["atoms"], **freq_job_kwargs)
+        if run_freq
+        else None
+    )
+    relax_summary["freq_job"] = freq_summary
+    relax_summary["irc_job"] = irc_summary
+
+    return relax_summary
+"""
+
+
+def sella_wrapper(
+    atoms_object,
+    traj_file=None,
+    sella_order=0,
+    use_internal=True,
+    traj_log_interval=2,
+    fmax_cutoff=1e-3,
+    max_steps=1000,
+):
+    if traj_file:
+        traj = Trajectory(traj_file, "w", atoms_object)
+    qn = Sella(atoms_object, order=sella_order, internal=use_internal)
+    if traj_file:
+        qn.attach(traj.write, interval=traj_log_interval)
+    qn.run(fmax=fmax_cutoff, steps=max_steps)
+    if traj_file:
+        traj.close()
+
+
+def geodesic_interpolate_wrapper(
+    r_p_atoms: Atoms,
+    nimages: int = 17,
+    sweep: bool | None = None,
+    output: str = "interpolated.xyz",
+    tol: float = 2e-3,
+    maxiter: int = 15,
+    microiter: int = 20,
+    scaling: float = 1.7,
+    friction: float = 1e-2,
+    dist_cutoff: float = 3,
+    save_raw: str | None = None,
+):
+    """
+    Interpolates between two geometries and optimizes the path.
+
+    Parameters:
+    filename (str): XYZ file containing geometries.
+    nimages (int): Number of images. Default is 17.
+    sweep (bool): Sweep across the path optimizing one image at a time.
+                  Default is to perform sweeping updates if there are more than 35 atoms.
+    output (str): Output filename. Default is "interpolated.xyz".
+    tol (float): Convergence tolerance. Default is 2e-3.
+    maxiter (int): Maximum number of minimization iterations. Default is 15.
+    microiter (int): Maximum number of micro iterations for sweeping algorithm. Default is 20.
+    scaling (float): Exponential parameter for Morse potential. Default is 1.7.
+    friction (float): Size of friction term used to prevent very large change of geometry. Default is 1e-2.
+    dist_cutoff (float): Cut-off value for the distance between a pair of atoms to be included in the coordinate system. Default is 3.
+    save_raw (str): When specified, save the raw path after bisections but before smoothing. Default is None.
+    """
+    # Read the initial geometries.
+    symbols = r_p_atoms[0].get_chemical_symbols()
+
+    X = [conf.get_positions() for conf in r_p_atoms]
+
+    if len(X) < 2:
+        raise ValueError("Need at least two initial geometries.")
+
+    # First redistribute number of images. Perform interpolation if too few and subsampling if too many images are given
+    raw = redistribute(symbols, X, nimages, tol=tol * 5)
+    if save_raw is not None:
+        write_xyz(save_raw, symbols, raw)
+
+    # Perform smoothing by minimizing distance in Cartesian coordinates with redundant internal metric
+    # to find the appropriate geodesic curve on the hyperspace.
+    smoother = Geodesic(symbols, raw, scaling, threshold=dist_cutoff, friction=friction)
+    if sweep is None:
+        sweep = len(symbols) > 35
+    try:
+        if sweep:
+            smoother.sweep(tol=tol, max_iter=maxiter, micro_iter=microiter)
+        else:
+            smoother.smooth(tol=tol, max_iter=maxiter)
+    finally:
+        # Save the smoothed path to output file. try block is to ensure output is saved if one ^C the process, or there is an error
+        write_xyz(output, symbols, smoother.path)
+    return symbols, smoother.path
+
+
+def setup_images(logdir: str, xyz_r_p: str, n_intermediate: int = 40):
+    """
+    Sets up intermediate images for NEB calculations between reactant and product states.
+
+    Parameters:
+    logdir (str): Directory to save the intermediate files.
+    xyz_r_p (str): Path to the XYZ file containing reactant and product structures.
+    n_intermediate (int): Number of intermediate images to generate.
+
+    Returns:
+    List: List of ASE Atoms objects with calculated energies and forces.
+    """
+    calc_defaults = {
+        "model_path": SETTINGS.NEWTONNET_MODEL_PATH,
+        "settings_path": SETTINGS.NEWTONNET_CONFIG_PATH,
+    }
+
+    calc_flags = recursive_dict_merge(calc_defaults, {})
+
+    try:
+        # Ensure the log directory exists
+        os.makedirs(logdir, exist_ok=True)
+
+        # Read reactant and product structures
+        reactant = read(xyz_r_p, index="0")
+        product = read(xyz_r_p, index="1")
+
+        # Optimize reactant and product structures using sella
+        for atom, name in zip([reactant, product], ["reactant", "product"]):
+            # atom.calc = calc()
+            atom.calc = NewtonNet(**calc_flags)
+            traj_file = os.path.join(logdir, f"{name}_opt.traj")
+            sella_wrapper(atom, traj_file=traj_file, sella_order=0)
+        # Save optimized reactant and product structures
+        r_p_path = os.path.join(logdir, "r_p.xyz")
+        write(r_p_path, [reactant.copy(), product.copy()])
+
+        # Generate intermediate images using geodesic interpolation
+        symbols, smoother_path = geodesic_interpolate_wrapper(
+            [reactant.copy(), product.copy()]
+        )
+        images = [Atoms(symbols=symbols, positions=conf) for conf in smoother_path]
+
+        # Calculate energies and forces for each intermediate image
+        for image in images:
+            # image.calc = calc()
+            # ml_calculator = calc()
+            image.calc = NewtonNet(**calc_flags)
+            ml_calculator = NewtonNet(**calc_flags)
+            ml_calculator.calculate(image)
+
+            energy = ml_calculator.results["energy"]
+            forces = ml_calculator.results["forces"]
+
+            image.info["energy"] = energy
+            image.arrays["forces"] = forces
+
+        # Save the geodesic path
+        geodesic_path = os.path.join(logdir, "geodesic_path.xyz")
+        write(geodesic_path, images)
+
+        return images
+
+    except Exception:
+        return []
+
+
+def run_neb_method(
+    method: str,
+    optimizer: Optimizer | None = NEBOptimizer,
+    opt_method: str | None = "aseneb",
+    precon: str | None = None,
+    logdir: str | None = None,
+    xyz_r_p: str | None = None,
+    n_intermediate: int | None = 20,
+    k: float | None = 0.1,
+    max_steps: int | None = 1000,
+    fmax_cutoff: float | None = 1e-2,
+) -> None:
+    """
+    Run NEB method.
+
+    Args:
+        method (str): NEB method.
+        optimizer (Optimizer, Optional): NEB path Optimizer function, Defaults to NEBOptimizer.
+        precon (str, optional): Preconditioner method. Defaults to None.
+        opt_method (str, Optimizer): Optimization method. Defaults to aseneb.
+        logdir (str, optional): Directory to save logs. Defaults to None.
+        xyz_r_p (str, optional): Path to reactant and product XYZ files. Defaults to None.
+        n_intermediate (int, optional): Number of intermediate images. Defaults to 20.
+        k (float, optional): force constant for the springs in NEB. Defaults to 0.1.
+        max_steps (int, optional): maximum number of optimization steps allowed. Defaults to 1000.
+        fmax_cutoff (float: optional): convergence cut-off criteria for the NEB optimization. Defaults to 1e-2.
+    """
+    images = setup_images(logdir, xyz_r_p, n_intermediate=n_intermediate)
+
+    mep = NEB(
+        images,
+        k=k,
+        method=method,
+        climb=True,
+        precon=precon,
+        remove_rotation_and_translation=True,
+        parallel=True,
+    )
+
+    os.makedirs(logdir, exist_ok=True)
+    log_filename = f"neb_band_{method}_{optimizer.__name__}_{precon}.txt"
+
+    logfile_path = os.path.join(logdir, log_filename)
+
+    opt = optimizer(mep, method=opt_method, logfile=logfile_path, verbose=2)
+
+    opt.run(fmax=fmax_cutoff, steps=max_steps)
+
+    # The following was written because of some error in writing the xyz file below
+    images_copy = []
+    for image in images:
+        image_copy = Atoms(symbols=image.symbols, positions=image.positions)
+        image_copy.info["energy"] = image.get_potential_energy()
+        images_copy.append(image_copy)
+
+    write(
+        f"{logdir}/optimized_path_{method}_{optimizer.__name__}_{precon}.xyz",
+        images_copy,
+    )
+    return images