Add steps option for NEB

ElliottKasoar · Feb 6, 2025 · 78883a1 · 78883a1
1 parent 3549f9e
commit 78883a1
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/janus_core/calculations/neb.py b/janus_core/calculations/neb.py
@@ -101,6 +101,8 @@ class NEB(BaseCalculation):
         Optimizer to apply to NEB object. Default is NEBOptimizer.
     fmax
         Maximum force for NEB optimizer. Default is 0.1.
+    steps
+        Maximum number of steps to optimize NEB. Default is 100.
     optimizer_kwargs
         Keyword arguments to pass to neb_optimizer. Deault is {}.
     plot_band
@@ -148,6 +150,7 @@ def __init__(
         interpolation_kwargs: dict[str, Any] | None = None,
         neb_optimizer: Callable | str = NEBOptimizer,
         fmax: float = 0.1,
+        steps: int = 100,
         optimizer_kwargs: dict[str, Any] | None = None,
         plot_band: bool = False,
         minimize: bool = False,
@@ -227,6 +230,8 @@ def __init__(
             Optimizer to apply to NEB object. Default is NEBOptimizer.
         fmax
             Maximum force for NEB optimizer. Default is 0.1.
+        steps
+            Maximum number of steps to optimize NEB. Default is 100.
         optimizer_kwargs
             Keyword arguments to pass to neb_optimizer. Deault is {}.
         plot_band
@@ -266,6 +271,7 @@ def __init__(
         self.interpolation_kwargs = interpolation_kwargs
         self.neb_optimizer = neb_optimizer
         self.fmax = fmax
+        self.steps = steps
         self.optimizer_kwargs = optimizer_kwargs
         self.plot_band = plot_band
         self.minimize = minimize
@@ -513,7 +519,7 @@ def run(self) -> dict[str, float]:
         self.set_interpolator()
 
         optimizer = self.neb_optimizer(self.neb, **self.optimizer_kwargs)
-        optimizer.run(fmax=self.fmax)
+        optimizer.run(fmax=self.fmax, steps=self.steps)
 
         # Optionally write band images to file
         output_structs(

diff --git a/janus_core/cli/neb.py b/janus_core/cli/neb.py
@@ -61,6 +61,9 @@ def neb(
         Option(help="Name of ASE NEB optimizer to use."),
     ] = "NEBOptimizer",
     fmax: Annotated[float, Option(help="Maximum force for NEB optimizer.")] = 0.1,
+    steps: Annotated[
+        int, Option(help="Maximum number of steps for optimization.")
+    ] = 100,
     optimizer_kwargs: NebOptKwargs = None,
     plot_band: Annotated[
         bool,
@@ -131,6 +134,8 @@ def neb(
         Optimizer to apply to NEB object. Default is NEBOptimizer.
     fmax
         Maximum force for NEB optimizer. Default is 0.1.
+    steps
+        Maximum number of steps to optimize NEB. Default is 100.
     optimizer_kwargs
         Keyword arguments to pass to neb_optimizer. Deault is {}.
     plot_band
@@ -237,6 +242,7 @@ def neb(
         "interpolation_kwargs": interpolation_kwargs,
         "neb_optimizer": neb_optimizer,
         "fmax": fmax,
+        "steps": steps,
         "optimizer_kwargs": optimizer_kwargs,
         "plot_band": plot_band,
         "minimize": minimize,