Source code for pysisyphus.optimizers.gdiis

# [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 dataclasses import dataclass
from collections import namedtuple
import functools
import logging
import warnings

import autograd.numpy as anp
from autograd import grad
import numpy as np
from scipy.optimize import minimize

from pysisyphus.helpers_pure import log as hp_log


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")


[docs] @dataclass class DIISResult: coeffs: np.ndarray coords: np.ndarray forces: np.ndarray energy: float N: int prefix: str @property def type(self): return f"{self.prefix}DIIS"
[docs] def valid_diis_direction(diis_step, ref_step, use): 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)
[docs] def from_coeffs(vec, coeffs): return np.sum(coeffs[:, None] * vec[::-1][: len(coeffs)], axis=0)
[docs] 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), prefix=prefix, ) return diis_result
[docs] def gdiis( err_vecs, coords, forces, ref_step, max_vecs=5, test_direction=True, logger=None ): # Scale error vectors so the smallest norm is 1 norms = np.linalg.norm(err_vecs, axis=1) err_vecs = err_vecs / norms.min() log = functools.partial(hp_log, logger) valid_coeffs = None for use in range(2, min(max_vecs, len(err_vecs)) + 1): log(f"Trying GDIIS with {use} previous cycles.") use_vecs = np.array(err_vecs[::-1][:use]) A = np.einsum("ij,kj->ik", use_vecs, use_vecs) try: coeffs = np.linalg.solve(A, np.ones(use)) except np.linalg.LinAlgError: log("LinAlgError when solving GDIIS matrix.") break # Scale coeffs so that their sum equals 1 coeffs_norm = np.linalg.norm(coeffs) valid_coeffs_norm = coeffs_norm <= 1e8 log(f"\tError vectors are linearly independent: {valid_coeffs_norm}") coeffs /= np.sum(coeffs) coeffs_str = np.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] 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 result = diis_result(valid_coeffs, coords, forces, prefix="G") log(f"\tUsed {len(result.coeffs)} error vectors for {result.type}.\n") return result
[docs] def gediis(coords, energies, forces, hessian=None, max_vecs=3, logger=None): log = functools.partial(hp_log, logger) use = min(len(coords), max_vecs) 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() 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: 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) coeffs = None if res.success: coeffs = x2c(res.x) en_ = res.fun log(f"\tOptimization converged!") coeff_str = np.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]): warnings.warn( f"GEDIIS converged, but proposed energy is above current energy! Returning None" ) return None result = diis_result(coeffs, coords, forces, energy=en_, prefix="GE") log(f"\tUsed {len(result.coeffs)} error vectors for {result.type}.\n") return result