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

pysisyphus/optimizers/PreconLBFGS.py

@@ -0,0 +1,260 @@
+from collections import deque
+from functools import partial
+from typing import Literal, Optional
+
+import numpy as np
+from scipy.sparse.linalg import spsolve
+
+from pysisyphus.calculators import Dimer
+from pysisyphus.cos.GrowingNT import GrowingNT
+from pysisyphus.Geometry import Geometry
+from pysisyphus.line_searches import *
+from pysisyphus.optimizers.closures import bfgs_multiply
+from pysisyphus.optimizers.Optimizer import Optimizer
+from pysisyphus.optimizers.precon import precon_getter, PreconKind
+
+
+LineSearch = Literal["armijo", "armijo_fg", "strong_wolfe", "hz", None, False]
+
+
+class PreconLBFGS(Optimizer):
+    def __init__(
+        self,
+        geometry: Geometry,
+        alpha_init: float = 1.0,
+        history: int = 7,
+        precon: bool = True,
+        precon_update: int = 1,
+        precon_getter_update: Optional[int] = None,
+        precon_kind: PreconKind = "full",
+        max_step_element: Optional[float] = None,
+        line_search: LineSearch = "armijo",
+        c_stab: Optional[float] = None,
+        **kwargs,
+    ) -> None:
+        """Preconditioned limited-memory BFGS optimizer.
+
+        See pysisyphus.optimizers.precon for related references.
+
+        Parameters
+        ----------
+        geometry
+            Geometry to be optimized.
+        alpha_init
+            Initial scaling factor for the first trial step in the excplicit line search.
+        history
+            History size. Keep last 'history' steps and gradient differences.
+        precon
+            Wheter to use preconditioning or not.
+        precon_update
+            Recalculate preconditioner P in every n-th cycle with the same topology.
+        precon_getter_update
+            Recalculate topology for preconditioner P in every n-th cycle. It is usually
+            sufficient to only determine the topology once at the beginning.
+        precon_kind
+            What types of primitive internal coordinates to consider in the preconditioner.
+        max_step_element
+            Maximum component of the absolute step vector when no line search is carried out.
+        line_search
+            Whether to use explicit line searches and if so, which kind of line search.
+        c_stab
+            Regularization constant c in (H + cI)⁻¹ in atomic units.
+
+        Other Parameters
+        ----------------
+        **kwargs
+            Keyword arguments passed to the Optimizer baseclass.
+        """
+        if precon:
+            assert geometry.coord_type in (
+                "cart",
+                "cartesian",
+            ), "Preconditioning makes only sense with 'coord_type: cart|cartesian'"
+        # Set before calling the superclass constructor so we can use the value
+        # in _set_opt_restart_info.
+        self.history = history
+
+        super().__init__(geometry, **kwargs)
+
+        self.alpha_init = alpha_init
+        self.precon = precon
+        # Disable preconditioning for 1 atom species, e.g., analytical potentials
+        if self.precon and (self.geometry.cart_coords.size == 3):
+            self.precon = None
+        self.precon_update = precon_update
+        self.precon_getter_update = precon_getter_update
+        self.precon_kind = precon_kind
+
+        self.P = None
+        try:
+            is_dimer = isinstance(self.geometry.calculator, Dimer)
+        # COS objectes may not have a calculator
+        except AttributeError:
+            is_dimer = False
+        go_uphill = is_dimer or isinstance(self.geometry, GrowingNT)
+
+        if c_stab is None:
+            self.log("No c_stab specified.")
+            if go_uphill:
+                c_stab = 0.103  # 1 eV/Ų
+                self.log(
+                    "Found climbing calculation. Using higher regularization "
+                    f"c_stab={c_stab:.4f} au/bohr**2."
+                )
+            else:
+                c_stab = 0.0103  # 0.1 eV/Ų
+
+        self.c_stab = float(c_stab)
+        self.log(f"Using c_stab={self.c_stab:.6f}")
+
+        if max_step_element is None and go_uphill:
+            max_step_element = 0.25
+            self.log(
+                f"Found climbing calculation. Using max_step_element={max_step_element:.2f}"
+            )
+        elif max_step_element is None and line_search in (None, False):
+            max_step_element = 0.2
+        self.max_step_element = max_step_element
+
+        # Disable linesearch if max_step_element is set
+        if self.max_step_element is not None:
+            self.log(
+                f"max_step_element={max_step_element:.6f} given. Setting "
+                "line_search to 'None'."
+            )
+            line_search = None
+
+        self.line_search = line_search
+        ls_cls = {
+            "armijo": Backtracking,
+            "armijo_fg": partial(Backtracking, use_grad=True),
+            "strong_wolfe": StrongWolfe,
+            "hz": HagerZhang,
+            None: None,
+            False: None,
+        }
+        self.line_search_cls = ls_cls[self.line_search]
+
+        if not self.restarted:
+            self.grad_diffs = deque(maxlen=self.history)
+            self.steps_ = deque(maxlen=self.history)
+            self.f_evals = 0
+            self.df_evals = 0
+
+    def get_precon_getter(self):
+        return precon_getter(
+            self.geometry, c_stab=self.c_stab, kind=self.precon_kind, logger=self.logger
+        )
+
+    def prepare_opt(self):
+        if self.precon:
+            self.precon_getter = self.get_precon_getter()
+
+    def _get_opt_restart_info(self):
+        opt_restart_info = {
+            "c_stab": self.c_stab,
+            "grad_diffs": np.array(self.grad_diffs).tolist(),
+            "steps_": np.array(self.steps_).tolist(),
+            "precon_kind": self.precon_kind,
+            "f_evals": self.f_evals,
+            "df_evals": self.df_evals,
+        }
+        return opt_restart_info
+
+    def _set_opt_restart_info(self, opt_restart_info):
+        self.grad_diffs = deque(
+            [np.array(gd) for gd in opt_restart_info["grad_diffs"]], maxlen=self.history
+        )
+        self.steps_ = deque(
+            [np.array(cd) for cd in opt_restart_info["steps_"]], maxlen=self.history
+        )
+
+        self.c_stab = opt_restart_info["c_stab"]
+        self.precon_kind = opt_restart_info["precon_kind"]
+        self.precon_getter = self.get_precon_getter()
+        self.f_evals = opt_restart_info["f_evals"]
+        self.df_evals = opt_restart_info["df_evals"]
+
+    def scale_max_element(self, step, max_step_element):
+        max_element = np.abs(step).max()
+        if max_element > max_step_element:
+            step *= max_step_element / max_element
+        return step
+
+    def optimize(self):
+        forces = self.geometry.forces
+        energy = self.geometry.energy
+
+        self.forces.append(forces)
+        self.energies.append(energy)
+
+        norm = np.linalg.norm(forces)
+        if not self.is_cos:
+            fmt = " >12.6f"
+            self.log(f" Current energy={energy:{fmt}} au")
+            if self.cur_cycle > 0:
+                prev_energy = self.energies[-2]
+                self.log(f"Previous energy={prev_energy:{fmt}} au")
+                self.log(f"             ΔE={energy-prev_energy:{fmt}} au")
+        self.log(f"norm(forces)={norm:.6f}")
+
+        # Check for preconditioner getter update
+        if (
+            self.precon_getter_update
+            and self.cur_cycle > 0
+            and self.cur_cycle % self.precon_getter_update == 0
+        ):
+            self.precon_getter = self.get_precon_getter()
+            self.log("Calculated new preconditioner getter")
+
+        # If requested, construct preconditioner
+        if self.precon and (self.cur_cycle % self.precon_update == 0):
+            self.P = self.precon_getter(self.geometry.coords)
+            self.log("Calculated new preconditioner P")
+
+        # Determine step direction
+
+        # Preconditioned LBFGS in later cycles
+        if self.cur_cycle > 0:
+            self.grad_diffs.append(-forces - -self.forces[-2])
+            self.steps_.append(self.steps[-1])
+            step = bfgs_multiply(self.steps_, self.grad_diffs, forces, P=self.P)
+        # Preconditioned steepest descent in cycle 0
+        elif self.precon:
+            step = spsolve(self.P, forces)
+        # Steepest descent w/o preconditioner in cycle 0
+        else:
+            step = forces
+
+        step_norm = np.linalg.norm(step)
+        self.log(f"norm(unscaled step)={step_norm:.6f} au")
+
+        # Determine step length
+        if self.line_search_cls is not None:
+            kwargs = {
+                "geometry": self.geometry,
+                "p": step,
+                "f0": energy,
+                "g0": -forces,
+                "alpha_init": self.alpha_init,
+            }
+            line_search = self.line_search_cls(**kwargs)
+            line_search_result = line_search.run()
+            try:
+                alpha = line_search_result.alpha
+                step *= alpha
+                f_evals = line_search_result.f_evals
+                df_evals = line_search_result.df_evals
+                self.log(
+                    f"Evaluated {f_evals} energies and {df_evals} gradients in line search."
+                )
+                self.f_evals += f_evals
+                self.df_evals += df_evals
+            except TypeError:
+                self.log("Line search did not converge!")
+                step = self.scale_max_element(step, self.max_step_element)
+        else:
+            step = self.scale_max_element(step, self.max_step_element)
+        step_norm = np.linalg.norm(step)
+        self.log(f"norm(step)={step_norm:.6f} au")
+        return step

pysisyphus/optimizers/PreconSteepestDescent.py

@@ -0,0 +1,7 @@
+from pysisyphus.optimizers.PreconLBFGS import PreconLBFGS
+
+
+class PreconSteepestDescent(PreconLBFGS):
+
+    def __init__(self, geometry, alpha_init=0.5, **kwargs):
+        super().__init__(geometry, alpha_init=alpha_init, history=0, **kwargs)

pysisyphus/optimizers/QuickMin.py

@@ -0,0 +1,74 @@
+import numpy as np
+
+from pysisyphus.optimizers.Optimizer import Optimizer
+
+
+class QuickMin(Optimizer):
+    def __init__(self, geometry, dt=0.35, **kwargs):
+        super(QuickMin, self).__init__(geometry, **kwargs)
+
+        self.dt = dt
+
+    def _get_opt_restart_info(self):
+        opt_restart_info = {
+            "velocity": self.velocities[-1].tolist(),
+        }
+        return opt_restart_info
+
+    def _set_opt_restart_info(self, opt_restart_info):
+        velocity = np.array(opt_restart_info["velocity"], dtype=float)
+        self.velocities = [
+            velocity,
+        ]
+
+    def prepare_opt(self):
+        self.velocities = [
+            np.zeros_like(self.geometry.coords),
+        ]
+
+    def reset(self):
+        if self.coords[-1].size != self.coords[-2].size:
+            self.log(
+                "Number of coordinates changed. Adapting velocity vector "
+                "to new coordinate number."
+            )
+            self.prepare_opt()
+        self.log("Resetted optimizer")
+
+    def optimize(self):
+        if self.align and self.is_cos:
+            (self.velocities[-1],), _, _ = self.fit_rigid(
+                vectors=(self.velocities[-1],)
+            )
+
+        prev_velocities = self.velocities[-1]
+        cur_forces = self.geometry.forces
+        self.forces.append(cur_forces)
+        self.energies.append(self.geometry.energy)
+
+        norm = np.linalg.norm(cur_forces)
+        if not self.is_cos:
+            self.log(f"Current energy={self.energies[-1]:.6f}")
+        self.log(f"norm(forces)={norm:.6f}")
+
+        if self.cur_cycle == 0:
+            tmp_velocities = np.zeros_like(prev_velocities)
+        else:
+            overlap = prev_velocities.dot(cur_forces)
+            self.log(f"Overlap of previous and current forces: {overlap:.6f}")
+            if overlap > 0:
+                tmp_velocities = overlap * cur_forces / cur_forces.dot(cur_forces)
+            else:
+                tmp_velocities = np.zeros_like(prev_velocities)
+                self.log("Zeroed velocities")
+
+        accelerations = cur_forces / self.geometry.masses_rep
+        cur_velocities = tmp_velocities + self.dt * accelerations
+        steps = cur_velocities * self.dt + 1 / 2 * accelerations * self.dt**2
+        steps = self.scale_by_max_step(steps)
+        self.velocities.append(cur_velocities)
+        velo_norm = np.linalg.norm(cur_velocities)
+        acc_norm = np.linalg.norm(accelerations)
+        self.log(f"norm(v) = {velo_norm:.4f}, norm(a) = {acc_norm:.4f}")
+
+        return steps

pysisyphus/optimizers/RFOptimizer.py

@@ -0,0 +1,181 @@
+# [1] http://aip.scitation.org/doi/10.1063/1.1515483 Optimization review
+# [2] https://doi.org/10.1063/1.450914 Trust region method
+# [3] 10.1007/978-0-387-40065-5 Numerical optimization
+# [4] 10.1007/s00214-016-1847-3 Explorations of some refinements
+
+
+import numpy as np
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import rms
+from pysisyphus.optimizers.HessianOptimizer import HessianOptimizer
+from pysisyphus.optimizers.poly_fit import poly_line_search
+from pysisyphus.optimizers.gdiis import gdiis, gediis
+
+import torch
+
+class RFOptimizer(HessianOptimizer):
+    def __init__(
+        self,
+        geometry: Geometry,
+        line_search: bool = True,
+        gediis: bool = False,
+        gdiis: bool = True,
+        gdiis_thresh: float = 2.5e-3,
+        gediis_thresh: float = 1e-2,
+        gdiis_test_direction: bool = True,
+        max_micro_cycles: int = 25,
+        adapt_step_func: bool = False,
+        **kwargs,
+    ) -> None:
+        """
+        Rational function Optimizer.
+
+        Parameters
+        ----------
+        geometry
+            Geometry to be optimized.
+        line_search
+            Whether to carry out implicit line searches.
+        gediis
+            Whether to enable GEDIIS.
+        gdiis
+            Whether to enable GDIIS.
+        gdiis_thresh
+            Threshold for rms(forces) to enable GDIIS.
+        gediis_thresh
+            Threshold for rms(step) to enable GEDIIS.
+        gdiis_test_direction
+            Whether to the overlap of the RFO step and the GDIIS step.
+        max_micro_cycles
+            Number of restricted-step microcycles. Disabled by default.
+        adapt_step_func
+            Whether to switch between shifted Newton and RFO-steps.
+
+        Other Parameters
+        ----------------
+        **kwargs
+            Keyword arguments passed to the Optimizer/HessianOptimizer baseclass.
+        """
+        super().__init__(geometry, max_micro_cycles=max_micro_cycles, **kwargs)
+
+        self.line_search = line_search
+        self.gediis = gediis
+        self.gdiis = gdiis
+        self.gdiis_thresh = gdiis_thresh  # Will be compared to rms(step)
+        self.gediis_thresh = gediis_thresh  # Will be compared to rms(forces)
+        self.gdiis_test_direction = gdiis_test_direction
+        self.adapt_step_func = adapt_step_func
+
+        self.successful_gediis = 0
+        self.successful_gdiis = 0
+        self.successful_line_search = 0
+
+    def optimize(self):
+        energy, gradient, H, big_eigvals, big_eigvecs, resetted = self.housekeeping()
+        step_func, pred_func = self.get_step_func(big_eigvals, gradient)
+
+        forces_act = self.active_list(self.forces)
+        coords_act = self.active_list(self.coords)
+        steps_act = self.active_list(self.steps)
+
+        ref_gradient = gradient.copy() if isinstance(gradient, np.ndarray) else gradient.clone()
+        # Reference step, used for judging the proposed GDIIS step
+        ref_step = step_func(big_eigvals, big_eigvecs, gradient) # heavy-compute
+
+        # Right everything is in place to check for convergence.  If all values are below
+        # the thresholds, there is no need to do additional inter/extrapolations.
+        if self.check_convergence(ref_step)[0]:  # Drop conv_info
+            self.log("Convergence achieved! Skipping inter/extrapolation.")
+            return ref_step
+
+        # Try to interpolate an intermediate geometry, either from GDIIS or line search.
+        #
+        # Set some defaults
+        ip_gradient = None
+        ip_step = None
+        diis_result = None
+
+        # Check if we can do GDIIS or GEDIIS. If we (can) do a line search is decided
+        # after trying GDIIS.
+        rms_forces = rms(gradient)
+        rms_step = rms(ref_step)
+        can_diis = (rms_step <= self.gdiis_thresh) and (not resetted)
+        can_gediis = (rms_forces <= self.gediis_thresh) and (not resetted)
+
+        # GDIIS / GEDIIS, prefer GDIIS over GEDIIS
+        if self.gdiis and can_diis:
+            # Gradients as error vectors
+            if isinstance(ref_step, torch.Tensor):
+                err_vecs = -torch.from_numpy(np.array(forces_act)).to(ref_step.dtype).to(ref_step.device)
+            else:
+                err_vecs = -np.array(forces_act)
+            diis_result = gdiis(
+                err_vecs,
+                coords_act,
+                forces_act,
+                ref_step,
+                test_direction=self.gdiis_test_direction,
+            )
+            self.successful_gdiis += 1 if diis_result else 0
+        # Don't try GEDIIS if GDIIS failed. If GEDIIS should be tried after GDIIS failed
+        # comment the line below and uncomment the line following it.
+        elif self.gediis and can_gediis:
+            # if self.gediis and can_gediis and (diis_result == None):
+            diis_result = gediis(coords_act, self.energies, forces_act, hessian=H)
+            self.successful_gediis += 1 if diis_result else 0
+
+        try:
+            ip_coords = diis_result.coords
+            if isinstance(ip_coords, torch.Tensor):
+                ip_step = ip_coords - torch.from_numpy(coords_act[-1]).to(ip_coords.device, ip_coords.dtype)
+            else:
+                ip_step = ip_coords - coords_act[-1]
+            ip_gradient = -diis_result.forces
+        # When diis_result is None
+        except AttributeError:
+            self.log("GDIIS didn't succeed.")
+
+        # Try line search if GDIIS failed or not requested
+        if self.line_search and (diis_result is None) and (not resetted):
+            ip_energy, ip_gradient, ip_step = poly_line_search(
+                energy,
+                self.energies[-2],
+                gradient,
+                -forces_act[-2],
+                steps_act[-1],
+                cubic_max_x=-1,
+                quartic_max_x=2,
+                logger=self.logger,
+            )
+            self.successful_line_search += 1 if ip_gradient is not None else 0
+
+        # Use the interpolated gradient for the RFO step if interpolation succeeded
+        if (ip_gradient is not None) and (ip_step is not None):
+            gradient = ip_gradient
+            step = step_func(big_eigvals, big_eigvecs, gradient) # heavy-compute
+            # Form full step with interpolation offset.
+            if isinstance(ip_step, torch.Tensor):
+                ip_step = ip_step.detach().cpu().numpy()
+            step = step + ip_step
+        # Keep the original gradient when the interpolation failed; reuse ref_step.
+        else:
+            step = ref_step
+
+        # Use the original, actually calculated, gradient
+        prediction = pred_func(ref_gradient, H, step)
+        self.predicted_energy_changes.append(prediction)
+
+        step = self.full_from_active(step)
+        if isinstance(step, torch.Tensor):
+            step = step.cpu().numpy()
+        return step
+
+    def postprocess_opt(self):
+        msg = (
+            f"Successful invocations:\n"
+            f"\t     GEDIIS: {self.successful_gediis}\n"
+            f"\t      GDIIS: {self.successful_gdiis}\n"
+            f"\tLine Search: {self.successful_line_search}\n"
+        )
+        self.log(msg)

pysisyphus/optimizers/RSA.py

@@ -0,0 +1,99 @@
+from math import sqrt
+
+import numpy as np
+from scipy.optimize import root_scalar
+
+from pysisyphus.optimizers.HessianOptimizer import HessianOptimizer
+
+
+class RSA(HessianOptimizer):
+    """The Importance of Step Control in Optimization Methods, del Campo, 2009."""
+
+    def optimize(self):
+        energy, gradient, H, big_eigvals, big_eigvecs, resetted = self.housekeeping()
+
+        assert big_eigvals.argmin() == 0
+        min_eigval = big_eigvals[0]
+        pos_definite = min_eigval > 0.0
+        gradient_trans = big_eigvecs.T.dot(gradient)
+        # This will be also be True when we come close to a minimizer,
+        # but then the Hessian will also be positive definite and a
+        # simple Newton step will be used.
+        hard_case = abs(gradient_trans[0]) <= 1e-6
+        self.log(f"Smallest eigenvalue: {min_eigval:.6f}")
+        self.log(f"Positive definite Hessian: {pos_definite}")
+        self.log(f"Hard case: {hard_case}")
+
+        def get_step(lambda_):
+            return -gradient_trans / (big_eigvals + lambda_)
+
+        # Unshifted Newton step
+        newton_step = get_step(0.0)
+        newton_norm = np.linalg.norm(newton_step)
+
+        # def on_trust_radius(step, thresh=1e-3):
+        # return abs(self.trust_radius - np.linalg.norm(step)) <= thresh
+
+        def on_trust_radius_lin(step):
+            return 1 / self.trust_radius - 1 / np.linalg.norm(step)
+
+        def finalize_step(lambda_):
+            step = get_step(lambda_)
+            step = big_eigvecs.dot(step)
+            predicted_change = step.dot(gradient) + 0.5 * step.dot(H).dot(step)
+            self.predicted_energy_changes.append(predicted_change)
+            return step
+
+        # Simplest case. Positive definite Hessian and predicted step is
+        # already in trust radius.
+        if pos_definite and newton_norm <= self.trust_radius:
+            lambda_ = 0.0
+            self.log("Using unshifted Newton step.")
+            return finalize_step(lambda_)
+
+        # If the Hessian is not positive definite or if the step is too
+        # long we have to determine the shift parameter lambda.
+        rs_kwargs = {
+            "f": lambda lambda_: on_trust_radius_lin(get_step(lambda_)),
+            "xtol": 1e-3,
+            # Would otherwise be chosen automatically, but we set it
+            # here explicitly for verbosity.
+            "method": "brentq",
+        }
+
+        def root_search(bracket):
+            rs_kwargs.update(
+                {
+                    "bracket": bracket,
+                    "x0": bracket[0] + 1e-3,
+                }
+            )
+            res = root_scalar(**rs_kwargs)
+            return res
+
+        BRACKET_END = 1e10
+        if not hard_case:
+            bracket_start = 0.0 if pos_definite else -min_eigval - 1e-3
+            bracket = (bracket_start, BRACKET_END)
+            res = root_search(bracket)
+            assert res.converged
+            return finalize_step(res.root)
+
+        # Hard case.
+        # First we try the bracket (-b1, ∞)
+        bracket = (-min_eigval, BRACKET_END)
+        res = root_search(bracket)
+        if res.converged:
+            return finalize_step(res.root)
+
+        # Now we would try the bracket (-b2, -b1). The resulting step should have
+        # a suitable length, but the (shifted) Hessian would have an incorrect
+        # eigenvalue spectrum (not positive definite). To solve this we use a
+        # different formula to calculate the step.
+        without_first = gradient_trans[1:] / (big_eigvals[1:] - min_eigval)
+        tau = sqrt(self.trust_radius ** 2 - (without_first ** 2).sum())
+        step_trans = [tau] + -without_first.tolist()
+        step = big_eigvecs.dot(step_trans)
+        predicted_change = step.dot(gradient) + 0.5 * step.dot(H).dot(step)
+        self.predicted_energy_changes.append(predicted_change)
+        return step