diis.h

/*
 *  Copyright (C) 2004-2024 Edward F. Valeev
 *
 *  This file is part of Libint library.
 *
 *  Libint library is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  Libint library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with Libint library.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
 
#ifndef _libint2_src_lib_libint_diis_h_
#define _libint2_src_lib_libint_diis_h_
 
#include <libint2/util/cxxstd.h>
#if LIBINT2_CPLUSPLUS_STD < 2011
#error "libint2/diis.h requires C++11 support"
#endif
 
#include <deque>
 
#pragma GCC diagnostic push
#pragma GCC system_header
#include <Eigen/Core>
#include <Eigen/QR>
#pragma GCC diagnostic pop
 
namespace libint2 {
 
namespace diis {
 
template <typename D>
struct traits;
 
/*
template <typename D>
typename traits<D>::element_type
dot_product(const D& d1, const D& d2);
 
template <typename D>
void
zero(D& d);
 
template <typename D, typename Scalar>
void
axpy(const D& y, Scalar a, const D& x);
*/
 
template <typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows,
          int _MaxCols>
struct traits<
    Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>> {
  typedef Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols> D;
  typedef typename D::Scalar element_type;
};
 
template <typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows,
          int _MaxCols>
typename traits<Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows,
                              _MaxCols>>::element_type
dot_product(const Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows,
                                _MaxCols>& d1,
            const Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows,
                                _MaxCols>& d2) {
  return d1.cwiseProduct(d2).sum();
}
 
template <typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows,
          int _MaxCols>
void zero(
    Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>& d) {
  d.setZero(d.rows(), d.cols());
}
 
template <typename _Scalar, int _Rows, int _Cols, int _Options, int _MaxRows,
          int _MaxCols, typename Scalar>
void axpy(Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows, _MaxCols>& y,
          Scalar a,
          const Eigen::Matrix<_Scalar, _Rows, _Cols, _Options, _MaxRows,
                              _MaxCols>& x) {
  y += a * x;
}
 
}  // namespace diis
 
 
template <typename D>
class DIIS {
 public:
  typedef typename diis::traits<D>::element_type value_type;
 
 
  DIIS(unsigned int strt = 1, unsigned int ndi = 5, value_type dmp = 0,
       unsigned int ngr = 1, unsigned int ngrdiis = 1, value_type mf = 0)
      : error_(0),
        errorset_(false),
        start(strt),
        ndiis(ndi),
        iter(0),
        ngroup(ngr),
        ngroupdiis(ngrdiis),
        damping_factor(dmp),
        mixing_fraction(mf) {
    init();
  }
  ~DIIS() {
    x_.clear();
    errors_.clear();
    x_extrap_.clear();
  }
 
  void extrapolate(D& x, D& error, bool extrapolate_error = false) {
    using namespace ::libint2::diis;
 
    const value_type zero_determinant =
        std::numeric_limits<value_type>::epsilon();
    const value_type zero_norm = 1.0e-10;
 
    iter++;
 
    const bool do_mixing = (mixing_fraction != 0.0);
    const value_type scale = 1.0 + damping_factor;
 
    // if have ndiis vectors
    if (errors_.size() == ndiis) {  // holding max # of vectors already? drop
                                    // the least recent {x, error} pair
      x_.pop_front();
      errors_.pop_front();
      if (not x_extrap_.empty()) x_extrap_.pop_front();
      EigenMatrixX Bcrop = B_.bottomRightCorner(ndiis - 1, ndiis - 1);
      Bcrop.conservativeResize(ndiis, ndiis);
      B_ = Bcrop;
    }
 
    // push {x, error} to the set
    x_.push_back(x);
    errors_.push_back(error);
    const unsigned int nvec = errors_.size();
    assert(x_.size() == errors_.size());
 
    // and compute the most recent elements of B, B(i,j) = <ei|ej>
    for (unsigned int i = 0; i < nvec - 1; i++)
      B_(i, nvec - 1) = B_(nvec - 1, i) =
          dot_product(errors_[i], errors_[nvec - 1]);
    B_(nvec - 1, nvec - 1) = dot_product(errors_[nvec - 1], errors_[nvec - 1]);
 
    if (iter == 1) {  // the first iteration
      if (not x_extrap_.empty() && do_mixing) {
        zero(x);
        axpy(x, (1.0 - mixing_fraction), x_[0]);
        axpy(x, mixing_fraction, x_extrap_[0]);
      }
    } else if (iter > start && (((iter - start) % ngroup) <
                                ngroupdiis)) {  // not the first iteration and
                                                // need to extrapolate?
 
      EigenVectorX c;
 
      value_type absdetA;
      unsigned int nskip = 0;  // how many oldest vectors to skip for the sake
                               // of conditioning? try zero
      // skip oldest vectors until found a numerically stable system
      do {
        const unsigned int rank = nvec - nskip + 1;  // size of matrix A
 
        // set up the DIIS linear system: A c = rhs
        EigenMatrixX A(rank, rank);
        A.col(0).setConstant(-1.0);
        A.row(0).setConstant(-1.0);
        A(0, 0) = 0.0;
        EigenVectorX rhs = EigenVectorX::Zero(rank);
        rhs[0] = -1.0;
 
        value_type norm = 1.0;
        if (B_(nskip, nskip) > zero_norm) norm = 1.0 / B_(nskip, nskip);
 
        A.block(1, 1, rank - 1, rank - 1) =
            B_.block(nskip, nskip, rank - 1, rank - 1) * norm;
        A.diagonal() *= scale;
        // for (unsigned int i=1; i < rank ; i++) {
        //   for (unsigned int j=1; j <= i ; j++) {
        //     A(i, j) = A(j, i) = B_(i+nskip-1, j+nskip-1) * norm;
        //     if (i==j) A(i, j) *= scale;
        //   }
        // }
 
#if 0
            std::cout << "DIIS: iter=" << iter << " nskip=" << nskip << " nvec=" << nvec << std::endl;
            std::cout << "DIIS: B=" << B_ << std::endl;
            std::cout << "DIIS: A=" << A << std::endl;
            std::cout << "DIIS: rhs=" << rhs << std::endl;
#endif
 
        // finally, solve the DIIS linear system
        Eigen::ColPivHouseholderQR<EigenMatrixX> A_QR = A.colPivHouseholderQr();
        c = A_QR.solve(rhs);
        absdetA = A_QR.absDeterminant();
 
        // std::cout << "DIIS: |A|=" << absdetA << " sol=" << c << std::endl;
 
        ++nskip;
 
      } while (absdetA < zero_determinant &&
               nskip < nvec);  // while (system is poorly conditioned)
 
      // failed?
      if (absdetA < zero_determinant) {
        std::ostringstream oss;
        oss << "DIIS::extrapolate: poorly-conditioned system, |A| = "
            << absdetA;
        throw std::domain_error(oss.str());
      }
      --nskip;  // undo the last ++ :-(
 
      {
        zero(x);
        for (unsigned int k = nskip, kk = 1; k < nvec; ++k, ++kk) {
          if (not do_mixing || x_extrap_.empty()) {
            // std::cout << "contrib " << k << " c=" << c[kk] << ":" <<
            // std::endl << x_[k] << std::endl;
            axpy(x, c[kk], x_[k]);
            if (extrapolate_error) axpy(error, c[kk], errors_[k]);
          } else {
            axpy(x, c[kk] * (1.0 - mixing_fraction), x_[k]);
            axpy(x, c[kk] * mixing_fraction, x_extrap_[k]);
          }
        }
      }
    }  // do DIIS
 
    // only need to keep extrapolated x if doing mixing
    if (do_mixing) x_extrap_.push_back(x);
  }
 
  void start_extrapolation() {
    if (start > iter) start = iter + 1;
  }
 
  void reinitialize(const D* data = 0) {
    iter = 0;
    if (data) {
      const bool do_mixing = (mixing_fraction != 0.0);
      if (do_mixing) x_extrap_.push_front(*data);
    }
  }
 
 private:
  value_type error_;
  bool errorset_;
 
  unsigned int start;
  unsigned int ndiis;
  unsigned int iter;
  unsigned int ngroup;
  unsigned int ngroupdiis;
  value_type damping_factor;
  value_type mixing_fraction;
 
  typedef Eigen::Matrix<value_type, Eigen::Dynamic, Eigen::Dynamic,
                        Eigen::RowMajor>
      EigenMatrixX;
  typedef Eigen::Matrix<value_type, Eigen::Dynamic, 1> EigenVectorX;
 
  EigenMatrixX B_;  
 
  std::deque<D>
      x_;  
  std::deque<D> errors_;    
  std::deque<D> x_extrap_;  
 
  void set_error(value_type e) {
    error_ = e;
    errorset_ = true;
  }
  value_type error() { return error_; }
 
  void init() {
    iter = 0;
 
    B_ = EigenMatrixX::Zero(ndiis, ndiis);
 
    x_.clear();
    errors_.clear();
    x_extrap_.clear();
    // x_.resize(ndiis);
    // errors_.resize(ndiis);
    //  x_extrap_ is bigger than the other because
    //  it must hold data associated with the next iteration
    // x_extrap_.resize(diis+1);
  }
 
};  // class DIIS
 
}  // namespace libint2
 
#include <libint2/engine.h>
 
#endif /* _libint2_src_lib_libint_diis_h_ */