LR-Grad exx force

Author: maki49

source/module_lr/Grad/esolver_lr_grad.cpp

@@ -4,6 +4,9 @@
 #include "module_lr/Grad/force/lr_force.h"
 #include "module_elecstate/module_dm/cal_dm_psi.h"
 #include "module_io/output_log.h"
+#ifdef __EXX
+#include <RI/ri/LRI_Cal_Aux.h>
+#endif
 
 template <typename Tstream>
 inline void print_force(const std::vector<ModuleBase::matrix>& force, Tstream& ofs)
@@ -59,6 +62,41 @@ inline void test_edm_H2(const T* const edm, const double* const eig_ks, const ps
     }
 }
 
+#ifdef __EXX
+// convert DensityMatrix to maps of RI::Tensors
+template <typename T>
+inline auto get_exx_Ds_spin1(const elecstate::DensityMatrix<T, T>& dm,
+    const UnitCell& ucell, const K_Vectors& kv, const Parallel_Orbitals& pmat)
+    -> std::map<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>>
+{
+    const int& nk = dm.get_DMK_nks();   // nks/nspin
+    std::vector<const std::vector<T>*> DMk_trans_pointer(nk);
+    for (int ik = 0;ik < nk;++ik) { DMk_trans_pointer[ik] = &dm.get_DMK_vector()[ik]; }
+    return RI_2D_Comm::split_m2D_ktoR<T>(ucell, kv, DMk_trans_pointer, pmat, /*nspin=*/1)[0];
+}
+template <typename T>
+inline auto get_exx_Ds_gs(elecstate::DensityMatrix<T, double>& dm,
+    const UnitCell& ucell, const K_Vectors& kv, const Parallel_Orbitals& pmat, const int nspin)
+    -> std::map<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>>
+{
+    const int& nk = dm.get_DMK_nks() / nspin;   // nks/nspin
+    std::vector<const std::vector<T>*> DMk_trans_pointer(nk);
+
+    std::map<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>> Ds_allspin;
+    for (int is = 0;is < nspin;++is)
+    {
+        for (int ik = 0;ik < nk;++ik)
+        {
+            DMk_trans_pointer[ik] = &dm.get_DMK_vector()[ik + is * nk];
+        }
+        std::map<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>>  Ds_tmp
+            = RI_2D_Comm::split_m2D_ktoR<T>(ucell, kv, DMk_trans_pointer, pmat, /*nspin=*/1)[0];
+        RI::LRI_Cal_Aux::add_Ds<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>>(std::move(Ds_tmp), Ds_allspin);
+    }
+    return Ds_allspin;
+}
+#endif
+
 template<typename T, typename TR>
 void LR::ESolver_LR<T, TR>::init_pot_groundstate(const Charge& chg_gs)
 {
@@ -129,7 +167,11 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
     const auto& c = LR_Util::get_psi_spin(*this->psi_ks, ispin, this->nk);   // wavefunction coefficients of ground state
 
     // calculate the force (the partial gradient of Lagrangian)
-    LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_);
+    LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_
+#ifdef __EXX
+        , std::weak_ptr<Exx_LRI<T>>(this->exx_lri), this->exx_info.info_global.hybrid_alpha
+#endif
+    );
     GlobalV::ofs_running << "Start to calculate excited-state force of " << this->spin_types[ispin] << std::endl;
     // ground state dm for currrent spin (only for test the correctness of the force)
     // elecstate::DensityMatrix<T, T> dm_gs(this->paraMat_, 1, this->kv.kvec_d, this->nk);
@@ -141,9 +183,12 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
         const int offset = istate * this->nloc_per_band;
         // The imag part will be cancelled in the force calculation, so we use double DM(R) to calculate force. 
         // But complex transition DM(R) is still used in energy density matrix calculation.
-        elecstate::DensityMatrix<T, double> dm_trans_real =   // D(X)
-            LR_Util::build_dm_from_dmk<T, double>(
-                cal_dm_trans_pblas(this->X[ispin].template data<T>() + offset, this->paraX_[ispin], c, this->paraC_, this->nbasis, this->nocc[ispin], this->nvirt[ispin], this->paraMat_),
+        const auto& dm_trans_k = cal_dm_trans_pblas(this->X[ispin].template data<T>() + offset, this->paraX_[ispin], c, this->paraC_, this->nbasis, this->nocc[ispin], this->nvirt[ispin], this->paraMat_);
+        const elecstate::DensityMatrix<T, double>& dm_trans_real =   // D(X), double (FIXME: not enough for periodic system!)
+            LR_Util::build_dm_from_dmk<T, double>(dm_trans_k,
+                this->paraMat_, this->nk, this->kv.kvec_d, this->ucell, this->gd, this->orb_cutoff_);
+        const elecstate::DensityMatrix<T, T>& dm_trans =   // D(X) complex
+            LR_Util::build_dm_from_dmk<T, T>(dm_trans_k,
                 this->paraMat_, this->nk, this->kv.kvec_d, this->ucell, this->gd, this->orb_cutoff_);
 
         elecstate::DensityMatrix<T, T> relaxed_diff_dm =    // T+D(Z), (R) can be complex
@@ -199,13 +244,35 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
         ModuleBase::matrix force_hxc_dmtrans = lr_force.cal_force_hxc_dmtrans(dm_trans_real, *this->pot[ispin]);
         std::cout << "Force (Hxc-DMTrans term) of state " << istate << ": " << std::endl;
         LR_Util::print_value(force_hxc_dmtrans.c, ucell.nat, 3);
+
         ModuleBase::matrix force_hamiltgs_relaxed_diff = lr_force.cal_force_hamilt_gs_dm_relaxed_diff(relaxed_diff_dm_real, *pot_gs, /*with_ewald=*/false);
         std::cout << "Force (GS-(T+Z) term) of state " << istate << ": " << std::endl;
         LR_Util::print_value(force_hamiltgs_relaxed_diff.c, ucell.nat, 3);
+
         ModuleBase::matrix force_overlap_edm = lr_force.cal_force_overlap_edm(edm_real);    // "-" sign has been included in the force factor
         std::cout << "Force (Overlap-EDM term) of state " << istate << ": " << std::endl;
         LR_Util::print_value(force_overlap_edm.c, ucell.nat, 3);
-        // remaining for EXX
+
+#ifdef __EXX
+        const double& alpha = this->exx_info.info_global.hybrid_alpha;
+
+        const auto& Ds_trans = get_exx_Ds_spin1(dm_trans, this->ucell, this->kv, this->paraMat_);
+        ModuleBase::matrix force_exx_dmtrans = lr_force.cal_force_exx_dm_trans(Ds_trans, alpha);
+        std::cout << "Force (EXX-DMTrans term) of state " << istate << ": " << std::endl;
+        LR_Util::print_value(force_exx_dmtrans.c, ucell.nat, 3);
+
+        elecstate::DensityMatrix<T, double> dm_gs(&this->paraMat_, this->nspin, this->kv.kvec_d, this->nk);   //DX
+        elecstate::cal_dm_psi(&this->paraMat_all_, this->wg_ks_all, *this->psi_ks_all, dm_gs);   // nbands is important here
+        const auto& Ds_gs = get_exx_Ds_gs(dm_gs, this->ucell, this->kv, this->paraMat_, this->nspin);
+        const auto& Ds_relaxed_diff = get_exx_Ds_spin1(relaxed_diff_dm, this->ucell, this->kv, this->paraMat_);
+        ModuleBase::matrix force_exx_gs_diff = lr_force.cal_force_exx_gs_dm_relaxed_diff(Ds_gs, Ds_relaxed_diff, alpha);
+        std::cout << "Force (EXX-GS-(T+Z) term) of state " << istate << ": " << std::endl;
+        LR_Util::print_value(force_exx_gs_diff.c, ucell.nat, 3);
+
+        // add exx force to the corresponding terms
+        force_hxc_dmtrans += force_exx_dmtrans;
+        force_hamiltgs_relaxed_diff += force_exx_gs_diff;
+#endif
         forces[istate] = force_hxc_dmtrans + force_hamiltgs_relaxed_diff + force_overlap_edm;
     }
     ModuleBase::timer::tick("ESolver_LR", "cal_force");
@@ -217,7 +284,11 @@ std::vector<ModuleBase::matrix> LR::ESolver_LR<T, TR>::cal_force(const int ispin
 template<typename T, typename TR>
 void LR::ESolver_LR<T, TR>::test_force()
 {
-    LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_);
+    LR_Force<T> lr_force(this->ucell, this->kv.kvec_d, this->paraMat_, *this->pw_rho, this->locpp, this->sf, this->gd, this->gint_, this->two_center_bundle_
+#ifdef __EXX
+        , std::weak_ptr<Exx_LRI<T>>(this->exx_lri), this->exx_info.info_global.hybrid_alpha
+#endif
+    );
 
     elecstate::DensityMatrix<T, double> dm_gs(&this->paraMat_, this->nspin, this->kv.kvec_d, this->nk);   //DX
     elecstate::cal_dm_psi(&this->paraMat_all_, this->wg_ks_all, *this->psi_ks_all, dm_gs);   // nbands is important here

source/module_lr/Grad/force/lr_force.cpp

@@ -112,6 +112,49 @@ namespace LR
         this->gint_->reset_DMRGint(1);
         return fvl_dphi;
     }
+
+#ifdef __EXX
+    template<typename TK>
+    ModuleBase::matrix LR_Force<TK>::cal_force_exx_dm_trans(
+        const std::map<int, std::map<TAC, RI::Tensor<TK>>>& dm_trans,
+        const double& alpha)
+    {
+        ModuleBase::matrix f_exx_dmtrans(this->ucell_.nat, 3);
+        auto& exx_lri_kernel = this->exx_lri_.lock()->get();
+        exx_lri_kernel.set_Ds(dm_trans, this->exx_lri_.lock()->get_info().dm_threshold, "0");
+        exx_lri_kernel.cal_Hs();
+        exx_lri_kernel.cal_force();// using dm_trans
+        for (std::size_t idim = 0; idim < 3; ++idim)
+            for (const auto& force_item : exx_lri_kernel.force[idim])
+                f_exx_dmtrans(force_item.first, idim) = std::real(force_item.second);
+        return f_exx_dmtrans * alpha;
+    }
+
+    template<typename TK>
+    ModuleBase::matrix LR_Force<TK>::cal_force_exx_gs_dm_relaxed_diff(
+        const std::map<int, std::map<TAC, RI::Tensor<TK>>>& dm_gs,
+        const std::map<int, std::map<TAC, RI::Tensor<TK>>>& relaxed_diff_dm,
+        const double& alpha)
+    {
+        ModuleBase::matrix f_exx_gs_diff(this->ucell_.nat, 3);
+        auto& exx_lri_kernel = this->exx_lri_.lock()->get();
+        exx_lri_kernel.set_Ds(dm_gs, this->exx_lri_.lock()->get_info().dm_threshold, "0");
+        exx_lri_kernel.cal_Hs();  // using dm_gs
+        // auto* lr_ptr = dynamic_cast<RI::LR<int, int, 3, TK>*>(&exx_lri_kernel);  // wrong: 
+        // assert(lr_ptr != nullptr);
+        RI::LR<int, int, 3, TK> lr_exx_kernel(std::move(exx_lri_kernel));
+        std::cout << "post_2D adress moved: " << &(lr_exx_kernel.post_2D) << std::endl;
+        std::cout << "begin RI::LR::cal_force" << std::endl;
+        lr_exx_kernel.cal_force(relaxed_diff_dm); // using relaxed_diff_dm
+        exx_lri_kernel = std::move(lr_exx_kernel);        // move back 
+        std::cout << "post_2D adress after move back: " << &(exx_lri_kernel.post_2D) << std::endl;
+        std::cout << "end RI::LR::cal_force" << std::endl;
+        for (std::size_t idim = 0; idim < 3; ++idim)
+            for (const auto& force_item : exx_lri_kernel.force[idim])
+                f_exx_gs_diff(force_item.first, idim) = std::real(force_item.second);
+        return f_exx_gs_diff * alpha;
+    }
+#endif
 }
 
 template class LR::LR_Force<double>;

source/module_lr/Grad/force/lr_force.h

@@ -2,7 +2,11 @@
 #include "module_lr/potentials/pot_hxc_lrtd.h"
 #include "module_lr/utils/gint_template.h"
 // free functions, usefull for both ground and excited state
-
+#ifdef __EXX
+#include "module_ri/Exx_LRI.h"
+#include <RI/physics/LR.h>
+using TAC = std::pair<int, std::array<int, 3>>;
+#endif
 namespace LR
 {
 
@@ -18,10 +22,19 @@ namespace LR
             const Structure_Factor& sf,
             const Grid_Driver& gd,
             typename TGint<TK>::type* gint, ///<  for grid integrals
-            const TwoCenterBundle& two_center_bundle)///<  for 2-center integrals
+            const TwoCenterBundle& two_center_bundle
+#ifdef __EXX
+            , std::weak_ptr<Exx_LRI<TK>> exx_lri_in,
+            const double& alpha
+#endif
+        )///<  for 2-center integrals
             : ucell_(ucell), kvec_d_(kvec_d), pv_(pv),
             rhopw_(rhopw), sf_(sf), locpp_(locpp), gd_(gd),
-            gint_(gint), two_center_bundle_(two_center_bundle) {
+            gint_(gint), two_center_bundle_(two_center_bundle)
+#ifdef __EXX
+            , exx_lri_(exx_lri_in), alpha_(alpha)
+#endif
+        {
         }
 
         /// 1. $Tr[H_{GS}^x * (T+D^Z)]$, where GS=groud state and $(T+D^Z)$ is the relaxed difference density matrix
@@ -35,11 +48,18 @@ namespace LR
         ModuleBase::matrix cal_force_hxc_dmtrans(const elecstate::DensityMatrix<TK, double>& dm_trans, const PotHxcLR& pot_hxc);
 
 #ifdef __EXX
+        // auto* lrexx_ptr = dynamic_cast<RI::LR<int, std::array<int, 3>, 3, TK>*>(&exx_lri_in.get());
         /// 4. $\alpha \sum_{mnkl}(mk|nl)^x *D^X *D^X$
-        // ModuleBase::matrix cal_force_exx_dm_diff(const elecstate::DensityMatrix<TK, double>& dm_trans);
+        ModuleBase::matrix cal_force_exx_dm_trans(
+            const std::map<int, std::map<TAC, RI::Tensor<TK>>>& dm_trans,
+            const double& alpha);
+        ModuleBase::matrix cal_force_exx_gs_dm_relaxed_diff(
+            const std::map<int, std::map<TAC, RI::Tensor<TK>>>& dm_gs,
+            const std::map<int, std::map<TAC, RI::Tensor<TK>>>& relaxed_diff_dm,
+            const double& alpha);
 #endif
 
-        /// test functions
+        // test functions
         /// reproduce the force of the ground state
         ModuleBase::matrix reproduce_force_gs(const elecstate::DensityMatrix<TK, double>& dm_gs,
             const elecstate::DensityMatrix<TK, double>& edm_gs,
@@ -58,6 +78,10 @@ namespace LR
         const Grid_Driver& gd_;
         typename TGint<TK>::type* gint_;
         const TwoCenterBundle& two_center_bundle_;
+#ifdef __EXX
+        std::weak_ptr<Exx_LRI<TK>> exx_lri_;
+        const double alpha_;
+#endif
 
         Charge dm_to_charge(const elecstate::DensityMatrix<TK, double>& dm);
 

source/module_lr/dm_band/dm_band.cpp

@@ -27,17 +27,18 @@ namespace LR
                         {
                             const int iat1 = ucell_.itia2iat(it1, ia1);
                             const int iat2 = ucell_.itia2iat(it2, ia2);
-                            auto& D2d = dm_band[iat1][std::make_pair(iat2, cell)];
-                            const int nw1 = ucell_.atoms[it1].nw;
-                            const int nw2 = ucell_.atoms[it2].nw;
+                            const std::size_t nw1 = ucell_.atoms[it1].nw;
+                            const std::size_t nw2 = ucell_.atoms[it2].nw;
+                            RI::Tensor<double> dm_tmp({ nw1, nw2 });
                             for (int iw1 = 0;iw1 < nw1;++iw1)
                                 for (int iw2 = 0;iw2 < nw2;++iw2)
                                 {
                                     const int iwt1 = use_nws1 ? nws1[it1] : ucell_.itiaiw2iwt(it1, ia1, iw1);
                                     const int iwt2 = use_nws2 ? nws2[it2] : ucell_.itiaiw2iwt(it2, ia2, iw2);
                                     if (pmat_.in_this_processor(iwt1, iwt2))
-                                        D2d(iw1, iw2) = fac * c1_(ik, iband1, iwt1) * c2_(ik, iband2, iwt2);
+                                        dm_tmp(iw1, iw2) = fac * c1_(ik, iband1, iwt1) * c2_(ik, iband2, iwt2);
                                 }
+                            dm_band[iat1][std::make_pair(iat2, cell)] = dm_tmp;
                         }
         }
     }
@@ -56,25 +57,26 @@ namespace LR
         for (auto cell : bvk_cells_)
         {
             std::complex<double> fac_phase = RI::Global_Func::convert<std::complex<double>>(std::exp(
-                -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kvec_c_.at(ik) * (RI_Util::array3_to_Vector3(cell) * ucell_.latvec))));
+                -ModuleBase::TWO_PI * ModuleBase::IMAG_UNIT * (kvec_c_.at(ik) * (RI_Util::array3_to_Vector3(cell) * ucell_.latvec)))) * fac;
             for (int it1 = 0;it1 < ucell_.ntype;++it1)
                 for (int ia1 = 0; ia1 < ucell_.atoms[it1].na; ++ia1)
                     for (int it2 = 0;it2 < ucell_.ntype;++it2)
                         for (int ia2 = 0;ia2 < ucell_.atoms[it2].na;++ia2)
                         {
-                            int iat1 = ucell_.itia2iat(it1, ia1);
-                            int iat2 = ucell_.itia2iat(it2, ia2);
-                            auto& D2d = dm_band[iat1][std::make_pair(iat2, cell)];
-                            const int nw1 = ucell_.atoms[it1].nw;
-                            const int nw2 = ucell_.atoms[it2].nw;
+                            const int iat1 = ucell_.itia2iat(it1, ia1);
+                            const int iat2 = ucell_.itia2iat(it2, ia2);
+                            const std::size_t nw1 = ucell_.atoms[it1].nw;
+                            const std::size_t nw2 = ucell_.atoms[it2].nw;
+                            RI::Tensor<std::complex<double>> dm_tmp({ nw1, nw2 });
                             for (int iw1 = 0;iw1 < nw1;++iw1)
                                 for (int iw2 = 0;iw2 < nw2;++iw2)
                                 {
                                     const int iwt1 = use_nws1 ? nws1[it1] : ucell_.itiaiw2iwt(it1, ia1, iw1);
                                     const int iwt2 = use_nws2 ? nws2[it2] : ucell_.itiaiw2iwt(it2, ia2, iw2);
                                     if (pmat_.in_this_processor(iwt1, iwt2))
-                                        D2d(iw1, iw2) = fac * c1_(ik, iband1, iwt1) * std::conj(c2_(ik, iband2, iwt2));
+                                        dm_tmp(iw1, iw2) = fac_phase * c1_(ik, iband1, iwt1) * std::conj(c2_(ik, iband2, iwt2));
                                 }
+                            dm_band[iat1][std::make_pair(iat2, cell)] = dm_tmp;
                         }
         }
     }

source/module_ri/Exx_LRI.h

@@ -53,6 +53,8 @@ class Exx_LRI
     Exx_LRI operator=(const Exx_LRI&) = delete;
     Exx_LRI operator=(Exx_LRI&&);
 
+    RI::Exx<TA, Tcell, Ndim, Tdata>& get() { return this->exx_lri; }
+    auto& get_info() const { return this->info; }
     void reset_Cs(const std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>& Cs_in) { this->exx_lri.set_Cs(Cs_in, this->info.C_threshold); }
     void reset_Vs(const std::map<TA, std::map<TAC, RI::Tensor<Tdata>>>& Vs_in) { this->exx_lri.set_Vs(Vs_in, this->info.V_threshold); }