Add average iterative diagonalization step output (#6795)

* Call output_iterInfo in hsolver

* Add average iter output for cg

* Fix average iter output

* Add comment

* Align format

* Update output format

* Make clear output format

* Make clear output format

Author: Cstandardlib

source/source_hsolver/diago_cg.cpp

@@ -171,6 +171,7 @@ void DiagoCG<T, Device>::diag_once(const ct::Tensor& prec_in,
         {
             ++this->notconv_;
         }
+        iter_band.push_back(iter);
         avg += static_cast<Real>(iter) + 1.00;
 
         // reorder eigenvalue if they are not in the right order
@@ -575,7 +576,7 @@ bool DiagoCG<T, Device>::test_exit_cond(const int& ntry, const int& notconv) con
 }
 
 template <typename T, typename Device>
-void DiagoCG<T, Device>::diag(const Func& hpsi_func,
+double DiagoCG<T, Device>::diag(const Func& hpsi_func,
                               const Func& spsi_func,
                               ct::Tensor& psi,
                               ct::Tensor& eigen,
@@ -626,6 +627,20 @@ void DiagoCG<T, Device>::diag(const Func& hpsi_func,
     psi.zero();
     // copy psi_temp to psi for 0 to npw.
     psi.sync(psi_temp);
+
+#ifdef __DEBUG
+// only output iter count for each band if DEBUG!
+// this should not be output in production log
+    std::cout << "\n DiagoCG::diag' avg_iter_ = " << avg_iter_;
+    std::cout << "\n DiagoCG::diag' iter_band = ";
+    for (auto iter_in_band : iter_band)
+    {
+        std::cout << iter_in_band << " ";
+    }
+    std::cout << "\n";
+#endif
+
+    return avg_iter_;
 }
 
 namespace hsolver
@@ -644,4 +659,4 @@ template class DiagoCG<double, base_device::DEVICE_CPU>;
 template class DiagoCG<double, base_device::DEVICE_GPU>;
 #endif
 #endif
-} // namespace hsolver
+} // namespace hsolver

source/source_hsolver/diago_cg.h

@@ -2,6 +2,7 @@
 #define MODULE_HSOLVER_DIAGO_CG_H_
 
 #include <functional>
+#include <vector>
 
 #include <source_base/macros.h>
 #include <source_base/kernels/math_kernel_op.h>
@@ -35,13 +36,14 @@ class DiagoCG final
         const Real& pw_diag_thr,
         const int& pw_diag_nmax,
         const int& nproc_in_pool);
-    
+
     ~DiagoCG();
 
     // virtual void init(){};
     // refactor hpsi_info
     // this is the diag() function for CG method
-    void diag(const Func& hpsi_func,
+    // returns avg_iter
+    double diag(const Func& hpsi_func,
               const Func& spsi_func,
               ct::Tensor& psi,
               ct::Tensor& eigen,
@@ -59,7 +61,9 @@ class DiagoCG final
     /// col size for input psi matrix
     int n_basis_ = 0;
     /// average iteration steps for cg diagonalization
-    int avg_iter_ = 0;
+    double avg_iter_ = 0;
+    /// std::vector for iter count of each band
+    std::vector<int> iter_band;
     /// threshold for cg diagonalization
     Real pw_diag_thr_ = 1e-5;
     /// maximum iteration steps for cg diagonalization
@@ -87,15 +91,15 @@ class DiagoCG final
         ct::Tensor& pphi);
 
     void orth_grad(
-        const ct::Tensor& psi, 
-        const int& m, 
-        ct::Tensor& grad, 
+        const ct::Tensor& psi,
+        const int& m,
+        ct::Tensor& grad,
         ct::Tensor& scg,
         ct::Tensor& lagrange);
 
     void calc_gamma_cg(
         const int& iter,
-        const Real& cg_norm, 
+        const Real& cg_norm,
         const Real& theta,
         const ct::Tensor& prec,
         const ct::Tensor& scg,
@@ -110,8 +114,8 @@ class DiagoCG final
         const ct::Tensor& cg,
         const ct::Tensor& scg,
         const double& ethreshold,
-        Real &cg_norm, 
-        Real &theta, 
+        Real &cg_norm,
+        Real &theta,
         Real &eigen,
         ct::Tensor& phi_m,
         ct::Tensor& sphi,
@@ -133,4 +137,4 @@ class DiagoCG final
 
 } // namespace hsolver
 
-#endif // MODULE_HSOLVER_DIAGO_CG_H_
+#endif // MODULE_HSOLVER_DIAGO_CG_H_

source/source_hsolver/hsolver_pw.cpp

@@ -105,7 +105,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
         for (int i = 0; i < this->wfc_basis->nks; ++i)
         {
             const int ik = k_order[i];
-            
+
             // update H(k) for each k point
             pHamilt->updateHk(ik);
 
@@ -142,13 +142,13 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
 
             if (skip_charge)
             {
-                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;
+                GlobalV::ofs_running << " Average iterative diagonalization steps for k-points " << ik
+                                    << " is " << DiagoIterAssist<T, Device>::avg_iter
+                                    << "\n current threshold of diagonalization is " << this->diag_thr << std::endl;
                 DiagoIterAssist<T, Device>::avg_iter = 0.0;
             }
         }
-    }
+    } // if (use_k_continuity)
     else {
         // Original code without k-point continuity
         for (int ik = 0; ik < this->wfc_basis->nks; ++ik)
@@ -182,17 +182,22 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             // solve eigenvector and eigenvalue for H(k)
             this->hamiltSolvePsiK(pHamilt, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks);
 
+            // output iteration information and reset avg_iter
             if (skip_charge)
             {
                 GlobalV::ofs_running << " k(" << ik+1 << "/" << pes->klist->get_nkstot()
                                      << ") Iter steps (avg)=" << DiagoIterAssist<T, Device>::avg_iter
                                      << " threshold=" << this->diag_thr << std::endl;
                 DiagoIterAssist<T, Device>::avg_iter = 0.0;
             }
+
             /// calculate the contribution of Psi for charge density rho
         }
-    }
-    
+    } // else (use_k_continuity)
+
+    // output average iteration information and reset avg_iter
+    this->output_iterInfo();
+
     count++;
     // END Loop over k points
 
@@ -341,7 +346,9 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                .to_device<ct_Device>()
                                .slice({0}, {psi.get_current_ngk()});
 
-        cg.diag(hpsi_func, spsi_func, psi_tensor, eigen_tensor, this->ethr_band, prec_tensor);
+        DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
+            cg.diag(hpsi_func, spsi_func, psi_tensor, eigen_tensor, this->ethr_band, prec_tensor)
+        );
         // TODO: Double check tensormap's potential problem
         // ct::TensorMap(psi.get_pointer(), psi_tensor, {psi.get_nbands(), psi.get_nbasis()}).sync(psi_tensor);
     }
@@ -519,9 +526,9 @@ void HSolverPW<T, Device>::output_iterInfo()
     // in PW base, average iteration steps for each band and k-point should be printing
     if (DiagoIterAssist<T, Device>::avg_iter > 0.0)
     {
-        GlobalV::ofs_running << "Average iterative diagonalization steps: "
+        GlobalV::ofs_running << " Average iterative diagonalization steps for k-points is "
                              << DiagoIterAssist<T, Device>::avg_iter / this->wfc_basis->nks
-                             << " ; where current threshold is: " << this->diag_thr << " . " << std::endl;
+                             << "\n current threshold of diagonalizaiton is " << this->diag_thr << std::endl;
         // reset avg_iter
         DiagoIterAssist<T, Device>::avg_iter = 0.0;
     }
@@ -533,39 +540,39 @@ void HSolverPW<T, Device>::build_k_neighbors() {
     kvecs_c.resize(nk);
     k_order.clear();
     k_order.reserve(nk);
-    
+
     // Store k-points and corresponding indices
     struct KPoint {
         ModuleBase::Vector3<double> kvec;
         int index;
         double norm;
-        
-        KPoint(const ModuleBase::Vector3<double>& v, int i) : 
+
+        KPoint(const ModuleBase::Vector3<double>& v, int i) :
             kvec(v), index(i), norm(v.norm()) {}
     };
-    
+
     // Build k-point list
     std::vector<KPoint> klist;
     for (int ik = 0; ik < nk; ++ik) {
         kvecs_c[ik] = this->wfc_basis->kvec_c[ik];
         klist.push_back(KPoint(kvecs_c[ik], ik));
     }
-    
+
     // Sort k-points by distance from origin
     std::sort(klist.begin(), klist.end(),
         [](const KPoint& a, const KPoint& b) {
             return a.norm < b.norm;
         });
-    
+
     // Build parent-child relationships
     k_order.push_back(klist[0].index);
-    
+
     // Find nearest processed k-point as parent for each k-point
     for (int i = 1; i < nk; ++i) {
         int current_k = klist[i].index;
         double min_dist = 1e10;
         int parent = -1;
-        
+
         // find the nearest k-point as parent
         for (int j = 0; j < k_order.size(); ++j) {
             int processed_k = k_order[j];
@@ -575,7 +582,7 @@ void HSolverPW<T, Device>::build_k_neighbors() {
                 parent = processed_k;
             }
         }
-        
+
         k_parent[current_k] = parent;
         k_order.push_back(current_k);
     }
@@ -585,34 +592,34 @@ template <typename T, typename Device>
 void HSolverPW<T, Device>::propagate_psi(psi::Psi<T, Device>& psi, const int from_ik, const int to_ik) {
     const int nbands = psi.get_nbands();
     const int npwk = this->wfc_basis->npwk[to_ik];
-    
+
     // Get k-point difference
     ModuleBase::Vector3<double> dk = kvecs_c[to_ik] - kvecs_c[from_ik];
-    
+
     // Allocate porter locally
     T* porter = nullptr;
     resmem_complex_op()(porter, this->wfc_basis->nmaxgr, "HSolverPW::porter");
-    
+
     // Process each band
     for (int ib = 0; ib < nbands; ib++)
     {
         // Fix current k-point and band
         // psi.fix_k(from_ik);
-        
+
         // FFT to real space
         // this->wfc_basis->recip_to_real(this->ctx, psi.get_pointer(ib), porter, from_ik);
         this->wfc_basis->recip_to_real(this->ctx, &psi(from_ik, ib, 0), porter, from_ik);
-        
+
         // Apply phase factor
         //     // TODO: Check how to get the r vector
         //     ModuleBase::Vector3<double> r = this->wfc_basis->get_ir2r(ir);
         //     double phase = this->wfc_basis->tpiba * (dk.x * r.x + dk.y * r.y + dk.z * r.z);
         //     psi_real[ir] *= std::exp(std::complex<double>(0.0, phase));
         // }
-        
+
         // Fix k-point for target
         // psi.fix_k(to_ik);
-        
+
         // FFT back to reciprocal space
         // this->wfc_basis->real_to_recip(this->ctx, porter, psi.get_pointer(ib), to_ik, true);
         this->wfc_basis->real_to_recip(this->ctx, porter, &psi(to_ik, ib, 0), to_ik);

source/source_hsolver/test/hsolver_pw_sup.h

@@ -92,7 +92,7 @@ DiagoCG<T, Device>::~DiagoCG() {
 }
 
 template <typename T, typename Device>
-void DiagoCG<T, Device>::diag(const Func& hpsi_func,
+double DiagoCG<T, Device>::diag(const Func& hpsi_func,
                               const Func& spsi_func,
                               ct::Tensor& psi,
                               ct::Tensor& eigen,
@@ -112,7 +112,7 @@ void DiagoCG<T, Device>::diag(const Func& hpsi_func,
         eigen_pack[ib] /= n_basis;
     }
     DiagoIterAssist<T, Device>::avg_iter += 1.0;
-    return;
+    return avg_iter_;
 }
 
 template class DiagoCG<std::complex<float>, base_device::DEVICE_CPU>;