diff --git a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp index 8ef539902e..112eb836e2 100644 --- a/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp +++ b/python/pyabacus/src/hsolver/py_diago_dav_subspace.hpp @@ -101,12 +101,12 @@ class PyDiagoDavSubspace int diag( std::function>(py::array_t>)> mm_op, - std::vector precond_vec, + std::vector& precond_vec, int dav_ndim, double tol, int max_iter, bool need_subspace, - std::vector diag_ethr, + std::vector& diag_ethr, bool scf_type, hsolver::diag_comm_info comm_info ) { @@ -141,7 +141,7 @@ class PyDiagoDavSubspace comm_info ); - return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr.data(), scf_type); + return obj->diag(hpsi_func, psi, nbasis, eigenvalue, diag_ethr, scf_type); } private: diff --git a/python/pyabacus/src/hsolver/py_diago_david.hpp b/python/pyabacus/src/hsolver/py_diago_david.hpp index 2008eb6b85..8a8d2c727e 100644 --- a/python/pyabacus/src/hsolver/py_diago_david.hpp +++ b/python/pyabacus/src/hsolver/py_diago_david.hpp @@ -101,9 +101,10 @@ class PyDiagoDavid int diag( std::function>(py::array_t>)> mm_op, - std::vector precond_vec, + std::vector& precond_vec, int dav_ndim, double tol, + std::vector& diag_ethr, int max_iter, bool use_paw, hsolver::diag_comm_info comm_info @@ -146,7 +147,7 @@ class PyDiagoDavid comm_info ); - return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, tol, max_iter); + return obj->diag(hpsi_func, spsi_func, nbasis, psi, eigenvalue, diag_ethr, max_iter); } private: diff --git a/python/pyabacus/src/hsolver/py_hsolver.cpp b/python/pyabacus/src/hsolver/py_hsolver.cpp index 7e76d1f874..2ae23d0966 100644 --- a/python/pyabacus/src/hsolver/py_hsolver.cpp +++ b/python/pyabacus/src/hsolver/py_hsolver.cpp @@ -121,6 +121,8 @@ void bind_hsolver(py::module& m) eigenvectors to be calculated. tol : double The tolerance for the convergence. + diag_ethr: np.ndarray + The tolerance vector. max_iter : int The maximum number of iterations. use_paw : bool @@ -130,6 +132,7 @@ void bind_hsolver(py::module& m) "precond_vec"_a, "dav_ndim"_a, "tol"_a, + "diag_ethr"_a, "max_iter"_a, "use_paw"_a, "comm_info"_a) diff --git a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py index d1cebad5dd..56f3ce532b 100644 --- a/python/pyabacus/src/pyabacus/hsolver/_hsolver.py +++ b/python/pyabacus/src/pyabacus/hsolver/_hsolver.py @@ -118,6 +118,7 @@ def davidson( dav_ndim: int = 2, tol: float = 1e-2, max_iter: int = 1000, + diag_ethr: Union[List[float], None] = None, use_paw: bool = False, # scf_type: bool = False ) -> Tuple[NDArray[np.float64], NDArray[np.complex128]]: @@ -143,6 +144,8 @@ def davidson( The tolerance for the convergence, by default 1e-2. max_iter : int, optional The maximum number of iterations, by default 1000. + diag_ethr : List[float] | None, optional + The list of thresholds of bands, by default None. use_paw : bool, optional Whether to use projector augmented wave (PAW) method, by default False. @@ -164,12 +167,16 @@ def davidson( _diago_obj_david.init_eigenvalue() comm_info = diag_comm_info(0, 1) + + if diag_ethr is None: + diag_ethr = [tol] * num_eigs _ = _diago_obj_david.diag( mvv_op, precondition, dav_ndim, tol, + diag_ethr, max_iter, use_paw, comm_info diff --git a/source/module_hsolver/diago_dav_subspace.cpp b/source/module_hsolver/diago_dav_subspace.cpp index d11d7093f1..b180d72c13 100644 --- a/source/module_hsolver/diago_dav_subspace.cpp +++ b/source/module_hsolver/diago_dav_subspace.cpp @@ -88,7 +88,7 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, - const double* ethr_band) + const std::vector& ethr_band) { ModuleBase::timer::tick("Diago_DavSubspace", "diag_once"); @@ -726,7 +726,7 @@ int Diago_DavSubspace::diag(const HPsiFunc& hpsi_func, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, - const double* ethr_band, + const std::vector& ethr_band, const bool& scf_type) { /// record the times of trying iterative diagonalization diff --git a/source/module_hsolver/diago_dav_subspace.h b/source/module_hsolver/diago_dav_subspace.h index d93c252602..622b73638c 100644 --- a/source/module_hsolver/diago_dav_subspace.h +++ b/source/module_hsolver/diago_dav_subspace.h @@ -42,7 +42,7 @@ class Diago_DavSubspace T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, - const double* ethr_band, + const std::vector& ethr_band, const bool& scf_type); private: @@ -135,7 +135,7 @@ class Diago_DavSubspace T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, - const double* ethr_band); + const std::vector& ethr_band); bool test_exit_cond(const int& ntry, const int& notconv, const bool& scf); diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp index 8f404da135..cb9b63c3a4 100644 --- a/source/module_hsolver/diago_david.cpp +++ b/source/module_hsolver/diago_david.cpp @@ -156,7 +156,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, const int ld_psi, T *psi_in, Real* eigenvalue_in, - const Real david_diag_thr, + const std::vector& ethr_band, const int david_maxiter) { if (test_david == 1) @@ -273,7 +273,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, this->notconv = 0; for (int m = 0; m < nband; m++) { - convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < david_diag_thr); + convflag[m] = (std::abs(this->eigenvalue[m] - eigenvalue_in[m]) < ethr_band[m]); if (!convflag[m]) { unconv[this->notconv] = m; @@ -1177,7 +1177,7 @@ int DiagoDavid::diag(const HPsiFunc& hpsi_func, const int ld_psi, T *psi_in, Real* eigenvalue_in, - const Real david_diag_thr, + const std::vector& ethr_band, const int david_maxiter, const int ntry_max, const int notconv_max) @@ -1189,7 +1189,7 @@ int DiagoDavid::diag(const HPsiFunc& hpsi_func, int sum_dav_iter = 0; do { - sum_dav_iter += this->diag_once(hpsi_func, spsi_func, dim, nband, ld_psi, psi_in, eigenvalue_in, david_diag_thr, david_maxiter); + sum_dav_iter += this->diag_once(hpsi_func, spsi_func, dim, nband, ld_psi, psi_in, eigenvalue_in, ethr_band, david_maxiter); ++ntry; } while (!check_block_conv(ntry, this->notconv, ntry_max, notconv_max)); diff --git a/source/module_hsolver/diago_david.h b/source/module_hsolver/diago_david.h index 62ffc2654f..3a13c6b860 100644 --- a/source/module_hsolver/diago_david.h +++ b/source/module_hsolver/diago_david.h @@ -79,7 +79,7 @@ class DiagoDavid const int ld_psi, // Leading dimension of the psi input T *psi_in, // Pointer to eigenvectors Real* eigenvalue_in, // Pointer to store the resulting eigenvalues - const Real david_diag_thr, // Convergence threshold for the Davidson iteration + const std::vector& ethr_band, // Convergence threshold for the Davidson iteration const int david_maxiter, // Maximum allowed iterations for the Davidson method const int ntry_max = 5, // Maximum number of diagonalization attempts (5 by default) const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors @@ -134,7 +134,7 @@ class DiagoDavid const int ld_psi, T *psi_in, Real* eigenvalue_in, - const Real david_diag_thr, + const std::vector& ethr_band, const int david_maxiter); void cal_grad(const HPsiFunc& hpsi_func, diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp index 434fefe073..44f5571116 100644 --- a/source/module_hsolver/hsolver_pw.cpp +++ b/source/module_hsolver/hsolver_pw.cpp @@ -526,7 +526,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, comm_info); DiagoIterAssist::avg_iter += static_cast( - dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band.data(), scf)); + dav_subspace.diag(hpsi_func, psi.get_pointer(), psi.get_nbasis(), eigenvalue, this->ethr_band, scf)); } else if (this->method == "dav") { @@ -589,7 +589,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, ld_psi, psi.get_pointer(), eigenvalue, - david_diag_thr, + this->ethr_band, david_maxiter, ntry_max, notconv_max)); diff --git a/source/module_hsolver/test/diago_david_float_test.cpp b/source/module_hsolver/test/diago_david_float_test.cpp index f6d97f87c9..50b52593a3 100644 --- a/source/module_hsolver/test/diago_david_float_test.cpp +++ b/source/module_hsolver/test/diago_david_float_test.cpp @@ -119,10 +119,12 @@ class DiagoDavPrepare hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out); phm->ops->hPsi(info); }; - auto spsi_func = [phm](const std::complex* psi_in, std::complex* spsi_out,const int ld_psi, const int nbands){ - phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); - }; - dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter); + auto spsi_func = [phm](const std::complex* psi_in, + std::complex* spsi_out, + const int ld_psi, + const int nbands) { phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); }; + std::vector ethr_band(phi.get_nbands(), eps); + dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); diff --git a/source/module_hsolver/test/diago_david_real_test.cpp b/source/module_hsolver/test/diago_david_real_test.cpp index f6597fa78f..0914c322ea 100644 --- a/source/module_hsolver/test/diago_david_real_test.cpp +++ b/source/module_hsolver/test/diago_david_real_test.cpp @@ -118,10 +118,11 @@ class DiagoDavPrepare hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out); phm->ops->hPsi(info); }; - auto spsi_func = [phm](const double* psi_in, double* spsi_out,const int ld_psi, const int nbands){ + auto spsi_func = [phm](const double* psi_in, double* spsi_out, const int ld_psi, const int nbands) { phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); - }; - dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter); + }; + std::vector ethr_band(phi.get_nbands(), eps); + dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); diff --git a/source/module_hsolver/test/diago_david_test.cpp b/source/module_hsolver/test/diago_david_test.cpp index de88aadfe0..fa1a099207 100644 --- a/source/module_hsolver/test/diago_david_test.cpp +++ b/source/module_hsolver/test/diago_david_test.cpp @@ -121,10 +121,12 @@ class DiagoDavPrepare hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out); phm->ops->hPsi(info); }; - auto spsi_func = [phm](const std::complex* psi_in, std::complex* spsi_out,const int ld_psi, const int nbands){ - phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); - }; - dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter); + auto spsi_func = [phm](const std::complex* psi_in, + std::complex* spsi_out, + const int ld_psi, + const int nbands) { phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands); }; + std::vector ethr_band(phi.get_nbands(), eps); + dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); diff --git a/source/module_hsolver/test/hsolver_pw_sup.h b/source/module_hsolver/test/hsolver_pw_sup.h index cea4b6118c..0651fa89b3 100644 --- a/source/module_hsolver/test/hsolver_pw_sup.h +++ b/source/module_hsolver/test/hsolver_pw_sup.h @@ -153,9 +153,9 @@ template int DiagoDavid::diag(const std::function& hpsi_func, const std::function& spsi_func, const int ld_psi, - T *psi_in, + T* psi_in, Real* eigenvalue_in, - const Real david_diag_thr, + const std::vector& ethr_band, const int david_maxiter, const int ntry_max, const int notconv_max) { diff --git a/source/module_lr/hsolver_lrtd.hpp b/source/module_lr/hsolver_lrtd.hpp index 6c8be619eb..77d33a528c 100644 --- a/source/module_lr/hsolver_lrtd.hpp +++ b/source/module_lr/hsolver_lrtd.hpp @@ -82,9 +82,15 @@ namespace LR // converged. const int notconv_max = ("nscf" == PARAM.inp.calculation) ? 0 : 5; // do diag and add davidson iteration counts up to avg_iter - hsolver::DiagoDavid david(precondition.data(), nband, dim, PARAM.inp.pw_diag_ndim, PARAM.inp.use_paw, comm_info); + hsolver::DiagoDavid david(precondition.data(), + nband, + dim, + PARAM.inp.pw_diag_ndim, + PARAM.inp.use_paw, + comm_info); + std::vector ethr_band(nband, diag_ethr); hsolver::DiagoIterAssist::avg_iter += static_cast(david.diag(hpsi_func, spsi_func, - dim, psi, eigenvalue.data(), diag_ethr, maxiter, ntry_max, 0)); + dim, psi, eigenvalue.data(), ethr_band, maxiter, ntry_max, 0)); } else if (method == "dav_subspace") //need refactor { @@ -102,7 +108,7 @@ namespace LR hpsi_func, psi, dim, eigenvalue.data(), - ethr_band.data(), + ethr_band, false /*scf*/)); } else if (method == "cg")