diff --git a/source/module_lr/hamilt_casida.h b/source/module_lr/hamilt_casida.h index a005f976a1..81f456851d 100644 --- a/source/module_lr/hamilt_casida.h +++ b/source/module_lr/hamilt_casida.h @@ -1,9 +1,11 @@ #pragma once +#include #include "module_hamilt_general/hamilt.h" #include "module_elecstate/module_dm/density_matrix.h" #include "module_lr/operator_casida/operator_lr_diag.h" #include "module_lr/operator_casida/operator_lr_hxc.h" #include "module_basis/module_ao/parallel_orbitals.h" +#include "module_lr/dm_trans/dm_trans.h" #ifdef __EXX #include "module_lr/operator_casida/operator_lr_exx.h" #include "module_lr/ri_benchmark/operator_ri_hartree.h" @@ -42,8 +44,10 @@ namespace LR { ModuleBase::TITLE("HamiltLR", "HamiltLR"); if (ri_hartree_benchmark != "aims") { assert(aims_nbasis.empty()); } - this->DM_trans.resize(1); - this->DM_trans[0] = LR_Util::make_unique>(&pmat_in, nspin, kv_in.kvec_d, nk); + // always use nspin=1 for transition density matrix + this->DM_trans = LR_Util::make_unique>(&pmat_in, 1, kv_in.kvec_d, nk); + this->DM_trans->init_DMR(&gd_in, &ucell_in); + // add the diag operator (the first one) this->ops = new OperatorLRDiag(eig_ks.c, pX[0], nk, nocc[0], nvirt[0]); //add Hxc operator @@ -104,11 +108,24 @@ namespace LR hamilt::Operator* lr_exx = new OperatorLREXX(nspin, naos, nocc[0], nvirt[0], ucell_in, psi_ks_in, this->DM_trans, exx_lri_in, kv_in, pX_in[0], pc_in, pmat_in, xc_kernel == "hf" ? 1.0 : exx_alpha, //alpha - ri_hartree_benchmark != "none"/*whether to cal_dm_trans first here*/, aims_nbasis); this->ops->add(lr_exx); } #endif + + this->cal_dm_trans = [&, this](const int& is, const T* X)->void + { + const auto psi_ks_is = LR_Util::get_psi_spin(psi_ks_in, is, nk); +#ifdef __MPI + std::vector dm_trans_2d = cal_dm_trans_pblas(X, pX[is], psi_ks_is, pc_in, naos, nocc[is], nvirt[is], pmat_in); + if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos, pmat_in); +#else + std::vector dm_trans_2d = cal_dm_trans_blas(X, psi_ks_is, nocc[is], nvirt[is]); + if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos); +#endif + // tensor to vector, then set DMK + for (int ik = 0;ik < nk;++ik) { this->DM_trans->set_DMK_pointer(ik, dm_trans_2d[ik].data()); } + }; } ~HamiltLR() { @@ -123,11 +140,16 @@ namespace LR virtual void hPsi(const T* psi_in, T* hpsi, const int ld_psi, const int& nband) const { assert(ld_psi == nk * pX[0].get_local_size()); - hamilt::Operator* node(this->ops); - while (node != nullptr) + for (int ib = 0;ib < nband;++ib) { - node->act(nband, ld_psi, /*npol=*/1, psi_in, hpsi); - node = (hamilt::Operator*)(node->next_op); + const int offset = ib * ld_psi; + this->cal_dm_trans(0, psi_in + offset); // calculate transition density matrix here + hamilt::Operator* node(this->ops); + while (node != nullptr) + { + node->act(/*nband=*/1, ld_psi, /*npol=*/1, psi_in + offset, hpsi + offset); + node = (hamilt::Operator*)(node->next_op); + } } } @@ -136,14 +158,16 @@ namespace LR const std::vector& nvirt; const int nspin = 1; const int nk = 1; + const bool tdm_sym = false; ///< whether to symmetrize the transition density matrix const std::vector& pX; T one()const; /// transition density matrix in AO representation - /// Hxc only: size=1, calculate on the same address for each bands - /// Hxc+Exx: size=nbands, store the result of each bands for common use - std::vector>> DM_trans; + /// calculate on the same address for each bands, and commonly used by all the operators + std::unique_ptr> DM_trans; /// first node operator, add operations from each operators hamilt::Operator* ops = nullptr; + + std::function cal_dm_trans; }; } diff --git a/source/module_lr/hamilt_ulr.hpp b/source/module_lr/hamilt_ulr.hpp index b2317d6aa0..ec16047c94 100644 --- a/source/module_lr/hamilt_ulr.hpp +++ b/source/module_lr/hamilt_ulr.hpp @@ -38,8 +38,8 @@ namespace LR nloc_per_band(nk* pX[0].get_local_size() + nk * pX[1].get_local_size()) { ModuleBase::TITLE("HamiltULR", "HamiltULR"); - this->DM_trans.resize(1); - this->DM_trans[0] = LR_Util::make_unique>(&pmat_in, nspin, kv_in.kvec_d, nk); + this->DM_trans = LR_Util::make_unique>(&pmat_in, 1, kv_in.kvec_d, nk); + this->DM_trans->init_DMR(&gd_in, &ucell_in); // how to change the index of eig_ks and psi_ks_In? // modify the interface of opetators to support different left- and right- spin-pairs @@ -68,6 +68,20 @@ namespace LR } } #endif + + this->cal_dm_trans = [&, this](const int& is, const T* X)->void + { + const auto psi_ks_is = LR_Util::get_psi_spin(psi_ks_in, is, nk); +#ifdef __MPI + std::vector dm_trans_2d = cal_dm_trans_pblas(X, pX[is], psi_ks_is, pc_in, naos, nocc[is], nvirt[is], pmat_in); + if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos, pmat_in); +#else + std::vector dm_trans_2d = cal_dm_trans_blas(X, psi_ks_is, nocc[is], nvirt[is]); + if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos); +#endif + // tensor to vector, then set DMK + for (int ik = 0;ik < nk;++ik) { this->DM_trans->set_DMK_pointer(ik, dm_trans_2d[ik].data()); } + }; } ~HamiltULR() { @@ -85,6 +99,7 @@ namespace LR for (int is_bj : {0, 1}) { const int offset = offset_b + is_bj * xdim_is[0]; + cal_dm_trans(is_bj, psi_in + offset); // calculate transition density matrix here for (int is_ai : {0, 1}) { hamilt::Operator* node(this->ops[(is_ai << 1) + is_bj]); @@ -195,6 +210,9 @@ namespace LR /// transition density matrix in AO representation /// Hxc only: size=1, calculate on the same address for each bands /// Hxc+Exx: size=nbands, store the result of each bands for common use - std::vector>> DM_trans; + std::unique_ptr> DM_trans; + + std::function cal_dm_trans; + const bool tdm_sym = false; ///< whether to symmetrize the transition density matrix }; } \ No newline at end of file diff --git a/source/module_lr/lr_spectrum.cpp b/source/module_lr/lr_spectrum.cpp index 687a456df3..9616480bb4 100644 --- a/source/module_lr/lr_spectrum.cpp +++ b/source/module_lr/lr_spectrum.cpp @@ -33,7 +33,7 @@ void LR::LR_Spectrum::oscillator_strength() osc.resize(nstate, 0.0); // const int nspin0 = (this->nspin == 2) ? 2 : 1; use this in NSPIN=4 implementation double osc_tot = 0.0; - elecstate::DensityMatrix DM_trans(&this->pmat, this->nspin, this->kv.kvec_d, this->nk); + elecstate::DensityMatrix DM_trans(&this->pmat, 1, this->kv.kvec_d, this->nk); DM_trans.init_DMR(&GlobalC::GridD, &this->ucell); this->transition_dipole_.resize(nstate, ModuleBase::Vector3(0.0, 0.0, 0.0)); for (int istate = 0;istate < nstate;++istate) @@ -87,9 +87,9 @@ void LR::LR_Spectrum>::oscillator_strength() osc.resize(nstate, 0.0); // const int nspin0 = (this->nspin == 2) ? 2 : 1; use this in NSPIN=4 implementation double osc_tot = 0.0; - elecstate::DensityMatrix, std::complex> DM_trans(&this->pmat, this->nspin, this->kv.kvec_d, this->nk); + elecstate::DensityMatrix, std::complex> DM_trans(&this->pmat, 1, this->kv.kvec_d, this->nk); DM_trans.init_DMR(&GlobalC::GridD, &this->ucell); - elecstate::DensityMatrix, double> DM_trans_real_imag(&this->pmat, this->nspin, this->kv.kvec_d, this->nk); + elecstate::DensityMatrix, double> DM_trans_real_imag(&this->pmat, 1, this->kv.kvec_d, this->nk); DM_trans_real_imag.init_DMR(&GlobalC::GridD, &this->ucell); this->transition_dipole_.resize(nstate, ModuleBase::Vector3>(0.0, 0.0, 0.0)); diff --git a/source/module_lr/operator_casida/operator_lr_exx.cpp b/source/module_lr/operator_casida/operator_lr_exx.cpp index f3bf1f4e65..8f754952dc 100644 --- a/source/module_lr/operator_casida/operator_lr_exx.cpp +++ b/source/module_lr/operator_casida/operator_lr_exx.cpp @@ -89,25 +89,12 @@ namespace LR for (int ib = 0;ib < nbands;++ib) { const int xstart_b = ib * nk * pX.get_local_size(); - // suppose Cs，Vs, have already been calculated in the ion-step of ground state, - // DM_trans(k) and DM_trans(R) has already been calculated from psi_in in OperatorLRHxc::act - // but int RI_benchmark, DM_trans(k) should be first calculated here - if (cal_dm_trans) - { -#ifdef __MPI - std::vector dm_trans_2d = cal_dm_trans_pblas(psi_in + xstart_b, pX, psi_ks, pc, naos, nocc, nvirt, pmat); - if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos, pmat); -#else - std::vector dm_trans_2d = cal_dm_trans_blas(psi_in + xstart, psi_ks, nocc, nvirt); - if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos); -#endif - // tensor to vector, then set DMK - for (int ik = 0;ik < nk;++ik) { this->DM_trans[ib]->set_DMK_pointer(ik, dm_trans_2d[ik].data()); } - } + // suppose Cs，Vs, have already been calculated in the ion-step of ground state + // and DM_trans has been calculated in hPsi() outside. // 1. set_Ds (once) // convert to vector for the interface of RI_2D_Comm::split_m2D_ktoR (interface will be unified to ct::Tensor) - std::vector> DMk_trans_vector = this->DM_trans[ib]->get_DMK_vector(); + std::vector> DMk_trans_vector = this->DM_trans->get_DMK_vector(); // assert(DMk_trans_vector.size() == nk); std::vector*> DMk_trans_pointer(nk); for (int ik = 0;ik < nk;++ik) {DMk_trans_pointer[ik] = &DMk_trans_vector[ik];} diff --git a/source/module_lr/operator_casida/operator_lr_exx.h b/source/module_lr/operator_casida/operator_lr_exx.h index ab2089b200..0597706da5 100644 --- a/source/module_lr/operator_casida/operator_lr_exx.h +++ b/source/module_lr/operator_casida/operator_lr_exx.h @@ -23,7 +23,7 @@ namespace LR const int& nvirt, const UnitCell& ucell_in, const psi::Psi& psi_ks_in, - std::vector>>& DM_trans_in, + std::unique_ptr>& DM_trans_in, // HContainer* hR_in, std::weak_ptr> exx_lri_in, const K_Vectors& kv_in, @@ -31,11 +31,10 @@ namespace LR const Parallel_2D& pc_in, const Parallel_Orbitals& pmat_in, const double& alpha = 1.0, - const bool& cal_dm_trans = false, const std::vector& aims_nbasis = {}) : nspin(nspin), naos(naos), nocc(nocc), nvirt(nvirt), psi_ks(psi_ks_in), DM_trans(DM_trans_in), exx_lri(exx_lri_in), kv(kv_in), - pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha), cal_dm_trans(cal_dm_trans), + pX(pX_in), pc(pc_in), pmat(pmat_in), ucell(ucell_in), alpha(alpha), aims_nbasis(aims_nbasis) { ModuleBase::TITLE("OperatorLREXX", "OperatorLREXX"); @@ -75,7 +74,7 @@ namespace LR const std::vector aims_nbasis={}; ///< number of basis functions for each type of atom in FHI-aims /// transition density matrix - std::vector>>& DM_trans; + std::unique_ptr>& DM_trans; /// density matrix of a certain (i, a, k), with full naos*naos size for each key /// D^{iak}_{\mu\nu}(k): 1/N_k * c^*_{ak,\mu} c_{ik,\nu} diff --git a/source/module_lr/operator_casida/operator_lr_hxc.cpp b/source/module_lr/operator_casida/operator_lr_hxc.cpp index d97d364690..4ef80d983b 100644 --- a/source/module_lr/operator_casida/operator_lr_hxc.cpp +++ b/source/module_lr/operator_casida/operator_lr_hxc.cpp @@ -8,7 +8,6 @@ #include "module_lr/utils/lr_util_print.h" // #include "module_hamilt_lcao/hamilt_lcaodft/DM_gamma_2d_to_grid.h" #include "module_hamilt_lcao/module_hcontainer/hcontainer_funcs.h" -#include "module_lr/dm_trans/dm_trans.h" #include "module_lr/AX/AX.h" #include "module_hamilt_pw/hamilt_pwdft/global.h" @@ -22,41 +21,16 @@ namespace LR { ModuleBase::TITLE("OperatorLRHxc", "act"); const int& sl = ispin_ks[0]; - const int& sr = ispin_ks.size() == 1 ? sl : ispin_ks[1]; - const auto psir_ks = LR_Util::get_psi_spin(psi_ks, sr, nk); const auto psil_ks = LR_Util::get_psi_spin(psi_ks, sl, nk); - this->init_DM_trans(nbands, this->DM_trans); // initialize transion density matrix - const int& lgd = gint->gridt->lgd; for (int ib = 0;ib < nbands;++ib) { - // if Hxc-only, the memory of single-band DM_trans is enough. - // if followed by EXX, we need to allocate memory for all bands. - const int ib_dm = (this->next_op == nullptr) ? 0 : ib; const int xstart_b = ib * nbasis; - // 1. transition density matrix -#ifdef __MPI - std::vector dm_trans_2d = cal_dm_trans_pblas(psi_in + xstart_b, pX[sr], psir_ks, pc, naos, nocc[sr], nvirt[sr], pmat); - if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos, pmat); -#else - std::vector dm_trans_2d = cal_dm_trans_blas(psi_in + xstart_b, psir_ks, nocc[sr], nvirt[sr]); - if (this->tdm_sym) for (auto& t : dm_trans_2d) LR_Util::matsym(t.data(), naos); -#endif - // tensor to vector, then set DMK - for (int ik = 0;ik < nk;++ik) { this->DM_trans[ib_dm]->set_DMK_pointer(ik, dm_trans_2d[ik].data()); } - - // if (this->first_print) - // for (int ik = 0;ik < nk;++ik) - // LR_Util::print_tensor(dm_trans_2d[ik], "1.DMK[ik=" + std::to_string(ik) + "]", this->pmat); - - // use cal_DMR to get DMR form DMK by FT - this->DM_trans[ib_dm]->cal_DMR(); //DM_trans->get_DMR_vector() is 2d-block parallized - // LR_Util::print_DMR(*this->DM_trans[0], ucell.nat, "DM(R) (complex)"); - + this->DM_trans->cal_DMR(); //DM_trans->get_DMR_vector() is 2d-block parallized // ========================= begin grid calculation========================= - this->grid_calculation(nbands, ib_dm); //DM(R) to H(R) + this->grid_calculation(nbands); //DM(R) to H(R) // ========================= end grid calculation ========================= // V(R)->V(k) @@ -80,11 +54,11 @@ namespace LR template<> - void OperatorLRHxc::grid_calculation(const int& nbands, const int& iband_dm) const + void OperatorLRHxc::grid_calculation(const int& nbands) const { ModuleBase::TITLE("OperatorLRHxc", "grid_calculation(real)"); ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation"); - this->gint->transfer_DM2DtoGrid(this->DM_trans[iband_dm]->get_DMR_vector()); // 2d block to grid + this->gint->transfer_DM2DtoGrid(this->DM_trans->get_DMR_vector()); // 2d block to grid // 2. transition electron density // \f[ \tilde{\rho}(r)=\sum_{\mu_j, \mu_b}\tilde{\rho}_{\mu_j,\mu_b}\phi_{\mu_b}(r)\phi_{\mu_j}(r) \f] @@ -116,7 +90,7 @@ namespace LR } template<> - void OperatorLRHxc, base_device::DEVICE_CPU>::grid_calculation(const int& nbands, const int& iband_dm) const + void OperatorLRHxc, base_device::DEVICE_CPU>::grid_calculation(const int& nbands) const { ModuleBase::TITLE("OperatorLRHxc", "grid_calculation(complex)"); ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation"); @@ -126,9 +100,9 @@ namespace LR hamilt::HContainer HR_real_imag(GlobalC::ucell, &this->pmat); this->initialize_HR(HR_real_imag, ucell, gd); - auto dmR_to_hR = [&, this](const int& iband_dm, const char& type) -> void + auto dmR_to_hR = [&, this](const char& type) -> void { - LR_Util::get_DMR_real_imag_part(*this->DM_trans[iband_dm], DM_trans_real_imag, ucell.nat, type); + LR_Util::get_DMR_real_imag_part(*this->DM_trans, DM_trans_real_imag, ucell.nat, type); // if (this->first_print)LR_Util::print_DMR(DM_trans_real_imag, ucell.nat, "DMR(2d, real)"); this->gint->transfer_DM2DtoGrid(DM_trans_real_imag.get_DMR_vector()); @@ -166,8 +140,8 @@ namespace LR LR_Util::set_HR_real_imag_part(HR_real_imag, *this->hR, GlobalC::ucell.nat, type); }; this->hR->set_zero(); - dmR_to_hR(iband_dm, 'R'); //real - if (kv.get_nks() / this->nspin > 1) { dmR_to_hR(iband_dm, 'I'); } //imag for multi-k + dmR_to_hR('R'); //real + if (kv.get_nks() / this->nspin > 1) { dmR_to_hR('I'); } //imag for multi-k ModuleBase::timer::tick("OperatorLRHxc", "grid_calculation"); } diff --git a/source/module_lr/operator_casida/operator_lr_hxc.h b/source/module_lr/operator_casida/operator_lr_hxc.h index 16ff807f8c..0fa7919134 100644 --- a/source/module_lr/operator_casida/operator_lr_hxc.h +++ b/source/module_lr/operator_casida/operator_lr_hxc.h @@ -20,7 +20,7 @@ namespace LR const std::vector& nocc, const std::vector& nvirt, const psi::Psi& psi_ks_in, - std::vector>>& DM_trans_in, + std::unique_ptr>& DM_trans_in, typename TGint::type* gint_in, std::weak_ptr pot_in, const UnitCell& ucell_in, @@ -41,7 +41,6 @@ namespace LR this->is_first_node = true; this->hR = std::unique_ptr>(new hamilt::HContainer(&pmat_in)); this->initialize_HR(*this->hR, ucell_in, gd_in); - this->DM_trans[0]->init_DMR(*this->hR); }; ~OperatorLRHxc() { }; @@ -76,24 +75,8 @@ namespace LR hR.set_paraV(&pmat); if (std::is_same::value) { hR.fix_gamma(); } } - template - void init_DM_trans(const int& nbands, std::vector>>& DM_trans)const - { - // LR_Util::print_DMR(*this->DM_trans[0], ucell.nat, "DMR[ib=" + std::to_string(0) + "]"); - if (this->next_op != nullptr) - { - int prev_size = DM_trans.size(); - if (prev_size > nbands) { for (int ib = nbands;ib < prev_size;++ib) { DM_trans[ib].reset(); } } - DM_trans.resize(nbands); - for (int ib = prev_size;ib < nbands;++ib) - { - // the first dimenstion of DensityMatrix is nk=nks/nspin - DM_trans[ib] = LR_Util::make_unique>(&this->pmat, this->nspin, this->kv.kvec_d, this->kv.get_nks() / nspin); - DM_trans[ib]->init_DMR(*this->hR); - } - } - } - void grid_calculation(const int& nbands, const int& iband_dm)const; + + void grid_calculation(const int& nbands)const; //global sizes const int& nspin; @@ -109,7 +92,7 @@ namespace LR const psi::Psi& psi_ks = nullptr; /// transition density matrix - std::vector>>& DM_trans; + std::unique_ptr>& DM_trans; /// transition hamiltonian in AO representation std::unique_ptr> hR = nullptr; @@ -127,8 +110,6 @@ namespace LR std::vector orb_cutoff_; Grid_Driver& gd; - bool tdm_sym = false; ///< whether transition density matrix is symmetric - /// test mutable bool first_print = true; };