mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/tsoptimizers/TSHessianOptimizer.py

@@ -0,0 +1,463 @@
+from typing import List, Optional
+
+import h5py
+import numpy as np
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import log
+from pysisyphus.intcoords.augment_bonds import augment_bonds
+from pysisyphus.intcoords.PrimTypes import normalize_prim_input, normalize_prim_inputs
+from pysisyphus.optimizers import poly_fit
+from pysisyphus.optimizers.guess_hessians import ts_hessian, HessInit
+from pysisyphus.optimizers.HessianOptimizer import HessianOptimizer, HessUpdate
+from pysisyphus.optimizers.Optimizer import get_data_model, get_h5_group
+
+from pysisyphus.helpers import array2string
+import torch
+
+class TSHessianOptimizer(HessianOptimizer):
+    """Optimizer to find first-order saddle points."""
+
+    valid_updates = ("bofill", "ts_bfgs", "ts_bfgs_org", "ts_bfgs_rev")
+
+    def __init__(
+        self,
+        geometry: Geometry,
+        roots: Optional[List[int]] = None,
+        root: int = 0,
+        hessian_ref: Optional[str] = None,
+        rx_modes=None,
+        prim_coord=None,
+        rx_coords=None,
+        hessian_init: HessInit = "calc",
+        hessian_update: HessUpdate = "bofill",
+        hessian_recalc_reset: bool = True,
+        max_micro_cycles: int = 50,
+        trust_radius: float = 0.3,
+        trust_max: float = 0.5,
+        augment_bonds: bool = False,
+        min_line_search: bool = False,
+        max_line_search: bool = False,
+        assert_neg_eigval: bool = False,
+        **kwargs,
+    ) -> None:
+        """Baseclass for transition state optimizers utilizing Hessian information.
+
+        Several arguments expect a typed primitive or an iterable of typed primitives.
+        A typed primitive is specified as (PrimType, int, int, ...), e.g., for a bond
+        between atoms 0 and 1: (BOND, 0, 1) or for a bend between the atom triple 0, 1, 2
+        as (BEND, 0, 1, 2).
+
+        Parameters
+        ----------
+        geometry
+            Geometry to be optimized.
+        roots
+            Indices of modes to maximize along, e.g., to optimize saddle points of 2nd order.
+            Overrides 'root'.
+        root
+            Index of imaginary mode to maximize along. Shortcut for 'roots' with only one root.
+        hessian_ref
+            Filename pointing to a pysisyphus HDF5 Hessian.
+        rx_modes : iterable of (iterable of (typed_prim, phase_factor))
+            Select initial root(s) by overlap with a modes constructed from the given
+            typed primitives with respective phase factors.
+        prim_coord : typed_prim
+            Select initial root/imaginary mode by overlap with this internal coordinate.
+            Shortcut for 'rx_modes' with only one internal coordinate.
+        rx_coords : iterable of (typed_prim)
+            Construct imaginary mode comprising the given typed prims by modifying
+            a model Hessian.
+        hessian_init
+            Type of initial model Hessian.
+        hessian_update
+            Type of Hessian update. Defaults to BFGS for minimizations and Bofill
+            for saddle point searches.
+        hessian_recalc_reset
+            Whether to skip Hessian recalculation after reset. Undocumented.
+        max_micro_cycles
+            Maximum number of RS iterations.
+        trust_radius
+            Initial trust radius in whatever unit the optimization is carried out.
+        trust_max
+            Maximum trust radius.
+        augment_bonds
+            Try to derive additional streching coordinates from the imaginary mode.
+        min_line_search
+            Carry out line search along the imaginary mode.
+        max_line_search
+            Carry out line search in the subspace that is minimized.
+        assert_neg_eigval
+            Check for the existences for at least one significant negative eigenvalue.
+            If enabled and no negative eigenvalue is present the optimization will be
+            aborted.
+
+        Other Parameters
+        ----------------
+        **kwargs
+            Keyword arguments passed to the HessianOptimizer/Optimizer baseclass.
+        """
+
+        assert (
+            hessian_update in self.valid_updates
+        ), f"Invalid Hessian update. Please chose from: {self.valid_updates}!"
+
+        super().__init__(
+            geometry,
+            hessian_init=hessian_init,
+            hessian_update=hessian_update,
+            trust_radius=trust_radius,
+            trust_max=trust_max,
+            hessian_recalc_reset=hessian_recalc_reset,
+            **kwargs,
+        )
+
+        if (root is not None) and (roots is None):
+            roots = [
+                root,
+            ]
+        elif roots is None:
+            roots = list()
+        self.roots = roots
+        self.log(f"{self.roots=}")
+        self.hessian_ref = hessian_ref
+        try:
+            with h5py.File(self.hessian_ref, "r") as handle:
+                self.hessian_ref = handle["hessian"][:]
+            _ = self.geometry.coords.size
+            expected_shape = (_, _)
+            shape = self.hessian_ref.shape
+            # Hessian is not yet converted to the correct coordinate system if
+            # coord_type != cart.
+            assert (
+                self.geometry.coord_type == "cart"
+            ), "hessian_ref with internal coordinates are not yet supported."
+            assert shape == expected_shape, (
+                f"Shape of reference Hessian {shape} doesn't match the expected "
+                f"shape {expected_shape} of the Hessian for the current coordinates!"
+            )
+        except OSError as err:
+            self.log(
+                f"Tried to load reference Hessian from '{self.hessian_ref}' "
+                "but the file could not be found."
+            )
+            self.hessian_ref = None
+        except (ValueError, TypeError) as err:
+            self.log(f"No reference Hessian provided.")
+
+        # Select initial root according to highest contribution of 'prim_coord'
+        if prim_coord is not None:
+            self.log("'prim_coord' given. Overwriting/setting 'rx_modes'.")
+            rx_modes = [[[prim_coord, 1.0]]]
+        self.prim_coord = prim_coord
+
+        # Construct initial imaginary mode from the given inputs while respecting
+        # the given weight/phase factors.
+        self.rx_modes = rx_modes
+
+        # Construct initial imaginary mode according the primitive internals in
+        # 'rx_coords' by modifying a model Hessian.
+        if rx_coords is not None:
+            rx_coords = normalize_prim_inputs(rx_coords)
+        self.rx_coords = rx_coords
+
+        # Bond augmentation is only useful with calculated hessians
+        self.augment_bonds = augment_bonds and (self.hessian_init == "calc")
+        self.min_line_search = min_line_search
+        self.max_line_search = max_line_search
+        self.assert_neg_eigval = assert_neg_eigval
+
+        self.ts_modes = list()
+        self.max_micro_cycles = max_micro_cycles
+        self.prim_contrib_thresh = 0.05
+
+        self.alpha0 = 1
+
+    @property
+    def root(self):
+        return self.roots[0]
+
+    @root.setter
+    def root(self, root):
+        raise Exception("Setting 'self.root' is deprecated. Set 'self.roots' instead.")
+
+    @property
+    def roots(self):
+        return self._roots
+
+    @roots.setter
+    def roots(self, roots):
+        self._roots = np.array(roots, dtype=int)
+
+    def prepare_opt(self, *args, **kwargs):
+        if self.augment_bonds:
+            self.geometry = augment_bonds(self.geometry, root=self.root)
+            # Update data model and HD5 shapes, as the number of coordinates
+            # may have changed.
+            if self.dump:
+                self.data_model = get_data_model(
+                    self.geometry, self.is_cos, self.max_cycles
+                )
+                # self.h5_group = get_h5_group(
+                #     self.h5_fn, self.h5_group_name, self.data_model
+                # )
+
+        # Calculate/set initial hessian
+        super().prepare_opt(*args, **kwargs)
+
+        # Assume a guess hessian when not calculated. This hessian has to be
+        # modified according to the assumed reaction coordinates.
+        if self.hessian_init != "calc" and (self.rx_coords is not None):
+            assert self.geometry.coord_type != "cart", (
+                "Using a modified guess Hessian for TS-optimizations is "
+                "only supported in redundand internal coordinates "
+                "(coord_type=redund)"
+            )
+            prim_inds = [
+                self.geometry.internal.get_index_of_typed_prim(rxc)
+                for rxc in self.rx_coords
+            ]
+            missing_prim_inds = [
+                self.rx_coords[i] for i, _ in enumerate(prim_inds) if _ is None
+            ]
+            assert len(missing_prim_inds) == 0, (
+                "Some of the requested reaction coordinates are not defined: "
+                f"{missing_prim_inds}"
+            )
+            self.H = ts_hessian(self.H, coord_inds=prim_inds)
+
+        # Determiniation of initial mode either by using a provided
+        # reference hessian, or by using a supplied root.
+
+        if isinstance(self.H, torch.Tensor):
+            eigvals, eigvecs = torch.linalg.eigh(self.H)
+            eigvals = eigvals.cpu().numpy()
+        else:
+            eigvals, eigvecs = np.linalg.eigh(self.H)
+        neg_inds = eigvals < -self.small_eigval_thresh
+        self.log_negative_eigenvalues(eigvals, "Initial ")
+
+        self.log("Determining initial TS mode to follow uphill.")
+        # Select an initial TS-mode by highest overlap with eigenvectors from
+        # reference Hessian.
+        if self.hessian_ref is not None:
+            eigvals_ref, eigvecs_ref = np.linalg.eigh(self.hessian_ref)
+            self.log_negative_eigenvalues(eigvals_ref, "Reference ")
+            assert eigvals_ref[0] < -self.small_eigval_thresh
+            ref_mode = eigvecs_ref[:, 0]
+            overlaps = np.einsum("ij,j->i", eigvecs.T, ref_mode)
+            ovlp_str = array2string(overlaps, precision=4)
+            self.log(
+                "Overlaps between eigenvectors of current Hessian "
+                f"TS mode from reference Hessian:"
+            )
+            self.log(f"\t{ovlp_str}")
+
+            root = np.abs(overlaps).argmax()
+            self.roots = [
+                root,
+            ]
+            print(
+                "Highest overlap between reference TS mode and "
+                f"eigenvector {self.root:02d}."
+            )
+            used_str = "overlap with reference TS mode"
+        # Construct an approximate initial mode according to user input
+        # and calculate overlaps with the current eigenvectors.
+        elif self.rx_modes is not None:
+            self.log(f"Constructing reference mode, according to user input")
+            assert self.geometry.coord_type != "cart"
+            int_ = self.geometry.internal
+            modes = list()
+            for i, rx_mode in enumerate(self.rx_modes):
+                mode = np.zeros_like(self.geometry.coords)
+                for typed_prim, val in rx_mode:
+                    typed_prim = normalize_prim_input(typed_prim)[0]
+                    ind = int_.get_index_of_typed_prim(typed_prim)
+                    mode[ind] = val
+                    self.log(
+                        f"\tIndex {ind} (coordinate {typed_prim}) set to {val:.4f}"
+                    )
+                mode /= np.linalg.norm(mode)
+                modes.append(mode)
+                mode_str = array2string(mode, precision=2)
+                self.log(f"Normalized reference mode {i:02d}: {mode_str}")
+
+            # Calculate overlaps in non-redundant subspace by zeroing overlaps
+            # in the redundant subspace.
+            small_inds = np.abs(eigvals) < self.small_eigval_thresh
+            # Take absolute value, because sign of eigenvectors is ambiguous.
+            ovlps = np.abs(np.einsum("ij,ki->kj", eigvecs, modes))
+            ovlps[:, small_inds] = 0.0
+            self.roots = ovlps.argmax(axis=1)
+            used_str = "overlap with user-generated mode"
+        else:
+            used_str = f"root(s)={self.roots}"
+        self.log(f"Used {used_str} to select inital TS mode.")
+
+        # Below, some code is found, that checks if the chosen root(s) are a
+        # sensible choice, i.e., if they are negative. Currently, it is commented out,
+        # as we can also start from the convex region of the PES.
+        #
+        # Check if the selected mode (root) is a sensible choice.
+        #
+        # small_eigval_thresh is positive and we dont take the absolute value
+        # of the eigenvalues. So we multiply small_eigval_thresh to get a
+        # negative number.
+        # assert (eigvals[self.roots] < -self.small_eigval_thresh).all(), (
+        # "Expected negative eigenvalue(s)! Eigenvalues of selected TS-modes "
+        # f"are above the the threshold of self.small_eigval_thresh:.6e}!"
+        # )
+
+        # Select an initial TS-mode by root index. self.root may have been
+        # modified by using a reference hessian.
+        self.ts_modes = eigvecs[:, self.roots].T
+        self.ts_mode_eigvals = eigvals[self.roots]
+        self.log(
+            f"Using root(s) {self.roots} with eigenvalues "
+            f"{array2string(self.ts_mode_eigvals, precision=6)} as TS mode.\n"
+        )
+
+    def update_ts_mode(self, eigvals, eigvecs):
+        neg_eigval_inds = eigvals < -self.small_eigval_thresh
+        neg_num = neg_eigval_inds.sum()
+        self.log_negative_eigenvalues(eigvals)
+
+        # When we left the convex region of the PES we only compare to other
+        # imaginary modes ... is this a bad idea? Maybe we should use all modes
+        # for the overlaps?!
+        if (self.ts_mode_eigvals < 0).all() and neg_num > 0:
+            infix = "imaginary "
+            ovlp_eigvecs = eigvecs[:, :neg_num]
+            eigvals = eigvals[:neg_num]
+        # When the eigenvalue corresponding to the TS mode has been negative once,
+        # we should not lose all negative eigenvalues. If this happens something went
+        # wrong and we crash :)
+        elif self.assert_neg_eigval and neg_num == 0:
+            raise AssertionError(
+                "Need at least 1 negative eigenvalue for TS optimization."
+            )
+        # Use all eigenvectors for overlaps when the eigenvalue corresponding to the TS
+        # mode is still positive.
+        else:
+            infix = ""
+            ovlp_eigvecs = eigvecs
+
+        if (
+            self.using_active_dofs
+            and self.ts_modes is not None
+            and self.ts_modes.shape[1] != ovlp_eigvecs.shape[0]
+        ):
+            self.log("Projecting previous TS mode(s) onto active DOFs.")
+            if isinstance(self.ts_modes, torch.Tensor):
+                self.ts_modes = torch.stack(
+                    [self.active_from_full(mode) for mode in self.ts_modes]
+                )
+            else:
+                self.ts_modes = np.stack(
+                    [self.active_from_full(mode) for mode in self.ts_modes]
+                )
+
+        # Select new TS mode according to biggest overlap with previous TS mode.
+        self.log(f"Overlaps of previous TS mode with current {infix}mode(s):")
+        if isinstance(eigvecs, torch.Tensor):
+            ovlps = torch.abs(torch.einsum("ij,ki->kj", ovlp_eigvecs, self.ts_modes))
+            ovlps = ovlps.cpu().numpy()
+        else:
+            ovlps = np.abs(np.einsum("ij,ki->kj", ovlp_eigvecs, self.ts_modes))
+        for i, ovlp in enumerate(ovlps):
+            self.log(f"\tTS mode {i:02d}: {array2string(ovlp, precision=3)}")
+        max_ovlp_inds = np.argmax(ovlps, axis=1)
+        for i, _ in enumerate(self.ts_modes):
+            max_ovlp_ind = max_ovlp_inds[i].argmax()
+            self.log(
+                f"Mode {i}: highest overlap: {ovlps[i, max_ovlp_ind]:.6f} with mode "
+                f"{max_ovlp_ind}"
+            )
+        self.roots = max_ovlp_inds
+        self.ts_modes = ovlp_eigvecs.T[self.roots]
+        self.ts_mode_eigvals = eigvals[self.roots]
+
+    @staticmethod
+    def do_line_search(e0, e1, g0, g1, prev_step, maximize, logger=None):
+        poly_fit_kwargs = {
+            "e0": e0,
+            "e1": e1,
+            "g0": g0,
+            "g1": g1,
+            "maximize": maximize,
+        }
+        if not maximize:
+            poly_fit_kwargs.update(
+                {
+                    "g0": prev_step.dot(g0),
+                    "g1": prev_step.dot(g1),
+                }
+            )
+        prefix = "Max" if maximize else "Min"
+
+        fit_result = poly_fit.quartic_fit(**poly_fit_kwargs)
+        fit_energy = None
+        fit_grad = None
+        fit_step = None
+        if fit_result and (0.0 < fit_result.x <= 2.0):
+            x = fit_result.x
+            log(logger, f"{prefix}-subpsace interpolation succeeded: x={x:.6f}")
+            fit_energy = fit_result.y
+            fit_step = (1 - x) * -prev_step
+            fit_grad = (1 - x) * g0 + x * g1
+        return fit_energy, fit_grad, fit_step
+
+    def step_and_grad_from_line_search(
+        self,
+        energy,
+        gradient_trans,
+        eigvecs,
+        min_indices,
+        max_indices,
+    ):
+        ip_step = np.zeros_like(gradient_trans)
+        ip_gradient_trans = gradient_trans.copy()
+
+        if self.max_line_search and self.cur_cycle > 0:
+            prev_energy = self.energies[-2]
+            prev_gradient = -self.forces[-2]
+            try:
+                prev_gradient_trans = eigvecs.T.dot(prev_gradient)
+                prev_step = self.steps[-1]
+                prev_step_trans = eigvecs.T.dot(prev_step)
+            # Will be raised when coordinates were rebuilt and the array shapes differe.
+            except ValueError:
+                return ip_step, ip_gradient_trans
+
+            # Max subspace
+            # max_energy, max_gradient, max_step = self.do_max_line_search(
+            max_energy, max_gradient, max_step = self.do_line_search(
+                prev_energy,
+                energy,
+                prev_gradient_trans[max_indices],
+                gradient_trans[max_indices],
+                prev_step=prev_step_trans[max_indices],
+                maximize=True,
+                logger=self.logger,
+            )
+            if max_gradient is not None:
+                ip_gradient_trans[max_indices] = max_gradient
+                ip_step[max_indices] = max_step
+
+        if self.min_line_search and self.cur_cycle > 0:
+            # Min subspace
+            # min_energy, min_gradient, min_step = self.do_min_line_search(
+            min_energy, min_gradient, min_step = self.do_line_search(
+                prev_energy,
+                energy,
+                prev_gradient_trans[min_indices],
+                gradient_trans[min_indices],
+                prev_step=prev_step_trans[min_indices],
+                maximize=False,
+                logger=self.logger,
+            )
+            if min_gradient is not None:
+                ip_gradient_trans[min_indices] = min_gradient
+                ip_step[min_indices] = min_step
+        return ip_step, ip_gradient_trans

pysisyphus/tsoptimizers/__init__.py

@@ -0,0 +1,23 @@
+import logging
+
+from pysisyphus.tsoptimizers.RSPRFOptimizer import RSPRFOptimizer
+from pysisyphus.tsoptimizers.TRIM import TRIM
+from pysisyphus.tsoptimizers.RSIRFOptimizer import RSIRFOptimizer
+
+__all__ = [
+    "RSPRFOptimizer",
+    "TRIM",
+    "RSIRFOptimizer",
+    "TSHessianOptimizer",
+]
+
+
+logger = logging.getLogger("tsoptimizer")
+logger.setLevel(logging.DEBUG)
+# delay = True prevents creation of empty logfiles
+handler = logging.FileHandler("tsoptimizer.log", mode="w", delay=True)
+# fmt_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+fmt_str = "%(message)s"
+formatter = logging.Formatter(fmt_str)
+handler.setFormatter(formatter)
+logger.addHandler(handler)