Refactor the relevant code for psi initializer

haozhihan · haozhihan · commit 10d4ef72d635 · 2024-11-13T16:17:11.000+08:00
diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp
@@ -1,20 +1,19 @@
 #include "hsolver_pw.h"
 
-#include "module_parameter/parameter.h"
 #include "module_base/global_variable.h"
 #include "module_base/timer.h"
 #include "module_base/tool_quit.h"
 #include "module_elecstate/elecstate_pw.h"
 #include "module_hamilt_general/hamilt.h"
 #include "module_hamilt_pw/hamilt_pwdft/wavefunc.h"
-#include "module_psi/psi.h"
-
 #include "module_hsolver/diag_comm_info.h"
 #include "module_hsolver/diago_bpcg.h"
 #include "module_hsolver/diago_cg.h"
 #include "module_hsolver/diago_dav_subspace.h"
 #include "module_hsolver/diago_david.h"
 #include "module_hsolver/diago_iter_assist.h"
+#include "module_parameter/parameter.h"
+#include "module_psi/psi.h"
 
 #include <algorithm>
 #include <vector>
@@ -209,26 +208,30 @@ void HSolverPW<T, Device>::paw_func_after_kloop(psi::Psi<T, Device>& psi, elecst
 #endif
 
 template <typename T, typename Device>
-void HSolverPW<T, Device>::cal_ethr_band(const double& wk, const double* wg, const double& ethr, std::vector<double>& ethrs)
+void HSolverPW<T, Device>::cal_ethr_band(const double& wk,
+                                         const double* wg,
+                                         const double& ethr,
+                                         std::vector<double>& ethrs)
 {
     // threshold for classifying occupied and unoccupied bands
     const double occ_threshold = 1e-2;
     // diagonalization threshold limitation for unoccupied bands
     const double ethr_limit = 1e-5;
-    if(wk > 0.0)
+    if (wk > 0.0)
     {
         // Note: the idea of threshold for unoccupied bands (1e-5) comes from QE
-        // In ABACUS, We applied a smoothing process to this truncation to avoid abrupt changes in energy errors between different bands.
+        // In ABACUS, We applied a smoothing process to this truncation to avoid abrupt changes in energy errors between
+        // different bands.
         const double ethr_unocc = std::max(ethr_limit, ethr);
         for (int i = 0; i < ethrs.size(); i++)
         {
             double band_weight = wg[i] / wk;
             if (band_weight > occ_threshold)
             {
-                ethrs[i] = ethr; 
+                ethrs[i] = ethr;
             }
-            else if(band_weight > ethr_limit)
-            {// similar energy difference for different bands when band_weight in range [1e-5, 1e-2]
+            else if (band_weight > ethr_limit)
+            { // similar energy difference for different bands when band_weight in range [1e-5, 1e-2]
                 ethrs[i] = std::min(ethr_unocc, ethr / band_weight);
             }
             else
@@ -258,6 +261,27 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     ModuleBase::TITLE("HSolverPW", "solve");
     ModuleBase::timer::tick("HSolverPW", "solve");
 
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------for psi init guess!!!!--------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+    if (!PARAM.inp.psi_initializer && !this->initialed_psi && this->basis_type == "pw")
+    {
+        for (int ik = 0; ik < this->wfc_basis->nks; ++ik)
+        {
+            /// update H(k) for each k point
+            pHamilt->updateHk(ik);
+
+            /// update psi pointer for each k point
+            psi.fix_k(ik);
+
+            /// for psi init guess!!!!
+            hamilt::diago_PAO_in_pw_k2(this->ctx, ik, psi, this->wfc_basis, this->pwf, pHamilt);
+        }
+    }
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+
     this->rank_in_pool = rank_in_pool_in;
     this->nproc_in_pool = nproc_in_pool_in;
 
@@ -283,19 +307,20 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
         this->paw_func_in_kloop(ik);
 #endif
 
-        this->updatePsiK(pHamilt, psi, ik);
+        /// update psi pointer for each k point
+        psi.fix_k(ik);
 
         // template add precondition calculating here
         update_precondition(precondition, ik, this->wfc_basis->npwk[ik], Real(pes->pot->get_vl_of_0()));
-        
+
         // only dav_subspace method used smooth threshold for all bands now,
         // for other methods, this trick can be added in the future to accelerate calculation without accuracy loss.
-        if (this->method == "dav_subspace") 
+        if (this->method == "dav_subspace")
         {
             this->cal_ethr_band(pes->klist->wk[ik],
-                            &pes->wg(ik, 0),
-                            DiagoIterAssist<T, Device>::PW_DIAG_THR,
-                            ethr_band);
+                                &pes->wg(ik, 0),
+                                DiagoIterAssist<T, Device>::PW_DIAG_THR,
+                                ethr_band);
         }
 
 #ifdef USE_PAW
@@ -309,8 +334,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
         {
             GlobalV::ofs_running << "Average iterative diagonalization steps for k-points " << ik
                                  << " is: " << DiagoIterAssist<T, Device>::avg_iter
-                                 << " ; where current threshold is: " << this->diag_thr
-                                 << " . " << std::endl;
+                                 << " ; where current threshold is: " << this->diag_thr << " . " << std::endl;
             DiagoIterAssist<T, Device>::avg_iter = 0.0;
         }
         /// calculate the contribution of Psi for charge density rho
@@ -347,17 +371,6 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     }
 }
 
-template <typename T, typename Device>
-void HSolverPW<T, Device>::updatePsiK(hamilt::Hamilt<T, Device>* pHamilt, psi::Psi<T, Device>& psi, const int ik)
-{
-    psi.fix_k(ik);
-    if (!PARAM.inp.psi_initializer && !this->initialed_psi && this->basis_type == "pw")
-    {
-        hamilt::diago_PAO_in_pw_k2(this->ctx, ik, psi, this->wfc_basis, this->pwf, pHamilt);
-    }
-    /* lcao_in_pw now is based on newly implemented psi initializer, so it does not appear here*/
-}
-
 template <typename T, typename Device>
 void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                            psi::Psi<T, Device>& psi,
@@ -484,16 +497,13 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
     {
         auto ngk_pointer = psi.get_ngk_pointer();
         // hpsi_func (X, HX, ld, nvec) -> HX = H(X), X and HX blockvectors of size ld x nvec
-        auto hpsi_func = [hm, ngk_pointer](T *psi_in,
-                                           T *hpsi_out,
-                                           const int ld_psi,
-                                           const int nvec) {
+        auto hpsi_func = [hm, ngk_pointer](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) {
             ModuleBase::timer::tick("DavSubspace", "hpsi_func");
 
             // Convert "pointer data stucture" to a psi::Psi object
             auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ld_psi, ngk_pointer);
 
-            psi::Range bands_range(true, 0, 0, nvec-1);
+            psi::Range bands_range(true, 0, 0, nvec - 1);
 
             using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
             hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
@@ -513,8 +523,8 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                                   this->need_subspace,
                                                   comm_info);
 
-        DiagoIterAssist<T, Device>::avg_iter
-            += static_cast<double>(dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band.data(), scf));
+        DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
+            dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band.data(), scf));
     }
     else if (this->method == "dav")
     {
@@ -533,23 +543,20 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
         // dimensions of matrix to be solved
         const int dim = psi.get_current_nbas(); /// dimension of matrix
         const int nband = psi.get_nbands();     /// number of eigenpairs sought
-        const int ld_psi = psi.get_nbasis();     /// leading dimension of psi
+        const int ld_psi = psi.get_nbasis();    /// leading dimension of psi
 
         // Davidson matrix-blockvector functions
 
         auto ngk_pointer = psi.get_ngk_pointer();
         /// wrap hpsi into lambda function, Matrix \times blockvector
         // hpsi_func (X, HX, ld, nvec) -> HX = H(X), X and HX blockvectors of size ld x nvec
-        auto hpsi_func = [hm, ngk_pointer](T *psi_in,
-                                           T *hpsi_out,
-                                           const int ld_psi,
-                                           const int nvec) {
+        auto hpsi_func = [hm, ngk_pointer](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) {
             ModuleBase::timer::tick("David", "hpsi_func");
 
             // Convert pointer of psi_in to a psi::Psi object
             auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ld_psi, ngk_pointer);
 
-            psi::Range bands_range(true, 0, 0, nvec-1);
+            psi::Range bands_range(true, 0, 0, nvec - 1);
 
             using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
             hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
@@ -561,23 +568,19 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
         /// wrap spsi into lambda function, Matrix \times blockvector
         /// spsi(X, SX, ld, nvec)
         /// ld is leading dimension of psi and spsi
-        auto spsi_func = [hm](const T* psi_in, T* spsi_out,
-                               const int ld_psi,   // Leading dimension of psi and spsi.
-                               const int nvec      // Number of vectors(bands)
-                            ){
+        auto spsi_func = [hm](const T* psi_in,
+                              T* spsi_out,
+                              const int ld_psi, // Leading dimension of psi and spsi.
+                              const int nvec    // Number of vectors(bands)
+                         ) {
             ModuleBase::timer::tick("David", "spsi_func");
             // sPsi determines S=I or not by  PARAM.globalv.use_uspp inside
             // sPsi(psi, spsi, nrow, npw, nbands)
             hm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec);
             ModuleBase::timer::tick("David", "spsi_func");
         };
 
-        DiagoDavid<T, Device> david(pre_condition.data(),
-                                    nband,
-                                    dim,
-                                    PARAM.inp.pw_diag_ndim,
-                                    this->use_paw,
-                                    comm_info);
+        DiagoDavid<T, Device> david(pre_condition.data(), nband, dim, PARAM.inp.pw_diag_ndim, this->use_paw, comm_info);
         // do diag and add davidson iteration counts up to avg_iter
         DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(david.diag(hpsi_func,
                                                                                spsi_func,
@@ -593,7 +596,10 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
 }
 
 template <typename T, typename Device>
-void HSolverPW<T, Device>::update_precondition(std::vector<Real>& h_diag, const int ik, const int npw, const Real vl_of_0)
+void HSolverPW<T, Device>::update_precondition(std::vector<Real>& h_diag,
+                                               const int ik,
+                                               const int npw,
+                                               const Real vl_of_0)
 {
     h_diag.assign(h_diag.size(), 1.0);
     int precondition_type = 2;
@@ -646,8 +652,7 @@ void HSolverPW<T, Device>::output_iterInfo()
     {
         GlobalV::ofs_running << "Average iterative diagonalization steps: "
                              << DiagoIterAssist<T, Device>::avg_iter / this->wfc_basis->nks
-                             << " ; where current threshold is: " << this->diag_thr << " . "
-                             << std::endl;
+                             << " ; where current threshold is: " << this->diag_thr << " . " << std::endl;
         // reset avg_iter
         DiagoIterAssist<T, Device>::avg_iter = 0.0;
     }
diff --git a/source/module_hsolver/hsolver_pw.h b/source/module_hsolver/hsolver_pw.h
@@ -59,9 +59,6 @@ class HSolverPW
                          std::vector<Real>& pre_condition,
                          Real* eigenvalue);
 
-    // psi initializer && change k point in psi
-    void updatePsiK(hamilt::Hamilt<T, Device>* pHamilt, psi::Psi<T, Device>& psi, const int ik);
-
     // calculate the precondition array for diagonalization in PW base
     void update_precondition(std::vector<Real>& h_diag, const int ik, const int npw, const Real vl_of_0);
 
diff --git a/source/module_hsolver/hsolver_pw_sdft.cpp b/source/module_hsolver/hsolver_pw_sdft.cpp
@@ -1,9 +1,9 @@
 #include "hsolver_pw_sdft.h"
 
 #include "module_base/global_function.h"
+#include "module_base/parallel_device.h"
 #include "module_base/timer.h"
 #include "module_base/tool_title.h"
-#include "module_base/parallel_device.h"
 #include "module_elecstate/module_charge/symmetry_rho.h"
 
 #include <algorithm>
@@ -28,6 +28,30 @@ void HSolverPW_SDFT<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     const int nbands = psi.get_nbands();
     const int nks = psi.get_nk();
 
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------for psi init guess!!!!--------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+    if (!PARAM.inp.psi_initializer && !this->initialed_psi && this->basis_type == "pw")
+    {
+        for (int ik = 0; ik < nks; ++ik)
+        {
+            /// update H(k) for each k point
+            pHamilt->updateHk(ik);
+
+            if (nbands > 0 && GlobalV::MY_STOGROUP == 0)
+            {
+                /// update psi pointer for each k point
+                psi.fix_k(ik);
+
+                /// for psi init guess!!!!
+                hamilt::diago_PAO_in_pw_k2(this->ctx, ik, psi, this->wfc_basis, this->pwf, pHamilt);
+            }
+        }
+    }
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+    //---------------------------------------------------------------------------------------------------------------
+
     // prepare for the precondition of diagonalization
     std::vector<double> precondition(psi.get_nbasis(), 0.0);
 
@@ -44,8 +68,9 @@ void HSolverPW_SDFT<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
         pHamilt->updateHk(ik);
         if (nbands > 0 && GlobalV::MY_STOGROUP == 0)
         {
-            this->updatePsiK(pHamilt, psi, ik);
-            // template add precondition calculating here
+            /// update psi pointer for each k point
+            psi.fix_k(ik);
+            /// template add precondition calculating here
             this->update_precondition(precondition, ik, this->wfc_basis->npwk[ik], pes->pot->get_vl_of_0());
             /// solve eigenvector and eigenvalue for H(k)
             double* p_eigenvalues = &(pes->ekb(ik, 0));
@@ -105,7 +130,7 @@ void HSolverPW_SDFT<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     {
         for (int is = 0; is < this->nspin; is++)
         {
-            setmem_var_op()(this->ctx, pes_pw->rho[is], 0,  pes_pw->charge->nrxx);
+            setmem_var_op()(this->ctx, pes_pw->rho[is], 0, pes_pw->charge->nrxx);
         }
     }
     // calculate stochastic rho
