From 0131d7187acdcae5c4a629355fc8ecff339cf9fc Mon Sep 17 00:00:00 2001 From: Critsium-xy Date: Wed, 16 Sep 2026 16:10:09 +0800 Subject: [PATCH 1/3] Refactor: decouple hsolver from hamilt through HSOperator/HSMatrix interfaces The eigensolvers in source_hsolver used to see the Hamiltonian either as std::function callbacks built from hamilt::Hamilt (iterative PW solvers) or as hamilt::Hamilt* directly (HSolverPW/LCAO/LIP, DiagoIterAssist, Parallel_K2D). Both are replaced by two small abstract interfaces that carry only what the math needs: - hsolver::HSOperator: update_k / hpsi / spsi, plus two optional subspace hooks (used by lcao_in_pw EXX). Consumed by DiagoCG, DiagoDavid, Diago_DavSubspace, DiagoBPCG, DiagoIterAssist, HSolverPW and HSolverLIP. - hsolver::HSMatrix: hs_at_k(ik, hk, sk). Consumed by HSolverLCAO and Parallel_K2D (its HskFunc std::function is gone). hamilt::HamiltHSOperator / hamilt::HamiltHSMatrix (source_hamilt/ hamilt_hs_adapter.h) are the only place that wraps raw pointers into Psi/hpsi_info for the operator chain; HamiltLIPHSOperator adds the EXX subspace hooks that HSolverLIP used to reach through a dynamic_cast. LR-TDDFT gets its own LRHSOperator since HamiltLR is not a hamilt::Hamilt. hsolver_pw.h, hsolver_lcao.h, hsolver_lcaopw.h and diago_iter_assist.h no longer include source_hamilt/hamilt.h. HSolverPW_SDFT still takes a Hamilt (it depends on module_stodft) and is left for a follow-up. Tests: the iterative solver tests drive the solvers with an HSOperatorMock over the dense test matrix instead of a HamiltPW/OperatorMock, and no longer link operator.cpp/op_pw.cpp. The lcao_in_pw test previously exercised the "no operators allocated" fallback, which is now a hard error in the adapter; it now checks the subspace rotation with H = S = 1. Co-Authored-By: Claude Fable 5.1 --- source/source_esolver/esolver_ks_lcao.cpp | 5 +- .../source_esolver/esolver_ks_lcao_tddft.cpp | 5 +- source/source_esolver/esolver_ks_lcaopw.cpp | 11 +- source/source_esolver/esolver_ks_pw.cpp | 5 +- source/source_hamilt/hamilt_hs_adapter.h | 84 ++++++ source/source_hsolver/diago_bpcg.cpp | 30 +- source/source_hsolver/diago_bpcg.h | 30 +- source/source_hsolver/diago_cg.cpp | 53 ++-- source/source_hsolver/diago_cg.h | 36 +-- source/source_hsolver/diago_dav_subspace.cpp | 22 +- source/source_hsolver/diago_dav_subspace.h | 14 +- source/source_hsolver/diago_david.cpp | 26 +- source/source_hsolver/diago_david.h | 59 +--- source/source_hsolver/diago_iter_assist.cpp | 248 ++++++---------- source/source_hsolver/diago_iter_assist.h | 88 +++--- source/source_hsolver/hs_matrix.h | 30 ++ source/source_hsolver/hs_operator.h | 54 ++++ source/source_hsolver/hsolver_lcao.cpp | 42 ++- source/source_hsolver/hsolver_lcao.h | 14 +- source/source_hsolver/hsolver_lcaopw.cpp | 47 +-- source/source_hsolver/hsolver_lcaopw.h | 12 +- source/source_hsolver/hsolver_pw.cpp | 51 +--- source/source_hsolver/hsolver_pw.h | 9 +- source/source_hsolver/hsolver_pw_sdft.cpp | 8 +- source/source_hsolver/parallel_k2d.cpp | 4 +- source/source_hsolver/parallel_k2d.h | 13 +- source/source_hsolver/test/CMakeLists.txt | 14 - .../source_hsolver/test/diago_bpcg_test.cpp | 36 +-- .../test/diago_cg_float_test.cpp | 46 +-- .../test/diago_cg_real_test.cpp | 43 +-- source/source_hsolver/test/diago_cg_test.cpp | 46 +-- .../test/diago_david_float_test.cpp | 20 +- .../test/diago_david_real_test.cpp | 19 +- .../source_hsolver/test/diago_david_test.cpp | 20 +- source/source_hsolver/test/diago_mock.h | 269 ++++-------------- source/source_hsolver/test/hsolver_pw_sup.h | 23 +- .../source_hsolver/test/test_hsolver_pw.cpp | 90 ++++-- source/source_lcao/lcao_set.cpp | 4 +- .../module_deltaspin/cal_mw_from_lambda.cpp | 14 +- .../module_deltaspin/deltaspin_pw_mi.cpp | 7 +- source/source_lcao/module_lr/hsolver_lrtd.hpp | 69 ++--- source/source_psi/psi_prepare.cpp | 10 +- source/source_pw/module_pwdft/hamilt_lcaopw.h | 47 +++ 43 files changed, 722 insertions(+), 1055 deletions(-) create mode 100644 source/source_hamilt/hamilt_hs_adapter.h create mode 100644 source/source_hsolver/hs_matrix.h create mode 100644 source/source_hsolver/hs_operator.h diff --git a/source/source_esolver/esolver_ks_lcao.cpp b/source/source_esolver/esolver_ks_lcao.cpp index f28b0c98d98..e6747b42aa1 100644 --- a/source/source_esolver/esolver_ks_lcao.cpp +++ b/source/source_esolver/esolver_ks_lcao.cpp @@ -13,6 +13,7 @@ #include "source_hamilt/module_gint/gint.h" #include "source_estate/elecstate_lcao.h" #include "source_lcao/hamilt_lcao.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/hsolver_lcao.h" #ifdef __EXX #include "../source_lcao/module_ri/exx_opt_orb.h" @@ -480,7 +481,9 @@ void ESolver_KS_LCAO::hamilt2rho_single(UnitCell& ucell, int istep, int this->inp_->device == "gpu", GlobalV::NPROC, GlobalV::MY_RANK); - hsolver_lcao_obj.solve(static_cast*>(this->p_hamilt), this->psi[0], this->pelec, *this->dmat.dm, + // the eigensolvers only ever ask the Hamiltonian for H(k) and S(k) + hamilt::HamiltHSMatrix hs(static_cast*>(this->p_hamilt)); + hsolver_lcao_obj.solve(hs, this->psi[0], this->pelec, *this->dmat.dm, this->chr, this->inp_->nspin, skip_charge); } else diff --git a/source/source_esolver/esolver_ks_lcao_tddft.cpp b/source/source_esolver/esolver_ks_lcao_tddft.cpp index e74951f0e91..48b5d68ac1f 100644 --- a/source/source_esolver/esolver_ks_lcao_tddft.cpp +++ b/source/source_esolver/esolver_ks_lcao_tddft.cpp @@ -19,6 +19,7 @@ #include "source_estate/module_pot/h_tddft_pw.h" #include "source_estate/module_pot/potential_new.h" #include "source_estate/module_pot/td_field_manager.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/hsolver_lcao.h" #include "source_lcao/module_rt/evolve_elec.h" #include "source_lcao/rho_tau_lcao.h" @@ -365,7 +366,9 @@ void ESolver_KS_LCAO_TDDFT::hamilt2rho_single(UnitCell& ucell, const this->inp_->device == "gpu", GlobalV::NPROC, GlobalV::MY_RANK); - hsolver_lcao_obj.solve(static_cast>*>(this->p_hamilt), + hamilt::HamiltHSMatrix> hs( + static_cast>*>(this->p_hamilt)); + hsolver_lcao_obj.solve(hs, this->psi[0], this->pelec, *this->dmat.dm, diff --git a/source/source_esolver/esolver_ks_lcaopw.cpp b/source/source_esolver/esolver_ks_lcaopw.cpp index 49cb7ea6b1b..7afdb64cbcf 100644 --- a/source/source_esolver/esolver_ks_lcaopw.cpp +++ b/source/source_esolver/esolver_ks_lcaopw.cpp @@ -155,7 +155,13 @@ namespace ModuleESolver #else const hsolver::diag_comm_info diag_comm(0, 1); #endif - hsolver_lip_obj.solve(static_cast*>(this->p_hamilt), + // the solver sees the Hamiltonian only through this operator; the EXX + // subspace hooks live in the adapter, not in the solver + hamilt::HamiltLIPHSOperator op(dynamic_cast*>(this->p_hamilt), + this->pw_wfc, + this->general_exx_info_.cal_exx, + this->general_exx_info_.hybrid_alpha); + hsolver_lip_obj.solve(op, *this->stp.template get_psi_t(), this->pelec, *this->psi_local, @@ -163,8 +169,7 @@ namespace ModuleESolver GlobalV::ofs_running, skip_charge, ucell.tpiba, - ucell.nat, - this->general_exx_info_); + ucell.nat); // add exx #ifdef __EXX diff --git a/source/source_esolver/esolver_ks_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index 1c8dc1b798f..3a1bffbb854 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -4,6 +4,7 @@ #include "source_estate/module_charge/symm_rho.h" #include "source_hsolver/diago_iter_assist.h" #include "source_hsolver/diago_params.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/hsolver_pw.h" #include "source_io/module_parameter/parameter.h" #include "source_pw/module_pwdft/force_pw.h" @@ -264,7 +265,9 @@ void ESolver_KS_PW::hamilt2rho_single(UnitCell& ucell, const int iste this->inp_->nb2d, this->inp_->use_k_continuity); - hsolver_pw_obj.solve(static_cast*>(this->p_hamilt), + // the iterative eigensolvers see the Hamiltonian only through this operator + hamilt::HamiltHSOperator op(static_cast*>(this->p_hamilt), this->pw_wfc); + hsolver_pw_obj.solve(op, *this->stp.template get_psi_t(), this->pelec, this->pelec->ekb.c, diff --git a/source/source_hamilt/hamilt_hs_adapter.h b/source/source_hamilt/hamilt_hs_adapter.h new file mode 100644 index 00000000000..918226d6de5 --- /dev/null +++ b/source/source_hamilt/hamilt_hs_adapter.h @@ -0,0 +1,84 @@ +#ifndef HAMILT_HS_ADAPTER_H +#define HAMILT_HS_ADAPTER_H + +#include "source_base/tool_quit.h" +#include "source_basis/module_pw/pw_basis_k.h" +#include "source_hamilt/hamilt.h" +#include "source_hsolver/hs_matrix.h" +#include "source_hsolver/hs_operator.h" +#include "source_psi/psi.h" + +namespace hamilt +{ + +/** + * @brief Presents a Hamilt as the H/S block-vector operator the iterative + * eigensolvers work on (hsolver::HSOperator). + * + * This is the only place that wraps raw pointers into Psi / hpsi_info for the + * operator chain; hsolver itself never sees Hamilt. + */ +template +class HamiltHSOperator : public hsolver::HSOperator +{ + public: + HamiltHSOperator(Hamilt* hm, const ModulePW::PW_Basis_K* wfc_basis) : hm_(hm), wfc_basis_(wfc_basis) + { + } + + void update_k(const int ik) override + { + hm_->updateHk(ik); + ik_ = ik; + npw_ = wfc_basis_->npwk[ik]; + } + + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override + { + if (hm_->ops == nullptr) + { + ModuleBase::WARNING_QUIT("HamiltHSOperator::hpsi", "Operators in Hamilt are not allocated yet"); + } + // non-owning view of x: one k point, nvec bands, leading dimension ld, npw valid rows + psi::Psi x_view(const_cast(x), 1, nvec, ld, npw_); + typename Operator::hpsi_info info(&x_view, psi::Range(true, 0, 0, nvec - 1), hx); + hm_->ops->hPsi(info); + } + + void spsi(const T* x, T* sx, const int ld, const int nvec) const override + { + hm_->sPsi(x, sx, ld, npw_, nvec); + } + + protected: + Hamilt* hm_ = nullptr; + const ModulePW::PW_Basis_K* wfc_basis_ = nullptr; + int ik_ = 0; ///< k point set by the last update_k() + int npw_ = 0; ///< number of plane waves of that k point (without npol) +}; + +/** + * @brief Presents a Hamilt as the H(k)/S(k) matrix source the direct + * eigensolvers work on (hsolver::HSMatrix). + */ +template +class HamiltHSMatrix : public hsolver::HSMatrix +{ + public: + explicit HamiltHSMatrix(Hamilt* hm) : hm_(hm) + { + } + + void hs_at_k(const int ik, ModuleBase::MatrixBlock& hk, ModuleBase::MatrixBlock& sk) override + { + hm_->updateHk(ik); + hm_->matrix(hk, sk); + } + + private: + Hamilt* hm_ = nullptr; +}; + +} // namespace hamilt + +#endif // HAMILT_HS_ADAPTER_H diff --git a/source/source_hsolver/diago_bpcg.cpp b/source/source_hsolver/diago_bpcg.cpp index de4dfa09bfd..363dd5429cb 100644 --- a/source/source_hsolver/diago_bpcg.cpp +++ b/source/source_hsolver/diago_bpcg.cpp @@ -209,21 +209,21 @@ void DiagoBPCG::rotate_wf( template void DiagoBPCG::calc_hpsi_with_block( - const HPsiFunc& hpsi_func, - T *psi_in, + const HSOperator& op, + const T *psi_in, ct::Tensor& hpsi_out) { // calculate all-band hpsi - hpsi_func(psi_in, hpsi_out.data(), this->n_basis, this->n_band_l); + op.hpsi(psi_in, hpsi_out.data(), this->n_basis, this->n_band_l); } template void DiagoBPCG::calc_spsi_with_block( - const SPsiFunc& spsi_func, + const HSOperator& op, const T* psi_in, ct::Tensor& spsi_out) { - spsi_func(psi_in, spsi_out.data(), this->n_basis, this->n_band_l); + op.spsi(psi_in, spsi_out.data(), this->n_basis, this->n_band_l); } template @@ -244,8 +244,7 @@ void DiagoBPCG::diag_hsub( template void DiagoBPCG::calc_hsub_with_block( - const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + const HSOperator& op, T *psi_in, ct::Tensor& psi_out, ct::Tensor& hpsi_out, @@ -255,8 +254,8 @@ void DiagoBPCG::calc_hsub_with_block( ct::Tensor& eigenvalue_out) { // Apply the H operator to psi and obtain the hpsi matrix. - this->calc_hpsi_with_block(hpsi_func, psi_in, hpsi_out); - this->calc_spsi_with_block(spsi_func, psi_in, spsi_out); + this->calc_hpsi_with_block(op, psi_in, hpsi_out); + this->calc_spsi_with_block(op, psi_in, spsi_out); // Transform the generalized problem to an S-orthonormal subspace. this->orth_cholesky(workspace_in, psi_out, hpsi_out, spsi_out, hsub_out); @@ -293,8 +292,7 @@ void DiagoBPCG::calc_hsub_with_block_exit( } template -void DiagoBPCG::diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +void DiagoBPCG::diag(const HSOperator& op, T* psi_in, Real* eigenvalue_in, const std::vector& ethr_band) @@ -307,8 +305,7 @@ void DiagoBPCG::diag(const HPsiFunc& hpsi_func, this->calc_prec(); // Improving the initial guess of the wave function psi through a subspace diagonalization. - this->calc_hsub_with_block(hpsi_func, - spsi_func, + this->calc_hsub_with_block(op, psi_in, this->psi, this->hpsi, @@ -338,7 +335,7 @@ void DiagoBPCG::diag(const HPsiFunc& hpsi_func, this->psi, this->hpsi, this->spsi, this->grad, this->grad_old); // Apply S before projecting the search directions in the generalized metric. - this->calc_spsi_with_block(spsi_func, this->grad.template data(), this->sgrad); + this->calc_spsi_with_block(op, this->grad.template data(), this->sgrad); // Orthogonalize column vectors g_i in matrix grad to column vectors p_j in matrix psi // for all 'j less or equal to i'. @@ -349,7 +346,7 @@ void DiagoBPCG::diag(const HPsiFunc& hpsi_func, syncmem_complex_op()(this->grad_old.template data(), this->grad.template data(), n_basis * n_band_l); // Calculate H|grad> matrix - this->calc_hpsi_with_block(hpsi_func, this->grad.template data(), /*this->grad_wrapper[0],*/ this->hgrad); + this->calc_hpsi_with_block(op, this->grad.template data(), /*this->grad_wrapper[0],*/ this->hgrad); // optimize psi as well as the hpsi // 1. normalize grad @@ -361,8 +358,7 @@ void DiagoBPCG::diag(const HPsiFunc& hpsi_func, this->orth_cholesky(this->work, this->psi, this->hpsi, this->spsi, this->hsub); if (current_scf_iter == 1 && ntry % this->nline == 0) { - this->calc_hsub_with_block(hpsi_func, - spsi_func, + this->calc_hsub_with_block(op, psi_in, this->psi, this->hpsi, diff --git a/source/source_hsolver/diago_bpcg.h b/source/source_hsolver/diago_bpcg.h index caad91e3e83..66489e795c5 100644 --- a/source/source_hsolver/diago_bpcg.h +++ b/source/source_hsolver/diago_bpcg.h @@ -1,12 +1,11 @@ #ifndef DIAGO_BPCG_H_ #define DIAGO_BPCG_H_ -#include - #include "source_base/kernels/math_kernel_op.h" #include "source_base/module_device/memory_op.h" #include "source_base/module_device/types.h" #include "source_base/para_gemm.h" +#include "source_hsolver/hs_operator.h" #include "source_hsolver/kernels/hegvd_op.h" #include "source_hsolver/para_lin_tf.h" @@ -57,23 +56,16 @@ class DiagoBPCG */ void init_iter(const int nband, const int nband_l, const int nbasis, const int ndim); - using HPsiFunc = std::function; - using SPsiFunc = std::function; - /** * @brief Diagonalize the Hamiltonian using the BPCG method. * * This function is called by the HsolverPW::solve() function. * - * @param hpsi_func A function computing the product of the Hamiltonian matrix H - * and a wavefunction blockvector X. - * @param spsi_func A function computing the product of the overlap matrix S - * and a wavefunction blockvector X. + * @param op The H and S block-vector operator, see hsolver::HSOperator. * @param psi_in Pointer to input wavefunction psi matrix with [dim: n_basis x n_band, column major]. * @param eigenvalue_in Pointer to the eigen array with [dim: n_band, column major]. */ - void diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + void diag(const HSOperator& op, T* psi_in, Real* eigenvalue_in, const std::vector& ethr_band); @@ -163,24 +155,23 @@ class DiagoBPCG * psi_in[dim: n_basis x n_band, column major, lda = n_basis_max], * hpsi_out[dim: n_basis x n_band, column major, lda = n_basis_max]. * - * @param hpsi_func A function computing the product of the Hamiltonian matrix H - * and a wavefunction blockvector X. + * @param op The H and S block-vector operator. * @param psi_in The input wavefunction psi. * @param hpsi_out Pointer to the array where the resulting hpsi matrix will be stored. */ void calc_hpsi_with_block( - const HPsiFunc& hpsi_func, - T *psi_in, + const HSOperator& op, + const T *psi_in, ct::Tensor& hpsi_out); /** * @brief Apply the overlap operator to a wavefunction block. * - * @param spsi_func A function computing the product of the overlap matrix S and a wavefunction blockvector X. + * @param op The H and S block-vector operator. * @param psi_in The input wavefunction block. * @param spsi_out The resulting S|psi> block. */ - void calc_spsi_with_block(const SPsiFunc& spsi_func, const T* psi_in, ct::Tensor& spsi_out); + void calc_spsi_with_block(const HSOperator& op, const T* psi_in, ct::Tensor& spsi_out); /** * @brief Diagonalization of the subspace matrix. @@ -258,7 +249,7 @@ class DiagoBPCG * hsub_out[dim: n_band x n_band, column major, lda = n_band], * eigenvalue_out[dim: n_basis_max, column major]. * - * @param hpsi_func A function computing the product of matrix H and wavefunction blockvector X. + * @param op The H and S block-vector operator. * @param psi_in Input wavefunction pointer. * @param psi_out Output wavefunction. * @param hpsi_out Product of psi_out and Hamiltonian. @@ -266,8 +257,7 @@ class DiagoBPCG * @param eigenvalue_out Computed eigen. */ void calc_hsub_with_block( - const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + const HSOperator& op, T *psi_in, ct::Tensor& psi_out, ct::Tensor& hpsi_out, ct::Tensor& spsi_out, ct::Tensor& hsub_out, ct::Tensor& workspace_in, diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index 3b833f60fa8..d02a1051362 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -11,11 +11,17 @@ #include "source_base/tool_title.h" // ModuleBase::TITLE #include "source_base/global_function.h" // ModuleBase::GlobalFunc::NOTE #include "source_hsolver/diago_cg.h" +#include "source_hsolver/diago_iter_assist.h" using namespace hsolver; template DiagoCG::DiagoCG(const std::string& basis_type, const std::string& calculation) +#ifdef __MPI + : diag_comm_(MPI_COMM_SELF, 0, 1) +#else + : diag_comm_(0, 1) +#endif { basis_type_ = basis_type; calculation_ = calculation; @@ -28,18 +34,16 @@ template DiagoCG::DiagoCG(const std::string& basis_type, const std::string& calculation, const bool& need_subspace, - const SubspaceFunc& subspace_func, + const diag_comm_info& diag_comm, const Real& pw_diag_thr, - const int& pw_diag_nmax, - const int& nproc_in_pool) + const int& pw_diag_nmax) + : diag_comm_(diag_comm) { basis_type_ = basis_type; calculation_ = calculation; need_subspace_ = need_subspace; - subspace_func_ = subspace_func; pw_diag_thr_ = pw_diag_thr; pw_diag_nmax_ = pw_diag_nmax; - nproc_in_pool_ = nproc_in_pool; this->one_ = new T(static_cast(1.0)); this->zero_ = new T(static_cast(0.0)); this->neg_one_ = new T(static_cast(-1.0)); @@ -122,10 +126,10 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, { phi_m.sync(psi[m]); // copy psi_in into internal psi, m=0 has been done in Constructor - this->spsi_func_(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> + this->op_->spsi(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> this->schmit_orth(m, psi, sphi, phi_m); - this->spsi_func_(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> - this->hpsi_func_(phi_m.data(), hphi.data(), this->n_basis_, 1); // hphi = H|psi(m)> + this->op_->spsi(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> + this->op_->hpsi(phi_m.data(), hphi.data(), this->n_basis_, 1); // hphi = H|psi(m)> eigen_pack[m] = dot_real_op()(this->n_basis_, phi_m.data(), hphi.data()); @@ -150,8 +154,8 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, g0, cg); // Tensor& - this->hpsi_func_(cg.data(), pphi.data(), this->n_basis_, 1); - this->spsi_func_(cg.data(), scg.data(), this->n_basis_, 1); + this->op_->hpsi(cg.data(), pphi.data(), this->n_basis_, 1); + this->op_->spsi(cg.data(), scg.data(), this->n_basis_, 1); converged = this->update_psi(pphi, cg, @@ -264,7 +268,7 @@ void DiagoCG::orth_grad(const ct::Tensor& psi, ct::Tensor& scg, ct::Tensor& lagrange) { - this->spsi_func_(grad.data(), scg.data(), this->n_basis_, 1); // scg = S|grad> + this->op_->spsi(grad.data(), scg.data(), this->n_basis_, 1); // scg = S|grad> ModuleBase::gemv_op()('C', this->n_basis_, m, @@ -564,6 +568,15 @@ void DiagoCG::schmit_orth(const int& m, const ct::Tensor& psi, const // ModuleBase::timer::end("DiagoCG","schmit_orth"); } +template +void DiagoCG::diag_subspace(const T* psi_in, T* psi_out, const int dim, const int nband, const bool S_orth) +{ + // subspace diagonalization of the current nband vectors, packed with leading dimension dim; + // the eigenvalues it produces are not needed, CG recomputes them + std::vector eigen(nband, 0.0); + DiagoIterAssist::diag_subspace(*op_, psi_in, psi_out, nband, nband, dim, dim, eigen.data(), diag_comm_, S_orth); +} + template bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) const { @@ -579,8 +592,7 @@ bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) con } template -double DiagoCG::diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +double DiagoCG::diag(const HSOperator& op, const int ld_psi, const int nband, const int dim, @@ -615,8 +627,7 @@ double DiagoCG::diag(const HPsiFunc& hpsi_func, /// record the times of trying iterative diagonalization int ntry = 0; this->notconv_ = 0; - hpsi_func_ = hpsi_func; - spsi_func_ = spsi_func; + op_ = &op; // create a new slice of psi to do cg diagonalization ct::Tensor psi_temp = psi.slice({0, 0}, {nband, dim}); @@ -630,22 +641,14 @@ double DiagoCG::diag(const HPsiFunc& hpsi_func, { ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp); const bool assume_S_orthogonal = true; - this->subspace_func_(psi_temp.data(), - psi_map.data(), - dim, - nband, - assume_S_orthogonal); + this->diag_subspace(psi_temp.data(), psi_map.data(), dim, nband, assume_S_orthogonal); psi_temp.sync(psi_map); } else if (need_subspace_) { ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp); const bool assume_S_orthogonal = false; - this->subspace_func_(psi_temp.data(), - psi_map.data(), - dim, - nband, - assume_S_orthogonal); + this->diag_subspace(psi_temp.data(), psi_map.data(), dim, nband, assume_S_orthogonal); psi_temp.sync(psi_map); } diff --git a/source/source_hsolver/diago_cg.h b/source/source_hsolver/diago_cg.h index 99d9369a0a3..c5109cfb0cf 100644 --- a/source/source_hsolver/diago_cg.h +++ b/source/source_hsolver/diago_cg.h @@ -1,12 +1,14 @@ #ifndef MODULE_HSOLVER_DIAGO_CG_H_ #define MODULE_HSOLVER_DIAGO_CG_H_ -#include #include #include #include +#include "source_hsolver/diag_comm_info.h" +#include "source_hsolver/hs_operator.h" + #include #include @@ -22,30 +24,23 @@ class DiagoCG final using Real = typename GetTypeReal::type; using ct_Device = typename ct::PsiToContainer::type; public: - using HPsiFunc = std::function; - using SPsiFunc = std::function; - using SubspaceFunc = std::function; // Constructor need: - // 1. temporary mock of Hamiltonian "Hamilt_PW" - // 2. precondition pointer should point to place of precondition array. + // 1. basis type and calculation type of ABACUS + // 2. diag_comm: the pool over which the subspace matrices are reduced DiagoCG(const std::string& basis_type, const std::string& calculation); DiagoCG( const std::string& basis_type, const std::string& calculation, const bool& need_subspace, - const SubspaceFunc& subspace_func, + const diag_comm_info& diag_comm, const Real& pw_diag_thr, - const int& pw_diag_nmax, - const int& nproc_in_pool); + const int& pw_diag_nmax); ~DiagoCG(); - // virtual void init(){}; - // refactor hpsi_info // this is the diag() function for CG method // returns avg_iter - double diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + double diag(const HSOperator& op, const int ld_psi, const int nband, const int dim, @@ -72,20 +67,16 @@ class DiagoCG final Real pw_diag_thr_ = 1e-5; /// maximum iteration steps for cg diagonalization int pw_diag_nmax_ = 0; - /// number of processors in a node - int nproc_in_pool_ = 0; + /// communicator of the pool sharing the plane waves + const diag_comm_info diag_comm_; /// basis_type of psi std::string basis_type_ = {}; /// calculation type of ABACUS std::string calculation_ = {}; bool need_subspace_ = false; - /// A function object that performs the hPsi calculation. - HPsiFunc hpsi_func_ = nullptr; - /// A function object that performs the sPsi calculation. - SPsiFunc spsi_func_ = nullptr; - /// A function object that performs the subspace calculation. - SubspaceFunc subspace_func_ = nullptr; + /// The H and S operator being diagonalized, set for the duration of diag(). + const HSOperator* op_ = nullptr; void calc_grad( const ct::Tensor& prec, @@ -135,6 +126,9 @@ class DiagoCG final bool test_exit_cond(const int& ntry, const int& notconv) const; + /// subspace rotation of the current nband vectors (packed, leading dimension dim) + void diag_subspace(const T* psi_in, T* psi_out, const int dim, const int nband, const bool S_orth); + using dot_real_op = ModuleBase::dot_real_op; const T * one_ = nullptr, * zero_ = nullptr, * neg_one_ = nullptr; }; diff --git a/source/source_hsolver/diago_dav_subspace.cpp b/source/source_hsolver/diago_dav_subspace.cpp index 2dfbcf73de4..331eb9857e4 100644 --- a/source/source_hsolver/diago_dav_subspace.cpp +++ b/source/source_hsolver/diago_dav_subspace.cpp @@ -108,8 +108,7 @@ Diago_DavSubspace::~Diago_DavSubspace() } template -int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, +int Diago_DavSubspace::diag_once(const HSOperator& op, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, @@ -143,11 +142,11 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, // compute h*psi_in_iter // NOTE: bands after the first n_band should yield zero // hpsi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x] - hpsi_func(this->psi_in_iter, this->hpsi, this->dim, this->notconv); + op.hpsi(this->psi_in_iter, this->hpsi, this->dim, this->notconv); // compute s*psi_in_iter // spsi[:, 0:nbase_x] = S * psi_in_iter[:, 0:nbase_x] - spsi_func(this->psi_in_iter, this->spsi, this->dim, this->notconv); + op.spsi(this->psi_in_iter, this->spsi, this->dim, this->notconv); this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->spsi, this->hpsi, this->hcc, this->scc); @@ -166,8 +165,7 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, { dav_iter++; - this->cal_grad(hpsi_func, - spsi_func, + this->cal_grad(op, this->dim, nbase, this->notconv, @@ -274,8 +272,7 @@ int Diago_DavSubspace::diag_once(const HPsiFunc& hpsi_func, } template -void Diago_DavSubspace::cal_grad(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, +void Diago_DavSubspace::cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int& notconv, @@ -479,8 +476,8 @@ void Diago_DavSubspace::cal_grad(const HPsiFunc& hpsi_func, // update hpsi[:, nbase:nbase+notconv] // hpsi[:, nbase:nbase+notconv] = H * psi_iter[:, nbase:nbase+notconv] - hpsi_func(psi_iter + nbase * dim, hpsi + nbase * this->dim, this->dim, notconv); - spsi_func(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, notconv); + op.hpsi(psi_iter + nbase * dim, hpsi + nbase * this->dim, this->dim, notconv); + op.spsi(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, notconv); ModuleBase::timer::end("Diago_DavSubspace", "cal_grad"); return; @@ -815,8 +812,7 @@ void Diago_DavSubspace::refresh(const int& dim, } template -int Diago_DavSubspace::diag(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, +int Diago_DavSubspace::diag(const HSOperator& op, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in_hsolver, @@ -832,7 +828,7 @@ int Diago_DavSubspace::diag(const HPsiFunc& hpsi_func, do { - sum_iter += this->diag_once(hpsi_func, spsi_func, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band); + sum_iter += this->diag_once(op, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band); ++ntry; diff --git a/source/source_hsolver/diago_dav_subspace.h b/source/source_hsolver/diago_dav_subspace.h index 43e0304729c..4e1e97fbbbc 100644 --- a/source/source_hsolver/diago_dav_subspace.h +++ b/source/source_hsolver/diago_dav_subspace.h @@ -8,9 +8,9 @@ #include "source_base/module_container/ATen/kernels/lapack.h" #include "source_hsolver/diag_comm_info.h" +#include "source_hsolver/hs_operator.h" #include -#include namespace hsolver { @@ -37,11 +37,7 @@ class Diago_DavSubspace ~Diago_DavSubspace(); - // See diago_david.h for information on the HPsiFunc function type - using HPsiFunc = std::function; - - int diag(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, + int diag(const HSOperator& op, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, @@ -99,8 +95,7 @@ class Diago_DavSubspace base_device::DEVICE_CPU* cpu_ctx = {}; base_device::AbacusDevice_t device = {}; - void cal_grad(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, + void cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int& notconv, @@ -147,8 +142,7 @@ class Diago_DavSubspace std::vector* eigenvalue_iter, T* vcc); - int diag_once(const HPsiFunc& hpsi_func, - const HPsiFunc& spsi_func, + int diag_once(const HSOperator& op, T* psi_in, const int psi_in_dmax, Real* eigenvalue_in, diff --git a/source/source_hsolver/diago_david.cpp b/source/source_hsolver/diago_david.cpp index b7667692dc4..1455d4138f4 100644 --- a/source/source_hsolver/diago_david.cpp +++ b/source/source_hsolver/diago_david.cpp @@ -115,8 +115,7 @@ DiagoDavid::~DiagoDavid() } template -int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +int DiagoDavid::diag_once(const HSOperator& op, const int dim, const int nband, const int ld_psi, @@ -157,7 +156,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, { { // phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1); - spsi_func(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1); + op.spsi(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1); } } // begin SchmidtOrth @@ -174,7 +173,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, pre_matrix_mv_m[m]); { // phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1); - spsi_func(basis + dim*m, &this->spsi[m * dim], dim, 1); + op.spsi(basis + dim*m, &this->spsi[m * dim], dim, 1); } } @@ -183,7 +182,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, // phm_in->ops->hPsi(dav_hpsi_in); // hpsi[:, 0:nband] = H basis[:, 0:nband] // slice index in this piece of code is in C manner. i.e. 0:id stands for [0,id) - hpsi_func(basis, hpsi, dim, nband); + op.hpsi(basis, hpsi, dim, nband); this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc); @@ -201,8 +200,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, { dav_iter++; - this->cal_grad(hpsi_func, - spsi_func, + this->cal_grad(op, dim, nbase, nbase_x, @@ -292,8 +290,7 @@ int DiagoDavid::diag_once(const HPsiFunc& hpsi_func, template -void DiagoDavid::cal_grad(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +void DiagoDavid::cal_grad(const HSOperator& op, const int& dim, const int& nbase, // current dimension of the reduced basis const int nbase_x, // maximum dimension of the reduced basis set @@ -509,7 +506,7 @@ void DiagoDavid::cal_grad(const HPsiFunc& hpsi_func, { { // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); } } // first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix @@ -558,7 +555,7 @@ void DiagoDavid::cal_grad(const HPsiFunc& hpsi_func, pre_matrix_mv_m[m]); { // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); } } // calculate H|psi> for not convergence bands @@ -567,7 +564,7 @@ void DiagoDavid::cal_grad(const HPsiFunc& hpsi_func, // &hpsi[nbase * dim]); // &hp(nbase, 0) // phm_in->ops->hPsi(dav_hpsi_in); // hpsi[:, nbase:nbase+notcnv] = H basis[:, nbase:nbase+notcnv] - hpsi_func(basis + nbase * dim, hpsi + nbase * dim, dim, notconv); + op.hpsi(basis + nbase * dim, hpsi + nbase * dim, dim, notconv); delmem_complex_op()(lagrange); delmem_complex_op()(vc_ev_vector); @@ -1006,8 +1003,7 @@ void DiagoDavid::planSchmidtOrth(const int nband, std::vector& p template -int DiagoDavid::diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +int DiagoDavid::diag(const HSOperator& op, const int ld_psi, T *psi_in, Real* eigenvalue_in, @@ -1023,7 +1019,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, ethr_band, david_maxiter); + sum_dav_iter += this->diag_once(op, 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/source_hsolver/diago_david.h b/source/source_hsolver/diago_david.h index e9ee3a50fde..5dd29805378 100644 --- a/source/source_hsolver/diago_david.h +++ b/source/source_hsolver/diago_david.h @@ -8,10 +8,10 @@ #include "source_base/module_container/ATen/kernels/lapack.h" // container::kernels #include "source_hsolver/diag_comm_info.h" +#include "source_hsolver/hs_operator.h" #include "source_hsolver/kernels/hegvd_op.h" #include -#include namespace hsolver { @@ -68,54 +68,12 @@ class DiagoDavid ~DiagoDavid(); - // declare type of matrix-blockvector functions. - // the function type is defined as a std::function object. - /** - * @brief A function type representing the HX function. - * - * This function type is used to define a matrix-blockvector operator H. - * For eigenvalue problem HX = λX or generalized eigenvalue problem HX = λSX, - * this function computes the product of the Hamiltonian matrix H and a blockvector X. - * - * Called as follows: - * hpsi(X, HX, ld, nvec) where X and HX are (ld, nvec)-shaped blockvectors. - * Result HX = H * X is stored in HX. - * - * @param[out] X Head address of input blockvector of type `T*`. - * @param[in] HX Head address of output blockvector of type `T*`. - * @param[in] ld Leading dimension of blockvector. - * @param[in] nvec Number of vectors in a block. - * - * @warning X and HX are the exact address to read input X and store output H*X, - * @warning both of size ld * nvec. - */ - using HPsiFunc = std::function; - - /** - * @brief A function type representing the SX function. - * - * nrow is leading dimension of spsi, npw is leading dimension of psi, nbands is number of vecs - * - * This function type is used to define a matrix-blockvector operator S. - * For generalized eigenvalue problem HX = λSX, - * this function computes the product of the overlap matrix S and a blockvector X. - * - * @param[in] X Pointer to the input blockvector. - * @param[out] SX Pointer to the output blockvector. - * @param[in] ld_psi Leading dimension of psi and spsi. Dimension of X&SX: ld * nvec. - * @param[in] nvec Number of vectors. - */ - using SPsiFunc = std::function; - /** * @brief Performs iterative diagonalization using the David algorithm. * - * @warning Please see docs of `HPsiFunc` for more information about the hpsi mat-vec interface. - * * @tparam T The type of the elements in the matrix. * @tparam Device The device type (CPU or GPU). - * @param hpsi_func The function object that computes the matrix-blockvector product H * psi. - * @param spsi_func The function object that computes the matrix-blockvector product overlap S * psi. + * @param op The H and S block-vector operator, see hsolver::HSOperator. * @param ld_psi The leading dimension of the psi_in array. * @param psi_in The input wavefunction. * @param eigenvalue_in The array to store the eigenvalues. @@ -129,8 +87,7 @@ class DiagoDavid * notconv_max is determined by the accuracy required for the calculation, default 0 */ int diag( - const HPsiFunc& hpsi_func, // function void hpsi(T*, T*, const int, const int) - const SPsiFunc& spsi_func, // function void spsi(T*, T*, const int, const int, const int) + const HSOperator& op, // applies H and S to block vectors 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 @@ -179,8 +136,7 @@ class DiagoDavid base_device::DEVICE_CPU* cpu_ctx = {}; base_device::AbacusDevice_t device = {}; - int diag_once(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + int diag_once(const HSOperator& op, const int dim, const int nband, const int ld_psi, @@ -192,8 +148,7 @@ class DiagoDavid /** * Calculates the preconditioned gradient of the eigenvectors in Davidson method. * - * @param hpsi_func The function to calculate the matrix-blockvector product H * psi. - * @param spsi_func The function to calculate the matrix-blockvector product overlap S * psi. + * @param op The H and S block-vector operator. * @param dim The dimension of the blockvector. * @param nbase The current dimension of the reduced basis. * @param nbase_x The maximum dimension of the reduced basis set. @@ -204,8 +159,7 @@ class DiagoDavid * @param unconv The array of indices for the unconverged eigenpairs. * @param eigenvalue The array of eigenvalues. */ - void cal_grad(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, + void cal_grad(const HSOperator& op, const int& dim, const int& nbase, const int nbase_x, @@ -340,7 +294,6 @@ class DiagoDavid // Note that ct_Device is different from base_device! using ct_Device = typename ct::PsiToContainer::type; - // using hpsi_info = typename hamilt::Operator::hpsi_info; // Dependence of hpsi removed const T *one = nullptr, *zero = nullptr, *neg_one = nullptr; const T one_ = static_cast(1.0), zero_ = static_cast(0.0), neg_one_ = static_cast(-1.0); diff --git a/source/source_hsolver/diago_iter_assist.cpp b/source/source_hsolver/diago_iter_assist.cpp index efbfe4c9695..12e8fce057f 100644 --- a/source/source_hsolver/diago_iter_assist.cpp +++ b/source/source_hsolver/diago_iter_assist.cpp @@ -1,6 +1,5 @@ #include "diago_iter_assist.h" -#include "source_base/complexmatrix.h" #include "source_base/constants.h" #include "source_base/global_function.h" #include "source_base/kernels/math_kernel_op.h" @@ -10,38 +9,31 @@ #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/kernels/hegvd_op.h" +#include + namespace hsolver { //---------------------------------------------------------------------- // Hamiltonian diagonalization in the subspace spanned -// by nstart states psi (atomic or random wavefunctions). +// by nstart vectors psi (atomic or random wavefunctions). // Produces on output n_band eigenvectors (n_band <= nstart) in evc. //---------------------------------------------------------------------- template -void DiagoIterAssist::diag_subspace( - const hamilt::Hamilt* const pHamilt, // hamiltonian operator carrier - const psi::Psi& psi, // [in] wavefunction - psi::Psi& evc, // [out] wavefunction, eigenvectors - Real* en, // [out] eigenvalues - const diag_comm_info& diag_comm, - int n_band, // [in] number of bands to be calculated, also number of rows - // of evc, if set to 0, n_band = nstart, default 0 - const bool S_orth // [in] if true, psi is assumed to be already S-orthogonalized -) +void DiagoIterAssist::diag_subspace(const HSOperator& op, + const T* psi, + T* evc, + const int nstart, + const int n_band, + const int dmin, + const int dmax, + Real* en, + const diag_comm_info& diag_comm, + const bool S_orth) { ModuleBase::TITLE("DiagoIterAssist", "diag_subspace"); ModuleBase::timer::start("DiagoIterAssist", "diag_subspace"); - // two case: - // 1. pw base: nstart = n_band, psi(nbands * npwx) - // 2. lcao_in_pw base: nstart >= n_band, psi(NLOCAL * npwx) - const int nstart = psi.get_nbands(); - // n_band = 0 means default, set n_band = nstart - if (n_band == 0) - { - n_band = nstart; - } assert(n_band <= nstart); // scc is overlap (optional, only needed if input is not s-orthogonal) @@ -61,31 +53,15 @@ void DiagoIterAssist::diag_subspace( resmem_complex_op()(vcc, nstart * nstart, "DiagSub::vcc"); setmem_complex_op()(vcc, 0, nstart * nstart); - // dmin is the active number of plane waves or atomic orbitals - // dmax is the leading dimension of psi - const int dmin = psi.get_current_ngk(); - const int dmax = psi.get_nbasis(); - - T *temp = nullptr; /// temporary array for calculation of evc - bool in_place = false; ///< if temp and evc share the same memory - if (psi.get_pointer() != evc.get_pointer() && psi.get_nbands() == evc.get_nbands()) - { // use memory of evc as temp - temp = evc.get_pointer(); - in_place = true; - } - else - { - resmem_complex_op()(temp, nstart * dmax, "DiagSub::temp"); - } + // temp holds H|psi>, then S|psi>, then the rotated vectors; it is separate + // from evc so that evc may alias psi + T* temp = nullptr; + resmem_complex_op()(temp, nstart * dmax, "DiagSub::temp"); + setmem_complex_op()(temp, 0, nstart * dmax); { // code block to calculate hcc and scc - setmem_complex_op()(temp, 0, nstart * dmax); - T *hpsi = temp; - // do hPsi for all bands - psi::Range all_bands_range(1, psi.get_current_k(), 0, nstart - 1); - hpsi_info hpsi_in(&psi, all_bands_range, hpsi); - pHamilt->ops->hPsi(hpsi_in); + op.hpsi(psi, hpsi, dmax, nstart); ModuleBase::gemm_op()('C', 'N', @@ -93,7 +69,7 @@ void DiagoIterAssist::diag_subspace( nstart, dmin, &one, - psi.get_pointer(), + psi, dmax, hpsi, dmax, @@ -104,8 +80,7 @@ void DiagoIterAssist::diag_subspace( if(!S_orth){ // Only calculate S_sub if not orthogonal T *spsi = temp; - // do sPsi for all bands - pHamilt->sPsi(psi.get_pointer(), spsi, dmax, dmin, nstart); + op.spsi(psi, spsi, dmax, nstart); ModuleBase::gemm_op()('C', 'N', @@ -113,7 +88,7 @@ void DiagoIterAssist::diag_subspace( nstart, dmin, &one, - psi.get_pointer(), + psi, dmax, spsi, dmax, @@ -142,8 +117,6 @@ void DiagoIterAssist::diag_subspace( DiagoIterAssist::diag_hegvd(nstart, n_band, hcc, scc, nstart, en, vcc); } - const int ld_temp = in_place ? dmax : dmin; - { // code block to calculate evc ModuleBase::gemm_op()('N', 'N', @@ -151,20 +124,18 @@ void DiagoIterAssist::diag_subspace( n_band, nstart, &one, - psi.get_pointer(), // dmin * nstart + psi, // dmin * nstart dmax, vcc, // nstart * n_band nstart, &zero, temp, - ld_temp); + dmin); } - if (!in_place) - { - ModuleBase::matrixCopy()(n_band, ld_temp, temp, ld_temp, evc.get_pointer(), dmax); - delmem_complex_op()(temp); - } + ModuleBase::matrixCopy()(n_band, dmin, temp, dmin, evc, dmax); + + delmem_complex_op()(temp); delmem_complex_op()(hcc); if(!S_orth){ delmem_complex_op()(scc); @@ -175,18 +146,50 @@ void DiagoIterAssist::diag_subspace( } template -void DiagoIterAssist::diag_subspace_init( - hamilt::Hamilt* pHamilt, - const T* psi, - int psi_nr, - int psi_nc, - psi::Psi& evc, - Real* en, - const std::string& basis_type, - const std::string& calculation, - const diag_comm_info& diag_comm, - const std::function& add_to_hcc, - const std::function& export_vcc) +void DiagoIterAssist::diag_subspace(const HSOperator& op, + const psi::Psi& psi, // [in] wavefunction + psi::Psi& evc, // [out] wavefunction, eigenvectors + Real* en, // [out] eigenvalues + const diag_comm_info& diag_comm, + int n_band, // [in] number of bands to be calculated, also number of rows + // of evc, if set to 0, n_band = nstart, default 0 + const bool S_orth // [in] if true, psi is assumed to be already S-orthogonalized +) +{ + // two case: + // 1. pw base: nstart = n_band, psi(nbands * npwx) + // 2. lcao_in_pw base: nstart >= n_band, psi(NLOCAL * npwx) + const int nstart = psi.get_nbands(); + // n_band = 0 means default, set n_band = nstart + if (n_band == 0) + { + n_band = nstart; + } + + // dmin is the active number of plane waves or atomic orbitals + // dmax is the leading dimension of psi + diag_subspace(op, + psi.get_pointer(), + evc.get_pointer(), + nstart, + n_band, + psi.get_current_ngk(), + psi.get_nbasis(), + en, + diag_comm, + S_orth); +} + +template +void DiagoIterAssist::diag_subspace_init(const HSOperator& op, + const T* psi, + int psi_nr, + int psi_nc, + psi::Psi& evc, + Real* en, + const std::string& basis_type, + const std::string& calculation, + const diag_comm_info& diag_comm) { ModuleBase::TITLE("DiagoIterAssist", "diag_subspace_init"); ModuleBase::timer::start("DiagoIterAssist", "diag_subspace_init"); @@ -200,27 +203,6 @@ void DiagoIterAssist::diag_subspace_init( const int dmax = evc.get_nbasis(); const int dmin = evc.get_current_ngk(); - // skip the diagonalization if the operators are not allocated - if (pHamilt->ops == nullptr) - { - ModuleBase::WARNING( - "DiagoIterAssist::diag_subspace_init", - "Severe warning: Operators in Hamilt are not allocated yet, will return value of psi to evc directly\n"); - for (int iband = 0; iband < n_band; iband++) - { - for (int ig = 0; ig < dmax; ig++) - { - evc(iband, ig) = psi[iband * dmax + ig]; - } - en[iband] = 0.0; - } - ModuleBase::timer::end("DiagoIterAssist", "diag_subspace_init"); - return; - } - - // ModuleBase::ComplexMatrix hc(nstart, nstart); - // ModuleBase::ComplexMatrix sc(nstart, nstart); - // ModuleBase::ComplexMatrix hvec(nstart, n_band); T *hcc = nullptr, *scc = nullptr, *vcc = nullptr; resmem_complex_op()(hcc, nstart * nstart, "DiagSub::hcc"); resmem_complex_op()(scc, nstart * nstart, "DiagSub::scc"); @@ -231,38 +213,25 @@ void DiagoIterAssist::diag_subspace_init( if (base_device::get_device_type(ctx) == base_device::GpuDevice) { - /// NOTE: current_nbasis must be npw (WITHOUT npol) for Nonlocal::act's - /// gemm K (vkb only has npw rows). See CPU branch comment above. - psi::Psi psi_temp(1, 1, psi_nc, evc.get_current_nbas(), true); - - T* ppsi = psi_temp.get_pointer(); - // hpsi and spsi share the temp space + // band by band on the GPU: the scratch buffer holds one vector only T* temp = nullptr; resmem_complex_op()(temp, psi_nc, "DiagSub::temp"); setmem_complex_op()(temp, 0, psi_nc); T* hpsi = temp; - // do hPsi band by band for (int i = 0; i < nstart; i++) { - // psi_temp is one band psi, psi is all bands psi, the range always is 1 for the only band in psi_temp - syncmem_complex_op()(ppsi, psi + i * psi_nc, psi_nc); - psi::Range band_by_band_range(true, 0, 0, 0); - hpsi_info hpsi_in(&psi_temp, band_by_band_range, hpsi); - // H|Psi> to get hpsi for target band - pHamilt->ops->hPsi(hpsi_in); + op.hpsi(psi + i * psi_nc, hpsi, psi_nc, 1); // calculate the related elements in hcc ModuleBase::gemv_op()('C', psi_nc, nstart, &one, psi, psi_nc, hpsi, 1, &zero, hcc + i * nstart, 1); } T* spsi = temp; - // do sPsi band by band for (int i = 0; i < nstart; i++) { - syncmem_complex_op()(ppsi, psi + i * psi_nc, psi_nc); - pHamilt->sPsi(ppsi, spsi, dmin, dmin, 1); + op.spsi(psi + i * psi_nc, spsi, psi_nc, 1); ModuleBase::gemv_op()('C', psi_nc, @@ -280,37 +249,26 @@ void DiagoIterAssist::diag_subspace_init( } else if (base_device::get_device_type(ctx) == base_device::CpuDevice) { - /// NOTE: the 4th arg (current_nbasis) must be npw (WITHOUT npol), - /// NOT dmin (= nbasis = npol*npwx in SOC). Nonlocal::act uses - /// psi_temp.get_current_nbas() as gemm K, but vkb only has npw rows. - /// dmin (still = nbasis) is kept for hcc/scc gemm K which needs npol. - psi::Psi psi_temp(1, nstart, psi_nc, evc.get_current_nbas(), true); - - T* ppsi = psi_temp.get_pointer(); - syncmem_complex_op()(ppsi, psi, psi_temp.size()); // hpsi and spsi share the temp space T* temp = nullptr; resmem_complex_op()(temp, nstart * psi_nc, "DiagSub::temp"); setmem_complex_op()(temp, 0, nstart * psi_nc); T* hpsi = temp; - // do hPsi for all bands - psi::Range all_bands_range(true, 0, 0, nstart - 1); - hpsi_info hpsi_in(&psi_temp, all_bands_range, hpsi); - pHamilt->ops->hPsi(hpsi_in); + op.hpsi(psi, hpsi, psi_nc, nstart); - ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, ppsi, dmax, hpsi, dmax, &zero, hcc, nstart); + ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, psi, psi_nc, hpsi, psi_nc, &zero, hcc, nstart); T* spsi = temp; - // do sPsi for all bands - pHamilt->sPsi(ppsi, spsi, psi_temp.get_nbasis(), psi_temp.get_nbasis(), psi_temp.get_nbands()); + op.spsi(psi, spsi, psi_nc, nstart); - ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, ppsi, dmax, spsi, dmax, &zero, scc, nstart); + ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, psi, psi_nc, spsi, psi_nc, &zero, scc, nstart); delmem_complex_op()(temp); - - add_to_hcc(hcc, nstart); } + // a Hamiltonian may carry a term hpsi() does not cover (EXX in lcao_in_pw) + op.add_to_subspace_h(hcc, nstart); + if (diag_comm.nproc > 1) { #ifdef __MPI @@ -320,22 +278,9 @@ void DiagoIterAssist::diag_subspace_init( } // after generation of H and S matrix, diag them - /// this part only for test, eigenvector would have different phase caused by micro numerical perturbation - /// set 8 bit effective accuracy would help for debugging - /*for(int i=0;i(double(int(hc(i,j).real()*100000000))/100000000, 0); - if(std::norm(sc(i,j))<1e-10) sc(i,j) = ModuleBase::ZERO; - else sc(i,j) = std::complex(double(int(sc(i,j).real()*100000000))/100000000, 0); - } - }*/ - DiagoIterAssist::diag_hegvd(nstart, n_band, hcc, scc, nstart, en, vcc); - export_vcc(vcc, nstart, n_band); + op.export_subspace_vec(vcc, nstart, n_band); //======================= // diagonize the H-matrix @@ -369,10 +314,6 @@ void DiagoIterAssist::diag_subspace_init( { assert(psi != evc.get_pointer()); - // T* evctemp = nullptr; - // resmem_complex_op()(ctx, evctemp, n_band * dmin, "DiagSub::evctemp"); - // setmem_complex_op()(ctx, evctemp, 0, n_band * dmin); - ModuleBase::gemm_op()('N', 'N', dmin, @@ -386,10 +327,6 @@ void DiagoIterAssist::diag_subspace_init( &zero, evc.get_pointer(), dmax); - - // matrixCopy()(ctx, n_band, evctemp, dmin, evc.get_pointer(), dmax); - - // delmem_complex_op()(ctx, evctemp); } delmem_complex_op()(hcc); @@ -487,12 +424,11 @@ void DiagoIterAssist::diag_hegvd(const int nstart, } template -void DiagoIterAssist::cal_hs_subspace( - const hamilt::Hamilt* pHamilt, // hamiltonian operator carrier - const psi::Psi& psi, // [in] wavefunction - T* hcc, - T* scc, - const diag_comm_info& diag_comm) +void DiagoIterAssist::cal_hs_subspace(const HSOperator& op, + const psi::Psi& psi, // [in] wavefunction + T* hcc, + T* scc, + const diag_comm_info& diag_comm) { const int nstart = psi.get_nbands(); @@ -507,13 +443,8 @@ void DiagoIterAssist::cal_hs_subspace( setmem_complex_op()(temp, 0, nstart * dmax); { // code block to calculate hcc and scc - setmem_complex_op()(temp, 0, nstart * dmax); - T* hpsi = temp; - // do hPsi for all bands - psi::Range all_bands_range(1, psi.get_current_k(), 0, nstart - 1); - hpsi_info hpsi_in(&psi, all_bands_range, hpsi); - pHamilt->ops->hPsi(hpsi_in); + op.hpsi(psi.get_pointer(), hpsi, dmax, nstart); ModuleBase::gemm_op()('C', 'N', @@ -530,8 +461,7 @@ void DiagoIterAssist::cal_hs_subspace( nstart); T* spsi = temp; - // do sPsi for all bands - pHamilt->sPsi(psi.get_pointer(), spsi, dmax, dmin, nstart); + op.spsi(psi.get_pointer(), spsi, dmax, nstart); ModuleBase::gemm_op()('C', 'N', diff --git a/source/source_hsolver/diago_iter_assist.h b/source/source_hsolver/diago_iter_assist.h index 225e8b89277..ef89de59bee 100644 --- a/source/source_hsolver/diago_iter_assist.h +++ b/source/source_hsolver/diago_iter_assist.h @@ -1,12 +1,11 @@ #ifndef DIAGOITERASSIST_H #define DIAGOITERASSIST_H -#include "source_base/complexmatrix.h" #include "source_base/macros.h" -#include "source_hamilt/hamilt.h" +#include "source_base/module_device/memory_op.h" +#include "source_hsolver/hs_operator.h" #include "source_psi/psi.h" -#include #include namespace hsolver @@ -30,24 +29,38 @@ class DiagoIterAssist static int SCF_ITER; - // for psi::Psi structure - /** - * @brief Diagonalizes the Hamiltonian in a subspace defined by the given wavefunction. + /** + * @brief Diagonalizes H in the subspace spanned by nstart vectors. * - * This static function computes the eigenvalues and eigenvectors of the Hamiltonian - * within the subspace spanned by the provided wavefunction `psi`. The resulting eigenvectors - * are stored in `evc`, and the corresponding eigenvalues are written to `en`. + * Builds the nstart*nstart matrices (and unless the + * input is S-orthogonal), solves the small eigenproblem and rotates psi + * into the lowest n_band eigenvectors, written to evc. * - * @tparam T Data type for computation (e.g., float, double). - * @tparam Device Device type for computation (e.g., CPU, GPU). - * @param pHamilt Pointer to the Hamiltonian object. - * @param psi Input wavefunction defining the subspace. - * @param evc Output container for computed eigenvectors. - * @param en Output array for computed eigenvalues. - * @param n_band Number of bands (eigenvalues/eigenvectors) to compute. Default is 0 (all). - * @param is_S_orthogonal If true, assumes the input wavefunction is already orthogonalized. + * @param op applies H and S to block vectors + * @param psi [in] nstart vectors, leading dimension dmax + * @param evc [out] n_band vectors, leading dimension dmax; may alias psi + * @param nstart number of input vectors + * @param n_band number of eigenvectors wanted (<= nstart) + * @param dmin active length of each vector + * @param dmax leading dimension of psi and evc + * @param en [out] n_band eigenvalues (host memory) + * @param is_S_orthogonal if true, psi is already S-orthonormal and the + * standard eigenproblem is solved instead of the generalized one */ - static void diag_subspace(const hamilt::Hamilt* const pHamilt, + static void diag_subspace(const HSOperator& op, + const T* psi, + T* evc, + const int nstart, + const int n_band, + const int dmin, + const int dmax, + Real* en, + const diag_comm_info& diag_comm, + const bool is_S_orthogonal = false); + + /// psi::Psi flavour of diag_subspace(): nstart = psi.get_nbands(), + /// n_band = 0 means all of them, dimensions taken from psi. + static void diag_subspace(const HSOperator& op, const psi::Psi& psi, psi::Psi& evc, Real* en, @@ -55,9 +68,10 @@ class DiagoIterAssist int n_band = 0, const bool is_S_orthogonal = false); - /// @brief use LAPACK to diagonalize the Hamiltonian matrix - /// @param pHamilt interface to hamiltonian - /// @param psi wavefunction to diagonalize + /// @brief subspace diagonalization used to build the starting wavefunction + /// @param op interface to H and S; op.add_to_subspace_h() and + /// op.export_subspace_vec() are called around the small eigenproblem + /// @param psi vectors spanning the subspace /// @param psi_nr number of rows (nbands) /// @param psi_nc number of columns (nbasis) /// @param evc new wavefunction @@ -65,21 +79,15 @@ class DiagoIterAssist /// @param basis_type "lcao", "lcao_in_pw" or "pw"; together with calculation it selects /// how the rotation matrix is applied to psi /// @param calculation "scf", "nscf", "md", "relax", ... - /// @note exception handle: if there is no operator initialized in Hamilt, will directly copy value from psi to evc, - /// and return all - zero eigenenergies. - static void diag_subspace_init( - hamilt::Hamilt* pHamilt, - const T* psi, - int psi_nr, - int psi_nc, - psi::Psi& evc, - Real* en, - const std::string& basis_type, - const std::string& calculation, - const diag_comm_info& diag_comm, - const std::function& add_to_hcc = [](T* null, const int n) {}, - const std::function& export_vcc - = [](const T* null, const int n, const int m) {}); + static void diag_subspace_init(const HSOperator& op, + const T* psi, + int psi_nr, + int psi_nc, + psi::Psi& evc, + Real* en, + const std::string& basis_type, + const std::string& calculation, + const diag_comm_info& diag_comm); static void diag_heevx(const int nstart, const int nbands, @@ -96,12 +104,12 @@ class DiagoIterAssist T *vcc); /// @brief calculate Hamiltonian and overlap matrix in subspace spanned by nstart states psi - /// @param pHamilt : hamiltonian operator carrier + /// @param op : applies H and S /// @param psi : wavefunction /// @param hcc : Hamiltonian matrix /// @param scc : overlap matrix - static void cal_hs_subspace(const hamilt::Hamilt* pHamilt, // hamiltonian operator carrier - const psi::Psi& psi, // [in] wavefunction + static void cal_hs_subspace(const HSOperator& op, + const psi::Psi& psi, // [in] wavefunction T* hcc, T* scc, const diag_comm_info& diag_comm); @@ -132,8 +140,6 @@ class DiagoIterAssist private: constexpr static const Device* ctx = {}; - using hpsi_info = typename hamilt::Operator::hpsi_info; - using setmem_var_op = base_device::memory::set_memory_op; using resmem_var_op = base_device::memory::resize_memory_op; using delmem_var_op = base_device::memory::delete_memory_op; diff --git a/source/source_hsolver/hs_matrix.h b/source/source_hsolver/hs_matrix.h new file mode 100644 index 00000000000..46591bf71bc --- /dev/null +++ b/source/source_hsolver/hs_matrix.h @@ -0,0 +1,30 @@ +#ifndef HSOLVER_HS_MATRIX_H +#define HSOLVER_HS_MATRIX_H + +#include "source_base/matrix_block.h" + +namespace hsolver +{ + +/** + * @brief What a direct (dense) eigensolver needs from the generalized + * eigenproblem H x = e S x: the matrices H(k) and S(k) themselves. + * + * The code that owns the Hamiltonian implements this interface + * (hamilt::HamiltHSMatrix in source_hamilt/hamilt_hs_adapter.h); HSolverLCAO + * and Parallel_K2D only ever ask it for the two matrix views. + */ +template +class HSMatrix +{ + public: + virtual ~HSMatrix() = default; + + /// H(k) and S(k) of k point ik as non-owning views. The memory they point + /// to belongs to the implementer and stays valid until the next call. + virtual void hs_at_k(const int ik, ModuleBase::MatrixBlock& hk, ModuleBase::MatrixBlock& sk) = 0; +}; + +} // namespace hsolver + +#endif // HSOLVER_HS_MATRIX_H diff --git a/source/source_hsolver/hs_operator.h b/source/source_hsolver/hs_operator.h new file mode 100644 index 00000000000..30547b2336b --- /dev/null +++ b/source/source_hsolver/hs_operator.h @@ -0,0 +1,54 @@ +#ifndef HSOLVER_HS_OPERATOR_H +#define HSOLVER_HS_OPERATOR_H + +#include "source_base/module_device/types.h" + +namespace hsolver +{ + +/** + * @brief What an iterative eigensolver needs from the generalized eigenproblem + * H x = e S x: the ability to apply H and S to a block of vectors, and + * nothing else. + * + * The solvers never see how H is built. The code that owns the Hamiltonian + * implements this interface (hamilt::HamiltHSOperator in + * source_hamilt/hamilt_hs_adapter.h) and the solvers only call hpsi()/spsi(). + * Unit tests implement it with a dense matrix. + * + * Block vectors are column major: vector i occupies x[i*ld, i*ld + ld). Only + * the first npw rows carry data, ld is the leading dimension. + */ +template +class HSOperator +{ + public: + virtual ~HSOperator() = default; + + /// switch H and S to k point ik; must precede hpsi()/spsi() for that k + virtual void update_k(const int ik) = 0; + + /// hx[:, 0:nvec) = H * x[:, 0:nvec), both with leading dimension ld + virtual void hpsi(const T* x, T* hx, const int ld, const int nvec) const = 0; + + /// sx[:, 0:nvec) = S * x[:, 0:nvec), both with leading dimension ld + virtual void spsi(const T* x, T* sx, const int ld, const int nvec) const = 0; + + /// Hook used by DiagoIterAssist::diag_subspace_init: hcc is the n*n + /// subspace Hamiltonian (column major, ld n). A Hamiltonian carrying a + /// term that hpsi() does not cover (EXX in lcao_in_pw) adds it here. + virtual void add_to_subspace_h(T* hcc, const int n) const + { + } + + /// Hook used by DiagoIterAssist::diag_subspace_init: vcc is the n*nband + /// matrix of subspace eigenvectors, handed out right after the subspace + /// diagonalization. + virtual void export_subspace_vec(const T* vcc, const int n, const int nband) const + { + } +}; + +} // namespace hsolver + +#endif // HSOLVER_HS_OPERATOR_H diff --git a/source/source_hsolver/hsolver_lcao.cpp b/source/source_hsolver/hsolver_lcao.cpp index 79751179d77..9abcbbf0b7e 100644 --- a/source/source_hsolver/hsolver_lcao.cpp +++ b/source/source_hsolver/hsolver_lcao.cpp @@ -39,7 +39,7 @@ namespace hsolver { template -void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, +void HSolverLCAO::solve(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes, elecstate::DensityMatrix& dm, // mohan add 2025-11-03 @@ -56,13 +56,13 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, #ifdef __CUDA if (this->method == "cusolver" && this->world_nproc > 1) { - this->parakSolve_cusolver(pHamilt, psi, pes); + this->parakSolve_cusolver(hs, psi, pes); }else #endif if (this->kpar_lcao > 1 && (this->method == "genelpa" || this->method == "elpa" || this->method == "scalapack_gvx" || this->method == "lapack")) { - this->parakSolve(pHamilt, psi, pes, this->kpar_lcao, nspin); + this->parakSolve(hs, psi, pes, this->kpar_lcao, nspin); } else #endif if (this->kpar_lcao == 1) @@ -70,14 +70,15 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, /// Loop over k points for solve Hamiltonian to eigenpairs(eigenvalues and eigenvectors). for (int ik = 0; ik < psi.get_nk(); ++ik) { - /// update H(k) for each k point - pHamilt->updateHk(ik); + /// H(k) and S(k) are all the eigensolvers need from the Hamiltonian + ModuleBase::MatrixBlock hk, sk; + hs.hs_at_k(ik, hk, sk); /// find psi pointer for each k point psi.fix_k(ik); /// solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(pHamilt, psi, &(pes->ekb(ik, 0))); + this->hamiltSolvePsiK(hk, sk, psi, &(pes->ekb(ik, 0))); } } else @@ -117,11 +118,9 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, DiagoPexsi pe(ParaV, nspin, this->nlocal, this->nelec, this->world_nproc); for (int ik = 0; ik < psi.get_nk(); ++ik) { - /// update H(k) for each k point - pHamilt->updateHk(ik); psi.fix_k(ik); ModuleBase::MatrixBlock hk, sk; - pHamilt->matrix(hk, sk); + hs.hs_at_k(ik, hk, sk); // solve eigenvector and eigenvalue for H(k) pe.diag(hk, sk, psi, nullptr); } @@ -137,16 +136,14 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, } template -void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue) +void HSolverLCAO::hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, + ModuleBase::MatrixBlock& sk, + psi::Psi& psi, + double* eigenvalue) { ModuleBase::TITLE("HSolverLCAO", "hamiltSolvePsiK"); ModuleBase::timer::start("HSolverLCAO", "hamiltSolvePsiK"); - // H(k) and S(k) are all the eigensolvers need from the Hamiltonian, so - // fetch them once here rather than once inside each solver. - ModuleBase::MatrixBlock hk, sk; - hm->matrix(hk, sk); - if (this->method == "scalapack_gvx") { #ifdef __MPI @@ -195,7 +192,7 @@ void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, do } template -void HSolverLCAO::parakSolve(hamilt::Hamilt* pHamilt, +void HSolverLCAO::parakSolve(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes, const int kpar, @@ -215,12 +212,6 @@ void HSolverLCAO::parakSolve(hamilt::Hamilt* pHamilt, int coord_col = k2d.get_p2D_pool()->get_coord_col(); int ncol_bands_pool = numroc_(&(nbands), &(nb2d), &coord_col, &zero, &(k2d.get_p2D_pool()->dim1)); - /// Parallel_K2D only redistributes H(k)/S(k); updating the Hamiltonian - /// for a given k point stays here, where the Hamiltonian is known. - auto get_hsk = [pHamilt](int ik, ModuleBase::MatrixBlock& hk, ModuleBase::MatrixBlock& sk) { - pHamilt->updateHk(ik); - pHamilt->matrix(hk, sk); - }; /// Loop over k points for solve Hamiltonian to charge density for (int ik = 0; ik < k2d.get_pKpoints()->get_max_nks_pool(); ++ik) { @@ -246,7 +237,7 @@ void HSolverLCAO::parakSolve(hamilt::Hamilt* pHamilt, ik_kpar[i] = ik + k2d.get_pKpoints()->startk_pool[i]; } } - k2d.distribute_hsk(get_hsk, ik_kpar, nrow); + k2d.distribute_hsk(hs, ik_kpar, nrow); /// global index of k point int ik_global = ik + k2d.get_pKpoints()->startk_pool[k2d.get_my_pool()]; auto psi_pool = psi::Psi(1, ncol_bands_pool, k2d.get_p2D_pool()->nrow, k2d.get_p2D_pool()->nrow, true); @@ -327,7 +318,7 @@ void HSolverLCAO::parakSolve(hamilt::Hamilt* pHamilt, #if defined (__MPI) && defined (__CUDA) template -void HSolverLCAO::parakSolve_cusolver(hamilt::Hamilt* pHamilt, +void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes) { @@ -414,9 +405,8 @@ void HSolverLCAO::parakSolve_cusolver(hamilt::Hamilt* pHamilt, hk_mat.resize(nrow * ncol); sk_mat.resize(nrow * ncol); } - pHamilt->updateHk(ik); ModuleBase::MatrixBlock hk_2D, sk_2D; - pHamilt->matrix(hk_2D, sk_2D); + hs.hs_at_k(ik, hk_2D, sk_2D); int desc_tmp[9]; T* hk_local_ptr = hk_mat.data(); T* sk_local_ptr = sk_mat.data(); diff --git a/source/source_hsolver/hsolver_lcao.h b/source/source_hsolver/hsolver_lcao.h index 9fb67fbd530..dc44a7fe3b5 100644 --- a/source/source_hsolver/hsolver_lcao.h +++ b/source/source_hsolver/hsolver_lcao.h @@ -2,7 +2,7 @@ #define HSOLVERLCAO_H #include "source_estate/elecstate.h" -#include "source_hamilt/hamilt.h" +#include "source_hsolver/hs_matrix.h" #include "source_basis/module_ao/parallel_orbitals.h" #include "source_estate/module_charge/charge.h" // mohan add 20251024 @@ -27,7 +27,8 @@ class HSolverLCAO : ParaV(ParaV_in), method(method_in), kpar_lcao(kpar_lcao_in), nlocal(nlocal_in), nbands(nbands_in), nelec(nelec_in), use_gpu(use_gpu_in), world_nproc(world_nproc_in), world_rank(world_rank_in){}; - void solve(hamilt::Hamilt* pHamilt, + /// @param hs supplies H(k) and S(k) of the Hamiltonian for each k point + void solve(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes, elecstate::DensityMatrix& dm, // mohan add 2025-11-03 @@ -36,16 +37,19 @@ class HSolverLCAO const bool skip_charge); private: - void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue); // for kpar_lcao == 1 + void hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, + ModuleBase::MatrixBlock& sk, + psi::Psi& psi, + double* eigenvalue); // for kpar_lcao == 1 - void parakSolve(hamilt::Hamilt* pHamilt, + void parakSolve(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes, const int kpar, const int nspin); // for kpar_lcao > 1 // The solving algorithm using cusolver is different from others, so a separate function is needed - void parakSolve_cusolver(hamilt::Hamilt* pHamilt, + void parakSolve_cusolver(HSMatrix& hs, psi::Psi& psi, elecstate::ElecState* pes); diff --git a/source/source_hsolver/hsolver_lcaopw.cpp b/source/source_hsolver/hsolver_lcaopw.cpp index 4dbc31b56a4..7937124b82c 100644 --- a/source/source_hsolver/hsolver_lcaopw.cpp +++ b/source/source_hsolver/hsolver_lcaopw.cpp @@ -5,17 +5,11 @@ #include "source_base/tool_quit.h" #include "source_estate/elecstate_pw.h" #include "source_estate/elecstate_tools.h" -#include "source_hamilt/module_xc/general_exx_info.h" #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/diago_iter_assist.h" -#include "source_pw/module_pwdft/hamilt_pw.h" #include -#ifdef __EXX -#include "source_pw/module_pwdft/hamilt_lcaopw.h" -#endif - namespace hsolver { @@ -23,7 +17,7 @@ namespace hsolver lcao_in_pw */ template -void HSolverLIP::solve(hamilt::Hamilt* pHamilt, // ESolver_KS_PW::p_hamilt +void HSolverLIP::solve(HSOperator& op, // ESolver_KS_PW::p_hamilt behind the operator interface psi::Psi& psi, // ESolver_KS_PW::kspw_psi elecstate::ElecState* pes, // ESolver_KS_PW::pes psi::Psi& transform, @@ -31,8 +25,7 @@ void HSolverLIP::solve(hamilt::Hamilt* pHamilt, // ESolver_KS_PW::p_hamilt std::ostream& log, const bool skip_charge, const double tpiba, - const int nat, - const General_Exx_Info& exx_info) + const int nat) { ModuleBase::TITLE("HSolverLIP", "solve"); ModuleBase::timer::start("HSolverLIP", "solve"); @@ -40,37 +33,13 @@ void HSolverLIP::solve(hamilt::Hamilt* pHamilt, // ESolver_KS_PW::p_hamilt for (int ik = 0; ik < this->wfc_basis->nks; ++ik) { /// update H(k) for each k point - pHamilt->updateHk(ik); + op.update_k(ik); psi.fix_k(ik); transform.fix_k(ik); -#ifdef __EXX - auto& exx_lip = dynamic_cast*>(pHamilt)->exx_lip; - bool cal_exx = exx_info.cal_exx; - double hybrid_alpha = exx_info.hybrid_alpha; - auto add_exx_to_subspace_hamilt = [&ik, &exx_lip, cal_exx, hybrid_alpha](T* hcc, const int naos) -> void { - if (cal_exx) - { - for (int n = 0; n < naos; ++n) - { - for (int m = 0; m < naos; ++m) - { - hcc[n * naos + m] - += (T)hybrid_alpha * exx_lip.get_exx_matrix()[ik][m][n]; - } - } - } - }; - auto set_exxlip_lcaowfc = [&ik, &exx_lip, cal_exx](const T* const vcc, const int naos, const int nbands) -> void { - if (cal_exx) - { - exx_lip.set_hvec(ik, vcc, naos, nbands); - } - }; -#endif /// solve eigenvector and eigenvalue for H(k) - hsolver::DiagoIterAssist::diag_subspace_init(pHamilt, // interface to hamilt + hsolver::DiagoIterAssist::diag_subspace_init(op, transform.get_pointer(), // transform matrix between lcao and pw transform.get_nbands(), transform.get_nbasis(), @@ -78,13 +47,7 @@ void HSolverLIP::solve(hamilt::Hamilt* pHamilt, // ESolver_KS_PW::p_hamilt eigenvalues.data() + ik * pes->ekb.nc, // eigenvalues this->basis_type, this->calculation, - diag_comm -#ifdef __EXX - , - add_exx_to_subspace_hamilt, - set_exxlip_lcaowfc -#endif - ); + diag_comm); if (skip_charge) { diff --git a/source/source_hsolver/hsolver_lcaopw.h b/source/source_hsolver/hsolver_lcaopw.h index b9569156985..6c2dfc3dd2f 100644 --- a/source/source_hsolver/hsolver_lcaopw.h +++ b/source/source_hsolver/hsolver_lcaopw.h @@ -3,12 +3,9 @@ #include "source_base/macros.h" #include "source_estate/elecstate.h" -#include "source_hamilt/hamilt.h" +#include "source_hsolver/hs_operator.h" #include -/// General_Exx_Info forward declaration, full definition in general_exx_info.h -struct General_Exx_Info; - namespace hsolver { @@ -34,12 +31,12 @@ class HSolverLIP global_nbands(global_nbands_in) {}; /// @brief solve function for lcao_in_pw - /// @param pHamilt interface to hamilt + /// @param op the H and S operator of the Hamiltonian; its subspace hooks carry the EXX term /// @param psi reference to psi /// @param pes interface to elecstate /// @param transform transformation matrix between lcao and pw /// @param skip_charge - void solve(hamilt::Hamilt* pHamilt, + void solve(HSOperator& op, psi::Psi& psi, elecstate::ElecState* pes, psi::Psi& transform, @@ -47,8 +44,7 @@ class HSolverLIP std::ostream& log, const bool skip_charge, const double tpiba, - const int nat, - const General_Exx_Info& exx_info); + const int nat); private: ModulePW::PW_Basis_K* wfc_basis = nullptr; diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index d365734475e..a48aff5a031 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -5,7 +5,6 @@ #include "source_base/tool_quit.h" #include "source_estate/elecstate_pw.h" #include "source_estate/elecstate_tools.h" -#include "source_hamilt/hamilt.h" #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/diago_bpcg.h" #include "source_hsolver/diago_cg.h" @@ -64,7 +63,7 @@ void HSolverPW::cal_smooth_ethr(const double& wk, } template -void HSolverPW::solve(hamilt::Hamilt* pHamilt, +void HSolverPW::solve(HSOperator& op, psi::Psi& psi, elecstate::ElecState* pes, double* out_eigenvalues, @@ -107,7 +106,7 @@ void HSolverPW::solve(hamilt::Hamilt* pHamilt, const int ik = k_order[i]; // update H(k) for each k point - pHamilt->updateHk(ik); + op.update_k(ik); @@ -134,7 +133,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); + this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); if (skip_charge) { @@ -150,7 +149,7 @@ void HSolverPW::solve(hamilt::Hamilt* pHamilt, for (int ik = 0; ik < this->wfc_basis->nks; ++ik) { // update H(k) for each k point - pHamilt->updateHk(ik); + op.update_k(ik); @@ -172,7 +171,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); + this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); // output iteration information and reset avg_iter if (skip_charge) @@ -227,7 +226,7 @@ void HSolverPW::solve(hamilt::Hamilt* pHamilt, } template -void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, +void HSolverPW::hamiltSolvePsiK(const HSOperator& op, psi::Psi& psi, std::vector& pre_condition, Real* eigenvalue, @@ -265,41 +264,17 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, ModuleBase::WARNING_QUIT("HSolverPW::hamiltSolvePsiK", msg); } - // Shared matrix-blockvector operators used by all iterative solvers. - auto hpsi_func = [hm, cur_nbasis](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) { - auto psi_wrapper = psi::Psi(psi_in, 1, nvec, ld_psi, cur_nbasis); - psi::Range bands_range(true, 0, 0, nvec - 1); - using hpsi_info = typename hamilt::Operator::hpsi_info; - hpsi_info info(&psi_wrapper, bands_range, hpsi_out); - hm->ops->hPsi(info); - }; - auto spsi_func = [hm, cur_nbasis](const T* psi_in, T* spsi_out, const int ld_psi, const int nvec) { - hm->sPsi(psi_in, spsi_out, ld_psi, cur_nbasis, nvec); - }; - if (this->method == "cg") { - // wrap the subspace_func into a lambda function - // if S_orth is true, then assume psi is S-orthogonal, solve standard eigenproblem - // otherwise, solve generalized eigenproblem - auto subspace_func = - [hm, cur_nbasis, &comm_info](T* psi_in, T* psi_out, const int ld_psi, const int nband, const bool S_orth) { - auto psi_in_wrapper = psi::Psi(psi_in, 1, nband, ld_psi, cur_nbasis); - auto psi_out_wrapper = psi::Psi(psi_out, 1, nband, ld_psi, cur_nbasis); - std::vector eigen(nband, 0.0); - DiagoIterAssist::diag_subspace(hm, psi_in_wrapper, psi_out_wrapper, eigen.data(), comm_info); - }; DiagoCG cg(this->basis_type, this->calculation_type, this->need_subspace, - subspace_func, + comm_info, this->diag_thr, - this->diag_iter_max, - this->nproc_in_pool); + this->diag_iter_max); DiagoIterAssist::avg_iter += static_cast( - cg.diag(hpsi_func, - spsi_func, + cg.diag(op, psi.get_nbasis(), psi.get_nbands(), psi.get_current_ngk(), @@ -318,7 +293,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, const int ndim = psi.get_current_ngk(); DiagoBPCG bpcg(pre_condition.data()); bpcg.init_iter(this->nbands, nband_l, nbasis, ndim); - bpcg.diag(hpsi_func, spsi_func, psi.get_pointer(), eigenvalue, this->ethr_band); + bpcg.diag(op, psi.get_pointer(), eigenvalue, this->ethr_band); } else if (this->method == "dav_subspace") { @@ -336,8 +311,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, this->nb2d); DiagoIterAssist::avg_iter += static_cast( - dav_subspace.diag(hpsi_func, - spsi_func, + dav_subspace.diag(op, psi.get_pointer(), psi.get_nbasis(), eigenvalue, @@ -366,8 +340,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, DiagoDavid david(pre_condition.data(), nband, dim, this->pw_diag_ndim, comm_info); // do diag and add davidson iteration counts up to avg_iter DiagoIterAssist::avg_iter += static_cast( - david.diag(hpsi_func, - spsi_func, + david.diag(op, ld_psi, psi.get_pointer(), eigenvalue, diff --git a/source/source_hsolver/hsolver_pw.h b/source/source_hsolver/hsolver_pw.h index a44053e3403..12d3418fcd4 100644 --- a/source/source_hsolver/hsolver_pw.h +++ b/source/source_hsolver/hsolver_pw.h @@ -4,7 +4,7 @@ #include "source_base/macros.h" #include "source_basis/module_pw/pw_basis_k.h" #include "source_estate/elecstate.h" -#include "source_hamilt/hamilt.h" +#include "source_hsolver/hs_operator.h" #include #include @@ -55,12 +55,11 @@ class HSolverPW use_k_continuity(use_k_continuity_in) {}; /// @brief solve function for pw - /// @param pHamilt interface to hamilt + /// @param op the H and S operator of the Hamiltonian, switched to each k point in turn /// @param psi reference to psi /// @param pes interface to elecstate - /// @param method_in dav or cg /// @param skip_charge - void solve(hamilt::Hamilt* pHamilt, + void solve(HSOperator& op, psi::Psi& psi, elecstate::ElecState* pes, double* out_eigenvalues, @@ -73,7 +72,7 @@ class HSolverPW protected: // diago caller - void hamiltSolvePsiK(hamilt::Hamilt* hm, + void hamiltSolvePsiK(const HSOperator& op, psi::Psi& psi, std::vector& pre_condition, Real* eigenvalue, diff --git a/source/source_hsolver/hsolver_pw_sdft.cpp b/source/source_hsolver/hsolver_pw_sdft.cpp index db075375515..c87064afcee 100644 --- a/source/source_hsolver/hsolver_pw_sdft.cpp +++ b/source/source_hsolver/hsolver_pw_sdft.cpp @@ -7,6 +7,7 @@ #include "source_base/tool_title.h" #include "source_estate/module_charge/symm_rho.h" #include "source_estate/elecstate_tools.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include @@ -52,11 +53,14 @@ void HSolverPW_SDFT::solve(const UnitCell& ucell, ModuleBase::WARNING_QUIT("HSolverPW::solve", "This type of eigensolver is not supported!"); } + // the iterative eigensolvers see the Hamiltonian only through this operator + hamilt::HamiltHSOperator op(pHamilt, wfc_basis); + // part of KSDFT to get KS orbitals for (int ik = 0; ik < nks; ++ik) { ModuleBase::timer::start("HSolverPW_SDFT", "solve_KS"); - pHamilt->updateHk(ik); + op.update_k(ik); if (nbands > 0 && this->ks_run) { /// update psi pointer for each k point @@ -65,7 +69,7 @@ void HSolverPW_SDFT::solve(const UnitCell& ucell, 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)); - this->hamiltSolvePsiK(pHamilt, psi, precondition, p_eigenvalues, nks); + this->hamiltSolvePsiK(op, psi, precondition, p_eigenvalues, nks); } #ifdef __MPI diff --git a/source/source_hsolver/parallel_k2d.cpp b/source/source_hsolver/parallel_k2d.cpp index a80205440df..098d6458674 100644 --- a/source/source_hsolver/parallel_k2d.cpp +++ b/source/source_hsolver/parallel_k2d.cpp @@ -36,7 +36,7 @@ void Parallel_K2D::set_para_env(int nks, } template -void Parallel_K2D::distribute_hsk(const typename Parallel_K2D::HskFunc& get_hsk, +void Parallel_K2D::distribute_hsk(hsolver::HSMatrix& hs, const std::vector& ik_kpar, const int& nw) { #ifdef __MPI @@ -44,7 +44,7 @@ void Parallel_K2D::distribute_hsk(const typename Parallel_K2D::HskFunc& for (int ipool = 0; ipool < ik_kpar.size(); ++ipool) { ModuleBase::MatrixBlock HK_global, SK_global; - get_hsk(ik_kpar[ipool], HK_global, SK_global); + hs.hs_at_k(ik_kpar[ipool], HK_global, SK_global); if (this->MY_POOL == this->Pkpoints->whichpool[ik_kpar[ipool]]) { this->hk_pool.resize(this->P2D_pool->get_local_size(), 0.0); this->sk_pool.resize(this->P2D_pool->get_local_size(), 0.0); diff --git a/source/source_hsolver/parallel_k2d.h b/source/source_hsolver/parallel_k2d.h index fb4ee104c94..3989aa83623 100644 --- a/source/source_hsolver/parallel_k2d.h +++ b/source/source_hsolver/parallel_k2d.h @@ -4,11 +4,11 @@ #include "source_base/matrix_block.h" #include "source_base/parallel_2d.h" #include "source_cell/parallel_kpoints.h" +#include "source_hsolver/hs_matrix.h" #ifdef __MPI #include "mpi.h" #endif -#include #include /*** @@ -35,14 +35,9 @@ class Parallel_K2D { const int& my_rank, const int& nspin); - /// Supplies H(k) and S(k) for one k point. The caller owns whatever has - /// to happen before the blocks are valid (updating the Hamiltonian for - /// that k point, for instance); this class only redistributes them. - using HskFunc = std::function< - void(int ik, ModuleBase::MatrixBlock& hk, ModuleBase::MatrixBlock& sk)>; - - /// this function distributes the Hk and Sk matrices to hk_pool and sk_pool - void distribute_hsk(const HskFunc& get_hsk, + /// this function distributes the Hk and Sk matrices to hk_pool and sk_pool; + /// hs supplies H(k) and S(k) for one k point, this class only redistributes them + void distribute_hsk(hsolver::HSMatrix& hs, const std::vector& ik_kpar, const int& nw); diff --git a/source/source_hsolver/test/CMakeLists.txt b/source/source_hsolver/test/CMakeLists.txt index 043c0596c1f..e387be12a7f 100644 --- a/source/source_hsolver/test/CMakeLists.txt +++ b/source/source_hsolver/test/CMakeLists.txt @@ -13,40 +13,30 @@ if (ENABLE_MPI) LIBS parameter base psi device container SOURCES diago_bpcg_test.cpp ../diago_bpcg.cpp ../para_lin_tf.cpp ../diago_iter_assist.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) AddTest( TARGET MODULE_HSOLVER_cg LIBS parameter base psi device container SOURCES diago_cg_test.cpp ../diago_cg.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) AddTest( TARGET MODULE_HSOLVER_cg_float LIBS parameter base psi device container SOURCES diago_cg_float_test.cpp ../diago_cg.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) AddTest( TARGET MODULE_HSOLVER_dav LIBS parameter base psi device SOURCES diago_david_test.cpp ../diago_david.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) AddTest( TARGET MODULE_HSOLVER_dav_float LIBS parameter base psi device SOURCES diago_david_float_test.cpp ../diago_david.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) if(ENABLE_LCAO) AddTest( @@ -54,16 +44,12 @@ if (ENABLE_MPI) LIBS parameter base psi device container SOURCES diago_cg_float_test.cpp ../diago_cg.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) AddTest( TARGET MODULE_HSOLVER_dav_real LIBS parameter base psi device SOURCES diago_david_real_test.cpp ../diago_david.cpp ../diago_iter_assist.cpp ../diag_const_nums.cpp ../../source_basis/module_pw/test/test_tool.cpp - ../../source_hamilt/operator.cpp - ../../source_pw/module_pwdft/op_pw.cpp ) endif() diff --git a/source/source_hsolver/test/diago_bpcg_test.cpp b/source/source_hsolver/test/diago_bpcg_test.cpp index 8a6a1558e16..34bed8b4735 100644 --- a/source/source_hsolver/test/diago_bpcg_test.cpp +++ b/source/source_hsolver/test/diago_bpcg_test.cpp @@ -2,8 +2,6 @@ #include "source_base/module_external/lapack_connector.h" #include "source_base/parallel_comm.h" #include "source_psi/psi.h" -#include "source_hamilt/hamilt.h" -#include "source_pw/module_pwdft/hamilt_pw.h" #include "../diago_iter_assist.h" #include "../diago_bpcg.h" #include "diago_mock.h" @@ -99,8 +97,7 @@ class DiagoBPCGPrepare //====================================================================== double *en = new double[npw]; int ik = 1; - hamilt::Hamilt>* ha; - ha =new hamilt::HamiltPW>(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock> ha; int* ngk = new int [1]; //psi::Psi> psi(ngk,ik,nband,npw); psi::Psi> psi; @@ -133,35 +130,13 @@ class DiagoBPCGPrepare psi_local.fix_k(0); double start, end; start = MPI_Wtime(); - using T = std::complex; - const int dim = DIAGOTEST::npw; - const std::vector &h_mat = DIAGOTEST::hmatrix_local; - auto hpsi_func = [h_mat, dim](T *psi_in, T *hpsi_out, - const int ld_psi, const int nvec) { - const T one(1.0); - const T zero(0.0); - - base_device::DEVICE_CPU *ctx = {}; - // hpsi_out(dim * nvec) = h_mat(dim * dim) * psi_in(dim * nvec) - ModuleBase::gemm_op()( - 'N', 'N', - dim, nvec, dim, - &one, - h_mat.data(), dim, - psi_in, ld_psi, - &zero, - hpsi_out, ld_psi); - }; - auto spsi_func = [](const T* psi_in, T* spsi_out, const int ld_psi, const int nvec) { - std::copy(psi_in, psi_in + ld_psi * nvec, spsi_out); - }; const int ndim = psi_local.get_current_ngk(); bpcg.init_iter(nband, nband, npw, ndim); std::vector ethr_band(nband, 1e-5); - bpcg.diag(hpsi_func, spsi_func, psi_local.get_pointer(), en, ethr_band); - bpcg.diag(hpsi_func, spsi_func, psi_local.get_pointer(), en, ethr_band); - bpcg.diag(hpsi_func, spsi_func, psi_local.get_pointer(), en, ethr_band); - bpcg.diag(hpsi_func, spsi_func, psi_local.get_pointer(), en, ethr_band); + bpcg.diag(ha, psi_local.get_pointer(), en, ethr_band); + bpcg.diag(ha, psi_local.get_pointer(), en, ethr_band); + bpcg.diag(ha, psi_local.get_pointer(), en, ethr_band); + bpcg.diag(ha, psi_local.get_pointer(), en, ethr_band); end = MPI_Wtime(); //if(mypnum == 0) printf("diago time:%7.3f\n",end-start); delete [] DIAGOTEST::npw_local; @@ -174,7 +149,6 @@ class DiagoBPCGPrepare delete[] en; delete[] e_lapack; - delete ha; } }; diff --git a/source/source_hsolver/test/diago_cg_float_test.cpp b/source/source_hsolver/test/diago_cg_float_test.cpp index d085a900bec..0651daca593 100644 --- a/source/source_hsolver/test/diago_cg_float_test.cpp +++ b/source/source_hsolver/test/diago_cg_float_test.cpp @@ -7,9 +7,7 @@ #include "source_base/inverse_matrix.h" #include "source_base/module_external/lapack_connector.h" #include "source_basis/module_pw/test/test_tool.h" -#include "source_hamilt/hamilt.h" #include "source_psi/psi.h" -#include "source_pw/module_pwdft/hamilt_pw.h" #include "gtest/gtest.h" #include @@ -104,8 +102,7 @@ class DiagoCGPrepare //====================================================================== float *en = new float[npw]; int ik = 1; - hamilt::Hamilt>* ha; - ha =new hamilt::HamiltPW>(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock> ha; psi::Psi> psi; psi.resize(ik,nband,npw); //psi.fix_k(0); @@ -143,54 +140,20 @@ class DiagoCGPrepare #else const hsolver::diag_comm_info diag_comm(mypnum, nprocs); #endif - // warp the subspace_func into a lambda function - auto subspace_func = [ha, &diag_comm](std::complex* psi_in, - std::complex* psi_out, - const int ld_psi, - const int nband, - const bool S_orth) { - auto psi_in_wrapper = psi::Psi>(psi_in, 1, nband, ld_psi, true); - auto psi_out_wrapper = psi::Psi>(psi_out, 1, nband, ld_psi, true); - std::vector eigen(nband, 0.0f); - hsolver::DiagoIterAssist>::diag_subspace(ha, - psi_in_wrapper, - psi_out_wrapper, - eigen.data(), - diag_comm); - }; hsolver::DiagoCG> cg("pw", "scf", hsolver::DiagoIterAssist>::need_subspace, - subspace_func, + diag_comm, hsolver::DiagoIterAssist>::PW_DIAG_THR, - hsolver::DiagoIterAssist>::PW_DIAG_NMAX, - nprocs); + hsolver::DiagoIterAssist>::PW_DIAG_NMAX); // hsolver::DiagoCG> cg(precondition_local); psi_local.fix_k(0); float start, end; start = MPI_Wtime(); - auto hpsi_func = [ha](std::complex* psi_in, - std::complex* hpsi_out, - const int ld_psi, - const int nvec) { - auto psi_wrapper = psi::Psi>(psi_in, 1, nvec, ld_psi, true); - psi::Range all_bands_range(true, 0, 0, nvec - 1); - using hpsi_info = typename hamilt::Operator>::hpsi_info; - hpsi_info info(&psi_wrapper, all_bands_range, hpsi_out); - ha->ops->hPsi(info); - }; - auto spsi_func = [ha](std::complex* psi_in, - std::complex* spsi_out, - const int ld_psi, - const int nvec) { - ha->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec); - }; std::vector ethr_band(nband, 1e-5); - cg.diag(hpsi_func, - spsi_func, - psi_local.get_nbasis(), + cg.diag(ha, psi_local.get_nbasis(), psi_local.get_nbands(), psi_local.get_current_ngk(), psi_local.get_pointer(), @@ -211,7 +174,6 @@ class DiagoCGPrepare delete[] en; delete[] e_lapack; - delete ha; } }; diff --git a/source/source_hsolver/test/diago_cg_real_test.cpp b/source/source_hsolver/test/diago_cg_real_test.cpp index 7d5e364e892..f1afbd99a45 100644 --- a/source/source_hsolver/test/diago_cg_real_test.cpp +++ b/source/source_hsolver/test/diago_cg_real_test.cpp @@ -7,9 +7,7 @@ #include "source_base/inverse_matrix.h" #include "source_base/module_external/lapack_connector.h" #include "source_basis/module_pw/test/test_tool.h" -#include "source_hamilt/hamilt.h" #include "source_psi/psi.h" -#include "source_pw/module_pwdft/hamilt_pw.h" #include "gtest/gtest.h" #include @@ -105,8 +103,7 @@ class DiagoCGPrepare //====================================================================== double* en = new double[npw]; int ik = 1; - hamilt::Hamilt* ha; - ha = new hamilt::HamiltPW(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock ha; int* ngk = new int[1]; psi::Psi psi; psi.resize(ik, nband, npw); @@ -149,51 +146,20 @@ class DiagoCGPrepare #else const hsolver::diag_comm_info diag_comm(mypnum, nprocs); #endif - // warp the subspace_func into a lambda function - auto subspace_func - = [ha, &diag_comm](double* psi_in, double* psi_out, const int ld_psi, const int nband, const bool S_orth) { - auto psi_in_wrapper = psi::Psi(psi_in, 1, nband, ld_psi, true); - auto psi_out_wrapper = psi::Psi(psi_out, 1, nband, ld_psi, true); - std::vector eigen(nband, 0.0); - hsolver::DiagoIterAssist::diag_subspace(ha, - psi_in_wrapper, - psi_out_wrapper, - eigen.data(), - diag_comm); - }; hsolver::DiagoCG cg("pw", "scf", hsolver::DiagoIterAssist::need_subspace, - subspace_func, + diag_comm, hsolver::DiagoIterAssist::PW_DIAG_THR, - hsolver::DiagoIterAssist::PW_DIAG_NMAX, - nprocs); + hsolver::DiagoIterAssist::PW_DIAG_NMAX); // hsolver::DiagoCG cg(precondition_local); psi_local.fix_k(0); double start, end; start = MPI_Wtime(); - auto hpsi_func = [ha](double* psi_in, - double* hpsi_out, - const int ld_psi, - const int nvec) { - auto psi_wrapper = psi::Psi(psi_in, 1, nvec, ld_psi, true); - psi::Range all_bands_range(true, 0, 0, nvec - 1); - using hpsi_info = typename hamilt::Operator::hpsi_info; - hpsi_info info(&psi_wrapper, all_bands_range, hpsi_out); - ha->ops->hPsi(info); - }; - auto spsi_func = [ha](double* psi_in, - double* spsi_out, - const int ld_psi, - const int nvec) { - ha->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec); - }; std::vector ethr_band(nband, 1e-5); - cg.diag(hpsi_func, - spsi_func, - psi_local.get_nbasis(), + cg.diag(ha, psi_local.get_nbasis(), psi_local.get_nbands(), psi_local.get_current_ngk(), psi_local.get_pointer(), @@ -214,7 +180,6 @@ class DiagoCGPrepare delete[] en; delete[] e_lapack; - delete ha; } }; diff --git a/source/source_hsolver/test/diago_cg_test.cpp b/source/source_hsolver/test/diago_cg_test.cpp index 7f304947d78..853b51c9f71 100644 --- a/source/source_hsolver/test/diago_cg_test.cpp +++ b/source/source_hsolver/test/diago_cg_test.cpp @@ -8,9 +8,7 @@ #include "source_base/inverse_matrix.h" #include "source_base/module_external/lapack_connector.h" #include "source_basis/module_pw/test/test_tool.h" -#include "source_hamilt/hamilt.h" #include "source_psi/psi.h" -#include "source_pw/module_pwdft/hamilt_pw.h" #include "gtest/gtest.h" #include @@ -101,8 +99,7 @@ class DiagoCGPrepare //====================================================================== double *en = new double[npw]; int ik = 1; - hamilt::Hamilt>* ha; - ha =new hamilt::HamiltPW>(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock> ha; psi::Psi> psi; psi.resize(ik,nband,npw); //psi.fix_k(0); @@ -138,54 +135,20 @@ class DiagoCGPrepare #else const hsolver::diag_comm_info diag_comm(mypnum, nprocs); #endif - // warp the subspace_func into a lambda function - auto subspace_func = [ha, &diag_comm](std::complex* psi_in, - std::complex* psi_out, - const int ld_psi, - const int nband, - const bool S_orth) { - auto psi_in_wrapper = psi::Psi>(psi_in, 1, nband, ld_psi, true); - auto psi_out_wrapper = psi::Psi>(psi_out, 1, nband, ld_psi, true); - std::vector eigen(nband, 0.0); - hsolver::DiagoIterAssist>::diag_subspace(ha, - psi_in_wrapper, - psi_out_wrapper, - eigen.data(), - diag_comm); - }; hsolver::DiagoCG> cg("pw", "scf", hsolver::DiagoIterAssist>::need_subspace, - subspace_func, + diag_comm, hsolver::DiagoIterAssist>::PW_DIAG_THR, - hsolver::DiagoIterAssist>::PW_DIAG_NMAX, - nprocs); + hsolver::DiagoIterAssist>::PW_DIAG_NMAX); // hsolver::DiagoCG> cg(precondition_local); psi_local.fix_k(0); double start, end; start = MPI_Wtime(); - auto hpsi_func = [ha](std::complex* psi_in, - std::complex* hpsi_out, - const int ld_psi, - const int nvec) { - auto psi_wrapper = psi::Psi>(psi_in, 1, nvec, ld_psi, true); - psi::Range all_bands_range(true, 0, 0, nvec - 1); - using hpsi_info = typename hamilt::Operator>::hpsi_info; - hpsi_info info(&psi_wrapper, all_bands_range, hpsi_out); - ha->ops->hPsi(info); - }; - auto spsi_func = [ha](std::complex* psi_in, - std::complex* spsi_out, - const int ld_psi, - const int nvec) { - ha->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec); - }; std::vector ethr_band(nband, 1e-5); - cg.diag(hpsi_func, - spsi_func, - psi_local.get_nbasis(), + cg.diag(ha, psi_local.get_nbasis(), psi_local.get_nbands(), psi_local.get_current_ngk(), psi_local.get_pointer(), @@ -207,7 +170,6 @@ class DiagoCGPrepare delete[] en; delete[] e_lapack; - delete ha; } }; diff --git a/source/source_hsolver/test/diago_david_float_test.cpp b/source/source_hsolver/test/diago_david_float_test.cpp index a5ac77be1ae..930b4920cb6 100644 --- a/source/source_hsolver/test/diago_david_float_test.cpp +++ b/source/source_hsolver/test/diago_david_float_test.cpp @@ -2,7 +2,6 @@ #include "source_hsolver/diag_comm_info.h" #include"source_hsolver/diago_iter_assist.h" #include "source_base/parallel_comm.h" -#include"source_pw/module_pwdft/hamilt_pw.h" #include"diago_mock.h" #include "source_psi/psi.h" #include"gtest/gtest.h" @@ -83,8 +82,7 @@ class DiagoDavPrepare //do Diago_David::diag() float* en = new float[npw]; - hamilt::Hamilt> *phm; - phm = new hamilt::HamiltPW>(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock> phm; #ifdef __MPI const hsolver::diag_comm_info comm_info = {POOL_WORLD, mypnum, nprocs}; @@ -111,21 +109,8 @@ class DiagoDavPrepare #endif - auto hpsi_func = [phm](std::complex* psi_in,std::complex* hpsi_out, - const int ld_psi, const int nvec) - { - auto psi_iter_wrapper = psi::Psi>(psi_in, 1, nvec, ld_psi, true); - psi::Range bands_range(true, 0, 0, nvec-1); - using hpsi_info = typename hamilt::Operator>::hpsi_info; - 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); }; std::vector ethr_band(phi.get_nbands(), eps); - dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); + dav.diag(phm, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); @@ -145,7 +130,6 @@ class DiagoDavPrepare } } delete [] en; - delete phm; delete [] e_lapack; } }; diff --git a/source/source_hsolver/test/diago_david_real_test.cpp b/source/source_hsolver/test/diago_david_real_test.cpp index 00f6917d652..3ff4e33bc7b 100644 --- a/source/source_hsolver/test/diago_david_real_test.cpp +++ b/source/source_hsolver/test/diago_david_real_test.cpp @@ -2,7 +2,6 @@ #include "source_hsolver/diag_comm_info.h" #include"source_hsolver/diago_iter_assist.h" #include "source_base/parallel_comm.h" -#include"source_pw/module_pwdft/hamilt_pw.h" #include"diago_mock.h" #include "source_psi/psi.h" #include"gtest/gtest.h" @@ -82,8 +81,7 @@ class DiagoDavPrepare //do Diago_David::diag() double* en = new double[npw]; - hamilt::Hamilt* phm; - phm = new hamilt::HamiltPW(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock phm; #ifdef __MPI const hsolver::diag_comm_info comm_info = {POOL_WORLD, mypnum, nprocs}; @@ -110,20 +108,8 @@ class DiagoDavPrepare #endif - auto hpsi_func = [phm](double* psi_in,double* hpsi_out, - const int ld_psi, const int nvec) - { - auto psi_iter_wrapper = psi::Psi(psi_in, 1, nvec, ld_psi, true); - psi::Range bands_range(true, 0, 0, nvec-1); - using hpsi_info = typename hamilt::Operator::hpsi_info; - 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) { - 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); + dav.diag(phm, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); @@ -143,7 +129,6 @@ class DiagoDavPrepare } } delete[] en; - delete phm; delete[] e_lapack; } }; diff --git a/source/source_hsolver/test/diago_david_test.cpp b/source/source_hsolver/test/diago_david_test.cpp index 771e0ae489d..7d4b8cbfea0 100644 --- a/source/source_hsolver/test/diago_david_test.cpp +++ b/source/source_hsolver/test/diago_david_test.cpp @@ -2,7 +2,6 @@ #include "source_hsolver/diag_comm_info.h" #include"source_hsolver/diago_iter_assist.h" #include "source_base/parallel_comm.h" -#include"source_pw/module_pwdft/hamilt_pw.h" #include"diago_mock.h" #include "source_psi/psi.h" #include"gtest/gtest.h" @@ -87,8 +86,7 @@ class DiagoDavPrepare //do Diago_David::diag() double* en = new double[npw]; - hamilt::Hamilt> *phm; - phm = new hamilt::HamiltPW>(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + HSOperatorMock> phm; #ifdef __MPI const hsolver::diag_comm_info comm_info = {POOL_WORLD, mypnum, nprocs}; @@ -115,21 +113,8 @@ class DiagoDavPrepare #endif - auto hpsi_func = [phm](std::complex* psi_in,std::complex* hpsi_out, - const int ld_psi, const int nvec) - { - auto psi_iter_wrapper = psi::Psi>(psi_in, 1, nvec, ld_psi, true); - psi::Range bands_range(true, 0, 0, nvec-1); - using hpsi_info = typename hamilt::Operator>::hpsi_info; - 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); }; std::vector ethr_band(phi.get_nbands(), eps); - dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); + dav.diag(phm, ld_psi, phi.get_pointer(), en, ethr_band, maxiter); #ifdef __MPI end = MPI_Wtime(); @@ -149,7 +134,6 @@ class DiagoDavPrepare } } delete [] en; - delete phm; delete [] e_lapack; } }; diff --git a/source/source_hsolver/test/diago_mock.h b/source/source_hsolver/test/diago_mock.h index 279b5dc566e..cc4e0853e7b 100644 --- a/source/source_hsolver/test/diago_mock.h +++ b/source/source_hsolver/test/diago_mock.h @@ -1,7 +1,14 @@ #include #include "mpi.h" #include "source_base/parallel_reduce.h" -#include "source_pw/module_pwdft/hamilt_pw.h" +#include "source_hsolver/hs_operator.h" +#include "source_psi/psi.h" + +#include +#include +#include +#include +#include namespace DIAGOTEST { @@ -210,8 +217,6 @@ class HPsi //ModuleBase::ComplexMatrix psi() {return psimatrix;} psi::Psi psi() { - Structure_Factor* sf; - int* ngk = nullptr; psi::Psi psitmp(1, nband, npw, npw, true); for(int i=0;i; template class HPsi>; template class HPsi>; -//totally same as the original function -template <> -void hamilt::HamiltPW::sPsi(const double* psi_in, - double* spsi, - const int nrow, - const int npw, - const int nbands) const -{ - for (size_t i = 0; i < static_cast(nbands * nrow); i++) - { - spsi[i] = psi_in[i]; - } - return; -} -template <> -void hamilt::HamiltPW, base_device::DEVICE_CPU>::sPsi(const std::complex* psi_in, - std::complex* spsi, - const int nrow, - const int npw, - const int nbands) const +namespace DIAGOTEST { - for (size_t i = 0; i < static_cast(nbands * nrow); i++) - { - spsi[i] = psi_in[i]; - } - return; -} -template <> -void hamilt::HamiltPW, base_device::DEVICE_CPU>::sPsi(const std::complex* psi_in, - std::complex* spsi, - const int nrow, - const int npw, - const int nbands) const + /// the process-local slice of the test matrix for each scalar type + template std::vector& hmatrix_local_of(); + template <> std::vector& hmatrix_local_of() { return hmatrix_local_d; } + template <> std::vector>& hmatrix_local_of>() { return hmatrix_local; } + template <> std::vector>& hmatrix_local_of>() { return hmatrix_local_f; } +} + +/** + * The test matrix seen through hsolver::HSOperator, which is all the iterative + * eigensolvers need. H is the dense (distributed) matrix DIAGOTEST::hmatrix_local, + * S is the identity. + * + * Every process holds npw_local[mypnum] columns of H, so H*x is formed as a + * partial product, reduced over the pool and redistributed like psi. + */ +template +class HSOperatorMock : public hsolver::HSOperator { - for (size_t i = 0; i < static_cast(nbands * nrow); i++) + public: + void update_k(const int ik) override { - spsi[i] = psi_in[i]; } - return; -} -//Mock function h_psi -#include "source_pw/module_pwdft/op_pw.h" -template -class OperatorMock : public hamilt::Operator -{ - ~OperatorMock() + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override { - if(this->hpsi != nullptr) - { - delete this->hpsi; - this->hpsi = nullptr; - } - } - virtual void act - ( - const int nbands, - const int nbasis, - const int npol, - const T* tmpsi_in, - T* tmhpsi, - const int ngk_ik = 0, - const bool is_first_node = false)const; -}; -template<> -void OperatorMock::act( - const int nbands, - const int nbasis, - const int npol, - const double* tmpsi_in, - double* tmhpsi, - const int ngk_ik, - const bool is_first_node)const -{ - int nprocs = 1, mypnum = 0; -#ifdef __MPI - MPI_Comm_size(MPI_COMM_WORLD, &nprocs); - MPI_Comm_rank(MPI_COMM_WORLD, &mypnum); -#endif - - double* hpsi0 = new double[DIAGOTEST::npw]; - for (int m = 0; m < nbands; m++) - { - for (int i = 0;i < DIAGOTEST::npw;i++) - { - hpsi0[i] = 0.0; - for (int j = 0;j < (DIAGOTEST::npw_local[mypnum]);j++) - { - hpsi0[i] += DIAGOTEST::hmatrix_local_d[i * DIAGOTEST::h_nc + j] * tmpsi_in[j]; - } - } - Parallel_Reduce::reduce_pool(hpsi0, DIAGOTEST::npw); - DIAGOTEST::divide_psi(hpsi0, tmhpsi); - tmhpsi += nbasis; - tmpsi_in += nbasis; - } - delete[] hpsi0; -} -template<> -void OperatorMock>::act( - const int nbands, - const int nbasis, - const int npol, - const std::complex* tmpsi_in, - std::complex* tmhpsi, - const int ngk_ik, - const bool is_first_node)const -{ - int nprocs = 1, mypnum = 0; -#ifdef __MPI - MPI_Comm_size(MPI_COMM_WORLD, &nprocs); - MPI_Comm_rank(MPI_COMM_WORLD, &mypnum); -#endif - - std::complex* hpsi0 = new std::complex[DIAGOTEST::npw]; - for (int m = 0; m < nbands; m++) - { - for (int i = 0;i < DIAGOTEST::npw;i++) + int mypnum = 0; +#ifdef __MPI + MPI_Comm_rank(MPI_COMM_WORLD, &mypnum); +#endif + const std::vector& hmat = DIAGOTEST::hmatrix_local_of(); + const int npw = DIAGOTEST::npw; + const int ncol = DIAGOTEST::npw_local[mypnum]; + std::vector hx_full(npw); + for (int m = 0; m < nvec; m++) { - hpsi0[i] = 0.0; - for (int j = 0;j < (DIAGOTEST::npw_local[mypnum]);j++) + for (int i = 0; i < npw; i++) { - hpsi0[i] += DIAGOTEST::hmatrix_local[i * DIAGOTEST::h_nc + j] * tmpsi_in[j]; + hx_full[i] = T(0); + for (int j = 0; j < ncol; j++) + { + hx_full[i] += hmat[i * DIAGOTEST::h_nc + j] * x[m * ld + j]; + } } + Parallel_Reduce::reduce_pool(hx_full.data(), npw); + DIAGOTEST::divide_psi(hx_full.data(), hx + m * ld); } - Parallel_Reduce::reduce_pool(hpsi0, DIAGOTEST::npw); - DIAGOTEST::divide_psi>(hpsi0, tmhpsi); - tmhpsi += nbasis; - tmpsi_in += nbasis; } - delete[] hpsi0; -} -template<> -void OperatorMock>::act( - const int nbands, - const int nbasis, - const int npol, - const std::complex* tmpsi_in, - std::complex* tmhpsi, - const int ngk_ik, - const bool is_first_node)const -{ - int nprocs = 1, mypnum = 0; -#ifdef __MPI - MPI_Comm_size(MPI_COMM_WORLD, &nprocs); - MPI_Comm_rank(MPI_COMM_WORLD, &mypnum); -#endif - std::complex* hpsi0 = new std::complex[DIAGOTEST::npw]; - for (int m = 0; m < nbands; m++) + void spsi(const T* x, T* sx, const int ld, const int nvec) const override { - for (int i = 0;i < DIAGOTEST::npw;i++) - { - hpsi0[i] = 0.0; - for (int j = 0;j < (DIAGOTEST::npw_local[mypnum]);j++) - { - hpsi0[i] += DIAGOTEST::hmatrix_local_f[i * DIAGOTEST::h_nc + j] * tmpsi_in[j]; - } - } - Parallel_Reduce::reduce_pool(hpsi0, DIAGOTEST::npw); - DIAGOTEST::divide_psi>(hpsi0, tmhpsi); - tmhpsi += nbasis; - tmpsi_in += nbasis; + std::copy(x, x + static_cast(ld) * nvec, sx); } - delete[] hpsi0; -} -template<> void hamilt::HamiltPW::updateHk(const int ik) -{ - return; -} - -template<> hamilt::HamiltPW::HamiltPW( - elecstate::Potential* pot_in, - ModulePW::PW_Basis_K* wfc_basis, - K_Vectors* pkv, - pseudopot_cell_vnl* ppcell, - Plus_U_Base* p_dftu, - const UnitCell* ucell, - const General_Exx_Info* exx_info) -{ - this->ops = new OperatorMock; -} - -template<> hamilt::HamiltPW::~HamiltPW() -{ - delete this->ops; -} - -template<> void hamilt::HamiltPW>::updateHk(const int ik) -{ - return; -} - -template<> hamilt::HamiltPW>::HamiltPW( - elecstate::Potential* pot_in, - ModulePW::PW_Basis_K* wfc_basis, - K_Vectors* pkv, - pseudopot_cell_vnl* ppcell, - Plus_U_Base* p_dftu, - const UnitCell* ucell, - const General_Exx_Info* exx_info) -{ - this->ops = new OperatorMock>; -} - -template<> hamilt::HamiltPW>::~HamiltPW() -{ - delete this->ops; -} - -template<> void hamilt::HamiltPW>::updateHk(const int ik) -{ - return; -} - -template<> hamilt::HamiltPW>::HamiltPW( - elecstate::Potential* pot_in, - ModulePW::PW_Basis_K* wfc_basis, - K_Vectors* pkv, - pseudopot_cell_vnl* ppcell, - Plus_U_Base* p_dftu, - const UnitCell* ucell, - const General_Exx_Info* exx_info) -{ - this->ops = new OperatorMock>; -} - -template<> hamilt::HamiltPW>::~HamiltPW() -{ - delete this->ops; -} +}; diff --git a/source/source_hsolver/test/hsolver_pw_sup.h b/source/source_hsolver/test/hsolver_pw_sup.h index b41196c396d..75ec3a90500 100644 --- a/source/source_hsolver/test/hsolver_pw_sup.h +++ b/source/source_hsolver/test/hsolver_pw_sup.h @@ -49,6 +49,7 @@ double& PW_Basis_K::getgk2(const int ik, const int igl) const { #include "source_hsolver/diago_cg.h" #include "source_hsolver/diago_david.h" +#include "source_hsolver/diag_comm_info.h" #include "source_hsolver/diago_iter_assist.h" @@ -56,7 +57,13 @@ namespace hsolver { template DiagoCG::DiagoCG(const std::string& basis_type, - const std::string& calculation) { + const std::string& calculation) +#ifdef __MPI + : diag_comm_(MPI_COMM_SELF, 0, 1) +#else + : diag_comm_(0, 1) +#endif +{ basis_type_ = basis_type; calculation_ = calculation; this->one_ = new T(static_cast(1.0)); @@ -68,17 +75,15 @@ template DiagoCG::DiagoCG(const std::string& basis_type, const std::string& calculation, const bool& need_subspace, - const SubspaceFunc& subspace_func, + const diag_comm_info& diag_comm, const Real& pw_diag_thr, - const int& pw_diag_nmax, - const int& nproc_in_pool) { + const int& pw_diag_nmax) + : diag_comm_(diag_comm) { basis_type_ = basis_type; calculation_ = calculation; need_subspace_ = need_subspace; - subspace_func_ = subspace_func; pw_diag_thr_ = pw_diag_thr; pw_diag_nmax_ = pw_diag_nmax; - nproc_in_pool_ = nproc_in_pool; this->one_ = new T(static_cast(1.0)); this->zero_ = new T(static_cast(0.0)); this->neg_one_ = new T(static_cast(-1.0)); @@ -92,8 +97,7 @@ DiagoCG::~DiagoCG() { } template -double DiagoCG::diag(const HPsiFunc& hpsi_func, - const SPsiFunc& spsi_func, +double DiagoCG::diag(const HSOperator& op, const int ld_psi, const int nband, const int dim, @@ -146,8 +150,7 @@ DiagoDavid::~DiagoDavid() { } template -int DiagoDavid::diag(const std::function& hpsi_func, - const std::function& spsi_func, +int DiagoDavid::diag(const HSOperator& op, const int ld_psi, T* psi_in, Real* eigenvalue_in, diff --git a/source/source_hsolver/test/test_hsolver_pw.cpp b/source/source_hsolver/test/test_hsolver_pw.cpp index 0bd18093fb2..e687d8f330c 100644 --- a/source/source_hsolver/test/test_hsolver_pw.cpp +++ b/source/source_hsolver/test/test_hsolver_pw.cpp @@ -5,11 +5,32 @@ #include "hsolver_pw_sup.h" #include "hsolver_supplementary_mock.h" -#include "source_hamilt/module_xc/general_exx_info.h" // for General_Exx_Info type #include "source_hsolver/diag_comm_info.h" +#include "source_hsolver/hs_operator.h" #include "source_hsolver/hsolver_lcaopw.h" #include "source_hsolver/hsolver_pw.h" +#include +#include + +/// H = S = identity: the simplest operator the solvers can be handed +template +class IdentityHSOperator : public hsolver::HSOperator +{ + public: + void update_k(const int ik) override + { + } + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override + { + std::copy(x, x + static_cast(ld) * nvec, hx); + } + void spsi(const T* x, T* sx, const int ld, const int nvec) const override + { + std::copy(x, x + static_cast(ld) * nvec, sx); + } +}; + // Mock implementations for the template functions causing linking errors namespace ModulePW { // Mock implementation for recip_to_real @@ -152,13 +173,13 @@ class TestHSolverPW : public ::testing::Test { // the protected hamiltSolvePsiK() is routed through here. template static void hamiltSolvePsiK(hsolver::HSolverPW& hs, - hamilt::Hamilt* h, + const hsolver::HSOperator& op, psi::Psi& ps, std::vector::type>& pre, typename GetTypeReal::type* eig, const int ntry) { - hs.hamiltSolvePsiK(h, ps, pre, eig, ntry); + hs.hamiltSolvePsiK(op, ps, pre, eig, ntry); } ModulePW::PW_Basis_K pwbk; @@ -197,8 +218,8 @@ class TestHSolverPW : public ::testing::Test { 0, 0); - hamilt::Hamilt> hamilt_test_d; - hamilt::Hamilt> hamilt_test_f; + IdentityHSOperator> hamilt_test_d; + IdentityHSOperator> hamilt_test_f; psi::Psi> psi_test_cd; psi::Psi> psi_test_cf; @@ -364,26 +385,41 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { transform_test_cd.resize(1, 3, 3); transform_test_cf.resize(1, 3, 3); - std::complex psi_value_d = {0.0, 0.0}; - std::complex psi_value_f = {0.0, 0.0}; + // 1, 2, 3 / 4, 5, 6 / 7, 8, 9 would be rank deficient, so the diagonal is + // lifted to keep the three subspace vectors linearly independent for (int iband = 0; iband < transform_test_cd.get_nbands(); iband++) { for (int ibasis = 0; ibasis < transform_test_cd.get_nbasis(); ibasis++) { + const double value = iband * transform_test_cd.get_nbasis() + ibasis + 1 + (iband == ibasis ? 10.0 : 0.0); transform_test_cd .get_pointer()[iband * transform_test_cd.get_nbasis() + ibasis] - = psi_value_d; + = std::complex(value, 0.0); transform_test_cf .get_pointer()[iband * transform_test_cf.get_nbasis() + ibasis] - = psi_value_f; - psi_value_d += std::complex(1.0, 0.0); - psi_value_f += std::complex(1.0, 0.0); + = std::complex(value, 0.0); } } + // with H = S = 1 every subspace eigenvalue is 1 and the rotated psi must + // come out orthonormal + auto check_orthonormal = [](const auto& p, const double tol) { + const int nb = p.get_nbands(); + const int nbasis = p.get_nbasis(); + for (int i = 0; i < nb; i++) { + for (int j = 0; j < nb; j++) { + std::complex dot = 0.0; + for (int ig = 0; ig < nbasis; ig++) { + dot += std::conj(std::complex(p.get_pointer()[i * nbasis + ig])) + * std::complex(p.get_pointer()[j * nbasis + ig]); + } + EXPECT_NEAR(dot.real(), i == j ? 1.0 : 0.0, tol); + EXPECT_NEAR(dot.imag(), 0.0, tol); + } + } + }; // check solve() elecstate_test.ekb.c[0] = 1.0; elecstate_test.ekb.c[1] = 2.0; - General_Exx_Info exx_info_local; hsolver::HSolverLIP> hs_f_lip = hsolver::HSolverLIP>(&pwbk, false, "pw", "scf", elecstate_test.ekb.nc); hsolver::HSolverLIP> hs_d_lip @@ -394,7 +430,7 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { const hsolver::diag_comm_info diag_comm(0, 1); #endif std::ostringstream log; - hs_f_lip.solve(&hamilt_test_f, + hs_f_lip.solve(hamilt_test_f, psi_test_cf, &elecstate_test, transform_test_cf, @@ -402,20 +438,16 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { log, true, 0.0, - 0, - exx_info_local); + 0); EXPECT_NE(log.str().find("Average iterative diagonalization steps"), std::string::npos); EXPECT_DOUBLE_EQ(hsolver::DiagoIterAssist>::avg_iter, 0.0); - for (int i = 0; i < psi_test_cf.size(); i++) - { - EXPECT_DOUBLE_EQ(psi_test_cf.get_pointer()[i].real(), i); - } - EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[0], 0.0); - EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[1], 0.0); + check_orthonormal(psi_test_cf, 1e-5); + EXPECT_NEAR(elecstate_test.ekb.c[0], 1.0, 1e-5); + EXPECT_NEAR(elecstate_test.ekb.c[1], 1.0, 1e-5); elecstate_test.ekb.c[0] = 1.0; elecstate_test.ekb.c[1] = 2.0; - hs_d_lip.solve(&hamilt_test_d, + hs_d_lip.solve(hamilt_test_d, psi_test_cd, &elecstate_test, transform_test_cd, @@ -423,15 +455,11 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { log, true, 0.0, - 0, - exx_info_local); + 0); EXPECT_DOUBLE_EQ(hsolver::DiagoIterAssist>::avg_iter, 0.0); - for (int i = 0; i < psi_test_cd.size(); i++) - { - EXPECT_DOUBLE_EQ(psi_test_cd.get_pointer()[i].real(), i); - } - EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[0], 0.0); - EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[1], 0.0); + check_orthonormal(psi_test_cd, 1e-10); + EXPECT_NEAR(elecstate_test.ekb.c[0], 1.0, 1e-10); + EXPECT_NEAR(elecstate_test.ekb.c[1], 1.0, 1e-10); } // Test that the program exits with an error when npwx < nbands, @@ -444,7 +472,7 @@ TEST_F(TestHSolverPW, NpwxLessThanNbandsDeath) std::vector eigenvalues(5, 0.0); // Expect death from WARNING_QUIT due to npwx < nbands EXPECT_EXIT( - hamiltSolvePsiK(hs_d, &hamilt_test_d, psi_test_cd, precond, eigenvalues.data(), 1), + hamiltSolvePsiK(hs_d, hamilt_test_d, psi_test_cd, precond, eigenvalues.data(), 1), ::testing::ExitedWithCode(1), ".*" ); diff --git a/source/source_lcao/lcao_set.cpp b/source/source_lcao/lcao_set.cpp index b4ecedc8c2c..8d346ddd881 100644 --- a/source/source_lcao/lcao_set.cpp +++ b/source/source_lcao/lcao_set.cpp @@ -6,6 +6,7 @@ #include "source_hamilt/module_hcontainer/read_hcontainer.h" #include "source_lcao/rho_tau_lcao.h" // use dm2rho #include "source_lcao/hamilt_lcao.h" // use HamiltLCAO for init_chg_hr +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/hsolver_lcao.h" // use HSolverLCAO for init_chg_hr #include "source_pw/module_pwdft/dftu_base.h" // use Plus_U_Base for the DFT+U init @@ -252,7 +253,8 @@ void LCAO_domain::init_chg_hr( PARAM.inp.device == "gpu", GlobalV::NPROC, GlobalV::MY_RANK); - hsolver_lcao_obj.solve(p_hamilt, psi, pelec, dm, chr, nspin, 0); + hamilt::HamiltHSMatrix hs(p_hamilt); + hsolver_lcao_obj.solve(hs, psi, pelec, dm, chr, nspin, 0); } diff --git a/source/source_lcao/module_deltaspin/cal_mw_from_lambda.cpp b/source/source_lcao/module_deltaspin/cal_mw_from_lambda.cpp index a148a0b92fd..3c16185c5e2 100644 --- a/source/source_lcao/module_deltaspin/cal_mw_from_lambda.cpp +++ b/source/source_lcao/module_deltaspin/cal_mw_from_lambda.cpp @@ -7,6 +7,7 @@ #include "source_base/tool_title.h" #include "source_estate/elecstate_tools.h" #include "source_hsolver/diag_comm_info.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/diago_iter_assist.h" #include "source_hsolver/hsolver_lcao.h" #include "source_io/module_parameter/parameter.h" @@ -128,7 +129,8 @@ void spinconstrain::SpinConstrain>::cal_mw_from_lambda( ->update_lambda(); } // Diagonalization without updating charge density (last param = true means skip charge update) - hsolver_t.solve(hamilt_t, psi_t[0], this->pelec, *this->dm_, *this->pelec->charge, this->state_.nspin_, true); + hamilt::HamiltHSMatrix> hs(hamilt_t); + hsolver_t.solve(hs, psi_t[0], this->pelec, *this->dm_, *this->pelec->charge, this->state_.nspin_, true); // Note: although update_lambda() modifies lambda in-place above, // solve() unconditionally recomputes DM and DMR (via cal_dm_psi + // cal_DMR) from the psi obtained by diagonalizing with the new @@ -155,6 +157,7 @@ void spinconstrain::SpinConstrain>::cal_mw_from_lambda( // ============================================================= psi::Psi>* psi_t = static_cast>*>(this->psi); hamilt::Hamilt, base_device::DEVICE_CPU>* hamilt_t = static_cast, base_device::DEVICE_CPU>*>(this->p_hamilt); + hamilt::HamiltHSOperator, base_device::DEVICE_CPU> op(hamilt_t, this->pw_wfc_); auto* onsite_p = projectors::OnsiteProjector::get_instance(); nbands = psi_t->get_nbands(); npol = psi_t->get_npol(); @@ -183,8 +186,8 @@ void spinconstrain::SpinConstrain>::cal_mw_from_lambda( if(initial_hs) { /// Compute H(k) and extract subspace matrices for this k-point - hamilt_t->updateHk(ik); - hsolver::DiagoIterAssist>::cal_hs_subspace(hamilt_t, + op.update_k(ik); + hsolver::DiagoIterAssist>::cal_hs_subspace(op, psi_t[0], h_k, s_k, @@ -214,6 +217,7 @@ void spinconstrain::SpinConstrain>::cal_mw_from_lambda( // ============================================================= psi::Psi, base_device::DEVICE_GPU>* psi_t = static_cast, base_device::DEVICE_GPU>*>(this->psi); hamilt::Hamilt, base_device::DEVICE_GPU>* hamilt_t = static_cast, base_device::DEVICE_GPU>*>(this->p_hamilt); + hamilt::HamiltHSOperator, base_device::DEVICE_GPU> op(hamilt_t, this->pw_wfc_); auto* onsite_p = projectors::OnsiteProjector::get_instance(); nbands = psi_t->get_nbands(); npol = psi_t->get_npol(); @@ -244,9 +248,9 @@ void spinconstrain::SpinConstrain>::cal_mw_from_lambda( std::complex* becp_k = this->pw_cache_.becp_k(ik, size_becp); if(initial_hs) { - hamilt_t->updateHk(ik); + op.update_k(ik); hsolver::DiagoIterAssist, base_device::DEVICE_GPU>::cal_hs_subspace( - hamilt_t, + op, psi_t[0], h_k, s_k, diff --git a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp index 3a8c73d23e0..f1de44fe9ba 100644 --- a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp +++ b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp @@ -30,6 +30,7 @@ #include "mi_tools.h" #include "source_io/module_parameter/parameter.h" #include "source_hsolver/diago_iter_assist.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/hsolver_pw.h" #include "source_estate/elecstate.h" #include "source_estate/elecstate_pw.h" @@ -376,7 +377,8 @@ void update_psi_charge_pw_cpu(ScState& state, PARAM.inp.nb2d, PARAM.inp.use_k_continuity); - hsolver_pw_obj.solve(hamilt_t, + hamilt::HamiltHSOperator, base_device::DEVICE_CPU> op(hamilt_t, pw_wfc); + hsolver_pw_obj.solve(op, psi_t[0], pelec, pelec->ekb.c, @@ -490,7 +492,8 @@ void update_psi_charge_pw_gpu(ScState& state, PARAM.inp.nb2d, PARAM.inp.use_k_continuity); - hsolver_pw_obj.solve(hamilt_t, + hamilt::HamiltHSOperator, base_device::DEVICE_GPU> op(hamilt_t, pw_wfc); + hsolver_pw_obj.solve(op, psi_t[0], pelec, pelec->ekb.c, diff --git a/source/source_lcao/module_lr/hsolver_lrtd.hpp b/source/source_lcao/module_lr/hsolver_lrtd.hpp index 629bf2b9628..bbecbffefd7 100644 --- a/source/source_lcao/module_lr/hsolver_lrtd.hpp +++ b/source/source_lcao/module_lr/hsolver_lrtd.hpp @@ -16,6 +16,23 @@ namespace LR namespace HSolver { + /// The LR Hamiltonians (HamiltLR, HamiltULR) are not hamilt::Hamilt, so + /// they get their own hsolver::HSOperator view. S is the identity. + template + class LRHSOperator : public hsolver::HSOperator + { + public: + explicit LRHSOperator(const THamilt& hm) : hm_(hm) {} + void update_k(const int ik) override {} + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override { hm_.hPsi(x, hx, ld, nvec); } + void spsi(const T* x, T* sx, const int ld, const int nvec) const override + { + std::memcpy(sx, x, sizeof(T) * static_cast(ld) * static_cast(nvec)); + } + private: + const THamilt& hm_; + }; + template inline void print_eigs(const std::vector& eigs, const std::string& label = "", const double factor = 1.0) { @@ -78,12 +95,9 @@ namespace LR } else { - // 3. set maxiter and funcs + // 3. set maxiter and the operator const int maxiter = hsolver::DiagoIterAssist::PW_DIAG_NMAX; - - auto hpsi_func = [&hm](T* psi_in, T* hpsi, const int ld_psi, const int nvec) {hm.hPsi(psi_in, hpsi, ld_psi, nvec);}; - auto spsi_func = [&hm](const T* psi_in, T* spsi, const int ld_psi, const int nvec) - { std::memcpy(spsi, psi_in, sizeof(T) * ld_psi * nvec); }; + const LRHSOperator op(hm); if (method == "dav") { @@ -99,7 +113,7 @@ namespace LR PARAM.inp.pw_diag_ndim, comm_info); std::vector ethr_band(nband, diag_ethr); - hsolver::DiagoIterAssist::avg_iter += static_cast(david.diag(hpsi_func, spsi_func, + hsolver::DiagoIterAssist::avg_iter += static_cast(david.diag(op, dim, psi, eigenvalue.data(), ethr_band, maxiter, ntry_max, 0)); } else if (method == "dav_subspace") //need refactor @@ -115,49 +129,16 @@ namespace LR PARAM.inp.nb2d); std::vector ethr_band(nband, diag_ethr); hsolver::DiagoIterAssist::avg_iter += static_cast( - dav_subspace.diag(hpsi_func, spsi_func, psi, dim, eigenvalue.data(), ethr_band, false /*scf*/)); + dav_subspace.diag(op, psi, dim, eigenvalue.data(), ethr_band, false /*scf*/)); } else if (method == "cg") { - ////// `diagH_subspace` needs refactor: - ////// replace `Hamilt*` with `hpsi_func` - ////// or I cannot use `is_subspace=true` as my `HamiltLR` does not inherit `Hamilt`. - - // auto subspace_func = [&hm](const ct::Tensor& psi_in, ct::Tensor& psi_out) { - // const auto ndim = psi_in.shape().ndim(); - // REQUIRES_OK(ndim == 2, "dims of psi_in should be less than or equal to 2"); - // // Convert a Tensor object to a psi::Psi object - // auto psi_in_wrapper = psi::Psi(psi_in.data(), - // 1, - // psi_in.shape().dim_size(0), - // psi_in.shape().dim_size(1)); - // auto psi_out_wrapper = psi::Psi(psi_out.data(), - // 1, - // psi_out.shape().dim_size(0), - // psi_out.shape().dim_size(1)); - // auto eigen = ct::Tensor(ct::DataTypeToEnum>::value, - // ct::DeviceType::CpuDevice, - // ct::TensorShape({ psi_in.shape().dim_size(0) })); - // hsolver::DiagoIterAssist::diagH_subspace(hm, psi_in_wrapper, psi_out_wrapper, eigen.data>()); - // }; - - ////// why diago_cg depends on basis_type? - // hsolver::DiagoCG cg("lcao", "nscf", true, subspace_func, diag_ethr, maxiter, GlobalV::NPROC_IN_POOL); - - auto subspace_func = [](T* psi_in, T* psi_out, const int ld_psi, const int nband, const bool S_orth) { - }; - hsolver::DiagoCG cg("lcao", "nscf", false, subspace_func, diag_ethr, maxiter, GlobalV::NPROC_IN_POOL); - - auto hpsi_func = [&hm](T* psi_in, T* hpsi, const int ld_psi, const int nvec) { - hm.hPsi(psi_in, hpsi, ld_psi, nvec); - }; - auto spsi_func = [](T* psi_in, T* spsi, const int ld_psi, const int nvec) { - std::memcpy(spsi, psi_in, sizeof(T) * static_cast(ld_psi) * static_cast(nvec)); - }; + // the subspace rotation of DiagoCG now works on any HSOperator, so it could be + // switched on here; it is kept off to leave the LR results unchanged + hsolver::DiagoCG cg("lcao", "nscf", false, comm_info, diag_ethr, maxiter); std::vector ethr_band(nband, diag_ethr); - cg.diag(hpsi_func, - spsi_func, + cg.diag(op, dim, nband, dim, diff --git a/source/source_psi/psi_prepare.cpp b/source/source_psi/psi_prepare.cpp index a0664292eed..353e58adce7 100644 --- a/source/source_psi/psi_prepare.cpp +++ b/source/source_psi/psi_prepare.cpp @@ -11,6 +11,7 @@ #include "source_basis/module_pw/pw_basis_k.h" #include "source_cell/unitcell.h" #include "source_hsolver/diag_comm_info.h" +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_hsolver/diago_iter_assist.h" #include "source_io/module_parameter/parameter.h" #include "source_psi/psi_init_atomic.h" @@ -221,6 +222,9 @@ void PSIPrepare::initialize_psi(Psi>* psi, } } + // the subspace diagonalization sees the Hamiltonian only through this operator + hamilt::HamiltHSOperator op(p_hamilt, &this->pw_wfc); + // loop over kpoints, make it possible to only allocate memory for psig at the only one kpt // like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints for (int ik = 0; ik < this->pw_wfc.nks; ik++) @@ -231,7 +235,7 @@ void PSIPrepare::initialize_psi(Psi>* psi, kspw_psi->fix_k(ik); //! Update Hamiltonian from other kpoint to the given one - p_hamilt->updateHk(ik); + op.update_k(ik); if (fill) { //! initialize psi_cpu @@ -249,7 +253,7 @@ void PSIPrepare::initialize_psi(Psi>* psi, { // for diagH_subspace_init, psi_device->get_pointer() and kspw_psi->get_pointer() should be // different - hsolver::DiagoIterAssist::diag_subspace_init(p_hamilt, + hsolver::DiagoIterAssist::diag_subspace_init(op, psi_device->get_pointer(), nbands_start, nbasis, @@ -262,7 +266,7 @@ void PSIPrepare::initialize_psi(Psi>* psi, else { // for diagH_subspace, psi_device->get_pointer() and kspw_psi->get_pointer() can be the same - hsolver::DiagoIterAssist::diag_subspace(p_hamilt, + hsolver::DiagoIterAssist::diag_subspace(op, *psi_device, *kspw_psi, etatom.data(), diff --git a/source/source_pw/module_pwdft/hamilt_lcaopw.h b/source/source_pw/module_pwdft/hamilt_lcaopw.h index 414c5cddb95..24038ae3505 100644 --- a/source/source_pw/module_pwdft/hamilt_lcaopw.h +++ b/source/source_pw/module_pwdft/hamilt_lcaopw.h @@ -1,6 +1,7 @@ #ifndef HAMILTLIP_H #define HAMILTLIP_H +#include "source_hamilt/hamilt_hs_adapter.h" #include "source_pw/module_pwdft/hamilt_pw.h" #ifdef __EXX #include "source_lcao/module_ri/exx_lip.h" @@ -32,6 +33,52 @@ namespace hamilt #endif }; + /// HamiltLIP seen through hsolver::HSOperator. Besides H and S it feeds the + /// EXX term into the subspace Hamiltonian and hands the subspace + /// eigenvectors back to Exx_Lip, which is what HSolverLIP needs. + template + class HamiltLIPHSOperator : public HamiltHSOperator + { + public: + HamiltLIPHSOperator(HamiltLIP* hm, + const ModulePW::PW_Basis_K* wfc_basis, + const bool cal_exx, + const double hybrid_alpha) + : HamiltHSOperator(hm, wfc_basis), hm_lip_(hm), cal_exx_(cal_exx), + hybrid_alpha_(hybrid_alpha){}; + +#ifdef __EXX + void add_to_subspace_h(T* hcc, const int naos) const override + { + if (!cal_exx_) + { + return; + } + const int ik = this->ik_; + for (int n = 0; n < naos; ++n) + { + for (int m = 0; m < naos; ++m) + { + hcc[n * naos + m] += (T)hybrid_alpha_ * hm_lip_->exx_lip.get_exx_matrix()[ik][m][n]; + } + } + } + + void export_subspace_vec(const T* vcc, const int naos, const int nbands) const override + { + if (cal_exx_) + { + hm_lip_->exx_lip.set_hvec(this->ik_, vcc, naos, nbands); + } + } +#endif + + private: + HamiltLIP* hm_lip_ = nullptr; + const bool cal_exx_; + const double hybrid_alpha_; + }; + } // namespace hamilt #endif From 449af9f5226a5f57b9f4bfebbb36e41486dfffc7 Mon Sep 17 00:00:00 2001 From: Critsium-xy Date: Wed, 16 Sep 2026 22:12:43 +0800 Subject: [PATCH 2/3] Fix CI: keep DiagoCG's subspace step generalized, port pyabacus to HSOperator - DiagoCG: the old subspace_func callback in HSolverPW ignored the S_orth flag and always solved the generalized subspace problem (hegvd). Passing the flag through switched CG restarts to heevx, which changes eigenvector phases and broke the Wannier90 projections of 101_PW_W90. Always solve the generalized problem, as before. - pyabacus: the Davidson adapters still built std::function callbacks for DiagoDavid / Diago_DavSubspace. Replace them with PyHSOperator, an HSOperator over the Python matrix-vector callable (S = identity). Co-Authored-By: Claude Fable 5.1 --- python/pyabacus/src/hsolver/diago_adapter.hpp | 12 ++-- python/pyabacus/src/hsolver/diago_traits.hpp | 62 ++++++++++--------- source/source_hsolver/diago_cg.cpp | 8 ++- 3 files changed, 42 insertions(+), 40 deletions(-) diff --git a/python/pyabacus/src/hsolver/diago_adapter.hpp b/python/pyabacus/src/hsolver/diago_adapter.hpp index d57c8fd82a5..0493e474a69 100644 --- a/python/pyabacus/src/hsolver/diago_adapter.hpp +++ b/python/pyabacus/src/hsolver/diago_adapter.hpp @@ -86,8 +86,7 @@ class PyDiagoDavidAdapter int max_iter, ::hsolver::diag_comm_info comm_info) { - auto hpsi_func = make_hpsi_func_fstyle(mm_op); - auto spsi_func = make_spsi_func_identity(); + const PyHSOperator op(mm_op); solver_ = std::make_unique( precond_vec.data(), @@ -98,8 +97,7 @@ class PyDiagoDavidAdapter ); return solver_->diag( - hpsi_func, - spsi_func, + op, nbasis_, storage_.psi_ptr(), storage_.eigenvalue_ptr(), @@ -178,8 +176,7 @@ class PyDiagoDavSubspaceAdapter int diag_subspace, int nb2d) { - auto hpsi_func = make_hpsi_func_fstyle(mm_op); - auto spsi_func = make_spsi_func_identity(); + const PyHSOperator op(mm_op); solver_ = std::make_unique( precond_vec, @@ -194,8 +191,7 @@ class PyDiagoDavSubspaceAdapter ); return solver_->diag( - hpsi_func, - spsi_func, + op, storage_.psi_ptr(), nbasis_, storage_.eigenvalue_ptr(), diff --git a/python/pyabacus/src/hsolver/diago_traits.hpp b/python/pyabacus/src/hsolver/diago_traits.hpp index c23d9f3a9f2..c73b6f3bea9 100644 --- a/python/pyabacus/src/hsolver/diago_traits.hpp +++ b/python/pyabacus/src/hsolver/diago_traits.hpp @@ -22,6 +22,7 @@ #include #include "source_hsolver/diago_david.h" +#include "source_hsolver/hs_operator.h" #include "source_hsolver/diago_dav_subspace.h" #include "source_hsolver/diago_cg.h" #include "source_base/module_device/memory_op.h" @@ -384,52 +385,53 @@ struct DiagoCGTraits #endif // __ENABLE_ATEN // ============================================================================ -// Helper Functions for Creating HPsi/SPsi Lambdas +// HSOperator view of a Python callable, plus tensor helpers for the ATen CG path // ============================================================================ /** - * @brief Create hpsi_func lambda for raw pointer interface (F-style) + * @brief The Python matrix-vector callable seen through hsolver::HSOperator. * - * Wraps a Python callable to work with ABACUS raw pointer interface. - * Handles array layout conversion between Python (row-major) and - * ABACUS (column-major for Davidson methods). + * H is applied by the Python callable mm_op on a column-major (ld_psi x nvec) + * array; S is the identity. This is all the Davidson-type solvers need. */ template -auto make_hpsi_func_fstyle( - std::function(py::array_t)> mm_op) +class PyHSOperator : public ::hsolver::HSOperator { - return [mm_op](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) { +public: + explicit PyHSOperator(std::function(py::array_t)> mm_op) + : mm_op_(std::move(mm_op)) + { + } + + void update_k(const int ik) override + { + } + + void hpsi(const T* x, T* hx, const int ld, const int nvec) const override + { // Create F-style numpy array (column-major) - py::array_t psi({ld_psi, nvec}); + py::array_t psi({ld, nvec}); py::buffer_info buf = psi.request(); T* ptr = static_cast(buf.ptr); - std::copy(psi_in, psi_in + nvec * ld_psi, ptr); + std::copy(x, x + static_cast(nvec) * ld, ptr); // Call Python function - py::array_t hpsi = mm_op(psi); + py::array_t hpsi_arr = mm_op_(psi); // Copy result back - py::buffer_info hpsi_buf = hpsi.request(); - T* hpsi_ptr = static_cast(hpsi_buf.ptr); - std::copy(hpsi_ptr, hpsi_ptr + nvec * ld_psi, hpsi_out); - }; -} + py::buffer_info hpsi_buf = hpsi_arr.request(); + const T* hpsi_ptr = static_cast(hpsi_buf.ptr); + std::copy(hpsi_ptr, hpsi_ptr + static_cast(nvec) * ld, hx); + } -/** - * @brief Create spsi_func lambda for raw pointer interface (identity) - * - * For non-orthogonal basis, S*psi = psi (identity operation). - */ -template -auto make_spsi_func_identity() -{ - using T = typename Traits::T; - using syncmem_op = typename Traits::syncmem_op; + void spsi(const T* x, T* sx, const int ld, const int nvec) const override + { + std::copy(x, x + static_cast(nvec) * ld, sx); + } - return [](const T* psi_in, T* spsi_out, const int nrow, const int nbands) { - syncmem_op()(spsi_out, psi_in, static_cast(nbands * nrow)); - }; -} +private: + std::function(py::array_t)> mm_op_; +}; #ifdef __ENABLE_ATEN /** diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index d02a1051362..ca60dbcb7b5 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -572,9 +572,13 @@ template void DiagoCG::diag_subspace(const T* psi_in, T* psi_out, const int dim, const int nband, const bool S_orth) { // subspace diagonalization of the current nband vectors, packed with leading dimension dim; - // the eigenvalues it produces are not needed, CG recomputes them + // the eigenvalues it produces are not needed, CG recomputes them. + // The generalized problem is always solved: the S-orthogonal shortcut (heevx instead of + // hegvd) changes eigenvector phases and, for vectors that are only approximately + // S-orthonormal after a CG restart, the results; wavefunction-sensitive outputs such as + // the Wannier90 projections rely on the generalized path. std::vector eigen(nband, 0.0); - DiagoIterAssist::diag_subspace(*op_, psi_in, psi_out, nband, nband, dim, dim, eigen.data(), diag_comm_, S_orth); + DiagoIterAssist::diag_subspace(*op_, psi_in, psi_out, nband, nband, dim, dim, eigen.data(), diag_comm_, false); } template From 5afafff18a9e950a3d9f0f1c77a3bb4149d2877c Mon Sep 17 00:00:00 2001 From: Xiaoyang Zhang Date: Fri, 18 Sep 2026 13:34:37 +0800 Subject: [PATCH 3/3] Drop the dead tpiba/nat parameters from HSolverPW/HSolverLIP::solve Review feedback on #7974: `HSolverLIP::solve` still takes `tpiba` and `nat`, which no code in its body reads. The same two parameters are equally dead in `HSolverPW::solve`; both were left over from an earlier PW/EXX path. Remove them from the declarations, the definitions and every call site (`ESolver_KS_PW`, `ESolver_KS_LIP`, the CPU and GPU deltaspin PW solves, and the `SolveLcaoInPW` unit test). No behaviour change. Co-Authored-By: Claude Opus 5 (1M context) --- source/source_esolver/esolver_ks_lcaopw.cpp | 4 +--- source/source_esolver/esolver_ks_pw.cpp | 4 +--- source/source_hsolver/hsolver_lcaopw.cpp | 4 +--- source/source_hsolver/hsolver_lcaopw.h | 4 +--- source/source_hsolver/hsolver_pw.cpp | 4 +--- source/source_hsolver/hsolver_pw.h | 4 +--- source/source_hsolver/test/test_hsolver_pw.cpp | 8 ++------ source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp | 8 ++------ 8 files changed, 10 insertions(+), 30 deletions(-) diff --git a/source/source_esolver/esolver_ks_lcaopw.cpp b/source/source_esolver/esolver_ks_lcaopw.cpp index 7afdb64cbcf..587beb60f51 100644 --- a/source/source_esolver/esolver_ks_lcaopw.cpp +++ b/source/source_esolver/esolver_ks_lcaopw.cpp @@ -167,9 +167,7 @@ namespace ModuleESolver *this->psi_local, diag_comm, GlobalV::ofs_running, - skip_charge, - ucell.tpiba, - ucell.nat); + skip_charge); // add exx #ifdef __EXX diff --git a/source/source_esolver/esolver_ks_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index 3a1bffbb854..cce89b61005 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -274,9 +274,7 @@ void ESolver_KS_PW::hamilt2rho_single(UnitCell& ucell, const int iste GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, GlobalV::ofs_running, - skip_charge, - ucell.tpiba, - ucell.nat); + skip_charge); } // symmetrize the charge density diff --git a/source/source_hsolver/hsolver_lcaopw.cpp b/source/source_hsolver/hsolver_lcaopw.cpp index 7937124b82c..7c324c22296 100644 --- a/source/source_hsolver/hsolver_lcaopw.cpp +++ b/source/source_hsolver/hsolver_lcaopw.cpp @@ -23,9 +23,7 @@ void HSolverLIP::solve(HSOperator& op, // ESolver_KS_PW::p_hamilt be psi::Psi& transform, const diag_comm_info& diag_comm, std::ostream& log, - const bool skip_charge, - const double tpiba, - const int nat) + const bool skip_charge) { ModuleBase::TITLE("HSolverLIP", "solve"); ModuleBase::timer::start("HSolverLIP", "solve"); diff --git a/source/source_hsolver/hsolver_lcaopw.h b/source/source_hsolver/hsolver_lcaopw.h index 6c2dfc3dd2f..0e5fac0a9fc 100644 --- a/source/source_hsolver/hsolver_lcaopw.h +++ b/source/source_hsolver/hsolver_lcaopw.h @@ -42,9 +42,7 @@ class HSolverLIP psi::Psi& transform, const diag_comm_info& diag_comm, std::ostream& log, - const bool skip_charge, - const double tpiba, - const int nat); + const bool skip_charge); private: ModulePW::PW_Basis_K* wfc_basis = nullptr; diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index a48aff5a031..a11edfd2732 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -70,9 +70,7 @@ void HSolverPW::solve(HSOperator& op, const int rank_in_pool_in, const int nproc_in_pool_in, std::ostream& log, - const bool skip_charge, - const double tpiba, - const int nat) + const bool skip_charge) { ModuleBase::TITLE("HSolverPW", "solve"); ModuleBase::timer::start("HSolverPW", "solve"); diff --git a/source/source_hsolver/hsolver_pw.h b/source/source_hsolver/hsolver_pw.h index 12d3418fcd4..d5cbb772a8c 100644 --- a/source/source_hsolver/hsolver_pw.h +++ b/source/source_hsolver/hsolver_pw.h @@ -66,9 +66,7 @@ class HSolverPW const int rank_in_pool_in, const int nproc_in_pool_in, std::ostream& log, - const bool skip_charge, - const double tpiba, - const int nat); + const bool skip_charge); protected: // diago caller diff --git a/source/source_hsolver/test/test_hsolver_pw.cpp b/source/source_hsolver/test/test_hsolver_pw.cpp index e687d8f330c..694349948d8 100644 --- a/source/source_hsolver/test/test_hsolver_pw.cpp +++ b/source/source_hsolver/test/test_hsolver_pw.cpp @@ -436,9 +436,7 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { transform_test_cf, diag_comm, log, - true, - 0.0, - 0); + true); EXPECT_NE(log.str().find("Average iterative diagonalization steps"), std::string::npos); EXPECT_DOUBLE_EQ(hsolver::DiagoIterAssist>::avg_iter, 0.0); check_orthonormal(psi_test_cf, 1e-5); @@ -453,9 +451,7 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) { transform_test_cd, diag_comm, log, - true, - 0.0, - 0); + true); EXPECT_DOUBLE_EQ(hsolver::DiagoIterAssist>::avg_iter, 0.0); check_orthonormal(psi_test_cd, 1e-10); EXPECT_NEAR(elecstate_test.ekb.c[0], 1.0, 1e-10); diff --git a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp index f1de44fe9ba..15656672fd7 100644 --- a/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp +++ b/source/source_lcao/module_deltaspin/deltaspin_pw_mi.cpp @@ -385,9 +385,7 @@ void update_psi_charge_pw_cpu(ScState& state, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, GlobalV::ofs_running, - false, - state.tpiba, - state.get_nat()); + false); } else { @@ -500,9 +498,7 @@ void update_psi_charge_pw_gpu(ScState& state, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, GlobalV::ofs_running, - false, - state.tpiba, - state.get_nat()); + false); } else {