Add NEB convergence and restart (#685)

Author: ElliottKasoar

docs/source/tutorials/python/neb.ipynb

@@ -375,6 +375,66 @@
     "v2.avr.show_hydrogen_bonds = True\n",
     "v2"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Continuing NEBs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also continue a NEB to a tighter `fmax`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "neb_b.optimize(fmax=0.05)\n",
+    "neb_b.run_nebtools()\n",
+    "print(neb_b.results)\n",
+    "\n",
+    "fig = neb_b.nebtools.plot_band()\n",
+    "v1=WeasWidget()\n",
+    "v1.from_ase(neb_b.nebtools.images)\n",
+    "v1.avr.model_style = 1\n",
+    "v1.avr.show_hydrogen_bonds = True\n",
+    "v1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "It is also possible to change NEB parameters, accessed through the attributes of `NEB.neb`, such as to continue or retry a NEB using different methods.\n",
+    "\n",
+    "Note that although the \"success\" of the NEB optimisation can be accessed through `NEB.converged`, this is typically only set to `False` if a `OptimizerConvergenceError`, not when the maximum number of steps is reached. We therefore recommend checking `NEB.opt.nsteps` or `NEB.results[\"max_force\"]`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "neb_c.neb.climb = True\n",
+    "neb_c.run(fmax=0.05)\n",
+    "print(neb_c.results)\n",
+    "print(neb_c.opt.nsteps)\n",
+    "\n",
+    "fig = neb_c.nebtools.plot_band()\n",
+    "v1=WeasWidget()\n",
+    "v1.from_ase(neb_c.nebtools.images)\n",
+    "v1.avr.model_style = 1\n",
+    "v1.avr.show_hydrogen_bonds = True\n",
+    "v1"
+   ]
   }
  ],
  "metadata": {
@@ -387,7 +447,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "janus",
+   "display_name": "janus-core (3.12.8)",
    "language": "python",
    "name": "python3"
   },
@@ -401,7 +461,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.8"
   }
  },
  "nbformat": 4,

janus_core/calculations/neb.py

@@ -5,6 +5,7 @@
 from collections.abc import Callable, Sequence
 from copy import copy
 from typing import Any, get_args
+import warnings
 
 from ase import Atoms
 import ase.mep
@@ -475,15 +476,42 @@ def interpolate(self) -> None:
             case _:
                 raise ValueError("Invalid interpolator selected")
 
-    def optimize(self):
-        """Run NEB optimization."""
+    def optimize(self, fmax: float | None = None, steps: int | None = None) -> None:
+        """
+        Run NEB optimisation.
+
+        Parameters
+        ----------
+        fmax
+            Maximum force for NEB optimiser. Default is `self.fmax`.
+        steps
+            Maximum steps for NEB optimiser. Default is `self.steps`.
+        """
+        fmax = fmax if fmax is not None else self.fmax
+        steps = steps if steps is not None else self.steps
+
         if not hasattr(self, "neb"):
             self.interpolate()
 
-        optimizer = self.optimizer(self.neb, **self.optimizer_kwargs)
-        optimizer.run(fmax=self.fmax, steps=self.steps)
+        if not hasattr(self, "opt"):
+            self.opt = self.optimizer(self.neb, **self.optimizer_kwargs)
+
+        self.converged = self.opt.run(fmax=fmax, steps=steps)
+
+        if not self.converged:
+            message = f"NEB optimization has not converged after {steps} steps."
+            try:
+                message += (
+                    f" Current max force {self.neb.get_residual()} > target force "
+                    f"{fmax}"
+                )
+            except RuntimeError:
+                pass
+
+            warnings.warn(message, stacklevel=2)
+
         if self.logger:
-            self.logger.info("Optimization steps: %s", optimizer.nsteps)
+            self.logger.info("Optimization steps: %s", self.opt.nsteps)
 
         # Optionally write band images to file
         output_structs(
@@ -511,10 +539,19 @@ def run_nebtools(self):
                 print("#Barrier [eV] | delta E [eV] | Max force [eV/Å] ", file=out)
                 print(*self.results.values(), file=out)
 
-    def run(self) -> dict[str, float]:
+    def run(
+        self, fmax: float | None = None, steps: int | None = None
+    ) -> dict[str, float]:
         """
         Run Nudged Elastic Band method.
 
+        Parameters
+        ----------
+        fmax
+            Maximum force for NEB optimiser. Default is `self.fmax`.
+        steps
+            Maximum steps for NEB optimiser. Default is `self.steps`.
+
         Returns
         -------
         dict[str, float]
@@ -527,17 +564,25 @@ def run(self) -> dict[str, float]:
 
         self._set_info_units()
 
-        if self.minimize:
-            GeomOpt(self.init_struct, **self.minimize_kwargs).run()
+        fmax = fmax if fmax is not None else self.fmax
+        steps = steps if steps is not None else self.steps
 
-            # Change filename to be written
-            self.minimize_kwargs["write_kwargs"]["filename"] = (
-                self.final_struct_min_path
-            )
-            GeomOpt(self.final_struct, **self.minimize_kwargs).run()
+        if hasattr(self, "neb"):
+            if self.logger:
+                self.logger.info("Continuing from previous NEB optimisation")
+        else:
+            if self.minimize:
+                GeomOpt(self.init_struct, **self.minimize_kwargs).run()
+
+                # Change filename to be written
+                self.minimize_kwargs["write_kwargs"]["filename"] = (
+                    self.final_struct_min_path
+                )
+                GeomOpt(self.final_struct, **self.minimize_kwargs).run()
+
+            self.interpolate()
 
-        self.interpolate()
-        self.optimize()
+        self.optimize(fmax=fmax, steps=steps)
         self.run_nebtools()
         self.plot()
 

tests/test_neb.py

@@ -168,6 +168,67 @@ def test_neb_plot(tmp_path):
     assert neb.results["max_force"] == pytest.approx(1.5425684122118983)
 
 
+def test_converge_warning(tmp_path, LFPO_start_b, LFPO_end_b):
+    """Test warning raised if NEB does not converge."""
+    neb = NEB(
+        init_struct=LFPO_start_b,
+        final_struct=LFPO_end_b,
+        arch="mace_mp",
+        model=MODEL_PATH,
+        n_images=5,
+        steps=1,
+        fmax=0.1,
+        file_prefix=tmp_path / "LFPO",
+    )
+    with pytest.warns(UserWarning, match="NEB optimization has not converged"):
+        neb.run()
+    assert not neb.converged
+
+
+def test_restart(tmp_path):
+    """Test restarting NEB optimisation."""
+    neb = NEB(
+        neb_structs=DATA_PATH / "LiFePO4-neb-band.xyz",
+        arch="mace",
+        model=MODEL_PATH,
+        interpolator=None,
+        file_prefix=tmp_path / "LFPO",
+        fmax=1.3,
+    )
+    neb.optimize()
+    neb.run_nebtools()
+    init_force = neb.results["max_force"]
+
+    neb.optimize(fmax=1.2)
+
+    neb.run_nebtools()
+    final_force = neb.results["max_force"]
+
+    assert final_force < init_force
+
+
+def test_restart_update_climb(tmp_path):
+    """Test updating NEB climb setting when continuing optimization."""
+    results = {}
+
+    for label in ("climb", "no-climb"):
+        neb = NEB(
+            neb_structs=DATA_PATH / "LiFePO4-neb-band.xyz",
+            arch="mace",
+            model=MODEL_PATH,
+            interpolator=None,
+            fmax=1.3,
+            file_prefix=tmp_path / "LFPO",
+        )
+        neb.run()
+        neb.neb.climb = label == "climb"
+
+        neb.run(fmax=1.2)
+        results[label] = neb.results
+
+    assert results["climb"]["barrier"] != results["no-climb"]["barrier"]
+
+
 @pytest.mark.parametrize("n_images", (-5, 0, 1.5))
 def test_invalid_n_images(LFPO_start_b, LFPO_end_b, n_images):
     """Test setting invalid invalid n_images."""