diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index 1a72774303..564c36e74b 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -171,6 +171,7 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, { ++this->notconv_; } + iter_band.push_back(iter); avg += static_cast(iter) + 1.00; // reorder eigenvalue if they are not in the right order @@ -575,7 +576,7 @@ bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) con } template -void DiagoCG::diag(const Func& hpsi_func, +double DiagoCG::diag(const Func& hpsi_func, const Func& spsi_func, ct::Tensor& psi, ct::Tensor& eigen, @@ -626,6 +627,20 @@ void DiagoCG::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; template class DiagoCG; #endif #endif -} // namespace hsolver \ No newline at end of file +} // namespace hsolver diff --git a/source/source_hsolver/diago_cg.h b/source/source_hsolver/diago_cg.h index 20a5890552..bf03cb5850 100644 --- a/source/source_hsolver/diago_cg.h +++ b/source/source_hsolver/diago_cg.h @@ -2,6 +2,7 @@ #define MODULE_HSOLVER_DIAGO_CG_H_ #include +#include #include #include @@ -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 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_ \ No newline at end of file +#endif // MODULE_HSOLVER_DIAGO_CG_H_ diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index f3e5a7ca60..7d178f699b 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -105,7 +105,7 @@ void HSolverPW::solve(hamilt::Hamilt* 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::solve(hamilt::Hamilt* pHamilt, if (skip_charge) { - GlobalV::ofs_running << "Average iterative diagonalization steps for k-points " << ik - << " is: " << DiagoIterAssist::avg_iter - << " ; where current threshold is: " << this->diag_thr << " . " << std::endl; + GlobalV::ofs_running << " Average iterative diagonalization steps for k-points " << ik + << " is " << DiagoIterAssist::avg_iter + << "\n current threshold of diagonalization is " << this->diag_thr << std::endl; DiagoIterAssist::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,6 +182,7 @@ void HSolverPW::solve(hamilt::Hamilt* 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() @@ -189,10 +190,14 @@ void HSolverPW::solve(hamilt::Hamilt* pHamilt, << " threshold=" << this->diag_thr << std::endl; DiagoIterAssist::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::hamiltSolvePsiK(hamilt::Hamilt* hm, .to_device() .slice({0}, {psi.get_current_ngk()}); - cg.diag(hpsi_func, spsi_func, psi_tensor, eigen_tensor, this->ethr_band, prec_tensor); + DiagoIterAssist::avg_iter += static_cast( + 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::output_iterInfo() // in PW base, average iteration steps for each band and k-point should be printing if (DiagoIterAssist::avg_iter > 0.0) { - GlobalV::ofs_running << "Average iterative diagonalization steps: " + GlobalV::ofs_running << " Average iterative diagonalization steps for k-points is " << DiagoIterAssist::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::avg_iter = 0.0; } @@ -533,39 +540,39 @@ void HSolverPW::build_k_neighbors() { kvecs_c.resize(nk); k_order.clear(); k_order.reserve(nk); - + // Store k-points and corresponding indices struct KPoint { ModuleBase::Vector3 kvec; int index; double norm; - - KPoint(const ModuleBase::Vector3& v, int i) : + + KPoint(const ModuleBase::Vector3& v, int i) : kvec(v), index(i), norm(v.norm()) {} }; - + // Build k-point list std::vector 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::build_k_neighbors() { parent = processed_k; } } - + k_parent[current_k] = parent; k_order.push_back(current_k); } @@ -585,34 +592,34 @@ template void HSolverPW::propagate_psi(psi::Psi& 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 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 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(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); diff --git a/source/source_hsolver/test/hsolver_pw_sup.h b/source/source_hsolver/test/hsolver_pw_sup.h index 07cc650013..fb3757a08b 100644 --- a/source/source_hsolver/test/hsolver_pw_sup.h +++ b/source/source_hsolver/test/hsolver_pw_sup.h @@ -92,7 +92,7 @@ DiagoCG::~DiagoCG() { } template -void DiagoCG::diag(const Func& hpsi_func, +double DiagoCG::diag(const Func& hpsi_func, const Func& spsi_func, ct::Tensor& psi, ct::Tensor& eigen, @@ -112,7 +112,7 @@ void DiagoCG::diag(const Func& hpsi_func, eigen_pack[ib] /= n_basis; } DiagoIterAssist::avg_iter += 1.0; - return; + return avg_iter_; } template class DiagoCG, base_device::DEVICE_CPU>;