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

pysisyphus/tsoptimizers/RSIRFOptimizer.py

@@ -0,0 +1,56 @@
+# [1] https://doi.org/10.1007/s002140050387
+#     Bofill, 1998
+
+
+import numpy as np
+
+from pysisyphus.tsoptimizers.TSHessianOptimizer import TSHessianOptimizer
+
+import torch
+
+class RSIRFOptimizer(TSHessianOptimizer):
+    def optimize(self):
+        energy, gradient, H, eigvals, eigvecs, resetted = self.housekeeping()
+        self.update_ts_mode(eigvals, eigvecs)
+
+        self.log(
+            "Using projection to construct image potential gradient "
+            f"and hessian for root(s) {self.roots}."
+        )
+        # Ensure gradient is in the same subspace as eigvecs (active DOFs)
+        if isinstance(H, torch.Tensor):
+            if gradient.size(0) != eigvecs.size(0):
+                gradient = self.active_from_full(gradient)
+        else:
+            if gradient.size != eigvecs.shape[0]:
+                gradient = self.active_from_full(gradient)
+        # Projection matrix to construct g* and H*
+        if isinstance(H, torch.Tensor):
+            P = torch.eye(gradient.size(0), device=H.device, dtype=H.dtype)
+            for root in self.roots:
+                trans_vec = eigvecs[:, root]
+                P -= 2 * torch.outer(trans_vec, trans_vec)
+            H_star = P @ H
+            eigvals_, eigvecs_ = torch.linalg.eigh(H_star)
+        else:
+            P = np.eye(gradient.size)
+            for root in self.roots:
+                trans_vec = eigvecs[:, root]
+                P -= 2 * np.outer(trans_vec, trans_vec)
+            H_star = P.dot(H)
+            eigvals_, eigvecs_ = np.linalg.eigh(H_star)
+        # Neglect small eigenvalues
+        eigvals_, eigvecs_ = self.filter_small_eigvals(eigvals_, eigvecs_)
+
+        if isinstance(H, torch.Tensor):
+            grad_star = P @ gradient
+        else:
+            grad_star = P.dot(gradient)
+        step = self.get_rs_step(eigvals_, eigvecs_, grad_star, name="RS-I-RFO")
+
+        self.predicted_energy_changes.append(self.rfo_model(gradient, self.cur_H, step))
+
+        step = self.full_from_active(step)
+        if isinstance(step, torch.Tensor):
+            step = step.cpu().numpy()
+        return step

pysisyphus/tsoptimizers/RSPRFOptimizer.py

@@ -0,0 +1,182 @@
+# See [1] https://pubs.acs.org/doi/pdf/10.1021/j100247a015
+#         Banerjee, 1985
+#     [2] https://aip.scitation.org/doi/abs/10.1063/1.2104507
+#         Heyden, 2005
+#     [3] https://onlinelibrary.wiley.com/doi/abs/10.1002/jcc.540070402
+#         Baker, 1985
+#     [4] https://link.springer.com/article/10.1007/s002140050387
+#         Besalu, 1998
+
+
+import numpy as np
+
+from pysisyphus.tsoptimizers.TSHessianOptimizer import TSHessianOptimizer
+
+
+class RSPRFOptimizer(TSHessianOptimizer):
+    def optimize(self):
+        energy, gradient, H, eigvals, eigvecs, resetted = self.housekeeping()
+        self.update_ts_mode(eigvals, eigvecs)
+
+        # Transform gradient to eigensystem of hessian
+        gradient_trans = eigvecs.T.dot(gradient)
+        # Minimize energy along all modes, except the TS-mode
+        min_indices = [i for i in range(gradient_trans.size) if i not in self.roots]
+        # Maximize energy along all requested modes.
+        max_indices = [i for i in range(gradient_trans.size) if i in self.roots]
+        # Get line search steps, if requested.
+        ip_step_trans, gradient_trans = self.step_and_grad_from_line_search(
+            energy,
+            gradient_trans,
+            eigvecs,
+            min_indices,
+            max_indices,
+        )
+
+        """In the RS-(P)RFO method we have to scale the matrices with alpha.
+        Unscaled matrix (Eq. 8) in [1]:
+            (H  g) (x)          (S 0) (x)
+                       = lambda
+            (g+ 0) (1)          (0 1) (1)
+        with
+            S = alpha * Identity matrix
+        and multiplying from the left with the inverse of the scaling matrix
+            (1/alpha 0)
+
+            (0       1)
+        we get
+            (1/alpha 0) (H  g) (x)          (x)
+                                   = lambda
+            (0       1) (g+ 0) (1)          (1)
+        eventually leading to the scaled matrix:
+            (H/alpha  g/alpha) (x)          (x)
+                                   = lambda     .
+            (g+             0) (1)          (1)
+        """
+
+        alpha = self.alpha0
+        for mu in range(self.max_micro_cycles):
+            self.log(f"RS-PRFO micro cycle {mu:02d}, alpha={alpha:.6f}")
+
+            # Maximize energy along the chosen TS mode. The matrix is hardcoded
+            # as 2x2, so only first-order saddle point searches are supported.
+            # Try secular equation solver first (O(N) vs O(N^3))
+            secular_max = self.solve_rfo_secular(
+                eigvals[max_indices], gradient_trans[max_indices], alpha,
+                kind="max", prev_eigvec=self.prev_eigvec_max,
+            )
+            if secular_max is not None:
+                step_max, eigval_max, nu_max, self.prev_eigvec_max = secular_max
+            else:
+                H_aug_max = self.get_augmented_hessian(
+                    eigvals[max_indices], gradient_trans[max_indices], alpha
+                )
+                step_max, eigval_max, nu_max, self.prev_eigvec_max = self.solve_rfo(
+                    H_aug_max, "max", prev_eigvec=self.prev_eigvec_max
+                )
+
+            # Minimize energy along all modes, but the TS mode.
+            secular_min = self.solve_rfo_secular(
+                eigvals[min_indices], gradient_trans[min_indices], alpha,
+                kind="min", prev_eigvec=self.prev_eigvec_min,
+            )
+            if secular_min is not None:
+                step_min, eigval_min, nu_min, self.prev_eigvec_min = secular_min
+            else:
+                H_aug_min = self.get_augmented_hessian(
+                    eigvals[min_indices], gradient_trans[min_indices], alpha
+                )
+                step_min, eigval_min, nu_min, self.prev_eigvec_min = self.solve_rfo(
+                    H_aug_min, "min", prev_eigvec=self.prev_eigvec_min
+                )
+
+            # Calculate overlap between directions over the course of the micro cycles
+            # if mu == 0:
+            # TODO: convert back to original space
+            # ref_step_max = step_max.copy()
+            # ref_step_min = step_min.copy()
+            min_norm = np.linalg.norm(step_min)
+            max_norm = np.linalg.norm(step_max)
+            self.log(f"norm(step_max)={max_norm:.6f}")
+            self.log(f"norm(step_min)={min_norm:.6f}")
+            self.log(f"norm(step_max)/norm(step_min)={max_norm/min_norm:.2%}")
+            # Calculate overlaps with originally proposed step in mu == 0
+            # TODO: convert back to original space
+            # max_ovlp = ref_step_max @ step_max
+            # min_ovlp = ref_step_min @ step_min
+
+            # As of Eq. (8a) of [4] max_eigval and min_eigval also
+            # correspond to:
+            # max_eigval = -forces_trans[max_indices].dot(max_step)
+            # min_eigval = -forces_trans[min_indices].dot(min_step)
+
+            # Create the full PRFO step
+            step = np.zeros_like(gradient_trans)
+            step[max_indices] = step_max
+            step[min_indices] = step_min
+            step_norm = np.linalg.norm(step)
+            self.log(f"norm(step)={step_norm:.6f}")
+
+            inside_trust = step_norm <= self.trust_radius
+            if inside_trust:
+                self.log(
+                    "Restricted step satisfies trust radius of "
+                    f"{self.trust_radius:.6f}"
+                )
+                self.log(
+                    f"Micro-cycles converged in cycle {mu:02d} with "
+                    f"alpha={alpha:.6f}!"
+                )
+                break
+
+            # Derivative of the squared step w.r.t. alpha
+            # max subspace
+            dstep2_dalpha_max = (
+                2
+                * eigval_max
+                / (1 + step_max.dot(step_max) ** 2 * alpha)
+                * np.sum(
+                    gradient_trans[max_indices] ** 2
+                    / (eigvals[max_indices] - eigval_max * alpha) ** 3
+                )
+            )
+            # min subspace
+            dstep2_dalpha_min = (
+                2
+                * eigval_min
+                / (1 + step_min.dot(step_min) * alpha)
+                * np.sum(
+                    gradient_trans[min_indices] ** 2
+                    / (eigvals[min_indices] - eigval_min * alpha) ** 3
+                )
+            )
+            dstep2_dalpha = dstep2_dalpha_max + dstep2_dalpha_min
+            # Update alpha
+            alpha_step = (
+                2 * (self.trust_radius * step_norm - step_norm**2) / dstep2_dalpha
+            )
+            alpha += alpha_step
+
+        # Right now the step is still given in the Hessians eigensystem. We
+        # transform it back now.
+        step += ip_step_trans
+        step = eigvecs.dot(step)
+        step_norm = np.linalg.norm(step)
+
+        # With max_micro_cycles = 1 the RS part is disabled and the step
+        # probably isn't scaled correctly in the one micro cycle.
+        # In this case we use a naive scaling if the step is too big.
+        if (self.max_micro_cycles == 1) and (step_norm > self.trust_radius):
+            step = step / step_norm * self.trust_radius
+        self.log(f"norm(step)={np.linalg.norm(step):.6f}")
+
+        # Eq. (6) from [4] seems erronous ... the prediction is usually only ~50%
+        # of the actual change ...
+        # predicted_energy_change = 1/2 * (eigval_max / nu_max**2 + eigval_min / nu_min**2)
+        # self.predicted_energy_changes.append(predicted_energy_change)
+
+        self.predicted_energy_changes.append(self.rfo_model(gradient, self.cur_H, step))
+
+        self.log("")
+        step = self.full_from_active(step)
+        return step

pysisyphus/tsoptimizers/TRIM.py

@@ -0,0 +1,59 @@
+# [1] https://doi.org/10.1016/0009-2614(91)90115-P
+#     Helgaker, 1991
+
+
+import numpy as np
+
+from scipy.optimize import newton
+from pysisyphus.tsoptimizers.TSHessianOptimizer import TSHessianOptimizer
+
+
+class TRIM(TSHessianOptimizer):
+
+    def optimize(self):
+        energy, gradient, H, eigvals, eigvecs, resetted = self.housekeeping()
+        self.update_ts_mode(eigvals, eigvecs)
+
+        self.log(f"Signs of eigenvalue and -vector of root(s) {self.roots} "
+                  "will be reversed!")
+        # Transform gradient to basis of eigenvectors
+        gradient_ = eigvecs.T.dot(gradient)
+
+        # Construct image function by inverting the signs of the eigenvalue and
+        # -vector of the mode to follow uphill.
+        eigvals_ = eigvals.copy()
+        eigvals_[self.roots] *= -1
+        gradient_ = gradient_.copy()
+        gradient_[self.roots] *= -1
+
+        def get_step(mu):
+            zetas = -gradient_ / (eigvals_ - mu)
+            # Replace nan with 0.
+            zetas = np.nan_to_num(zetas)
+            # Transform to original basis
+            step = eigvecs * zetas
+            step = step.sum(axis=1)
+            return step
+
+        def get_step_norm(mu):
+            return np.linalg.norm(get_step(mu))
+
+        def func(mu):
+            return get_step_norm(mu) - self.trust_radius
+
+        mu = 0
+        norm0 = get_step_norm(mu)
+        if norm0 > self.trust_radius:
+            mu, res = newton(func, x0=mu, full_output=True)
+            assert res.converged
+            self.log(f"Using levelshift of μ={mu:.4f}")
+        else:
+            self.log("Took pure newton step without levelshift")
+
+        step = get_step(mu)
+        step_norm = np.linalg.norm(step)
+        self.log(f"norm(step)={step_norm:.6f}")
+
+        self.predicted_energy_changes.append(self.quadratic_model(gradient, self.H, step))
+
+        return step