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

pysisyphus/optimizers/closures.py

@@ -0,0 +1,301 @@
+# [1] https://arxiv.org/pdf/2101.04413.pdf
+#     A Regularized Limited Memory BFGS method for Large-Scale Unconstrained
+#     Optimization and itsefficient Implementations
+#     Tankaria, Sugimoto, Yamashita, 2021
+# [2] Regularization of Limited Memory Quasi-Newton Methods for Large-Scale
+#     Nonconvex Minimization
+#     https://arxiv.org/pdf/1911.04584.pdf
+#     Kanzow, Steck 2021
+
+from collections import deque
+
+import numpy as np
+from scipy.sparse.linalg import spsolve
+
+
+def get_update_mu_reg(
+    mu_min=1e-3, gamma_1=0.1, gamma_2=5.0, eta_1=0.01, eta_2=0.9, logger=None
+):
+    """See 5.1 in [1]"""
+    assert 0.0 < mu_min
+    assert 0.0 < gamma_1 <= 1.0 < gamma_2
+    assert 0.0 < eta_1 < eta_2 <= 1.0
+
+    def log(msg):
+        if logger is not None:
+            logger.debug(msg)
+
+    def update_mu_reg(mu, cur_energy, trial_energy, cur_grad, step):
+        """Update regularization parameter μ_reg according to ratio r.
+        See 2.3 in [1].
+
+        r = actual_reduction / predicted_reduction
+        r = (f(x) - f(x + step)) / (f(x) - model_func(step, mu))
+
+        mode_func(step, mu) = f(x) + df(x)^T step(mu) + 1/2 step(mu)^T H^-1 step(mu)
+
+        f(x) - model_func(step, mu)
+        = f(x) - (f(x) + df(x)^T step(mu) + 1/2 step(mu)^T H(mu)^-1 step(mu))
+        = -df(x)^T step(mu) - 1/2 step(mu)^T H(mu)^-1 step(mu)
+        = -df(x)^T step(mu) + 1/2 step(mu)^T df(x)
+        = -1/2 df(x)^T step(mu)
+
+        r = (f(x) - f(x + step) / (-1/2 df(x)^T step(mu))
+        r = -2 * (f(x) - f(x + step) / (df(x)^T step(mu))
+
+        """
+        log(f"  Cur energy: {cur_energy:.8f}")
+        log(f"Trial energy: {trial_energy:.8f}")
+        act_change = cur_energy - trial_energy
+        log(f"   Actual change: {act_change: .8f}")
+        pred_change = -1 / 2 * cur_grad.dot(step)
+        log(f"Predicted change: {pred_change: .8f}")
+        r = act_change / pred_change
+        log(f"Ratio r={r:.8f}")
+
+        # Default case for eta_1 <= r < eta_2. Keep mu and accept step.
+        mu_updated = mu
+        recompute_step = False
+        # Poor actual reduction. Increase shift parameter and request new step.
+        if r < eta_1:
+            mu_updated *= gamma_2
+            recompute_step = True
+            log(f"Increased μ_reg to {mu_updated:.6f}.")
+        # Significant actual reduction. Reduce shift parameter und accept step.
+        elif r >= eta_2:
+            mu_updated = max(mu_min, mu_updated * gamma_1)
+            log(f"Decreased μ_reg to {mu_updated:.6f}.")
+        return mu_updated, recompute_step
+
+    return update_mu_reg
+
+
+def bfgs_multiply(
+    s_list,
+    y_list,
+    vector,
+    beta=1,
+    P=None,
+    logger=None,
+    gamma_mult=True,
+    mu_reg=None,
+    inds=None,
+    cur_size=None,
+):
+    """Matrix-vector product H·v.
+
+    Multiplies given vector with inverse Hessian, obtained
+    from repeated BFGS updates calculated from steps in 's_list'
+    and gradient differences in 'y_list'.
+
+    Based on algorithm 7.4 Nocedal, Num. Opt., p. 178."""
+
+    assert len(s_list) == len(y_list), (
+        "lengths of step list 's_list' and gradient list 'y_list' differ!"
+    )
+
+    cycles = len(s_list)
+    q = vector.copy()
+    alphas = list()
+    rhos = list()
+
+    if mu_reg is not None:
+        # Regularized L-BFGS with ŷ = y + μs, see [1]
+        assert mu_reg > 0.0
+        y_list_reg = list()
+        for i, si in enumerate(s_list):
+            yi = y_list[i]
+            y_hat = yi + mu_reg * si
+            if y_hat.dot(si) <= 0:
+                # See 2 in [1]
+                y_hat = yi + (max(0, -si.dot(yi) / si.dot(si)) + mu_reg) + si
+            y_list_reg.append(y_hat)
+        y_list = y_list_reg
+
+    # Store rho and alphas as they are also needed in the second loop
+    for i in reversed(range(cycles)):
+        s = s_list[i]
+        y = y_list[i]
+        rho = 1 / y.dot(s)
+        rhos.append(rho)
+        try:
+            alpha = rho * s.dot(q)
+            q -= alpha * y
+        except ValueError:
+            inds_i = inds[i]
+            q_ = q.reshape(cur_size, -1)
+            alpha = rho * s.dot(q_[inds_i].flatten())
+            # This also modifies q!
+            q_[inds_i] -= alpha * y.reshape(len(inds_i), -1)
+        alphas.append(alpha)
+
+    # Restore original order, so that rho[i] = 1/s_list[i].dot(y_list[i]) etc.
+    alphas = alphas[::-1]
+    rhos = rhos[::-1]
+
+    if P is not None:
+        r = spsolve(P, q)
+        msg = "preconditioner"
+    elif gamma_mult and (cycles > 0):
+        s = s_list[-1]
+        y = y_list[-1]
+        gamma = s.dot(y) / y.dot(y)
+        r = gamma * q
+        msg = f"gamma={gamma:.4f}"
+    else:
+        r = beta * q
+        msg = f"beta={beta:.4f}"
+
+    if mu_reg is not None:
+        msg += f" and μ_reg={mu_reg:.6f}"
+
+    if logger is not None:
+        msg = f"BFGS multiply using {cycles} previous cycles with {msg}."
+        if len(s_list) == 0:
+            msg += "\nProduced simple SD step."
+        logger.debug(msg)
+
+    for i in range(cycles):
+        s = s_list[i]
+        y = y_list[i]
+        try:
+            beta = rhos[i] * y.dot(r)
+            r += s * (alphas[i] - beta)
+        except ValueError:
+            inds_i = inds[i]
+            r_ = r.reshape(cur_size, -1)
+            beta = rhos[i] * y.dot(r_[inds_i].flatten())
+            # This also modifies r!
+            r_[inds_i] += s.reshape(len(inds_i), -1) * (alphas[i] - beta)
+
+    return r
+
+
+def lbfgs_closure(force_getter, M=10, beta=1, restrict_step=None):
+    x_list = list()
+    s_list = list()
+    y_list = list()
+    force_list = list()
+    cur_cycle = 0
+
+    if restrict_step is None:
+        restrict_step = lambda x, dx: dx
+
+    def lbfgs(x, *getter_args):
+        nonlocal cur_cycle
+        nonlocal s_list
+        nonlocal y_list
+
+        force = force_getter(x, *getter_args)
+        if cur_cycle > 0:
+            prev_x = x_list[-1]
+            s = x - prev_x
+            s_list.append(s)
+            prev_force = force_list[-1]
+            y = prev_force - force
+            y_list.append(y)
+        x_list.append(x)
+        force_list.append(force)
+
+        step = bfgs_multiply(s_list, y_list, force, beta=beta)
+        step = restrict_step(x, step)
+        # Only keep last m cycles
+        s_list = s_list[-M:]
+        y_list = y_list[-M:]
+        cur_cycle += 1
+        return step, force
+
+    return lbfgs
+
+
+def modified_broyden_closure(force_getter, M=5, beta=1, restrict_step=None):
+    """https://doi.org/10.1006/jcph.1996.0059
+    F corresponds to the residual gradient, so we after calling
+    force_getter we multiply the force by -1 to get the gradient."""
+
+    dxs = list()
+    dFs = list()
+    F = None
+    # The next line is used in SciPy, but then beta strongly depends
+    # on the magnitude of x, which is quite weird. Disregarding the
+    # analytical potentials the electronic energy is usually
+    # invariant under translation and rotation and doesn't depend on
+    # the magnitude of x.
+    # beta = 0.5 * max(np.linalg.norm(x), 1) / np.linalg.norm(F)
+    a = None
+
+    if restrict_step is None:
+        restrict_step = lambda x, dx: dx
+
+    def modified_broyden(x, *getter_args):
+        nonlocal dxs
+        nonlocal dFs
+        nonlocal F
+        nonlocal a
+
+        F_new = -force_getter(x, *getter_args)
+        if F is not None:
+            dF = F_new - F
+            dFs.append(dF)
+            dFs = dFs[-M:]
+            # Overlap matrix
+            a = np.zeros((len(dFs), len(dFs)))
+            for k, dF_k in enumerate(dFs):
+                for m, dF_m in enumerate(dFs):
+                    a[k, m] = dF_k.dot(dF_m)
+        F = F_new
+        dx = -beta * F
+
+        if len(dxs) > 0:
+            # Calculate gamma
+            dF_F = [dF_k.dot(F) for dF_k in dFs]
+            # Use least squares to avoid crashes with singular matrices
+            gammas, *_ = np.linalg.lstsq(a, dF_F, rcond=None)
+            gammas = gammas[:, None]
+            _ = np.array(dxs) - beta * np.array(dFs)
+            # Substract step correction
+            dx = dx - np.sum(gammas * _, axis=0)
+        dx = restrict_step(x, dx)
+        dxs.append(dx)
+
+        # Keep only informations of the last M cycles
+        dxs = dxs[-M:]
+
+        return dx, -F
+
+    return modified_broyden
+
+
+def small_lbfgs_closure(history=5, gamma_mult=True):
+    """Compact LBFGS closure.
+
+    The returned function takes two arguments: forces and prev_step.
+    forces are the forces at the current iterate and prev_step is the
+    previous step that lead us to the current iterate. In this way
+    step restriction/line search can be done outisde of the lbfgs function.
+    """
+
+    prev_forces = None  # lgtm [py/unused-local-variable]
+    grad_diffs = deque(maxlen=history)
+    steps = deque(maxlen=history)
+    cur_cycle = 0
+
+    def lbfgs(forces, prev_step=None):
+        nonlocal cur_cycle
+        nonlocal prev_forces
+
+        if prev_step is not None:
+            steps.append(prev_step)
+
+        # Steepest descent in the first cycle
+        step = forces
+        # LBFGS in the following cycles
+        if cur_cycle > 0:
+            grad_diffs.append(-forces - -prev_forces)
+            step = bfgs_multiply(steps, grad_diffs, forces, gamma_mult=gamma_mult)
+
+        prev_forces = forces
+        cur_cycle += 1
+        return step
+
+    return lbfgs

pysisyphus/optimizers/cls_map.py

@@ -0,0 +1,58 @@
+from pysisyphus.optimizers import (
+    BFGS,
+    CubicNewton,
+    ConjugateGradient,
+    FIRE,
+    LayerOpt,
+    LBFGS,
+    MicroOptimizer,
+    NCOptimizer,
+    PreconLBFGS,
+    PreconSteepestDescent,
+    QuickMin,
+    RFOptimizer,
+    SteepestDescent,
+    StabilizedQNMethod,
+    StringOptimizer,
+)
+from pysisyphus.tsoptimizers import (
+    RSPRFOptimizer,
+    TRIM,
+    RSIRFOptimizer,
+)
+
+
+OPT_DICT = {
+    "bfgs": BFGS.BFGS,
+    "cg": ConjugateGradient.ConjugateGradient,
+    "cnewton": CubicNewton,
+    "fire": FIRE.FIRE,
+    "layers": LayerOpt.LayerOpt,
+    "lbfgs": LBFGS.LBFGS,
+    "micro": MicroOptimizer,
+    "nc": NCOptimizer.NCOptimizer,
+    "plbfgs": PreconLBFGS.PreconLBFGS,
+    "psd": PreconSteepestDescent.PreconSteepestDescent,
+    "qm": QuickMin.QuickMin,
+    "rfo": RFOptimizer.RFOptimizer,
+    "sd": SteepestDescent.SteepestDescent,
+    "sqnm": StabilizedQNMethod.StabilizedQNMethod,
+    "string": StringOptimizer.StringOptimizer,
+}
+
+TSOPT_DICT = {
+    "rsprfo": RSPRFOptimizer,
+    "prfo": RSPRFOptimizer,  # Shortcut
+    "trim": TRIM,
+    "rsirfo": RSIRFOptimizer,
+    "irfo": RSIRFOptimizer,  # Shortcut
+}
+OPT_DICT.update(TSOPT_DICT)
+
+
+def get_opt_cls(opt_key):
+    return OPT_DICT[opt_key]
+
+
+def key_is_tsopt(opt_key):
+    return opt_key in TSOPT_DICT

pysisyphus/optimizers/exceptions.py

@@ -0,0 +1,6 @@
+class ZeroStepLength(Exception):
+    pass
+
+
+class OptimizationError(Exception):
+    pass

pysisyphus/optimizers/gdiis.py

@@ -0,0 +1,280 @@
+# [1] https://doi.org/10.1016/S0022-2860(84)87198-7
+#     Pulay, 1984
+# [2] https://pubs.rsc.org/en/content/articlehtml/2002/cp/b108658h
+#     Stabilized GDIIS
+#     Farkas, Schlegel, 2002
+# [3] https://pubs.acs.org/doi/abs/10.1021/ct050275a
+#     GEDIIS/Hybrid method
+#     Li, Frisch, 2006
+# [4] https://aip.scitation.org/doi/10.1063/1.2977735
+#     Sim-GEDIIS using hessian information
+#     Moss, Li, 2008
+
+from collections import namedtuple
+import logging
+
+import autograd.numpy as anp
+from autograd import grad
+import numpy as np
+from scipy.optimize import minimize
+
+from pysisyphus.helpers import array2string
+import torch
+
+COS_CUTOFFS = {
+    # Looser cutoffs
+    2: 0.80,
+    3: 0.75,
+    # Original cutoffs, as published in [2]
+    # 2: 0.97,
+    # 3: 0.84,
+    4: 0.71,
+    5: 0.67,
+    6: 0.62,
+    7: 0.56,
+    8: 0.49,
+    9: 0.41,
+}
+DIISResult = namedtuple("DIISResult", "coeffs coords forces energy N type")
+
+
+logger = logging.getLogger("optimizer")
+
+
+def log(msg):
+    logger.debug(msg)
+
+
+def valid_diis_direction(diis_step, ref_step, use):
+    if isinstance(diis_step, torch.Tensor):
+        ref_direction = ref_step / torch.linalg.norm(ref_step)
+        diis_direction = diis_step / torch.linalg.norm(diis_step)
+        cos = torch.dot(diis_direction, ref_direction)
+    else:
+        ref_direction = ref_step / np.linalg.norm(ref_step)
+        diis_direction = diis_step / np.linalg.norm(diis_step)
+        cos = diis_direction @ ref_direction
+    return (cos >= COS_CUTOFFS[use]) and (cos >= 0)
+
+
+def from_coeffs(vec, coeffs):
+    if isinstance(vec, torch.Tensor):
+        return torch.sum(coeffs[:, None] * vec.flip(0)[: len(coeffs)], dim=0)
+    else:
+        return np.sum(coeffs[:, None] * vec[::-1][: len(coeffs)], axis=0)
+
+
+def diis_result(coeffs, coords, forces, energy=None, prefix=""):
+    diis_coords = from_coeffs(coords, coeffs)
+    diis_forces = from_coeffs(forces, coeffs)
+    diis_result = DIISResult(
+        coeffs=coeffs,
+        coords=diis_coords,
+        forces=diis_forces,
+        energy=energy,
+        N=len(coeffs),
+        type=f"{prefix}DIIS",
+    )
+    log(f"\tUsed {len(coeffs)} error vectors for {prefix}DIIS.")
+    log("")
+    return diis_result
+
+
+def gdiis(err_vecs, coords, forces, ref_step, max_vecs=5, test_direction=True):
+    # Scale error vectors so the smallest norm is 1
+    if isinstance(err_vecs, torch.Tensor):
+        coords = torch.from_numpy(np.array(coords)).to(err_vecs.device, dtype=err_vecs.dtype)
+        forces = torch.from_numpy(np.array(forces)).to(err_vecs.device, dtype=err_vecs.dtype)
+        norms = torch.linalg.norm(err_vecs, dim=1)
+    else:
+        norms = np.linalg.norm(err_vecs, axis=1)
+    err_vecs = err_vecs / norms.min()
+
+    valid_coeffs = None
+    for use in range(2, min(max_vecs, len(err_vecs)) + 1):
+        log(f"Trying GDIIS with {use} previous cycles.")
+        if isinstance(err_vecs, torch.Tensor):
+            use_vecs = err_vecs.flip(0)[:use]
+            A = torch.einsum("ij,kj->ik", use_vecs, use_vecs)
+        else:
+            use_vecs = np.array(err_vecs[::-1][:use])
+            A = np.einsum("ij,kj->ik", use_vecs, use_vecs)
+        try:
+            if isinstance(err_vecs, torch.Tensor):
+                coeffs = torch.linalg.solve(A, torch.ones(use, device=err_vecs.device))
+            else:
+                coeffs = np.linalg.solve(A, np.ones(use))
+        except np.linalg.LinAlgError:
+            log("LinAlgError when solving GDIIS matrix.")
+            break
+        except torch.linalg.LinAlgError:
+            log("Torch LinAlgError when solving GDIIS matrix.")
+            break
+        # Scale coeffs so that their sum equals 1
+        if isinstance(err_vecs, torch.Tensor):
+            coeffs_norm = torch.linalg.norm(coeffs)
+        else:
+            coeffs_norm = np.linalg.norm(coeffs) 
+        valid_coeffs_norm = coeffs_norm <= 1e8
+        log(f"\tError vectors are linearly independent: {valid_coeffs_norm}")
+        if isinstance(err_vecs, torch.Tensor):
+            coeffs /= torch.sum(coeffs)
+        else:
+            coeffs /= np.sum(coeffs)
+        coeffs_str = array2string(coeffs, precision=4)
+        log(f"\tGDIIS coefficients: {coeffs_str}")
+
+        # Uncomment these lines and break here to only do the basic check
+        # for linear dependency above.
+        # valid_coeffs = coeffs
+        # break
+
+        # Check degree of extra- and interpolation.
+        pos_sum = abs(coeffs[coeffs > 0].sum())
+        neg_sum = abs(coeffs[coeffs < 0].sum())
+        valid_sums = (pos_sum <= 15) and (neg_sum <= 15)
+        log(f"\tSum of positive coefficients: {pos_sum:.2f}")
+        log(f"\tSum of negative coefficients: {neg_sum:.2f}")
+        log(f"\tSums are valid: {valid_sums}")
+
+        # Calculate GDIIS step for comparison to the reference step
+        diis_coords = from_coeffs(coords, coeffs)
+        diis_step = diis_coords - coords[-1]
+        if isinstance(err_vecs, torch.Tensor):
+            valid_length = torch.linalg.norm(diis_step) <= (10 * torch.linalg.norm(ref_step))
+        else:
+            valid_length = np.linalg.norm(diis_step) <= (10 * np.linalg.norm(ref_step))
+        log(f"\tGDIIS step has valid length: {valid_length}")
+
+        # Compare directions of GDIIS- and reference step
+        valid_direction = (
+            True
+            if (not test_direction)
+            else valid_diis_direction(diis_step, ref_step, use)
+        )
+        log(f"\tGDIIS step has valid direction: {valid_direction}")
+
+        gdiis_valid = (
+            valid_sums and valid_coeffs_norm and valid_direction and valid_length
+        )
+        log(f"\tGDIIS step is valid: {gdiis_valid}")
+        if not gdiis_valid:
+            break
+        # Update valid DIIS coefficients
+        valid_coeffs = coeffs
+        log("")
+
+    if valid_coeffs is None:
+        return None
+
+    # if len(valid_coeffs) is 2:
+    # print("GDIIS with only 2 cycles. Skipping! Return None")
+    # return None
+
+    return diis_result(valid_coeffs, coords, forces, prefix="G")
+
+
+def gediis(coords, energies, forces, hessian=None, max_vecs=3):
+    use = min(len(coords), max_vecs)
+    if isinstance(hessian, torch.Tensor):
+        coords = torch.from_numpy(np.array(coords)).to(device=hessian.device, dtype=hessian.dtype)
+        energies = torch.tensor(energies, device=hessian.device, dtype=hessian.dtype)
+        forces = torch.from_numpy(np.array(forces)).to(device=hessian.device, dtype=hessian.dtype)
+        R = coords.flip(0)[:use]
+        E = energies.flip(0)[:use].reshape(-1)
+        f = forces.flip(0)[:use]
+    else:
+        R = coords[::-1][:use]
+        E = np.ravel(energies[::-1][:use])
+        f = forces[::-1][:use]
+    assert len(R) == len(E) == len(f)
+    log(f"Trying GEDIIS with {use} previous cycles.")
+    # Precompute values so they can be reused in fun()
+    if isinstance(R, torch.Tensor):
+        Rifi = torch.einsum("ik,ik->i", R, f).cpu().numpy()
+        Rjfi = torch.einsum("jk,ik->ji", R, f).cpu().numpy()
+    else:
+        Rifi = np.einsum("ik,ik->i", R, f)
+        Rjfi = np.einsum("jk,ik->ji", R, f)
+
+    def x2c(x):
+        return x ** 2 / (x ** 2).sum()
+
+    # def fun(xs):
+    # """Naive implementation with loops."""
+    # cs = x2c(xs)
+    # first = (cs*E).sum()
+    # sec = 0.
+    # for i, ci in enumerate(cs):
+    # for j, cj in enumerate(cs):
+    # sec += ci * cj * (f[j] - f[i]) @ (R[i] - R[j])
+    # return first - 1/2 * sec
+
+    # def fun(xs):
+    # """Recalculation of all values in every call."""
+    # cs = x2c(xs)
+    # return anp.sum(cs*E) - anp.einsum("i,j,jk,ik", cs, cs, R, f) + anp.einsum("i,ij,ij", cs, R, f)
+
+    # Using precomputed values from above in 'fun()'
+    if hessian is None:
+
+        def fun(xs):
+            """Eq. (6) from [3]."""
+            cs = x2c(xs)
+            return (
+                anp.sum(cs * E) - anp.sum(anp.outer(cs, cs) * Rjfi) + (cs * Rifi).sum()
+            )
+
+    else:
+        if isinstance(hessian, torch.Tensor):
+            hessian_inv = torch.linalg.pinv(hessian, rcond=1e-6)
+            gHig = torch.einsum("ki,ji,ki->k", f, hessian_inv, f).cpu().numpy()
+        else:
+            hessian_inv = np.linalg.pinv(hessian, rcond=1e-6)
+            # It doesn't matter if we use forces or gradients, as the signs will cancel.
+            # gHig = 0.5 * np.einsum("ki,ji,ki->k", f, hessian_inv, f)
+            gHig = np.einsum("ki,ji,ki->k", f, hessian_inv, f)
+
+        def fun(xs):
+            """Eq. (5) from [4]."""
+            cs = x2c(xs)
+            # Consider the hessian in the first term
+            return (
+                0.5 * anp.sum(cs * gHig)
+                - anp.sum(anp.outer(cs, cs) * Rjfi)
+                + (cs * Rifi).sum()
+            )
+
+    # def fun(xs):
+    # cs = x2c(xs)
+    # cRjfi = anp.einsum("j,jk,ik->ji", cs, R, f).sum(axis=0)
+    # return anp.sum(
+    # cs * (E + cRjfi + Rifi)
+    # )
+
+    jac = grad(fun)
+
+    x0 = np.ones(use) / use
+    res = minimize(fun, x0=x0, jac=jac)  # , tol=1e-7)
+    # print(res)
+    # print("final x", res.x)
+    # x = res.x
+    # import pdb; pdb.set_trace()
+
+    coeffs = None
+    if res.success:
+        coeffs = x2c(res.x)
+        en_ = res.fun
+        if isinstance(hessian, torch.Tensor):
+            coeffs = torch.from_numpy(coeffs).to(device=hessian.device, dtype=hessian.dtype)
+    log(f"\tOptimization converged!")
+    coeff_str = array2string(coeffs, precision=4)
+    log(f"\tCoefficients: {coeff_str}")
+    # en_ = (E * coeffs).sum()
+    # import pdb; pdb.set_trace()
+    if (hessian is None) and (en_ >= E[0]):
+        print(
+            f"GEDIIS converged, but proposed energy is above current energy! Returning None"
+        )
+        return None
+    return diis_result(coeffs, coords, forces, energy=en_, prefix="GE")