From 833dcc33bf64516c7fd39106363c9060be854a21 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 1 Apr 2025 15:26:49 +0800 Subject: [PATCH 1/3] Use dm_r in LCAO_Deepks to avoid double counting in DeePKS. --- source/module_esolver/esolver_ks_lcao.cpp | 102 ++++++----- .../hamilt_lcaodft/FORCE_gamma.cpp | 6 +- .../hamilt_lcaodft/FORCE_k.cpp | 4 +- .../operator_lcao/deepks_lcao.cpp | 19 +- .../module_deepks/LCAO_deepks.cpp | 49 ++++++ .../module_deepks/LCAO_deepks.h | 9 + .../module_deepks/LCAO_deepks_interface.cpp | 23 ++- .../module_deepks/deepks_force.cpp | 56 +----- .../module_deepks/deepks_force.h | 2 +- .../module_deepks/deepks_fpre.cpp | 37 +--- .../module_deepks/deepks_fpre.h | 2 +- .../module_deepks/deepks_iterate.h | 3 +- .../module_deepks/deepks_pdm.cpp | 165 ++++++++++++------ .../module_deepks/deepks_pdm.h | 11 +- .../module_deepks/deepks_spre.cpp | 35 +--- .../module_deepks/deepks_spre.h | 2 +- .../module_deepks/test/LCAO_deepks_test.cpp | 84 +++++---- 17 files changed, 328 insertions(+), 281 deletions(-) diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp index 420fb040f74..941e4938be4 100644 --- a/source/module_esolver/esolver_ks_lcao.cpp +++ b/source/module_esolver/esolver_ks_lcao.cpp @@ -30,11 +30,11 @@ #include "module_io/write_wfc_nao.h" #include "module_parameter/parameter.h" -//be careful of hpp, there may be multiple definitions of functions, 20250302, mohan +// be careful of hpp, there may be multiple definitions of functions, 20250302, mohan +#include "module_hamilt_lcao/hamilt_lcaodft/hs_matrix_k.hpp" #include "module_io/write_eband_terms.hpp" #include "module_io/write_vxc.hpp" #include "module_io/write_vxc_r.hpp" -#include "module_hamilt_lcao/hamilt_lcaodft/hs_matrix_k.hpp" //--------------temporary---------------------------- #include "module_base/global_function.h" @@ -179,15 +179,13 @@ void ESolver_KS_LCAO::before_all_runners(UnitCell& ucell, const Input_pa // 7) initialize exact exchange calculations #ifdef __EXX - if (PARAM.inp.calculation == "scf" - || PARAM.inp.calculation == "relax" - || PARAM.inp.calculation == "cell-relax" + if (PARAM.inp.calculation == "scf" || PARAM.inp.calculation == "relax" || PARAM.inp.calculation == "cell-relax" || PARAM.inp.calculation == "md") { if (GlobalC::exx_info.info_global.cal_exx) { if (PARAM.inp.init_wfc != "file") - { // if init_wfc==file, directly enter the EXX loop + { // if init_wfc==file, directly enter the EXX loop XC_Functional::set_xc_first_loop(ucell); } @@ -307,7 +305,6 @@ void ESolver_KS_LCAO::before_all_runners(UnitCell& ucell, const Input_pa return; } - template double ESolver_KS_LCAO::cal_energy() { @@ -316,7 +313,6 @@ double ESolver_KS_LCAO::cal_energy() return this->pelec->f_en.etot; } - template void ESolver_KS_LCAO::cal_force(UnitCell& ucell, ModuleBase::matrix& force) { @@ -460,7 +456,7 @@ void ESolver_KS_LCAO::after_all_runners(UnitCell& ucell) this->pelec->ekb, this->kv); } - } + } // 4) write projected band structure by jiyy-2022-4-20 if (PARAM.inp.out_proj_band) @@ -489,49 +485,50 @@ void ESolver_KS_LCAO::after_all_runners(UnitCell& ucell) if (PARAM.inp.out_mat_xc) { ModuleIO::write_Vxc(PARAM.inp.nspin, - PARAM.globalv.nlocal, - GlobalV::DRANK, - &this->pv, - *this->psi, - ucell, - this->sf, - this->solvent, - *this->pw_rho, - *this->pw_rhod, - this->locpp.vloc, - this->chr, - this->GG, - this->GK, - this->kv, - orb_.cutoffs(), - this->pelec->wg, - this->gd + PARAM.globalv.nlocal, + GlobalV::DRANK, + &this->pv, + *this->psi, + ucell, + this->sf, + this->solvent, + *this->pw_rho, + *this->pw_rhod, + this->locpp.vloc, + this->chr, + this->GG, + this->GK, + this->kv, + orb_.cutoffs(), + this->pelec->wg, + this->gd #ifdef __EXX - , - this->exx_lri_double ? &this->exx_lri_double->Hexxs : nullptr, - this->exx_lri_complex ? &this->exx_lri_complex->Hexxs : nullptr + , + this->exx_lri_double ? &this->exx_lri_double->Hexxs : nullptr, + this->exx_lri_complex ? &this->exx_lri_complex->Hexxs : nullptr #endif ); } if (PARAM.inp.out_mat_xc2) { ModuleIO::write_Vxc_R(PARAM.inp.nspin, - &this->pv, - ucell, - this->sf, - this->solvent, - *this->pw_rho, - *this->pw_rhod, - this->locpp.vloc, - this->chr, - this->GG, - this->GK, - this->kv, - orb_.cutoffs(), - this->gd + &this->pv, + ucell, + this->sf, + this->solvent, + *this->pw_rho, + *this->pw_rhod, + this->locpp.vloc, + this->chr, + this->GG, + this->GK, + this->kv, + orb_.cutoffs(), + this->gd #ifdef __EXX - , this->exx_lri_double ? &this->exx_lri_double->Hexxs : nullptr, - this->exx_lri_complex ? &this->exx_lri_complex->Hexxs : nullptr + , + this->exx_lri_double ? &this->exx_lri_double->Hexxs : nullptr, + this->exx_lri_complex ? &this->exx_lri_complex->Hexxs : nullptr #endif ); } @@ -569,7 +566,6 @@ void ESolver_KS_LCAO::after_all_runners(UnitCell& ucell) ModuleBase::timer::tick("ESolver_KS_LCAO", "after_all_runners"); } - template void ESolver_KS_LCAO::iter_init(UnitCell& ucell, const int istep, const int iter) { @@ -639,11 +635,10 @@ void ESolver_KS_LCAO::iter_init(UnitCell& ucell, const int istep, const if (GlobalC::exx_info.info_global.cal_exx) { // the following steps are only needed in the first outer exx loop - exx_two_level_step = GlobalC::exx_info.info_ri.real_number ? - this->exd->two_level_step - : this->exc->two_level_step; + exx_two_level_step + = GlobalC::exx_info.info_ri.real_number ? this->exd->two_level_step : this->exc->two_level_step; } -#endif +#endif if (iter == 1 && exx_two_level_step == 0) { std::cout << " WAVEFUN -> CHARGE " << std::endl; @@ -742,6 +737,11 @@ void ESolver_KS_LCAO::iter_init(UnitCell& ucell, const int istep, const { this->p_hamilt->refresh(); } + if (iter == 1 && istep == 0) + { + // initialize DMR + this->ld.init_DMR(ucell, orb_, this->pv, this->gd); + } #endif if (PARAM.inp.vl_in_h) @@ -758,7 +758,6 @@ void ESolver_KS_LCAO::iter_init(UnitCell& ucell, const int istep, const } } - template void ESolver_KS_LCAO::hamilt2rho_single(UnitCell& ucell, int istep, int iter, double ethr) { @@ -822,7 +821,6 @@ void ESolver_KS_LCAO::hamilt2rho_single(UnitCell& ucell, int istep, int this->pelec->f_en.deband = this->pelec->cal_delta_eband(ucell); } - template void ESolver_KS_LCAO::update_pot(UnitCell& ucell, const int istep, const int iter, const bool conv_esolver) { @@ -840,7 +838,6 @@ void ESolver_KS_LCAO::update_pot(UnitCell& ucell, const int istep, const } } - template void ESolver_KS_LCAO::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver) { @@ -878,6 +875,7 @@ void ESolver_KS_LCAO::iter_finish(UnitCell& ucell, const int istep, int& = dynamic_cast*>(this->pelec)->get_DM()->get_DMK_vector(); ld.dpks_cal_e_delta_band(dm, this->kv.get_nks()); + DeePKS_domain::update_dmr(this->kv.kvec_d, dm, ucell, orb_, this->pv, this->gd, ld.dm_r); this->pelec->f_en.edeepks_scf = ld.E_delta - ld.e_delta_band; this->pelec->f_en.edeepks_delta = ld.E_delta; } diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp index 2c556ae11e3..ff16bbd81cf 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp @@ -37,7 +37,7 @@ void Force_LCAO::allocate(const UnitCell& ucell, // save the results in dense matrix by now. // pv.nloc: number of H elements in this proc. - assert(pv.nloc>0); + assert(pv.nloc > 0); fsr.DSloc_x = new double[pv.nloc]; fsr.DSloc_y = new double[pv.nloc]; fsr.DSloc_z = new double[pv.nloc]; @@ -230,11 +230,9 @@ void Force_LCAO::ftable(const bool isforce, #ifdef __DEEPKS if (PARAM.inp.deepks_scf) { - const std::vector>& dm_gamma = dm->get_DMK_vector(); - // No need to update E_delta here since it have been done in LCAO_Deepks_Interface in after_scf const int nks = 1; - DeePKS_domain::cal_f_delta(dm_gamma, + DeePKS_domain::cal_f_delta(ld.dm_r, ucell, orb, gd, diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp index 4b89aa93edf..7116c51c3bd 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp @@ -268,10 +268,8 @@ void Force_LCAO>::ftable(const bool isforce, #ifdef __DEEPKS if (PARAM.inp.deepks_scf) { - const std::vector>>& dm_k = dm->get_DMK_vector(); - // No need to update E_delta since it have been done in LCAO_Deepks_Interface in after_scf - DeePKS_domain::cal_f_delta>(dm_k, + DeePKS_domain::cal_f_delta>(ld.dm_r, ucell, orb, gd, diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp index af8455e856e..9ef5bc5e43e 100644 --- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp +++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp @@ -82,8 +82,8 @@ void hamilt::DeePKS>::initialize_HR(const Grid_Driv for (int iat0 = 0; iat0 < ucell->nat; iat0++) { auto tau0 = ucell->get_tau(iat0); - int T0=0; - int I0=0; + int T0 = 0; + int I0 = 0; ucell->iat2iait(iat0, &I0, &T0); AdjacentAtomInfo adjs; GridD->Find_atom(*ucell, tau0, T0, I0, &adjs); @@ -174,7 +174,7 @@ void hamilt::DeePKS>::contributeHR() this->ld->inl2l, this->ld->inl_index, this->kvec_d, - this->DM, + this->ld->dm_r, this->ld->phialpha, *this->ucell, *ptr_orb_, @@ -242,8 +242,8 @@ void hamilt::DeePKS>::pre_calculate_nlm( const Parallel_Orbitals* paraV = this->hR->get_paraV(); const int npol = this->ucell->get_npol(); auto tau0 = ucell->get_tau(iat0); - int T0=0; - int I0=0; + int T0 = 0; + int I0 = 0; ucell->iat2iait(iat0, &I0, &T0); AdjacentAtomInfo& adjs = this->adjs_all[iat0]; nlm_in.resize(adjs.adj_num + 1); @@ -307,8 +307,8 @@ void hamilt::DeePKS>::calculate_HR() for (int iat0 = 0; iat0 < this->ucell->nat; iat0++) { auto tau0 = ucell->get_tau(iat0); - int T0=0; - int I0=0; + int T0 = 0; + int I0 = 0; ucell->iat2iait(iat0, &I0, &T0); AdjacentAtomInfo& adjs = this->adjs_all[iat0]; @@ -370,8 +370,8 @@ void hamilt::DeePKS>::calculate_HR() this->pre_calculate_nlm(iat0, nlm_on_the_fly); } - std::vector>>& nlm_iat = - is_on_the_fly ? nlm_on_the_fly : nlm_tot[iat0]; + std::vector>>& nlm_iat + = is_on_the_fly ? nlm_on_the_fly : nlm_tot[iat0]; // 2. calculate D for each pair of atoms for (int ad1 = 0; ad1 < adjs.adj_num + 1; ++ad1) @@ -500,7 +500,6 @@ void hamilt::DeePKS>::cal_HR_IJR(const double* hr_i } } - // contributeHk() template void hamilt::DeePKS>::contributeHk(int ik) diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp index 1340f04e0c8..d3336355b0a 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp @@ -13,6 +13,7 @@ #ifdef __DEEPKS #include "LCAO_deepks.h" +#include "deepks_iterate.h" #include "module_hamilt_pw/hamilt_pwdft/global.h" // Constructor of the class @@ -206,6 +207,54 @@ void LCAO_Deepks::allocate_V_delta(const int nat, const int nks) return; } +template +void LCAO_Deepks::init_DMR(const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD) +{ + this->dm_r = new hamilt::HContainer(&pv); + DeePKS_domain::iterate_ad2( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt1, + const ModuleBase::Vector3& tau1, + const int start1, + const int nw1_tot, + ModuleBase::Vector3 dR1, + const int ibt2, + const ModuleBase::Vector3& tau2, + const int start2, + const int nw2_tot, + ModuleBase::Vector3 dR2) + { + auto row_indexes = pv.get_indexes_row(ibt1); + auto col_indexes = pv.get_indexes_col(ibt2); + if (row_indexes.size() * col_indexes.size() == 0) + { + return; // to next loop + } + + int dRx = 0; + int dRy = 0; + int dRz = 0; + if constexpr (std::is_same>::value) + { + dRx = (dR1 - dR2).x; + dRy = (dR1 - dR2).y; + dRz = (dR1 - dR2).z; + } + hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); + this->dm_r->insert_pair(dm_pair); + } + ); + this->dm_r->allocate(nullptr, true); +} + template void LCAO_Deepks::dpks_cal_e_delta_band(const std::vector>& dm, const int nks) { diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h index 27eec83ae00..e7bb86f1e82 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.h @@ -86,6 +86,9 @@ class LCAO_Deepks // index 0 for itself and index 1-3 for derivatives over x,y,z std::vector*> phialpha; + // density matrix in real space + hamilt::HContainer* dm_r = nullptr; + // projected density matrix // [tot_Inl][2l+1][2l+1], here l is corresponding to inl; // [nat][nlm*nlm] for equivariant version @@ -135,6 +138,12 @@ class LCAO_Deepks /// Allocate memory for correction to Hamiltonian void allocate_V_delta(const int nat, const int nks = 1); + /// Initialize the dm_r container + void init_DMR(const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD); + //! a temporary interface for cal_e_delta_band void dpks_cal_e_delta_band(const std::vector>& dm, const int nks); diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp index 6a3626712ce..0778e491676 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -53,6 +53,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, bool init_pdm = ld->init_pdm; double E_delta = ld->E_delta; double e_delta_band = ld->e_delta_band; + hamilt::HContainer* dmr = ld->dm_r; const int nspin = PARAM.inp.nspin; @@ -62,8 +63,10 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, { // this part is for integrated test of deepks // so it is printed no matter even if deepks_out_labels is not used + DeePKS_domain::update_dmr(kvec_d, dm->get_DMK_vector(), ucell, orb, *ParaV, GridD, dmr); + DeePKS_domain::cal_pdm< - TK>(init_pdm, inlmax, lmaxd, inl2l, inl_index, kvec_d, dm, phialpha, ucell, orb, GridD, *ParaV, pdm); + TK>(init_pdm, inlmax, lmaxd, inl2l, inl_index, kvec_d, dmr, phialpha, ucell, orb, GridD, *ParaV, pdm); DeePKS_domain::check_pdm(inlmax, inl2l, pdm); // print out the projected dm for NSCF calculaiton @@ -158,10 +161,9 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, if (PARAM.inp.deepks_scf && !PARAM.inp.deepks_equiv) // training with force label not supported by equivariant version now { - std::vector> dm_vec = dm->get_DMK_vector(); torch::Tensor gdmx; DeePKS_domain::cal_gdmx< - TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dm_vec, ucell, orb, *ParaV, GridD, gdmx); + TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dmr, ucell, orb, *ParaV, GridD, gdmx); torch::Tensor gvx; DeePKS_domain::cal_gvx(ucell.nat, inlmax, des_per_atom, inl2l, gevdm, gdmx, gvx, rank); @@ -182,10 +184,9 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, if (PARAM.inp.deepks_scf && !PARAM.inp.deepks_equiv) // training with stress label not supported by equivariant version now { - std::vector> dm_vec = dm->get_DMK_vector(); torch::Tensor gdmepsl; DeePKS_domain::cal_gdmepsl< - TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dm_vec, ucell, orb, *ParaV, GridD, gdmepsl); + TK>(lmaxd, inlmax, nks, kvec_d, phialpha, inl_index, dmr, ucell, orb, *ParaV, GridD, gdmepsl); torch::Tensor gvepsl; DeePKS_domain::cal_gvepsl(ucell.nat, inlmax, des_per_atom, inl2l, gevdm, gdmepsl, gvepsl, rank); @@ -316,7 +317,17 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, torch::Tensor phialpha_r_out; torch::Tensor R_query; - DeePKS_domain::prepare_phialpha_r(nlocal, lmaxd, inlmax, nat, phialpha, ucell, orb, *ParaV, GridD, phialpha_r_out, R_query); + DeePKS_domain::prepare_phialpha_r(nlocal, + lmaxd, + inlmax, + nat, + phialpha, + ucell, + orb, + *ParaV, + GridD, + phialpha_r_out, + R_query); const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy"; const std::string file_R_query = PARAM.globalv.global_out_dir + "deepks_R_query.npy"; LCAO_deepks_io::save_tensor2npy(file_phialpha_r, phialpha_r_out, rank); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp index bd392b6651d..f80f5ddf60d 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp @@ -11,7 +11,7 @@ #include "module_hamilt_lcao/module_hcontainer/atom_pair.h" template -void DeePKS_domain::cal_f_delta(const std::vector>& dm, +void DeePKS_domain::cal_f_delta(const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Grid_Driver& GridD, @@ -107,44 +107,7 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, dRz = dR1.z - dR2.z; } ModuleBase::Vector3 dR(dRx, dRy, dRz); - - hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); - - dm_pair.allocate(nullptr, true); - - if constexpr (std::is_same::value) // for gamma-only - { - for (int is = 0; is < dm.size(); is++) - { - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm[is].data(), pv.get_row_size(), 1.0, 1); - } - else - { - dm_pair.add_from_matrix(dm[is].data(), pv.get_col_size(), 1.0, 0); - } - } - } - else // for multi-k - { - for (int ik = 0; ik < nks; ik++) - { - const double arg = -(kvec_d[ik] * dR) * ModuleBase::TWO_PI; - double sinp = 0; - double cosp = 0; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - const std::complex kphase = std::complex(cosp, sinp); - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_row_size(), kphase, 1); - } - else - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_col_size(), kphase, 0); - } - } - } + const double* dm_current = dmr->find_matrix(ibt1, ibt2, dR.x, dR.y, dR.z)->get_pointer(); hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); @@ -162,7 +125,6 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, assert(overlap_1->get_col_size() == overlap_2->get_col_size()); - const double* dm_current = dm_pair.get_pointer(); for (int iw1 = 0; iw1 < row_indexes.size(); ++iw1) { for (int iw2 = 0; iw2 < col_indexes.size(); ++iw2) @@ -217,7 +179,7 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, } } } - + // HF term is minus, only one projector for each atom force. f_delta_local(iat, 0) -= 2.0 * *dm_current * nlm[0]; f_delta_local(iat, 1) -= 2.0 * *dm_current * nlm[1]; @@ -277,7 +239,7 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, } } } - + for (int ipol = 0; ipol < 3; ipol++) { for (int jpol = ipol; jpol < 3; jpol++) @@ -286,7 +248,6 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, += *dm_current * (nlm[ipol] * r2[jpol] + nlm_t[ipol] * r1[jpol]); } } - } dm_current++; } // iw2 @@ -294,15 +255,14 @@ void DeePKS_domain::cal_f_delta(const std::vector>& dm, } // ad2 } // ad1 } // iat - if(isstress) + if (isstress) { for (int ipol = 0; ipol < 3; ipol++) { for (int jpol = ipol; jpol < 3; jpol++) { #pragma omp atomic - svnl_dalpha(ipol, jpol) - += svnl_dalpha_local(ipol, jpol); + svnl_dalpha(ipol, jpol) += svnl_dalpha_local(ipol, jpol); } } } @@ -364,7 +324,7 @@ void DeePKS_domain::check_f_delta(const int nat, ModuleBase::matrix& f_delta, Mo return; } -template void DeePKS_domain::cal_f_delta(const std::vector>& dm, +template void DeePKS_domain::cal_f_delta(const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Grid_Driver& GridD, @@ -378,7 +338,7 @@ template void DeePKS_domain::cal_f_delta(const std::vector>(const std::vector>>& dm, +template void DeePKS_domain::cal_f_delta>(const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Grid_Driver& GridD, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_force.h b/source/module_hamilt_lcao/module_deepks/deepks_force.h index e42eaec21f0..bcf01087673 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_force.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_force.h @@ -26,7 +26,7 @@ namespace DeePKS_domain // 2. check_f_delta, which prints F_delta into F_delta.dat for checking template -void cal_f_delta(const std::vector>& dm, +void cal_f_delta(const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Grid_Driver& GridD, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp index 61274c3d0d7..c8879f74682 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp @@ -20,7 +20,7 @@ void DeePKS_domain::cal_gdmx(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, @@ -64,7 +64,6 @@ void DeePKS_domain::cal_gdmx(const int lmaxd, return; // to next loop } - double* dm_current = nullptr; int dRx = 0; int dRy = 0; int dRz = 0; @@ -75,35 +74,7 @@ void DeePKS_domain::cal_gdmx(const int lmaxd, dRz = (dR1 - dR2).z; } ModuleBase::Vector3 dR(dRx, dRy, dRz); - - hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); - dm_pair.allocate(nullptr, true); - // not support nspin = 2 now - for (int ik = 0; ik < nks; ik++) - { - TK kphase = TK(0); - if constexpr (std::is_same::value) - { - kphase = 1.0; - } - else - { - const double arg = -(kvec_d[ik] * dR) * ModuleBase::TWO_PI; - double sinp, cosp; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - kphase = TK(cosp, sinp); - } - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_row_size(), kphase, 1); - } - else - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_col_size(), kphase, 0); - } - } - - dm_current = dm_pair.get_pointer(); + const double* dm_current = dmr->find_matrix(ibt1, ibt2, dR.x, dR.y, dR.z)->get_pointer(); hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); if (overlap_1 == nullptr) @@ -353,7 +324,7 @@ template void DeePKS_domain::cal_gdmx(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, @@ -366,7 +337,7 @@ template void DeePKS_domain::cal_gdmx>(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_fpre.h b/source/module_hamilt_lcao/module_deepks/deepks_fpre.h index fcc10d13a94..4a78b8faff0 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_fpre.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_fpre.h @@ -38,7 +38,7 @@ void cal_gdmx(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_iterate.h b/source/module_hamilt_lcao/module_deepks/deepks_iterate.h index a2fba5b072f..0e622b4655d 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_iterate.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_iterate.h @@ -1,6 +1,8 @@ #ifndef DEEPKS_ITER_H #define DEEPKS_ITER_H +#include "module_parameter/parameter.h" + #ifdef __DEEPKS #include "module_base/complexmatrix.h" @@ -12,7 +14,6 @@ #include "module_basis/module_ao/parallel_orbitals.h" #include "module_cell/module_neighbor/sltk_grid_driver.h" #include "module_cell/unitcell.h" -#include "module_parameter/parameter.h" #include diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp index 54b23591da8..75f54b18ac4 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp @@ -15,6 +15,7 @@ #ifdef __DEEPKS +#include "deepks_iterate.h" #include "deepks_pdm.h" #include "module_base/constants.h" #include "module_base/libm/libm.h" @@ -85,6 +86,87 @@ void DeePKS_domain::read_pdm(bool read_pdm_file, } } +template +void DeePKS_domain::update_dmr(const std::vector>& kvec_d, + const std::vector>& dmk, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + hamilt::HContainer* dmr_deepks) +{ + dmr_deepks->set_zero(); + DeePKS_domain::iterate_ad2( + ucell, + GridD, + orb, + false, // no trace_alpha + [&](const int iat, + const ModuleBase::Vector3& tau0, + const int ibt1, + const ModuleBase::Vector3& tau1, + const int start1, + const int nw1_tot, + ModuleBase::Vector3 dR1, + const int ibt2, + const ModuleBase::Vector3& tau2, + const int start2, + const int nw2_tot, + ModuleBase::Vector3 dR2) + { + auto row_indexes = pv.get_indexes_row(ibt1); + auto col_indexes = pv.get_indexes_col(ibt2); + if (row_indexes.size() * col_indexes.size() == 0) + { + return; // to next loop + } + + hamilt::AtomPair dm_pair = dmr_deepks->get_atom_pair(ibt1, ibt2); + + int dRx = 0; + int dRy = 0; + int dRz = 0; + if constexpr (std::is_same>::value) + { + dRx = (dR1 - dR2).x; + dRy = (dR1 - dR2).y; + dRz = (dR1 - dR2).z; + } + ModuleBase::Vector3 dR(dRx, dRy, dRz); + + dm_pair.find_R(dR); + hamilt::BaseMatrix* dmr_ptr = dm_pair.find_matrix(dR); + dmr_ptr->set_zero(); // must reset to zero to avoid accumulation! + + for (int ik = 0; ik < dmk.size(); ik++) + { + TK kphase = TK(0); + if constexpr (std::is_same::value) + { + kphase = 1.0; + } + else + { + const double arg + = -(kvec_d[ik] * ModuleBase::Vector3(dR)) * ModuleBase::TWO_PI; + double sinp, cosp; + ModuleBase::libm::sincos(arg, &sinp, &cosp); + kphase = TK(cosp, sinp); + } + if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) + { + dm_pair.add_from_matrix(dmk[ik].data(), pv.get_row_size(), kphase, 1); + } + else + { + dm_pair.add_from_matrix(dmk[ik].data(), pv.get_col_size(), kphase, 0); + } + } + } + ); + return; +} + // this subroutine performs the calculation of projected density matrices // pdm_m,m'=\sum_{mu,nu} rho_{mu,nu} template @@ -94,7 +176,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm, const std::vector& inl2l, const ModuleBase::IntArray* inl_index, const std::vector>& kvec_d, - const elecstate::DensityMatrix* dm, + const hamilt::HContainer* dmr, const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, @@ -282,42 +364,8 @@ void DeePKS_domain::cal_pdm(bool& init_pdm, } ModuleBase::Vector3 dR(dRx, dRy, dRz); - hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); - dm_pair.allocate(nullptr, true); - auto dm_k = dm->get_DMK_vector(); - - if constexpr (std::is_same::value) // for gamma-only - { - for (int is = 0; is < dm_k.size(); is++) - { - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm_k[is].data(), pv.get_row_size(), 1.0, 1); - } - else - { - dm_pair.add_from_matrix(dm_k[is].data(), pv.get_col_size(), 1.0, 0); - } - } - } - else // for multi-k - { - for (int ik = 0; ik < dm_k.size(); ik++) - { - const double arg = -(kvec_d[ik] * dR) * ModuleBase::TWO_PI; - const std::complex kphase = std::complex(cos(arg), sin(arg)); - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm_k[ik].data(), pv.get_row_size(), kphase, 1); - } - else - { - dm_pair.add_from_matrix(dm_k[ik].data(), pv.get_col_size(), kphase, 0); - } - } - } + const double* dm_current = dmr->find_matrix(ibt1, ibt2, dR.x, dR.y, dR.z)->get_pointer(); - const double* dm_current = dm_pair.get_pointer(); // use s_2t and dm_current to get g_1dmt // dgemm_: C = alpha * A * B + beta * C // C = g_1dmt, A = dm_current, B = s_2t @@ -423,13 +471,29 @@ void DeePKS_domain::check_pdm(const int inlmax, const std::vector& inl2l, c } } +template void DeePKS_domain::update_dmr(const std::vector>& kvec_d, + const std::vector>& dmk, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + hamilt::HContainer* dmr_deepks); + +template void DeePKS_domain::update_dmr>(const std::vector>& kvec_d, + const std::vector>>& dmk, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + hamilt::HContainer* dmr_deepks); + template void DeePKS_domain::cal_pdm(bool& init_pdm, const int inlmax, const int lmaxd, const std::vector& inl2l, const ModuleBase::IntArray* inl_index, const std::vector>& kvec_d, - const elecstate::DensityMatrix* dm, + const hamilt::HContainer* dmr, const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, @@ -437,19 +501,18 @@ template void DeePKS_domain::cal_pdm(bool& init_pdm, const Parallel_Orbitals& pv, std::vector& pdm); -template void DeePKS_domain::cal_pdm>( - bool& init_pdm, - const int inlmax, - const int lmaxd, - const std::vector& inl2l, - const ModuleBase::IntArray* inl_index, - const std::vector>& kvec_d, - const elecstate::DensityMatrix, double>* dm, - const std::vector*> phialpha, - const UnitCell& ucell, - const LCAO_Orbitals& orb, - const Grid_Driver& GridD, - const Parallel_Orbitals& pv, - std::vector& pdm); +template void DeePKS_domain::cal_pdm>(bool& init_pdm, + const int inlmax, + const int lmaxd, + const std::vector& inl2l, + const ModuleBase::IntArray* inl_index, + const std::vector>& kvec_d, + const hamilt::HContainer* dmr, + const std::vector*> phialpha, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Grid_Driver& GridD, + const Parallel_Orbitals& pv, + std::vector& pdm); #endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.h b/source/module_hamilt_lcao/module_deepks/deepks_pdm.h index b89d12b6e35..d0165be8aa3 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_pdm.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.h @@ -39,6 +39,15 @@ void read_pdm(bool read_pdm_file, const Numerical_Orbital& alpha, std::vector& pdm); +template +void update_dmr(const std::vector>& kvec_d, + const std::vector>& dmk, + const UnitCell& ucell, + const LCAO_Orbitals& orb, + const Parallel_Orbitals& pv, + const Grid_Driver& GridD, + hamilt::HContainer* dmr_deepks); + // calculate projected density matrix: pdm = sum_i,occ // 3 cases to skip calculation of pdm: // - NSCF calculation of DeePKS, init_chg = file and pdm has been read @@ -51,7 +60,7 @@ void cal_pdm(bool& init_pdm, const std::vector& inl2l, const ModuleBase::IntArray* inl_index, const std::vector>& kvec_d, - const elecstate::DensityMatrix* dm, + const hamilt::HContainer* dmr, const std::vector*> phialpha, const UnitCell& ucell, const LCAO_Orbitals& orb, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp index 5ec34735d79..80d90ebc58c 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp @@ -21,7 +21,7 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, @@ -72,7 +72,6 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd, return; // to next loop } - double* dm_current = nullptr; int dRx = 0; int dRy = 0; int dRz = 0; @@ -84,33 +83,7 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd, } ModuleBase::Vector3 dR(dRx, dRy, dRz); - hamilt::AtomPair dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv); - dm_pair.allocate(nullptr, 1); - for (int ik = 0; ik < nks; ik++) - { - TK kphase = TK(0); - if constexpr (std::is_same::value) - { - kphase = 1.0; - } - else - { - const double arg = -(kvec_d[ik] * dR) * ModuleBase::TWO_PI; - double sinp, cosp; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - kphase = TK(cosp, sinp); - } - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_row_size(), kphase, 1); - } - else - { - dm_pair.add_from_matrix(dm[ik].data(), pv.get_col_size(), kphase, 0); - } - } - - dm_current = dm_pair.get_pointer(); + const double* dm_current = dmr->find_matrix(ibt1, ibt2, dR.x, dR.y, dR.z)->get_pointer(); hamilt::BaseMatrix* overlap_1 = phialpha[0]->find_matrix(iat, ibt1, dR1); hamilt::BaseMatrix* overlap_2 = phialpha[0]->find_matrix(iat, ibt2, dR2); @@ -322,7 +295,7 @@ template void DeePKS_domain::cal_gdmepsl(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, @@ -335,7 +308,7 @@ template void DeePKS_domain::cal_gdmepsl>(const int lmaxd, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_spre.h b/source/module_hamilt_lcao/module_deepks/deepks_spre.h index 674c6bbf3b2..8b7c7dae675 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_spre.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_spre.h @@ -40,7 +40,7 @@ void cal_gdmepsl( // const ModuleBase::matrix& dm, const std::vector>& kvec_d, std::vector*> phialpha, const ModuleBase::IntArray* inl_index, - const std::vector>& dm, + const hamilt::HContainer* dmr, const UnitCell& ucell, const LCAO_Orbitals& orb, const Parallel_Orbitals& pv, diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp index 2a7af4fefb9..cf6d6b10344 100644 --- a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp +++ b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.cpp @@ -137,19 +137,27 @@ void test_deepks::check_pdm() this->read_dm(kv.nkstot); this->set_dm_new(); this->set_p_elec_DM(); - DeePKS_domain::cal_pdm(this->ld.init_pdm, - this->ld.inlmax, - this->ld.lmaxd, - this->ld.inl2l, - this->ld.inl_index, - kv.kvec_d, - p_elec_DM, - this->ld.phialpha, - ucell, - ORB, - Test_Deepks::GridD, - ParaO, - this->ld.pdm); + this->ld.init_DMR(ucell, ORB, ParaO, Test_Deepks::GridD); + DeePKS_domain::update_dmr(kv.kvec_d, + p_elec_DM->get_DMK_vector(), + ucell, + ORB, + ParaO, + Test_Deepks::GridD, + this->ld.dm_r); + DeePKS_domain::cal_pdm(this->ld.init_pdm, + this->ld.inlmax, + this->ld.lmaxd, + this->ld.inl2l, + this->ld.inl_index, + kv.kvec_d, + this->ld.dm_r, + this->ld.phialpha, + ucell, + ORB, + Test_Deepks::GridD, + ParaO, + this->ld.pdm); DeePKS_domain::check_pdm(this->ld.inlmax, this->ld.inl2l, this->ld.pdm); this->compare_with_ref("pdm.dat", "pdm_ref.dat"); } @@ -157,18 +165,18 @@ void test_deepks::check_pdm() template void test_deepks::check_gdmx(torch::Tensor& gdmx) { - DeePKS_domain::cal_gdmx(this->ld.lmaxd, - this->ld.inlmax, - kv.nkstot, - kv.kvec_d, - this->ld.phialpha, - this->ld.inl_index, - dm_new, - ucell, - ORB, - ParaO, - Test_Deepks::GridD, - gdmx); + DeePKS_domain::cal_gdmx(this->ld.lmaxd, + this->ld.inlmax, + kv.nkstot, + kv.kvec_d, + this->ld.phialpha, + this->ld.inl_index, + this->ld.dm_r, + ucell, + ORB, + ParaO, + Test_Deepks::GridD, + gdmx); DeePKS_domain::check_gdmx(gdmx); for (int ia = 0; ia < ucell.nat; ia++) @@ -199,18 +207,18 @@ void test_deepks::check_gdmx(torch::Tensor& gdmx) template void test_deepks::check_gdmepsl(torch::Tensor& gdmepsl) { - DeePKS_domain::cal_gdmepsl(this->ld.lmaxd, - this->ld.inlmax, - kv.nkstot, - kv.kvec_d, - this->ld.phialpha, - this->ld.inl_index, - dm_new, - ucell, - ORB, - ParaO, - Test_Deepks::GridD, - gdmepsl); + DeePKS_domain::cal_gdmepsl(this->ld.lmaxd, + this->ld.inlmax, + kv.nkstot, + kv.kvec_d, + this->ld.phialpha, + this->ld.inl_index, + this->ld.dm_r, + ucell, + ORB, + ParaO, + Test_Deepks::GridD, + gdmepsl); DeePKS_domain::check_gdmepsl(gdmepsl); for (int i = 0; i < 6; i++) @@ -382,7 +390,7 @@ void test_deepks::check_f_delta_and_stress_delta() svnl_dalpha.create(3, 3); const int cal_stress = 1; const int nks = kv.nkstot; - DeePKS_domain::cal_f_delta(dm_new, + DeePKS_domain::cal_f_delta(this->ld.dm_r, ucell, ORB, Test_Deepks::GridD, From f05981e1a96066f6d0746ee85b3de33d197188c3 Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 1 Apr 2025 20:52:59 +0800 Subject: [PATCH 2/3] Remove 'if constexpr' usage in DeePKS. --- .../module_deepks/LCAO_deepks.cpp | 4 +- .../module_deepks/LCAO_deepks_interface.cpp | 34 ++---- .../module_deepks/LCAO_deepks_io.cpp | 15 +-- .../module_deepks/deepks_force.cpp | 2 +- .../module_deepks/deepks_fpre.cpp | 2 +- .../module_deepks/deepks_orbital.cpp | 35 ++---- .../module_deepks/deepks_orbital.h | 3 +- .../module_deepks/deepks_orbpre.cpp | 92 ++++++--------- .../module_deepks/deepks_pdm.cpp | 24 ++-- .../module_deepks/deepks_spre.cpp | 2 +- .../module_deepks/deepks_vdelta.cpp | 26 ++--- .../module_deepks/deepks_vdelta.h | 1 + .../module_deepks/deepks_vdpre.cpp | 108 +++++------------- 13 files changed, 114 insertions(+), 234 deletions(-) diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp index d3336355b0a..506dc4844bb 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks.cpp @@ -242,7 +242,7 @@ void LCAO_Deepks::init_DMR(const UnitCell& ucell, int dRx = 0; int dRy = 0; int dRz = 0; - if constexpr (std::is_same>::value) + if (std::is_same>::value) { dRx = (dR1 - dR2).x; dRy = (dR1 - dR2).y; @@ -258,7 +258,7 @@ void LCAO_Deepks::init_DMR(const UnitCell& ucell, template void LCAO_Deepks::dpks_cal_e_delta_band(const std::vector>& dm, const int nks) { - DeePKS_domain::cal_e_delta_band(dm, this->V_delta, nks, this->pv, this->e_delta_band); + DeePKS_domain::cal_e_delta_band(dm, this->V_delta, nks, PARAM.inp.nspin, this->pv, this->e_delta_band); } template class LCAO_Deepks; diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp index 0778e491676..ffdc80513c9 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp @@ -56,6 +56,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, hamilt::HContainer* dmr = ld->dm_r; const int nspin = PARAM.inp.nspin; + const int nk = nks / nspin; // Note : update PDM and all other quantities with the current dm // DeePKS PDM and descriptor @@ -221,30 +222,15 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, std::vector dm_bandgap; // Calculate O_delta - if constexpr (std::is_same::value) // for gamma only + wg_hl.create(nks, PARAM.inp.nbands); + dm_bandgap.resize(nks); + wg_hl.zero_out(); + for (int iks = 0; iks < nks; ++iks) { - wg_hl.create(nspin, PARAM.inp.nbands); - dm_bandgap.resize(nspin); - for (int is = 0; is < nspin; ++is) - { - wg_hl.zero_out(); - wg_hl(is, nocc - 1) = -1.0; - wg_hl(is, nocc) = 1.0; - elecstate::cal_dm(ParaV, wg_hl, psi, dm_bandgap); - } - } - else // for multi-k - { - wg_hl.create(nks, PARAM.inp.nbands); - dm_bandgap.resize(nks); - wg_hl.zero_out(); - for (int ik = 0; ik < nks; ik++) - { - wg_hl(ik, nocc - 1) = -1.0; - wg_hl(ik, nocc) = 1.0; - } - elecstate::cal_dm(ParaV, wg_hl, psi, dm_bandgap); + wg_hl(iks, nocc - 1) = -1.0; + wg_hl(iks, nocc) = 1.0; } + elecstate::cal_dm(ParaV, wg_hl, psi, dm_bandgap); ModuleBase::matrix o_delta(nks, 1); @@ -265,7 +251,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, *ParaV, GridD, orbital_precalc); - DeePKS_domain::cal_o_delta(dm_bandgap, *h_delta, o_delta, *ParaV, nks); + DeePKS_domain::cal_o_delta(dm_bandgap, *h_delta, o_delta, *ParaV, nks, nspin); // save obase and orbital_precalc const std::string file_orbpre = PARAM.globalv.global_out_dir + "deepks_orbpre.npy"; @@ -431,7 +417,7 @@ void LCAO_Deepks_Interface::out_deepks_labels(const double& etot, /// print out deepks information to the screen if (PARAM.inp.deepks_scf) { - DeePKS_domain::cal_e_delta_band(dm->get_DMK_vector(), *h_delta, nks, ParaV, e_delta_band); + DeePKS_domain::cal_e_delta_band(dm->get_DMK_vector(), *h_delta, nks, nspin, ParaV, e_delta_band); std::cout << "E_delta_band = " << std::setprecision(8) << e_delta_band << " Ry" << " = " << std::setprecision(8) << e_delta_band * ModuleBase::Ry_to_eV << " eV" << std::endl; std::cout << "E_delta_NN = " << std::setprecision(8) << E_delta << " Ry" diff --git a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp index 2be3a943f8c..2620b9a0d61 100644 --- a/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp +++ b/source/module_hamilt_lcao/module_deepks/LCAO_deepks_io.cpp @@ -266,6 +266,7 @@ void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch:: { return; } + using T_tensor = typename std::conditional>::value, c10::complex, T>::type; const int dim = tensor.dim(); std::vector shape(dim); for (int i = 0; i < dim; i++) @@ -275,17 +276,11 @@ void LCAO_deepks_io::save_tensor2npy(const std::string& file_name, const torch:: std::vector data(tensor.numel()); - if constexpr (std::is_same>::value) + T_tensor* data_ptr_tensor = tensor.data_ptr(); + T* data_ptr = reinterpret_cast(data_ptr_tensor); + for (size_t i = 0; i < tensor.numel(); ++i) { - auto tensor_data = tensor.data_ptr>(); - for (size_t i = 0; i < tensor.numel(); ++i) - { - data[i] = std::complex(tensor_data[i].real(), tensor_data[i].imag()); - } - } - else - { - std::memcpy(data.data(), tensor.data_ptr(), tensor.numel() * sizeof(T)); + data[i] = data_ptr[i]; } npy::SaveArrayAsNumpy(file_name, false, shape.size(), shape.data(), data); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp index f80f5ddf60d..d91ede9c0ad 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_force.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_force.cpp @@ -100,7 +100,7 @@ void DeePKS_domain::cal_f_delta(const hamilt::HContainer* dmr, int dRx = 0; int dRy = 0; int dRz = 0; - if constexpr (std::is_same>::value) // for multi-k + if (std::is_same>::value) // for multi-k { dRx = dR1.x - dR2.x; dRy = dR1.y - dR2.y; diff --git a/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp index c8879f74682..ef02b046baf 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_fpre.cpp @@ -67,7 +67,7 @@ void DeePKS_domain::cal_gdmx(const int lmaxd, int dRx = 0; int dRy = 0; int dRz = 0; - if constexpr (std::is_same>::value) + if (std::is_same>::value) { dRx = (dR1 - dR2).x; dRy = (dR1 - dR2).y; diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp index 8baf298a462..8ccddeda47f 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbital.cpp @@ -12,12 +12,13 @@ void DeePKS_domain::cal_o_delta(const std::vector& dm_hl, // std::vector& o_delta, ModuleBase::matrix& o_delta, const Parallel_Orbitals& pv, - const int nks) + const int nks, + const int nspin) { ModuleBase::TITLE("DeePKS_domain", "cal_o_delta"); ModuleBase::timer::tick("DeePKS_domain", "cal_o_delta"); - for (int ik = 0; ik < nks; ik++) + for (int ik = 0; ik < nks / nspin; ik++) { TK o_delta_tmp = TK(0.0); for (int i = 0; i < PARAM.globalv.nlocal; ++i) @@ -39,31 +40,17 @@ void DeePKS_domain::cal_o_delta(const std::vector& dm_hl, { iic = mu * pv.ncol + nu; } - if constexpr (std::is_same::value) + for (int is = 0; is < nspin; is++) { - for (int is = 0; is < PARAM.inp.nspin; ++is) - { - o_delta_tmp += dm_hl[is](nu, mu) * h_delta[ik][iic]; - } - } - else - { - o_delta_tmp += dm_hl[ik](nu, mu) * h_delta[ik][iic]; + o_delta_tmp += dm_hl[ik + is * nks / nspin](nu, mu) * h_delta[ik][iic]; } } } } Parallel_Reduce::reduce_all(o_delta_tmp); - if constexpr (std::is_same::value) - { - // o_delta[ik] = o_delta_tmp; - o_delta(ik, 0) = o_delta_tmp; - } - else - { - // o_delta[ik] = o_delta_tmp.real(); - o_delta(ik, 0) = o_delta_tmp.real(); - } + + const double* o_delta_ptr = reinterpret_cast(&o_delta_tmp); + o_delta(ik, 0) = o_delta_ptr[0]; // real part in complex case } ModuleBase::timer::tick("DeePKS_domain", "cal_o_delta"); return; @@ -74,7 +61,8 @@ template void DeePKS_domain::cal_o_delta(const std:: // std::vector& o_delta, ModuleBase::matrix& o_delta, const Parallel_Orbitals& pv, - const int nks); + const int nks, + const int nspin); template void DeePKS_domain::cal_o_delta, ModuleBase::ComplexMatrix>( const std::vector& dm_hl, @@ -82,6 +70,7 @@ template void DeePKS_domain::cal_o_delta, ModuleBase::Compl // std::vector& o_delta, ModuleBase::matrix& o_delta, const Parallel_Orbitals& pv, - const int nks); + const int nks, + const int nspin); #endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbital.h b/source/module_hamilt_lcao/module_deepks/deepks_orbital.h index 6c60fd1e96a..eb9e52360da 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbital.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbital.h @@ -29,7 +29,8 @@ void cal_o_delta(const std::vector& dm_hl, // std::vector& o_delta, ModuleBase::matrix& o_delta, const Parallel_Orbitals& pv, - const int nks); + const int nks, + const int nspin); } // namespace DeePKS_domain #endif diff --git a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp index 7d1631a92c1..e146a08dc61 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_orbpre.cpp @@ -126,12 +126,10 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, const int T2 = GridD.getType(ad2); const int I2 = GridD.getNatom(ad2); const int ibt2 = ucell.itia2iat(T2, I2); - if constexpr (std::is_same>::value) // Why only for multi-k? + // skip if ibt1 > ibt2 and set gemm_alpha = 2.0 later, for performance + if (ibt1 > ibt2) { - if (ibt1 > ibt2) - { - continue; - } + continue; } const ModuleBase::Vector3 tau2 = GridD.getAdjacentTau(ad2); const Atom* atom2 = &ucell.atoms[T2]; @@ -169,71 +167,51 @@ void DeePKS_domain::cal_orbital_precalc(const std::vector& dm_hl, } } - std::vector dm_array(row_size * nks * col_size, 0.0); + std::vector dm_array(row_size * dm_hl.size() * col_size, 0.0); - const int row_size_nks = row_size * nks; + const int row_size_nks = row_size * dm_hl.size(); - if constexpr (std::is_same::value) + int dRx = 0; + int dRy = 0; + int dRz = 0; + if (std::is_same>::value) { - for (int is = 0; is < PARAM.inp.nspin; is++) - { - hamilt::AtomPair dm_pair(ibt1, ibt2, 0, 0, 0, &pv); + dRx = dR1.x - dR2.x; + dRy = dR1.y - dR2.y; + dRz = dR1.z - dR2.z; + } + ModuleBase::Vector3 dR(dRx, dRy, dRz); - dm_pair.allocate(dm_array.data(), 0); + hamilt::AtomPair dm_pair(ibt1, ibt2, dR.x, dR.y, dR.z, &pv); - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm_hl[is].c, pv.get_row_size(), 1.0, 1); - } - else - { - dm_pair.add_from_matrix(dm_hl[is].c, pv.get_col_size(), 1.0, 0); - } - } // is - } - else + for (int ik = 0; ik < dm_hl.size(); ik++) { - for (int ik = 0; ik < nks; ik++) + dm_pair.allocate(&dm_array[ik * row_size * col_size], 0); + + std::complex kphase = std::complex(1, 0); + if (std::is_same>::value) { - hamilt::AtomPair dm_pair(ibt1, - ibt2, - (dR1 - dR2).x, - (dR1 - dR2).y, - (dR1 - dR2).z, - &pv); - - dm_pair.allocate(&dm_array[ik * row_size * col_size], 0); - - const double arg - = -(kvec_d[ik] * ModuleBase::Vector3(dR1 - dR2)) * ModuleBase::TWO_PI; - - double sinp, cosp; - - ModuleBase::libm::sincos(arg, &sinp, &cosp); - - const std::complex kphase = std::complex(cosp, sinp); + const double arg = -(kvec_d[ik] * ModuleBase::Vector3(dR1 - dR2)) * ModuleBase::TWO_PI; + kphase = std::complex(cos(arg), sin(arg)); + } + TK* kphase_ptr = reinterpret_cast(&kphase); - if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) - { - dm_pair.add_from_matrix(dm_hl[ik].c, pv.get_row_size(), kphase, 1); - } - else - { - dm_pair.add_from_matrix(dm_hl[ik].c, pv.get_col_size(), kphase, 0); - } - } // ik - } + if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) + { + dm_pair.add_from_matrix(dm_hl[ik].c, pv.get_row_size(), *kphase_ptr, 1); + } + else + { + dm_pair.add_from_matrix(dm_hl[ik].c, pv.get_col_size(), *kphase_ptr, 0); + } + } // ik // dgemm for s_2t and dm_array to get g_1dmt constexpr char transa = 'T', transb = 'N'; double gemm_alpha = 1.0, gemm_beta = 1.0; - - if constexpr (std::is_same>::value) + if (ibt1 < ibt2) { - if (ibt1 < ibt2) - { - gemm_alpha = 2.0; - } + gemm_alpha = 2.0; } dgemm_(&transa, diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp index 75f54b18ac4..a82cd10be09 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp @@ -126,7 +126,7 @@ void DeePKS_domain::update_dmr(const std::vector>& k int dRx = 0; int dRy = 0; int dRz = 0; - if constexpr (std::is_same>::value) + if (std::is_same>::value) { dRx = (dR1 - dR2).x; dRy = (dR1 - dR2).y; @@ -140,26 +140,20 @@ void DeePKS_domain::update_dmr(const std::vector>& k for (int ik = 0; ik < dmk.size(); ik++) { - TK kphase = TK(0); - if constexpr (std::is_same::value) + std::complex kphase = std::complex(1, 0); + if (std::is_same>::value) { - kphase = 1.0; - } - else - { - const double arg - = -(kvec_d[ik] * ModuleBase::Vector3(dR)) * ModuleBase::TWO_PI; - double sinp, cosp; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - kphase = TK(cosp, sinp); + const double arg = -(kvec_d[ik] * ModuleBase::Vector3(dR)) * ModuleBase::TWO_PI; + kphase = std::complex(cos(arg), sin(arg)); } + TK* kphase_ptr = reinterpret_cast(&kphase); if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER(PARAM.inp.ks_solver)) { - dm_pair.add_from_matrix(dmk[ik].data(), pv.get_row_size(), kphase, 1); + dm_pair.add_from_matrix(dmk[ik].data(), pv.get_row_size(), *kphase_ptr, 1); } else { - dm_pair.add_from_matrix(dmk[ik].data(), pv.get_col_size(), kphase, 0); + dm_pair.add_from_matrix(dmk[ik].data(), pv.get_col_size(), *kphase_ptr, 0); } } } @@ -356,7 +350,7 @@ void DeePKS_domain::cal_pdm(bool& init_pdm, } // prepare DM from DMR int dRx = 0, dRy = 0, dRz = 0; - if constexpr (std::is_same>::value) + if (std::is_same>::value) { dRx = dR1.x - dR2.x; dRy = dR1.y - dR2.y; diff --git a/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp index 80d90ebc58c..713fa05b12a 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_spre.cpp @@ -75,7 +75,7 @@ void DeePKS_domain::cal_gdmepsl(const int lmaxd, int dRx = 0; int dRy = 0; int dRz = 0; - if constexpr (std::is_same>::value) + if (std::is_same>::value) { dRx = (dR1 - dR2).x; dRy = (dR1 - dR2).y; diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp index f060ada40b6..335b7246bda 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.cpp @@ -21,6 +21,7 @@ template void DeePKS_domain::cal_e_delta_band(const std::vector>& dm, const std::vector>& V_delta, const int nks, + const int nspin, const Parallel_Orbitals* pv, double& e_delta_band) { @@ -45,31 +46,18 @@ void DeePKS_domain::cal_e_delta_band(const std::vector>& dm, { iic = mu * pv->ncol + nu; } - if constexpr (std::is_same::value) + for (int is = 0; is < nspin; is++) { - for (int is = 0; is < dm.size(); ++is) // dm.size() == PARAM.inp.nspin + for (int ik = 0; ik < nks / nspin; ik++) { - e_delta_band_tmp += dm[is][nu * pv->nrow + mu] * V_delta[0][iic]; - } - } - else - { - for (int ik = 0; ik < nks; ik++) - { - e_delta_band_tmp += dm[ik][nu * pv->nrow + mu] * V_delta[ik][iic]; + e_delta_band_tmp += dm[ik + is * nks / nspin][nu * pv->nrow + mu] * V_delta[ik][iic]; } } } } } - if constexpr (std::is_same::value) - { - e_delta_band = e_delta_band_tmp; - } - else - { - e_delta_band = e_delta_band_tmp.real(); - } + const double* e_delta_band_ptr = reinterpret_cast(&e_delta_band_tmp); + e_delta_band = e_delta_band_ptr[0]; // real part in complex case #ifdef __MPI Parallel_Reduce::reduce_all(e_delta_band); #endif @@ -147,12 +135,14 @@ void DeePKS_domain::collect_h_mat(const Parallel_Orbitals& pv, template void DeePKS_domain::cal_e_delta_band(const std::vector>& dm, const std::vector>& V_delta, const int nks, + const int nspin, const Parallel_Orbitals* pv, double& e_delta_band); template void DeePKS_domain::cal_e_delta_band>( const std::vector>>& dm, const std::vector>>& V_delta, const int nks, + const int nspin, const Parallel_Orbitals* pv, double& e_delta_band); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h index dab6bf2c13b..650119d0c32 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h +++ b/source/module_hamilt_lcao/module_deepks/deepks_vdelta.h @@ -21,6 +21,7 @@ template void cal_e_delta_band(const std::vector>& dm, const std::vector>& V_delta, const int nks, + const int nspin, const Parallel_Orbitals* pv, double& e_delta_band); diff --git a/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp b/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp index 6189c8acf6b..1cb89b1748e 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_vdpre.cpp @@ -44,6 +44,7 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal, // gettimeofday(&t_start,NULL); constexpr torch::Dtype dtype = std::is_same::value ? torch::kFloat64 : torch::kComplexDouble; + using TK_tensor = typename std::conditional>::value, c10::complex, TK>::type; torch::Tensor v_delta_pdm = torch::zeros({nks, nlocal, nlocal, inlmax, (2 * lmaxd + 1), (2 * lmaxd + 1)}, torch::dtype(dtype)); @@ -101,14 +102,13 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal, { int ib = 0; std::complex kphase = std::complex(1.0, 0.0); - if constexpr (std::is_same>::value) + if (std::is_same>::value) { const double arg = -(kvec_d[ik] * ModuleBase::Vector3(dR1 - dR2)) * ModuleBase::TWO_PI; - double sinp, cosp; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - kphase = std::complex(cosp, sinp); + kphase = std::complex(cos(arg), sin(arg)); } + TK_tensor* kpase_ptr = reinterpret_cast(&kphase); for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) { for (int N0 = 0; N0 < orb.Alpha[0].getNchi(L0); ++N0) @@ -119,19 +119,8 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal, { for (int m2 = 0; m2 < nm; ++m2) // nm = 1 for s, 3 for p, 5 for d { - if constexpr (std::is_same::value) - { - accessor[ik][iw1][iw2][inl][m1][m2] += overlap_1->get_value(iw1, ib + m1) - * overlap_2->get_value(iw2, ib + m2); - } - else - { - c10::complex tmp; - tmp = overlap_1->get_value(iw1, ib + m1) - * overlap_2->get_value(iw2, ib + m2) - * kphase; // from std::complex to c10::complex - accessor[ik][iw1][iw2][inl][m1][m2] += tmp; - } + TK_tensor tmp = overlap_1->get_value(iw1, ib + m1) * overlap_2->get_value(iw2, ib + m2) * *kpase_ptr; + accessor[ik][iw1][iw2][inl][m1][m2] += tmp; } } ib += nm; @@ -144,16 +133,8 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal, ); #ifdef __MPI const int size = nks * nlocal * nlocal * inlmax * (2 * lmaxd + 1) * (2 * lmaxd + 1); - TK* data_ptr; - if constexpr (std::is_same::value) - { - data_ptr = v_delta_pdm.data_ptr(); - } - else - { - c10::complex* c10_ptr = v_delta_pdm.data_ptr>(); - data_ptr = reinterpret_cast(c10_ptr); - } + TK_tensor* data_tensor_ptr = v_delta_pdm.data_ptr(); + TK* data_ptr = reinterpret_cast(data_tensor_ptr); Parallel_Reduce::reduce_all(data_ptr, size); #endif @@ -183,16 +164,9 @@ void DeePKS_domain::cal_v_delta_precalc(const int nlocal, { for (int m2 = 0; m2 < nm; ++m2) // m1 = 1 for s, 3 for p, 5 for d { - if constexpr (std::is_same::value) - { - mmv.push_back(accessor[iks][mu][nu][inl][m1][m2]); - } - else - { - c10::complex tmp_c10 = accessor[iks][mu][nu][inl][m1][m2]; - std::complex tmp = std::complex(tmp_c10.real(), tmp_c10.imag()); - mmv.push_back(tmp); - } + TK_tensor tmp = accessor[iks][mu][nu][inl][m1][m2]; + TK* tmp_ptr = reinterpret_cast(&tmp); + mmv.push_back(*tmp_ptr); } } torch::Tensor mm = torch::from_blob(mmv.data(), @@ -241,6 +215,7 @@ void DeePKS_domain::check_v_delta_precalc(const int nat, const int des_per_atom, const torch::Tensor& v_delta_precalc) { + using TK_tensor = typename std::conditional>::value, c10::complex, TK>::type; std::ofstream ofs("v_delta_precalc.dat"); ofs << std::setprecision(10); auto accessor @@ -256,16 +231,9 @@ void DeePKS_domain::check_v_delta_precalc(const int nat, { for (int p = 0; p < des_per_atom; ++p) { - if constexpr (std::is_same::value) - { - ofs << accessor[iks][mu][nu][iat][p] << " "; - } - else - { - c10::complex tmp_c10 = accessor[iks][mu][nu][iat][p]; - std::complex tmp = std::complex(tmp_c10.real(), tmp_c10.imag()); - ofs << tmp << " "; - } + TK_tensor tmp = accessor[iks][mu][nu][iat][p]; + TK* tmp_ptr = reinterpret_cast(&tmp); + ofs << *tmp_ptr << " "; } } ofs << std::endl; @@ -293,6 +261,7 @@ void DeePKS_domain::prepare_phialpha(const int nlocal, ModuleBase::TITLE("DeePKS_domain", "prepare_phialpha"); ModuleBase::timer::tick("DeePKS_domain", "prepare_phialpha"); constexpr torch::Dtype dtype = std::is_same::value ? torch::kFloat64 : torch::kComplexDouble; + using TK_tensor = typename std::conditional>::value, c10::complex, TK>::type; int nlmax = inlmax / nat; int mmax = 2 * lmaxd + 1; phialpha_out = torch::zeros({nat, nlmax, nks, nlocal, mmax}, dtype); @@ -332,13 +301,12 @@ void DeePKS_domain::prepare_phialpha(const int nlocal, for (int ik = 0; ik < nks; ik++) { std::complex kphase = std::complex(1.0, 0.0); - if constexpr (std::is_same>::value) + if (std::is_same>::value) { const double arg = -(kvec_d[ik] * ModuleBase::Vector3(dR)) * ModuleBase::TWO_PI; - double sinp, cosp; - ModuleBase::libm::sincos(arg, &sinp, &cosp); - kphase = std::complex(cosp, sinp); + kphase = std::complex(cos(arg), sin(arg)); } + TK_tensor* kpase_ptr = reinterpret_cast(&kphase); int ib = 0; int nl = 0; for (int L0 = 0; L0 <= orb.Alpha[0].getLmax(); ++L0) @@ -348,16 +316,8 @@ void DeePKS_domain::prepare_phialpha(const int nlocal, const int nm = 2 * L0 + 1; for (int m1 = 0; m1 < nm; ++m1) // nm = 1 for s, 3 for p, 5 for d { - if constexpr (std::is_same::value) - { - accessor[iat][nl][ik][iw1_all][m1] = overlap->get_value(iw1, ib + m1); - } - else - { - c10::complex tmp; - tmp = overlap->get_value(iw1, ib + m1) * kphase; - accessor[iat][nl][ik][iw1_all][m1] += tmp; - } + TK_tensor tmp = overlap->get_value(iw1, ib + m1) * *kpase_ptr; + accessor[iat][nl][ik][iw1_local][m1] += tmp; } ib += nm; nl++; @@ -370,16 +330,8 @@ void DeePKS_domain::prepare_phialpha(const int nlocal, #ifdef __MPI int size = nat * nlmax * nks * nlocal * mmax; - TK* data_ptr; - if constexpr (std::is_same::value) - { - data_ptr = phialpha_out.data_ptr(); - } - else - { - c10::complex* c10_ptr = phialpha_out.data_ptr>(); - data_ptr = reinterpret_cast(c10_ptr); - } + TK_tensor* data_tensor_ptr = phialpha_out.data_ptr(); + TK* data_ptr = reinterpret_cast(data_tensor_ptr); Parallel_Reduce::reduce_all(data_ptr, size); #endif @@ -396,6 +348,7 @@ void DeePKS_domain::check_vdp_phialpha(const int nat, const int lmaxd, const torch::Tensor& phialpha_out) { + using TK_tensor = typename std::conditional>::value, c10::complex, TK>::type; std::ofstream ofs("vdp_phialpha.dat"); ofs << std::setprecision(10); auto accessor @@ -413,16 +366,9 @@ void DeePKS_domain::check_vdp_phialpha(const int nat, { for (int m = 0; m < mmax; m++) { - if constexpr (std::is_same::value) - { - ofs << accessor[iat][nl][iks][mu][m] << " "; - } - else - { - c10::complex tmp_c10 = accessor[iat][nl][iks][mu][m]; - std::complex tmp = std::complex(tmp_c10.real(), tmp_c10.imag()); - ofs << tmp << " "; - } + TK_tensor tmp = accessor[iat][nl][iks][mu][m]; + TK* tmp_ptr = reinterpret_cast(&tmp); + ofs << *tmp_ptr << " "; } } } From a44c255771cf3db1dea881d94aa1b751076206cf Mon Sep 17 00:00:00 2001 From: ErjieWu Date: Tue, 1 Apr 2025 21:14:56 +0800 Subject: [PATCH 3/3] Fix a merge bug. --- source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp index a82cd10be09..55975ddc9f1 100644 --- a/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp +++ b/source/module_hamilt_lcao/module_deepks/deepks_pdm.cpp @@ -121,7 +121,7 @@ void DeePKS_domain::update_dmr(const std::vector>& k return; // to next loop } - hamilt::AtomPair dm_pair = dmr_deepks->get_atom_pair(ibt1, ibt2); + hamilt::AtomPair dm_pair = dmr_deepks->get_atom_pair(ibt1, ibt2); int dRx = 0; int dRy = 0;