mod: generalize update_image_step in Optimizer

pysisyphus/optimizers/Optimizer.py

@@ -27,15 +27,17 @@
 from pysisyphus.TablePrinter import TablePrinter
 
 
-def get_optimizer(geom, index, opt_kwargs):
+def get_optimizer(geom, index, opt_kwargs, prefix):
     # Import here, to avoid cyclic import error
     from pysisyphus.optimizers.cls_map import get_opt_cls
 
     opt_kwargs = opt_kwargs.copy()
     opt_type = opt_kwargs.pop("type")
 
+    prefix = f"{prefix}_img_{index:02d}"
     _opt_kwargs = {
-        "prefix": f"ts_img_{index:02d}",
+        "prefix": prefix,
+        "h5_group_name": f"{prefix}_opt",
         # Child optimizers don't print to STDOUT/pysisyphus.log by default,
         # only to their respective log files.
         "logging_level": logging.DEBUG,
@@ -359,6 +361,7 @@ def __init__(
         # Don't use prefix for this fn, as different optimizations
         # can be distinguished according to their group in the HDF5 file.
         self.h5_fn = self.get_path_for_fn(h5_fn, with_prefix=False)
+        # TODO: make this also depend on 'prefix'?!
         self.h5_group_name = h5_group_name
 
         current_fn = "current_geometries.trj" if self.is_cos else "current_geometry.xyz"
@@ -528,8 +531,8 @@ def check_convergence(self, step=None, multiple=1.0, overachieve_factor=None):
         if overachieve_factor is None:
             overachieve_factor = self.overachieve_factor
 
-        # When using a ChainOfStates method we are only interested
-        # in optimizing the forces perpendicular to the MEP.
+        # When using a ChainOfStates method convergence depends on the
+        # perpendicular component of the force along the MEP.
         if self.is_cos:
             forces = self.geometry.perpendicular_forces
         elif len(self.modified_forces) == len(self.forces):
@@ -723,49 +726,34 @@ def postprocess_opt(self):
     def optimize(self):
         pass
 
-    """
-    def update_step_for_ts_image(self, image_index, prev_image_step=None):
-        geom = self.geometry.images[image_index]
-        try:
-            ts_opt = self.ts_opts[image_index]
-            cur_cycle = self.ts_opt_cycles[image_index]
-        except KeyError:
-            ts_opt = get_ts_optimizer(geom, image_index)
-            self.ts_opts[image_index] = ts_opt
-            cur_cycle = 0
-            self.ts_opt_cycles[image_index] = cur_cycle
-            ts_opt.prepare_opt()
-
-        if prev_image_step is not None:
-            ts_opt.steps.append(prev_image_step)
-        ts_opt.coords.append(geom.coords.copy())
-        ts_opt.cart_coords.append(geom.cart_coords.copy())
-
-        # Calculate one step
-        updated_step = ts_opt.optimize()
-        cur_cycle = cur_cycle + 1
-        ts_opt.cur_cycle = cur_cycle
-        self.ts_opt_cycles[image_index] = cur_cycle
-        return updated_step
-    """
-
     def update_image_step(self, image_index, prefix=""):
+        """Calculate a new step for an image in a COS using another Optimizer.
+
+        This method is useful to update an image step, e.g, by a TS optimizer.
+        """
         if prefix != "":
             prefix = f"{prefix}-"
         geom = self.geometry.images[image_index]
+        # Try to look up the Optimizer ...
         try:
-            ts_opt = self.image_opts[image_index]
+            opt = self.image_opts[image_index]
+        # ... if it is not already present we create it and store it.
         except KeyError:
-            ts_opt = get_optimizer(geom, image_index, self.image_opt_kwargs)
-            self.image_opts[image_index] = ts_opt
+            opt = get_optimizer(
+                geom, image_index, opt_kwargs=self.image_opt_kwargs, prefix="ts"
+            )
+            self.image_opts[image_index] = opt
             self.print(f"Created new {prefix}optimizer for image {image_index}")
 
         try:
-            next(ts_opt.run_generator)
+            # Advance the generator one cycle and fetch the new step
+            next(opt.run_generator)
+            updated_step = opt.steps[-1]
         except StopIteration:
-            return np.zeros_like(self.geometry.images[0].coords)
+            # When the optimzation already converged the generator is exhausted
+            # and we return a zero-step.
+            updated_step = np.zeros_like(self.geometry.images[0].coords)
 
-        updated_step = self.image_opts[image_index].steps[-1]
         return updated_step
 
     def write_to_out_dir(self, out_fn, content, mode="w"):
@@ -877,6 +865,7 @@ def get_run_generator(self):
 
         if not self.restarted:
             prep_start_time = time.time()
+            # Set up the optimizer in prepare_opt, e.g., prepare initial Hessiane etc.
             self.prepare_opt()
             self.log(f"{self.geometry.coords.size} degrees of freedom.")
             prep_end_time = time.time()
@@ -921,7 +910,7 @@ def get_run_generator(self):
                     self.coords[-1], self.energies[-1], self.forces[-1]
                 )
                 if messages:
-                    self.log("\n".join(messages))
+                    self.print("\n".join(messages))
             # Reset when number of coordinates changed, compared to the previous cycle.
             elif self.cur_cycle > 0:
                 reset_flag = reset_flag or (

tests/test_cos/test_cos.py

@@ -9,12 +9,13 @@
 from pysisyphus.cos.SimpleZTS import SimpleZTS
 from pysisyphus.Geometry import Geometry
 from pysisyphus.interpolate.Interpolator import Interpolator
-from pysisyphus.optimizers.LBFGS import LBFGS
 from pysisyphus.optimizers.ConjugateGradient import ConjugateGradient
-from pysisyphus.optimizers.QuickMin import QuickMin
 from pysisyphus.optimizers.FIRE import FIRE
-from pysisyphus.optimizers.SteepestDescent import SteepestDescent
+from pysisyphus.optimizers.guess_hessians import ts_hessian_from_cos
+from pysisyphus.optimizers.LBFGS import LBFGS
+from pysisyphus.optimizers.QuickMin import QuickMin
 from pysisyphus.optimizers.RFOptimizer import RFOptimizer
+from pysisyphus.optimizers.SteepestDescent import SteepestDescent
 from pysisyphus.plotters.AnimPlot import AnimPlot
 from pysisyphus.run import run_from_dict
 from pysisyphus.testing import using
@@ -103,12 +104,12 @@ def test_anapot_szts(between, param, ref_cycle):
     assert_cos_opt(opt, ref_cycle)
 
 
-def animate(opt):
+def animate(opt, show=False):
     xlim = (-2, 2.5)
     ylim = (0, 5)
     levels = (-3, 4, 80)
     ap = AnimPlot(AnaPot(), opt, xlim=xlim, ylim=ylim, levels=levels)
-    ap.animate()
+    ap.animate(show=show)
     return ap
 
 
@@ -257,3 +258,47 @@ def test_neb_climb(climb, ref_cycles):
     # plt.show()
     assert opt.is_converged
     assert opt.cur_cycle == ref_cycles
+
+
+def test_cos_image_ts_opt():
+    initial = AnaPot.get_geom((-1.05274, 1.02776, 0))
+    final = AnaPot.get_geom((1.94101, 3.85427, 0))
+    images = (initial, final)
+    interpol = Interpolator(images, between=5)
+    images = interpol.interpolate_all()
+
+    for image in images:
+        image.set_calculator(AnaPot())
+
+    cos_kwargs = {
+        "k_min": 0.01,
+        "climb": "one",
+        "ts_opt": True,
+        "ts_opt_rms": 0.0045,
+    }
+    cos = NEB(images, **cos_kwargs)
+
+    opt_kwargs = {
+        "dump": True,
+        "image_opt_kwargs": {
+            "type": "rsprfo",
+            "thresh": "gau_vtight",
+            "dump": True,
+        },
+        # Level 10 equals logging.DEBUG and is the default for child optimizers
+        # "logging_level": 10,
+    }
+    opt = SteepestDescent(cos, **opt_kwargs)
+    opt.run()
+
+    hei_coords = cos.images[cos.get_hei_index()].coords
+    ref_ts_coords = (0.6117310435507957, 1.4929729492717998, 0.0)
+
+    H = ts_hessian_from_cos(cos, cos.get_hei_index())
+    w, _ = np.linalg.eigh(H)
+    w_min = w[0]
+    assert w_min == pytest.approx(-1.46736345)
+
+    # animate(opt, show=True)
+
+    np.testing.assert_allclose(hei_coords, ref_ts_coords, atol=2e-6)