Refactor: Remove C++17 syntax and optimize DMR calculation in DeePKS. (#6094)

* Use dm_r in LCAO_Deepks to avoid double counting in DeePKS.

* Remove 'if constexpr' usage in DeePKS.

* Fix a merge bug.

Author: ErjieWu

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<TK, TR>::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<TK, TR>::before_all_runners(UnitCell& ucell, const Input_pa
     return;
 }
 
-
 template <typename TK, typename TR>
 double ESolver_KS_LCAO<TK, TR>::cal_energy()
 {
@@ -316,7 +313,6 @@ double ESolver_KS_LCAO<TK, TR>::cal_energy()
     return this->pelec->f_en.etot;
 }
 
-
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::cal_force(UnitCell& ucell, ModuleBase::matrix& force)
 {
@@ -460,7 +456,7 @@ void ESolver_KS_LCAO<TK, TR>::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<TK, TR>::after_all_runners(UnitCell& ucell)
     if (PARAM.inp.out_mat_xc)
     {
         ModuleIO::write_Vxc<TK, TR>(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<TK, TR>(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<TK, TR>::after_all_runners(UnitCell& ucell)
     ModuleBase::timer::tick("ESolver_KS_LCAO", "after_all_runners");
 }
 
-
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::iter_init(UnitCell& ucell, const int istep, const int iter)
 {
@@ -639,11 +635,10 @@ void ESolver_KS_LCAO<TK, TR>::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<TK, TR>::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<TK, TR>::iter_init(UnitCell& ucell, const int istep, const
     }
 }
 
-
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::hamilt2rho_single(UnitCell& ucell, int istep, int iter, double ethr)
 {
@@ -822,7 +821,6 @@ void ESolver_KS_LCAO<TK, TR>::hamilt2rho_single(UnitCell& ucell, int istep, int
     this->pelec->f_en.deband = this->pelec->cal_delta_eband(ucell);
 }
 
-
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::update_pot(UnitCell& ucell, const int istep, const int iter, const bool conv_esolver)
 {
@@ -840,7 +838,6 @@ void ESolver_KS_LCAO<TK, TR>::update_pot(UnitCell& ucell, const int istep, const
     }
 }
 
-
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver)
 {
@@ -878,6 +875,7 @@ void ESolver_KS_LCAO<TK, TR>::iter_finish(UnitCell& ucell, const int istep, int&
             = dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(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;
     }

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_gamma.cpp

@@ -37,7 +37,7 @@ void Force_LCAO<double>::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<double>::ftable(const bool isforce,
 #ifdef __DEEPKS
     if (PARAM.inp.deepks_scf)
     {
-        const std::vector<std::vector<double>>& 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<double>(dm_gamma,
+        DeePKS_domain::cal_f_delta<double>(ld.dm_r,
                                            ucell,
                                            orb,
                                            gd,

source/module_hamilt_lcao/hamilt_lcaodft/FORCE_k.cpp

@@ -268,10 +268,8 @@ void Force_LCAO<std::complex<double>>::ftable(const bool isforce,
 #ifdef __DEEPKS
     if (PARAM.inp.deepks_scf)
     {
-        const std::vector<std::vector<std::complex<double>>>& 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<std::complex<double>>(dm_k,
+        DeePKS_domain::cal_f_delta<std::complex<double>>(ld.dm_r,
                                                          ucell,
                                                          orb,
                                                          gd,

source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/deepks_lcao.cpp

@@ -82,8 +82,8 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::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<hamilt::OperatorLCAO<TK, TR>>::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<hamilt::OperatorLCAO<TK, TR>>::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<hamilt::OperatorLCAO<TK, TR>>::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<hamilt::OperatorLCAO<TK, TR>>::calculate_HR()
             this->pre_calculate_nlm(iat0, nlm_on_the_fly);
         }
 
-        std::vector<std::unordered_map<int, std::vector<double>>>& nlm_iat = 
-          is_on_the_fly ? nlm_on_the_fly : nlm_tot[iat0];
+        std::vector<std::unordered_map<int, std::vector<double>>>& nlm_iat
+            = is_on_the_fly ? nlm_on_the_fly : nlm_tot[iat0];
 
         // 2. calculate <phi_I|beta>D<beta|phi_{J,R}> for each pair of <IJR> atoms
         for (int ad1 = 0; ad1 < adjs.adj_num + 1; ++ad1)
@@ -500,7 +500,6 @@ void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::cal_HR_IJR(const double* hr_i
     }
 }
 
-
 // contributeHk()
 template <typename TK, typename TR>
 void hamilt::DeePKS<hamilt::OperatorLCAO<TK, TR>>::contributeHk(int ik)

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,10 +207,58 @@ void LCAO_Deepks<T>::allocate_V_delta(const int nat, const int nks)
     return;
 }
 
+template <typename T>
+void LCAO_Deepks<T>::init_DMR(const UnitCell& ucell,
+                              const LCAO_Orbitals& orb,
+                              const Parallel_Orbitals& pv,
+                              const Grid_Driver& GridD)
+{
+    this->dm_r = new hamilt::HContainer<double>(&pv);
+    DeePKS_domain::iterate_ad2(
+        ucell,
+        GridD,
+        orb,
+        false, // no trace_alpha
+        [&](const int iat,
+            const ModuleBase::Vector3<double>& tau0,
+            const int ibt1,
+            const ModuleBase::Vector3<double>& tau1,
+            const int start1,
+            const int nw1_tot,
+            ModuleBase::Vector3<int> dR1,
+            const int ibt2,
+            const ModuleBase::Vector3<double>& tau2,
+            const int start2,
+            const int nw2_tot,
+            ModuleBase::Vector3<int> 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 (std::is_same<T, std::complex<double>>::value)
+            {
+                dRx = (dR1 - dR2).x;
+                dRy = (dR1 - dR2).y;
+                dRz = (dR1 - dR2).z;
+            }
+            hamilt::AtomPair<double> dm_pair(ibt1, ibt2, dRx, dRy, dRz, &pv);
+            this->dm_r->insert_pair(dm_pair);
+        }
+    );
+    this->dm_r->allocate(nullptr, true);
+}
+
 template <typename T>
 void LCAO_Deepks<T>::dpks_cal_e_delta_band(const std::vector<std::vector<T>>& 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<double>;

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<hamilt::HContainer<double>*> phialpha;
 
+    // density matrix in real space
+    hamilt::HContainer<double>* 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<std::vector<T>>& dm, const int nks);