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

pysisyphus/optimizers/StabilizedQNMethod.py

@@ -0,0 +1,248 @@
+# [1] https://aip.scitation.org/doi/pdf/10.1063/1.4905665?class=pdf
+
+import numpy as np
+
+from pysisyphus.optimizers.Optimizer import Optimizer
+from pysisyphus.intcoords.setup import get_bond_mat
+from pysisyphus.optimizers.restrict_step import scale_by_max_step
+
+
+class StabilizedQNMethod(Optimizer):
+    def __init__(
+        self,
+        geometry,
+        alpha=0.5,
+        alpha_max=1,
+        alpha_stretch=0.5,
+        alpha_stretch_max=1,
+        eps=1e-4,
+        hist_max=10,
+        E_thresh=1e-6,
+        bio=True,
+        trust_radius=0.1,
+        linesearch=True,
+        **kwargs,
+    ):
+        super().__init__(geometry, **kwargs)
+
+        self.alpha = alpha
+        self.alpha_max = alpha_max
+        self.alpha_stretch = alpha_stretch
+        self.alpha_stretch_max = alpha_stretch_max
+        self.eps = eps
+        self.hist_max = int(hist_max)
+        self.E_thresh = E_thresh
+        self.bio = bio
+        self.trust_radius = trust_radius
+        self.linesearch = linesearch
+
+        self.log(
+            f"Keeping at most information from {self.hist_max} " "previous cycles."
+        )
+
+        self.alpha_start = self.alpha
+        self.alpha_stretch_start = self.alpha_stretch
+
+        self.gradients_for_precon = list()
+        self.coords_for_precon = list()
+        self.stretch_proj_signs = list()
+
+    def prepare_opt(self):
+        pass
+
+    @property
+    def n_hist(self):
+        return len(self.steps_normed)
+
+    def bio_mode(self, gradient):
+        bond_mat = get_bond_mat(self.geometry)
+        atom_num = len(self.geometry.atoms)
+        bond_vec_empty = np.zeros((atom_num, 3))
+        c3d = self.geometry.coords3d
+        bond_vectors = list()
+        unique_bonds = set(
+            [(min(m, k), max(m, k)) for m, k in zip(*np.where(bond_mat == True))]
+        )
+        for m, k in unique_bonds:
+            bm = bond_vec_empty.copy()
+            bm[k] = c3d[k] - c3d[m]
+            bm[m] = -bm[k]
+            bond_vectors.append(bm)
+        # The 3d array is reshaped to 2d, so its last dimension matches
+        # the gradient's dimension. LHS and RHS are accessible by simple
+        # dot products (see [1] (27)).
+        bond_vectors = np.array(bond_vectors).reshape(-1, gradient.size)
+        lhs = bond_vectors.dot(gradient)
+        self.stretch_proj_signs.append(np.sign(lhs).astype(int))
+        # Bond-vector overlap matrix.
+        rhs = bond_vectors.dot(bond_vectors.T)
+        coeffs = np.linalg.solve(rhs, lhs)
+        stretch_gradient = np.einsum("i,ij->j", coeffs, bond_vectors)
+        remainder_gradient = gradient - stretch_gradient
+        return stretch_gradient, remainder_gradient
+
+    def adjust_alpha_stretch(self):
+        try:
+            # -1 when signs changed, 1 when sign is constant
+            mult = self.stretch_proj_signs[-1] * self.stretch_proj_signs[-2]
+        except IndexError:
+            self.log("Can't update alpha_stretch yet!")
+            return
+        except ValueError:
+            self.log(
+                "Number of bonds found changed between cycles! Skipping "
+                "adjustment of alpha_stretch."
+            )
+            return
+        constant_signs = np.sum(mult == 1)
+        fraction = constant_signs / mult.size
+        if fraction >= (2 / 3):
+            self.alpha_stretch = min(self.alpha_stretch_max, 1.1 * self.alpha_stretch)
+            msg = "Increased"
+        else:
+            self.alpha_stretch *= 1 / 1.1
+            msg = "Decreased"
+        self.log(f"{msg} alpha_stretch to {self.alpha_stretch:.6f}")
+
+    def adjust_alpha(self, gradient, precon_gradient):
+        # Overlap between full and preconditioned gradient.
+        grad_ovlp = (
+            gradient.dot(precon_gradient)
+            / np.linalg.norm(gradient)
+            / np.linalg.norm(precon_gradient)
+        )
+        if grad_ovlp > 0.2:
+            self.alpha = min(self.alpha_max, 1.1 * self.alpha)
+        else:
+            self.alpha *= 0.85
+        self.log(
+            "Overlap between full gradient and preconditioned gradient "
+            f"is {grad_ovlp:.4f}. New alpha={self.alpha:.4f}."
+        )
+
+    def precondition_gradient(self, gradient, steps, grad_diffs, eps):
+        # Construct overlap matrix S between normalized steps
+        #
+        step_norms = np.linalg.norm(steps, axis=1)
+        steps_normed = steps / step_norms[:, None]
+        S = np.einsum("ij,jk->ik", steps_normed, steps_normed.T)
+        # Determine significant subspace by checking the eigenvalues of
+        # the overlap matrix.
+        w, v = np.linalg.eigh(S)
+        # Significant subspace indices
+        significant = (w / w.max()) > eps
+        ndim = np.sum(significant)
+        self.log(f"ndim = {ndim}")
+        # Continue only with eigenvalues and -vectors in the significant
+        # subspace.
+        eigvals = w[significant]
+        eigvecs = v[:, significant]
+
+        # Transform steps and gradient differences to significant subspace
+        norm_fac = 1 / eigvals ** (1 / 2)
+        # Eq. (8) in [1]
+        steps_normed_sub = norm_fac * np.einsum("ki,kj->ji", eigvecs, steps_normed)
+
+        # Eq. (11) in [1]
+        eigvecs_weighted = eigvecs / np.array(step_norms)[:, None]
+        grad_diffs_sub = norm_fac * np.einsum("ki,kj->ji", eigvecs_weighted, grad_diffs)
+
+        # Assemble approximate hessian
+        hess_approx = np.einsum("ij,ik->jk", grad_diffs_sub, steps_normed_sub)
+        hess_approx = (hess_approx + hess_approx.T) / 2
+        hess_w, hess_v = np.linalg.eigh(hess_approx)
+
+        # Calculate step directions and gradient difference in the original
+        # space. Eq. (15)
+        proj_v = np.einsum("ki,jk->ji", hess_v, steps_normed_sub)
+        proj_dg = np.einsum("ki,jk->ji", hess_v, grad_diffs_sub)
+
+        residuals = np.linalg.norm(proj_dg - hess_w * proj_v, axis=0)
+        eigvals_mod = np.sqrt(hess_w ** 2 + residuals ** 2)
+
+        # precon_grad = np.einsum("i,j,ij,ij->i", cur_grad, 1/eigvals_mod, proj_v, proj_v)
+        precon_grad = np.einsum(
+            "i,j,ij,ij->i", gradient, 1 / eigvals_mod, proj_v, proj_v
+        )
+
+        projector = proj_v.dot(proj_v.T)
+        perp_projector = np.eye(gradient.size) - projector
+        perp_grad = perp_projector.dot(gradient)
+        # Alternative formulation
+        # perp_grad_ = gradient - (proj_v.T.dot(gradient) * proj_v).sum(axis=1)
+        # np.testing.assert_allclose(perp_grad_, perp_grad, atol=1e-16)
+
+        self.adjust_alpha(gradient, precon_grad)
+        tot_precon_gradient = precon_grad + self.alpha * perp_grad
+        return tot_precon_gradient
+
+    def optimize(self):
+        gradient = self.geometry.gradient
+        energy = self.geometry.energy
+        self.forces.append(-gradient)
+        self.energies.append(energy)
+        self.log(f"norm(forces)={np.linalg.norm(gradient):.4e}")
+
+        if self.bio:
+            stretch_gradient, remainder_gradient = self.bio_mode(gradient)
+            self.adjust_alpha_stretch()
+            # Steepest descent against the stretch_gradient
+            stretch_step = -self.alpha_stretch * stretch_gradient
+            new_coords = self.geometry.coords + stretch_step
+            self.geometry.coords = new_coords
+            # Use only the remaining gradient for the rest of this method
+            gradient = remainder_gradient
+
+        self.gradients_for_precon.append(gradient)
+        self.coords_for_precon.append(self.geometry.coords.copy())
+
+        if len(self.coords_for_precon) > 2:
+            steps = np.diff(self.coords_for_precon, axis=0)[-self.hist_max :]
+            grad_diffs = np.diff(self.gradients_for_precon, axis=0)[-self.hist_max :]
+
+            self.log(
+                "Preconditioning gradient with information from "
+                f"{len(steps)+1} previous cycles."
+            )
+            precon_grad = self.precondition_gradient(
+                gradient, steps, grad_diffs, self.eps
+            )
+            step = -precon_grad
+        else:
+            self.log("Took pure steepest descent step.")
+            step = self.alpha * -gradient
+            scale_by_max_step(step, self.trust_radius)
+
+        if self.linesearch:
+            # Calculate energy at new geometry
+            new_coords = self.geometry.coords + step
+            _ = self.geometry.get_energy_and_forces_at(new_coords)
+            new_energy = _["energy"]
+            delta_energy = new_energy - energy
+            self.log(
+                f"Current energy is {energy:.6f} au. New energy is "
+                f"{new_energy:.6f} au, ΔE={delta_energy:.6f} au."
+            )
+
+            energy_rise_too_big = new_energy > (energy + self.E_thresh)
+            alpha_still_big_enough = self.alpha > (self.alpha_start / 10)
+            alpha_stre_still_big_enough = self.alpha_stretch > (
+                self.alpha_stretch_start / 10
+            )
+            if (
+                energy_rise_too_big
+                and alpha_still_big_enough
+                and alpha_stre_still_big_enough
+            ):
+                self.log(f"Energy increased by {delta_energy:.6f} au")
+                self.gradients_for_precon = self.gradients_for_precon[-2:-1]
+                self.coords_for_precon = self.coords_for_precon[-2:-1]
+                self.log("Resetted history.")
+                self.alpha /= 2
+                self.alpha_stretch /= 2
+
+                self.log(f"Decreased alpha to {self.alpha}")
+                self.log("Reverting bad step")
+                return None
+
+        return step

pysisyphus/optimizers/SteepestDescent.py

@@ -0,0 +1,23 @@
+from pysisyphus.optimizers.BacktrackingOptimizer import BacktrackingOptimizer
+
+
+class SteepestDescent(BacktrackingOptimizer):
+    def __init__(self, geometry, alpha=0.1, **kwargs):
+        super().__init__(geometry, alpha=alpha, **kwargs)
+
+    def prepare_opt(self):
+        self.log("no backtracking in cycle 0")
+
+    def optimize(self):
+        if self.is_cos and self.align:
+            self.procrustes()
+
+        self.forces.append(self.geometry.forces)
+        self.energies.append(self.geometry.energy)
+
+        if self.cur_cycle > 0:
+            self.skip = self.backtrack(self.forces[-1], self.forces[-2])
+
+        step = self.alpha * self.forces[-1]
+        step = self.scale_by_max_step(step)
+        return step

pysisyphus/optimizers/StringOptimizer.py

@@ -0,0 +1,173 @@
+# [1] https://aip.scitation.org/doi/abs/10.1063/1.3664901
+#     Behn, 2011, Freezing string method
+# [2] https://aip.scitation.org/doi/pdf/10.1063/1.4804162
+#     Zimmerman, 2013, Growing string with interpolation and optimization
+#                      in internal coordiantes
+
+import numpy as np
+
+from pysisyphus.optimizers.hessian_updates import double_damp
+from pysisyphus.optimizers.closures import bfgs_multiply
+from pysisyphus.optimizers.Optimizer import Optimizer
+from pysisyphus.optimizers.restrict_step import scale_by_max_step
+
+
+class StringOptimizer(Optimizer):
+    def __init__(
+        self,
+        geometry,
+        max_step=0.1,
+        stop_in_when_full=-1,
+        keep_last=10,
+        lbfgs_when_full=True,
+        gamma_mult=False,
+        double_damp=True,
+        scale_step="global",
+        **kwargs,
+    ):
+        super().__init__(geometry, max_step=max_step, **kwargs)
+
+        assert (
+            self.is_cos
+        ), "StringOptimizer is only intended to be used with COS objects."
+
+        self.stop_in_when_full = int(stop_in_when_full)
+        self.keep_last = int(keep_last)
+        assert self.keep_last >= 0
+        self.lbfgs_when_full = lbfgs_when_full
+        if self.lbfgs_when_full and (self.keep_last == 0):
+            print("lbfgs_when_full is True, but keep_last is 0!")
+        self.gamma_mult = bool(gamma_mult)
+        self.double_damp = bool(double_damp)
+        self.scale_step = scale_step
+        assert self.scale_step in ("global", "per_image")
+
+        # Add one as we later subtract 1 before we check if this value is 0.
+        self.stop_in = self.stop_in_when_full + 1
+        self.is_cart_opt = self.geometry.coord_type == "cart"
+        self.s_list = list()
+        self.y_list = list()
+        self.inds = list()
+
+    def prepare_opt(self):
+        if self.align and self.is_cart_opt:
+            self.procrustes()
+
+    def reset(self):
+        pass
+
+    def restrict_step_components(self, steps):
+        too_big = np.abs(steps) > self.max_step
+        self.log(f"Found {np.sum(too_big)} big step components.")
+        signs = np.sign(steps[too_big])
+        steps[too_big] = signs * self.max_step
+        return steps
+
+    def check_convergence(self, *args, **kwargs):
+        # Normal convergence check with gradients etc.
+        converged, conv_info = super().check_convergence(*args, **kwargs)
+
+        if self.geometry.fully_grown:
+            # We only start decrementing the counter after the string is fully grown.
+            self.stop_in -= 1
+            # Don't print this message if stop_in was disabled in the first place (< 0).
+            if self.stop_in >= 0:
+                self.log(f"String is fully grown. Stopping in {self.stop_in} cycle(s).")
+
+        fully_grown = self.geometry.fully_grown
+        full_stop = fully_grown and (self.stop_in == 0)
+        # full_stop will take precedence when True.
+        return full_stop or (fully_grown and converged), conv_info
+
+    def optimize(self):
+        new_image_inds = self.geometry.new_image_inds
+        string_size_changed = len(new_image_inds) > 0
+
+        if self.align and string_size_changed and self.is_cart_opt:
+            self.procrustes()
+            self.log("Aligned string.")
+
+        forces = self.geometry.forces
+        self.energies.append(self.geometry.energy)
+        self.forces.append(forces)
+
+        cur_size = self.geometry.string_size
+        add_to_list = (
+            self.keep_last > 0  # Only add to s_list and y_list if we want to keep
+            and self.cur_cycle
+            > 0  # cycles and if we can actually calculate differences.
+            and (
+                not self.lbfgs_when_full  # Add when LBFGS is allowed before fully grown.
+                or self.lbfgs_when_full
+                and self.geometry.fully_grown
+                and not string_size_changed  # Don't add when string has to be fully grown
+                # but grew fully in this cycle.
+            )
+        )
+        if add_to_list:
+            inds = list(range(cur_size))
+            try:
+                y = self.forces[-2] - forces
+                s = self.coords[-1] - self.coords[-2]
+            # Will be raised when the string grew in the previous cycle.
+            except ValueError:
+                cur_forces = np.delete(
+                    forces.reshape(cur_size, -1), new_image_inds, axis=0
+                ).flatten()
+                y = self.forces[-2] - cur_forces
+                cur_coords = np.delete(
+                    self.coords[-1].reshape(cur_size, -1), new_image_inds, axis=0
+                ).flatten()
+                s = self.coords[-2] - cur_coords
+                inds = np.delete(inds, new_image_inds)
+
+            if self.double_damp:
+                s, y = double_damp(s, y, s_list=self.s_list, y_list=self.y_list)
+
+            self.s_list.append(s)
+            self.y_list.append(y)
+            self.inds.append(inds)
+            # Drop oldest vectors
+            self.s_list = self.s_list[-self.keep_last :]
+            self.y_list = self.y_list[-self.keep_last :]
+            self.inds = self.inds[-self.keep_last :]
+
+        # Results in steepest descent step for empty s_list & y_list
+        step = bfgs_multiply(
+            self.s_list,
+            self.y_list,
+            forces,
+            gamma_mult=self.gamma_mult,
+            inds=self.inds,
+            cur_size=cur_size,
+            logger=self.logger,
+        )
+
+        # When LBFGS is not yet enabled then s_list and y_list will
+        # be empty and the step from bfgs_multiply will be a simple SD step.
+        # We try to calculated an improved step it via conjugate gradient.
+        previous_step_with_same_size = (self.cur_cycle > 0) and (
+            not string_size_changed
+        )
+        lbfgs_lists_empty = (len(self.s_list) == 0) and (len(self.y_list) == 0)
+        if previous_step_with_same_size and lbfgs_lists_empty:
+            prev_forces = self.forces[-2]
+            # Fletcher-Reeves
+            kind = "Fletcher-Reeves"
+            beta = forces.dot(forces) / prev_forces.dot(prev_forces)
+            # Polak-Ribiere
+            # kind = "Polak-Ribiere"
+            # beta = forces.dot(forces - prev_forces) / prev_forces.dot(prev_forces)
+            beta = min(beta, 1)
+            step = forces + beta * self.steps[-1]
+            self.log(f"{kind} conjugate gradient correction, β={beta:.6f}")
+
+        if self.scale_step == "global":
+            step = scale_by_max_step(step, self.max_step)
+        elif self.scale_step == "per_image":
+            for image_step in step.reshape(len(self.geometry.images), -1):
+                scale_by_max_step(image_step, self.max_step)
+        else:
+            raise Exception("Invalid scale_step={self.scale_step}!")
+
+        return step

pysisyphus/optimizers/__init__.py

@@ -0,0 +1,41 @@
+import logging
+
+__all__ = [
+    "BFGS",
+    "ConjugateGradient",
+    "CubicNewton",
+    "FIRE",
+    "LayerOpt",
+    "LBFGS",
+    "MicroOptimizer",
+    "NCOptimizer",
+    "PreconLBFGS",
+    "PreconSteepestDescent",
+    "QuickMin",
+    "RFOptimizer",
+    "SteepestDescent",
+    "StringOptimizer",
+    "StabilizedQNMethod",
+]
+
+from pysisyphus.optimizers.CubicNewton import CubicNewton
+from pysisyphus.optimizers.MicroOptimizer import MicroOptimizer
+
+logger = logging.getLogger("optimizer")
+logger.setLevel(logging.DEBUG)
+# delay = True prevents creation of empty logfiles
+handler = logging.FileHandler("optimizer.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)
+
+# logger = logging.getLogger("gdiis")
+# logger.setLevel(logging.DEBUG)
+# # delay = True prevents creation of empty logfiles
+# handler = logging.FileHandler("gdiis.log", mode="w", delay=True)
+# fmt_str = "%(message)s"
+# formatter = logging.Formatter(fmt_str)
+# handler.setFormatter(formatter)
+# logger.addHandler(handler)