Adaptive Convergence Optimization for Eigenvalue Iterative Solvers (CG & David)

source/source_hsolver/diago_cg.cpp

@@ -12,6 +12,8 @@
 #include "source_base/global_function.h"        // ModuleBase::GlobalFunc::NOTE
 #include "source_hsolver/diago_cg.h"
 
+#include <algorithm>
+
 using namespace hsolver;
 
 template <typename T, typename Device>
@@ -53,6 +55,13 @@ DiagoCG<T, Device>::~DiagoCG()
     delete this->neg_one_;
 }
 
+template <typename T, typename Device>
+void DiagoCG<T, Device>::set_adaptive_cg(const bool use_pr_plus, const Real max_gamma)
+{
+    this->use_pr_plus_ = use_pr_plus;
+    this->max_cg_gamma_ = max_gamma;
+}
+
 template <typename T, typename Device>
 void DiagoCG<T, Device>::diag_once(const ct::Tensor& prec_in,
                                    ct::Tensor& psi,
@@ -358,7 +367,15 @@ void DiagoCG<T, Device>::calc_gamma_cg(const int& iter,
     {
         // (4) Update gamma !
         REQUIRES_OK(gg_last != 0.0, "DiagoCG_New::calc_gamma_cg: gg_last is zero, which is not allowed!");
-        const Real gamma = (gg_now - gg_inter) / gg_last;
+        Real gamma = (gg_now - gg_inter) / gg_last;
+        if (this->use_pr_plus_)
+        {
+            gamma = std::max(static_cast<Real>(0.0), gamma);
+        }
+        if (this->max_cg_gamma_ > static_cast<Real>(0.0))
+        {
+            gamma = std::min(gamma, this->max_cg_gamma_);
+        }
 
         // (5) Update gg_last !
         gg_last = gg_now;

source/source_hsolver/diago_cg.h

@@ -40,6 +40,8 @@ class DiagoCG final
 
     ~DiagoCG();
 
+    void set_adaptive_cg(const bool use_pr_plus, const Real max_gamma = static_cast<Real>(5.0));
+
     // virtual void init(){};
     // refactor hpsi_info
     // this is the diag() function for CG method
@@ -80,6 +82,8 @@ class DiagoCG final
     std::string calculation_ = {};
 
     bool need_subspace_ = false;
+    bool use_pr_plus_ = true;
+    Real max_cg_gamma_ = static_cast<Real>(5.0);
     /// A function object that performs the hPsi calculation.
     HPsiFunc hpsi_func_ = nullptr;
     /// A function object that performs the sPsi calculation.

source/source_hsolver/diago_david.cpp

@@ -8,6 +8,9 @@
 #include "source_base/kernels/math_kernel_op.h"
 #include "source_base/parallel_comm.h"
 
+#include <algorithm>
+#include <cmath>
+#include <limits>
 
 using namespace hsolver;
 
@@ -114,6 +117,40 @@ DiagoDavid<T, Device>::~DiagoDavid()
 #endif
 }
 
+template <typename T, typename Device>
+void DiagoDavid<T, Device>::set_preconditioner(const DiagoPreconditioner preconditioner_type_in)
+{
+    this->preconditioner_type = preconditioner_type_in;
+}
+
+template <typename T, typename Device>
+void DiagoDavid<T, Device>::set_adaptive_restart(const bool enabled, const Real fill_ratio)
+{
+    this->adaptive_restart = enabled;
+    this->adaptive_restart_fill_ratio = std::min(static_cast<Real>(1.0),
+                                                 std::max(static_cast<Real>(0.5), fill_ratio));
+}
+
+template <typename T, typename Device>
+void DiagoDavid<T, Device>::set_min_precondition(const Real min_precondition_in)
+{
+    this->min_precondition = std::max(min_precondition_in, std::numeric_limits<Real>::epsilon());
+}
+
+template <typename T, typename Device>
+typename DiagoDavid<T, Device>::Real DiagoDavid<T, Device>::shifted_precondition_denominator(
+    const Real diagonal,
+    const Real eigenvalue,
+    const Real min_precondition)
+{
+    const Real x = std::abs(diagonal - eigenvalue);
+    const Real denom = static_cast<Real>(0.5)
+                       * (static_cast<Real>(1.0) + x
+                          + std::sqrt(static_cast<Real>(1.0)
+                                      + (x - static_cast<Real>(1.0)) * (x - static_cast<Real>(1.0))));
+    return std::max(denom, std::max(min_precondition, std::numeric_limits<Real>::epsilon()));
+}
+
 template <typename T, typename Device>
 int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
                                      const SPsiFunc& spsi_func,
@@ -233,8 +270,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
         }
 
         ModuleBase::timer::end("DiagoDavid", "check_update");
-        if (!this->notconv || (nbase + this->notconv > nbase_x)
-            || (dav_iter == david_maxiter))
+        if (!this->notconv || this->should_refresh(nbase, this->notconv, nbase_x, dav_iter, david_maxiter))
         {
             ModuleBase::timer::start("DiagoDavid", "last");
 
@@ -462,37 +498,7 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     }
     //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
 
-    // Preconditioning
-    // basis[nbase] = T * basis[nbase] = T * (H - lambda * S) * psi
-    // where T, the preconditioner, is an approximate inverse of H
-    //          T is a diagonal stored in array `precondition`
-    // to do preconditioning, divide each column of basis by the corresponding element of precondition
-    for (int m = 0; m < notconv; m++)
-    {
-        //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-        if (this->device == base_device::GpuDevice)
-        {
-#if defined(__CUDA) || defined(__ROCM)
-            ModuleBase::vector_div_vector_op<T, Device>()(dim,
-                                                          basis + dim * (nbase + m),
-                                                          basis + dim * (nbase + m),
-                                                          this->d_precondition);
-#endif
-        }
-        else
-        {
-            ModuleBase::vector_div_vector_op<T, Device>()(dim,
-                                                          basis + dim * (nbase + m),
-                                                          basis + dim * (nbase + m),
-                                                          this->precondition);
-        }
-        //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-        // for (int ig = 0; ig < dim; ig++)
-        // {
-        //     ppsi[ig] /= this->precondition[ig];
-        // }
-        //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
-    }
+    this->apply_precondition(dim, nbase, notconv, basis, unconv, eigenvalue);
 
     // there is a nbase to nbase + notconv band orthogonalise
     // plan for SchmidtOrth
@@ -574,6 +580,95 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     return;
 }
 
+template <typename T, typename Device>
+void DiagoDavid<T, Device>::apply_precondition(const int& dim,
+                                               const int& nbase,
+                                               const int& notconv,
+                                               T* basis,
+                                               const int* unconv,
+                                               const Real* eigenvalue)
+{
+    ModuleBase::timer::start("DiagoDavid", "precondition");
+
+    for (int m = 0; m < notconv; m++)
+    {
+        if (this->preconditioner_type == DiagoPreconditioner::Diagonal)
+        {
+            if (this->device == base_device::GpuDevice)
+            {
+#if defined(__CUDA) || defined(__ROCM)
+                ModuleBase::vector_div_vector_op<T, Device>()(dim,
+                                                              basis + dim * (nbase + m),
+                                                              basis + dim * (nbase + m),
+                                                              this->d_precondition);
+#endif
+            }
+            else
+            {
+                ModuleBase::vector_div_vector_op<T, Device>()(dim,
+                                                              basis + dim * (nbase + m),
+                                                              basis + dim * (nbase + m),
+                                                              this->precondition);
+            }
+            continue;
+        }
+
+        std::vector<Real> shifted_precondition(dim);
+        for (int i = 0; i < dim; i++)
+        {
+            shifted_precondition[i] = shifted_precondition_denominator(this->precondition[i],
+                                                                       eigenvalue[unconv[m]],
+                                                                       this->min_precondition);
+        }
+
+        if (this->device == base_device::GpuDevice)
+        {
+#if defined(__CUDA) || defined(__ROCM)
+            Real* d_shifted_precondition = nullptr;
+            resmem_var_op()(d_shifted_precondition, dim);
+            syncmem_var_h2d_op()(d_shifted_precondition, shifted_precondition.data(), dim);
+            ModuleBase::vector_div_vector_op<T, Device>()(dim,
+                                                          basis + dim * (nbase + m),
+                                                          basis + dim * (nbase + m),
+                                                          d_shifted_precondition);
+            delmem_var_op()(d_shifted_precondition);
+#endif
+        }
+        else
+        {
+            ModuleBase::vector_div_vector_op<T, Device>()(dim,
+                                                          basis + dim * (nbase + m),
+                                                          basis + dim * (nbase + m),
+                                                          shifted_precondition.data());
+        }
+    }
+
+    ModuleBase::timer::end("DiagoDavid", "precondition");
+}
+
+template <typename T, typename Device>
+bool DiagoDavid<T, Device>::should_refresh(const int& nbase,
+                                           const int& notconv,
+                                           const int nbase_x,
+                                           const int& david_iter,
+                                           const int& david_maxiter) const
+{
+    if (notconv == 0 || david_iter == david_maxiter || nbase + notconv > nbase_x)
+    {
+        return true;
+    }
+    if (!this->adaptive_restart || david_iter < 2 || nbase <= nband)
+    {
+        return false;
+    }
+
+    const Real fill_ratio = static_cast<Real>(nbase) / static_cast<Real>(std::max(1, nbase_x));
+    const int free_subspace = nbase_x - nbase;
+    const bool subspace_nearly_full = fill_ratio >= this->adaptive_restart_fill_ratio
+                                      || free_subspace < std::max(1, nband / 2);
+    const bool many_unconverged = notconv > std::max(1, nband / 2);
+    return subspace_nearly_full && many_unconverged;
+}
 
 template <typename T, typename Device>
 void DiagoDavid<T, Device>::cal_elem(const int& dim,

source/source_hsolver/diago_david.h

@@ -15,6 +15,12 @@
 
 namespace hsolver
 {
+enum class DiagoPreconditioner
+{
+    Diagonal,
+    ShiftedDiagonal
+};
+
 /**
  * @class DiagoDavid
  * @brief A class that implements the block-Davidson algorithm for solving generalized eigenvalue problems.
@@ -67,6 +73,12 @@ class DiagoDavid
      */
     ~DiagoDavid();
 
+    void set_preconditioner(const DiagoPreconditioner preconditioner_type);
+    void set_adaptive_restart(const bool enabled, const Real fill_ratio = static_cast<Real>(0.85));
+    void set_min_precondition(const Real min_precondition);
+    static Real shifted_precondition_denominator(const Real diagonal,
+                                                 const Real eigenvalue,
+                                                 const Real min_precondition);
 
     // declare type of matrix-blockvector functions.
     // the function type is defined as a std::function object.
@@ -154,6 +166,10 @@ class DiagoDavid
     const int david_ndim = 4;
     /// number of unconverged eigenvalues
     int notconv = 0;
+    DiagoPreconditioner preconditioner_type = DiagoPreconditioner::ShiftedDiagonal;
+    bool adaptive_restart = true;
+    Real adaptive_restart_fill_ratio = static_cast<Real>(0.85);
+    Real min_precondition = static_cast<Real>(1.0e-8);
 
     /// precondition for diag, diagonal approximation of matrix A(i.e. Hamilt)
     const Real* precondition = nullptr;
@@ -216,6 +232,19 @@ class DiagoDavid
                   const int* unconv,
                   const Real* eigenvalue);
 
+    void apply_precondition(const int& dim,
+                            const int& nbase,
+                            const int& notconv,
+                            T* basis,
+                            const int* unconv,
+                            const Real* eigenvalue);
+
+    bool should_refresh(const int& nbase,
+                        const int& notconv,
+                        const int nbase_x,
+                        const int& david_iter,
+                        const int& david_maxiter) const;
+
     /**
      * Calculates the elements of the diagonalization matrix for the DiagoDavid class.
      *

source/source_hsolver/test/diago_david_test.cpp

@@ -61,8 +61,15 @@ void lapackEigen(int& npw, std::vector<std::complex<double>>& hm, double* e, boo
 class DiagoDavPrepare 
 {
 public:
-	DiagoDavPrepare(int nband, int npw, int sparsity, int order,double eps,int maxiter):
-		nband(nband),npw(npw),sparsity(sparsity),order(order),eps(eps),maxiter(maxiter) 
+	DiagoDavPrepare(int nband,
+                    int npw,
+                    int sparsity,
+                    int order,
+                    double eps,
+                    int maxiter,
+                    hsolver::DiagoPreconditioner preconditioner_type = hsolver::DiagoPreconditioner::ShiftedDiagonal):
+		nband(nband),npw(npw),sparsity(sparsity),order(order),eps(eps),maxiter(maxiter),
+        preconditioner_type(preconditioner_type)
 	{
 #ifdef __MPI	
 		MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
@@ -73,6 +80,7 @@ class DiagoDavPrepare
 	int nband, npw, sparsity, order, maxiter, notconv;
 	double eps, avg_iter;
 	int nprocs=1, mypnum=0;
+    hsolver::DiagoPreconditioner preconditioner_type;
 
 	void CompareEigen(psi::Psi<std::complex<double>> &phi, double *precondition)
 	{
@@ -99,6 +107,7 @@ class DiagoDavPrepare
 		const int nband = phi.get_nbands();
 		const int ld_psi = phi.get_nbasis();
 		hsolver::DiagoDavid<std::complex<double>> dav(precondition, nband, dim, order, comm_info);
+        dav.set_preconditioner(preconditioner_type);
 
 		hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_NMAX = maxiter;
 		hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_THR = eps;
@@ -193,11 +202,21 @@ TEST_P(DiagoDavTest,RandomHamilt)
 INSTANTIATE_TEST_SUITE_P(VerifyDiag,DiagoDavTest,::testing::Values(
 		//DiagoDavPrepare(int nband, int npw, int sparsity, int order,double eps,int maxiter)
         DiagoDavPrepare(10,100,0,4,1e-5,500),
-        DiagoDavPrepare(20,500,7,4,1e-5,500)
+        DiagoDavPrepare(20,500,7,4,1e-5,500),
+        DiagoDavPrepare(8,80,0,3,1e-5,300, hsolver::DiagoPreconditioner::Diagonal)
         //DiagoDavPrepare(50,1000,8,4,1e-5,500)
 		//DiagoDavPrepare(20,2000,8,4,1e-5,500)
 ));
 
+TEST(DiagoDavPreconditionerTest, ShiftedDiagonalDenominatorIsBounded)
+{
+    using Dav = hsolver::DiagoDavid<std::complex<double>>;
+    EXPECT_NEAR(Dav::shifted_precondition_denominator(3.0, 2.0, 1.0e-8), 1.5, 1.0e-12);
+    EXPECT_DOUBLE_EQ(Dav::shifted_precondition_denominator(1.0, 3.0, 1.0e-8),
+                     Dav::shifted_precondition_denominator(3.0, 1.0, 1.0e-8));
+    EXPECT_DOUBLE_EQ(Dav::shifted_precondition_denominator(2.0, 2.0, 2.0), 2.0);
+}
+
 TEST(DiagoDavRealSystemTest, dataH)
 {
 	std::vector<std::complex<double>> hmatrix;