WIP: Added ase Nudged Elastic Band in the NewtonNet recipe

src/quacc/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from importlib.metadata import version
 
 from ase.atoms import Atoms
@@ -37,6 +38,9 @@
 # Load the settings
 SETTINGS = QuaccSettings()
 
+# Ignore ASE config file
+os.environ["ASE_CONFIG_PATH"] = ""
+
 # Set logging info
 logging.basicConfig(level=logging.DEBUG if SETTINGS.DEBUG else logging.INFO)
 

src/quacc/recipes/newtonnet/ts.py

@@ -5,22 +5,25 @@
 from importlib.util import find_spec
 from typing import TYPE_CHECKING
 
+import numpy as np
 from monty.dev import requires
 
 from quacc import SETTINGS, change_settings, job, strip_decorator
 from quacc.recipes.newtonnet.core import _add_stdev_and_hess, freq_job, relax_job
-from quacc.runners.ase import Runner
-from quacc.schemas.ase import summarize_opt_run
+from quacc.runners.ase import Runner, run_neb
+from quacc.schemas.ase import summarize_neb_run, summarize_opt_run
 from quacc.utils.dicts import recursive_dict_merge
 
+has_geodesic_interpolate = bool(find_spec("geodesic_interpolate"))
 has_sella = bool(find_spec("sella"))
 has_newtonnet = bool(find_spec("newtonnet"))
 
 if has_sella:
     from sella import IRC, Sella
 if has_newtonnet:
     from newtonnet.utils.ase_interface import MLAseCalculator as NewtonNet
-
+if has_geodesic_interpolate:
+    from quacc.runners.ase import _geodesic_interpolate_wrapper
 
 if TYPE_CHECKING:
     from typing import Any, Literal
@@ -264,6 +267,246 @@ def quasi_irc_job(
     return relax_summary
 
 
+@job
+@requires(
+    has_newtonnet, "NewtonNet must be installed. Refer to the quacc documentation."
+)
+@requires(
+    has_geodesic_interpolate,
+    "geodesic-interpolate must be installed. Refer to the quacc documentation.",
+)
+def neb_job(
+    reactant_atoms: Atoms,
+    product_atoms: Atoms,
+    relax_job_kwargs: dict[str, Any] | None = None,
+    calc_kwargs: dict[str, Any] | None = None,
+    geodesic_interpolate_kwargs: dict[str, Any] | None = None,
+    neb_kwargs: dict[str, Any] | None = None,
+) -> dict:
+    """
+    Perform a quasi-IRC job using the given atoms object. The initial IRC job by default
+    is run with `max_steps: 5`.
+
+    Parameters
+    ----------
+    atoms
+        The atoms object representing the system
+    direction
+        The direction of the IRC calculation
+    relax_job_kwargs
+        Keyword arguments to use for the [quacc.recipes.newtonnet.core.relax_job][]
+
+    Returns
+    -------
+    A dictionary containing the neb optimization summary.
+            See the type-hint for the data structure.
+    """
+    relax_job_kwargs = relax_job_kwargs or {}
+    neb_kwargs = neb_kwargs or {}
+    geodesic_interpolate_kwargs = geodesic_interpolate_kwargs or {}
+
+    calc_defaults = {
+        "model_path": SETTINGS.NEWTONNET_MODEL_PATH,
+        "settings_path": SETTINGS.NEWTONNET_CONFIG_PATH,
+    }
+
+    geodesic_defaults = {"nimages": 20}
+
+    neb_defaults = {"method": "aseneb", "precon": None}
+    calc_flags = recursive_dict_merge(calc_defaults, calc_kwargs)
+    geodesic_interpolate_flags = recursive_dict_merge(
+        geodesic_defaults, geodesic_interpolate_kwargs
+    )
+    neb_flags = recursive_dict_merge(neb_defaults, neb_kwargs)
+
+    # Define calculator
+    reactant_atoms.calc = NewtonNet(**calc_flags)
+    product_atoms.calc = NewtonNet(**calc_flags)
+
+    # Run IRC
+    relax_summary_r = strip_decorator(relax_job)(reactant_atoms, **relax_job_kwargs)
+    relax_summary_p = strip_decorator(relax_job)(product_atoms, **relax_job_kwargs)
+
+    images = _geodesic_interpolate_wrapper(
+        relax_summary_r["atoms"].copy(),
+        relax_summary_p["atoms"].copy(),
+        **geodesic_interpolate_flags,
+    )
+
+    for image in images:
+        image.calc = NewtonNet(**calc_flags)
+
+    dyn = run_neb(images, neb_kwargs=neb_flags)
+
+    return {
+        "relax_reactant": relax_summary_r,
+        "relax_product": relax_summary_p,
+        "geodesic_results": images,
+        "neb_results": summarize_neb_run(dyn),
+    }
+
+
+@job
+@requires(
+    has_newtonnet, "NewtonNet must be installed. Refer to the quacc documentation."
+)
+@requires(
+    has_geodesic_interpolate,
+    "geodesic-interpolate must be installed. Refer to the quacc documentation.",
+)
+def neb_ts_job(
+    reactant_atoms: Atoms,
+    product_atoms: Atoms,
+    relax_job_kwargs: dict[str, Any] | None = None,
+    calc_kwargs: dict[str, Any] | None = None,
+    geodesic_interpolate_kwargs: dict[str, Any] | None = None,
+    neb_kwargs: dict[str, Any] | None = None,
+) -> dict:
+    """
+    Perform a quasi-IRC job using the given atoms object. The initial IRC job by default
+    is run with `max_steps: 5`.
+
+    Parameters
+    ----------
+    atoms
+        The atoms object representing the system
+    direction
+        The direction of the IRC calculation
+    relax_job_kwargs
+        Keyword arguments to use for the [quacc.recipes.newtonnet.core.relax_job][]
+
+    Returns
+    -------
+    A dictionary containing the neb optimization summary.
+            See the type-hint for the data structure.
+    """
+    relax_job_kwargs = relax_job_kwargs or {}
+    neb_kwargs = neb_kwargs or {}
+    geodesic_interpolate_kwargs = geodesic_interpolate_kwargs or {}
+    calc_kwargs = calc_kwargs or {}
+
+    calc_defaults = {
+        "model_path": SETTINGS.NEWTONNET_MODEL_PATH,
+        "settings_path": SETTINGS.NEWTONNET_CONFIG_PATH,
+    }
+
+    geodesic_defaults = {"nimages": 20}
+
+    neb_defaults = {"method": "aseneb", "precon": None}
+    calc_flags = recursive_dict_merge(calc_defaults, calc_kwargs)
+    geodesic_interpolate_flags = recursive_dict_merge(
+        geodesic_defaults, geodesic_interpolate_kwargs
+    )
+    neb_flags = recursive_dict_merge(neb_defaults, neb_kwargs)
+
+    neb_results = strip_decorator(neb_job)(
+        reactant_atoms,
+        product_atoms,
+        calc_kwargs=calc_flags,
+        geodesic_interpolate_kwargs=geodesic_interpolate_flags,
+        neb_kwargs=neb_flags,
+        relax_job_kwargs=relax_job_kwargs,
+    )
+
+    traj = neb_results["neb_results"]["trajectory"]
+    traj_results = neb_results["neb_results"]["trajectory_results"]
+    n_images = len(neb_results["geodesic_results"])
+
+    ts_index = np.argmax([i["energy"] for i in traj_results[-(n_images - 1) : -1]]) + 1
+    ts_atoms = traj[-(n_images) + ts_index]
+
+    output = ts_job(ts_atoms)
+    neb_results["ts_results"] = output
+
+    return neb_results
+
+
+@job
+@requires(
+    has_newtonnet, "NewtonNet must be installed. Refer to the quacc documentation."
+)
+@requires(
+    has_geodesic_interpolate,
+    "geodesic-interpolate must be installed. Refer to the quacc documentation.",
+)
+def geodesic_ts_job(
+    reactant_atoms: Atoms,
+    product_atoms: Atoms,
+    relax_job_kwargs: dict[str, Any] | None = None,
+    calc_kwargs: dict[str, Any] | None = None,
+    geodesic_interpolate_kwargs: dict[str, Any] | None = None,
+) -> dict:
+    """
+    Perform a transition state search using geodesic interpolation between reactant and product states.
+
+    Parameters
+    ----------
+    reactant_atoms : Atoms
+        The atoms object representing the reactant state.
+    product_atoms : Atoms
+        The atoms object representing the product state.
+    relax_job_kwargs : dict[str, Any], optional
+        Keyword arguments to use for the relaxation job.
+    calc_kwargs : dict[str, Any], optional
+        Keyword arguments to configure the NewtonNet calculator.
+    geodesic_interpolate_kwargs : dict[str, Any], optional
+        Keyword arguments to configure the geodesic interpolation.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the NEB optimization summary, including:
+        - 'relax_reactant': Summary of the relaxation job for the reactant.
+        - 'relax_product': Summary of the relaxation job for the product.
+        - 'geodesic_results': The interpolated images between reactant and product.
+        - 'ts_results': The transition state optimization results.
+    """
+    relax_job_kwargs = relax_job_kwargs or {}
+    geodesic_interpolate_kwargs = geodesic_interpolate_kwargs or {}
+
+    calc_defaults = {
+        "model_path": SETTINGS.NEWTONNET_MODEL_PATH,
+        "settings_path": SETTINGS.NEWTONNET_CONFIG_PATH,
+    }
+
+    geodesic_defaults = {"nimages": 20}
+
+    calc_flags = recursive_dict_merge(calc_defaults, calc_kwargs)
+    geodesic_interpolate_flags = recursive_dict_merge(
+        geodesic_defaults, geodesic_interpolate_kwargs
+    )
+
+    # Define calculator
+    reactant_atoms.calc = NewtonNet(**calc_flags)
+    product_atoms.calc = NewtonNet(**calc_flags)
+
+    # Run IRC
+    relax_summary_r = strip_decorator(relax_job)(reactant_atoms, **relax_job_kwargs)
+    relax_summary_p = strip_decorator(relax_job)(product_atoms, **relax_job_kwargs)
+
+    images = _geodesic_interpolate_wrapper(
+        relax_summary_r["atoms"].copy(),
+        relax_summary_p["atoms"].copy(),
+        **geodesic_interpolate_flags,
+    )
+
+    potential_energies = []
+    for image in images:
+        image.calc = NewtonNet(**calc_flags)
+        potential_energies.append(image.get_potential_energy())
+
+    ts_index = np.argmax(potential_energies)
+    ts_atoms = images[ts_index]
+
+    output = ts_job(ts_atoms)
+    return {
+        "relax_reactant": relax_summary_r,
+        "relax_product": relax_summary_p,
+        "geodesic_results": images,
+        "ts_results": output,
+    }
+
+
 def _get_hessian(atoms: Atoms) -> NDArray:
     """
     Calculate and retrieve the Hessian matrix for the given molecular configuration.