Reconstruct: move exx_helper to hamilt_pwdft

Author: Flying-dragon-boxing

source/module_esolver/CMakeLists.txt

@@ -27,7 +27,8 @@ add_library(
     esolver
     OBJECT
     ${objects}
-        module_exx_helper/exx_helper.cpp
+        ../module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.cpp
+        ../module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.h
 )
 
 if(ENABLE_COVERAGE)

source/module_esolver/esolver_ks.cpp

@@ -661,22 +661,6 @@ void ESolver_KS<T, Device>::iter_finish(UnitCell& ucell, const int istep, int& i
     this->pelec->cal_energies(1);
     this->pelec->cal_energies(2);
 
-//    // for separate loop in hybrid functionals in pw
-//    // this gives correct energy for the first "pure" functional loop
-//    // and update the functional for the next loop
-//    if (PARAM.inp.basis_type == "pw")
-//    {
-//        auto p_esolver_ks_pw = dynamic_cast<ESolver_KS_PW<T, Device>*>(this);
-//
-//        if (conv_esolver && GlobalC::exx_info.info_global.cal_exx && GlobalC::exx_info.info_global.separate_loop && p_esolver_ks_pw->exx_helper.first_iter)
-//        {
-//            conv_esolver = false;
-//            XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);
-//            this->update_pot(ucell, istep, iter, conv_esolver);
-//            conv_esolver = true;
-//        }
-//    }
-
     if (iter == 1)
     {
         this->pelec->f_en.etot_old = this->pelec->f_en.etot;

source/module_esolver/esolver_ks_pw.cpp

@@ -513,18 +513,6 @@ void ESolver_KS_PW<T, Device>::hamilt2density_single(UnitCell& ucell,
         srho.begin(is, this->chr, this->pw_rhod, ucell.symm);
     }
 
-#ifdef __EXX
-    if (GlobalC::exx_info.info_global.cal_exx && !exx_helper.first_iter)
-    {
-        this->pelec->set_exx(exx_helper.cal_exx_energy(this->kspw_psi[0], this));
-    }
-#endif
-
-    // deband is calculated from "output" charge density calculated
-    // in sum_band
-    // need 'rho(out)' and 'vr (v_h(in) and v_xc(in))'
-    this->pelec->f_en.deband = this->pelec->cal_delta_eband(ucell);
-
     ModuleBase::timer::tick("ESolver_KS_PW", "hamilt2density_single");
 }
 
@@ -550,6 +538,18 @@ void ESolver_KS_PW<T, Device>::update_pot(UnitCell& ucell, const int istep, cons
 template <typename T, typename Device>
 void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver)
 {
+#ifdef __EXX
+    if (GlobalC::exx_info.info_global.cal_exx && !exx_helper.first_iter)
+    {
+        this->pelec->set_exx(exx_helper.cal_exx_energy(this->ctx, this->kspw_psi[0], this->pw_wfc, this->pw_rho, &ucell, &this->kv));
+    }
+#endif
+
+    // deband is calculated from "output" charge density calculated
+    // in sum_band
+    // need 'rho(out)' and 'vr (v_h(in) and v_xc(in))'
+    this->pelec->f_en.deband = this->pelec->cal_delta_eband(ucell);
+
     // 1) Call iter_finish() of ESolver_KS
     ESolver_KS<T, Device>::iter_finish(ucell, istep, iter, conv_esolver);
 

source/module_esolver/esolver_ks_pw.h

@@ -3,7 +3,7 @@
 #include "./esolver_ks.h"
 #include "module_hamilt_pw/hamilt_pwdft/operator_pw/velocity_pw.h"
 #include "module_psi/psi_init.h"
-
+#include "module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.h"
 #include "module_hamilt_pw/hamilt_pwdft/global.h"
 
 #include <memory>
@@ -33,58 +33,10 @@ class ESolver_KS_PW : public ESolver_KS<T, Device>
 
     void after_all_runners(UnitCell& ucell) override;
 
-#ifdef __EXX
-    struct Exx_Helper
-    {
-      public:
-        Exx_Helper() = default;
-        ModuleBase::matrix * wf_wg;
-        psi::Psi<T, base_device::DEVICE_CPU> psi;
-        static constexpr double DIV_UNDEFINED = 0x0d000721;
-        double div = DIV_UNDEFINED;
-        bool construct_ace = false;
-
-        bool exx_after_converge(int &iter)
-        {
-            if (first_iter)
-            {
-                first_iter = false;
-            }
-            else if (!GlobalC::exx_info.info_global.separate_loop)
-            {
-                return true;
-            }
-            else if (iter == 1)
-            {
-                return true;
-            }
-            GlobalV::ofs_running << "Updating EXX and rerun SCF" << std::endl;
-            iter = 0;
-            return false;
-
-        }
-
-        void set_psi(psi::Psi<T, Device> &psi_)
-        {
-            this->psi = psi_;
-            construct_ace = true;
-        }
-
-        void reset_div()
-        {
-            this->div = DIV_UNDEFINED;
-        }
-
-        double cal_exx_energy(psi::Psi<T, Device> &psi, ESolver_KS_PW<T, Device> *this_);
-
-        bool first_iter = false;
-    };
-#endif
-
     // EXX Todo: verify current implementation for after_converge
     // virtual bool do_after_converge(int &iter) override;
 #ifdef __EXX
-    Exx_Helper exx_helper;
+    Exx_Helper<T, Device> exx_helper;
 #endif
 
 

source/module_hamilt_pw/hamilt_pwdft/hamilt_pw.cpp

@@ -397,7 +397,7 @@ void HamiltPW<T, Device>::sPsi(const T* psi_in, // psi
 }
 
 template<typename T, typename Device>
-void HamiltPW<T, Device>::set_exx_helper(typename ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper *exx_helper)
+void HamiltPW<T, Device>::set_exx_helper(Exx_Helper<T, Device> *exx_helper)
 {
     auto op = this->ops;
     while (op != nullptr)

source/module_hamilt_pw/hamilt_pwdft/hamilt_pw.h

@@ -8,6 +8,7 @@
 #include "module_hamilt_general/hamilt.h"
 #include "module_hamilt_pw/hamilt_pwdft/VNL_in_pw.h"
 #include "module_base/kernels/math_kernel_op.h"
+#include "module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.h"
 
 namespace hamilt
 {
@@ -37,8 +38,8 @@ class HamiltPW : public Hamilt<T, Device>
     ) const override;
 
 #ifdef __EXX
-    void set_exx_helper(typename ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper* exx_helper_in);
-    typename ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper* p_exx_helper;
+    void set_exx_helper(Exx_Helper<T, Device>* exx_helper_in);
+    Exx_Helper<T, Device>* p_exx_helper;
 #endif
 
 protected:

source/module_esolver/module_exx_helper/exx_helper.cpp renamed to source/module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.cpp

@@ -1,25 +1,24 @@
-#include "module_esolver/esolver_ks_pw.h"
+#include "exx_helper.h"
 
 template <typename T, typename Device>
-double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::Psi<T, Device>& psi, ESolver_KS_PW<T, Device>* this_)
+double Exx_Helper<T, Device>::cal_exx_energy(const Device *ctx, psi::Psi<T, Device>& psi, ModulePW::PW_Basis_K* pw_wfc, ModulePW::PW_Basis* pw_rho, UnitCell* ucell, K_Vectors *kv)
 {
     ModuleBase::timer::tick("ESolver_KS_PW", "cal_exx_energy");
 
     using setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
-    T* psi_nk_real = new T[this_->pw_wfc->nrxx];
-    T* psi_mq_real = new T[this_->pw_wfc->nrxx];
-    T* h_psi_recip = new T[this_->pw_wfc->npwk_max];
-    T* h_psi_real = new T[this_->pw_wfc->nrxx];
-    T* density_real = new T[this_->pw_wfc->nrxx];
-    auto rhopw = this_->pw_rho;
+    T* psi_nk_real = new T[pw_wfc->nrxx];
+    T* psi_mq_real = new T[pw_wfc->nrxx];
+    T* h_psi_recip = new T[pw_wfc->npwk_max];
+    T* h_psi_real = new T[pw_wfc->nrxx];
+    T* density_real = new T[pw_wfc->nrxx];
+    auto rhopw = pw_rho;
     T* density_recip = new T[rhopw->npw];
-    auto *kv = &this_->kv;
 
     // lambda
     auto exx_divergence = [&]() -> double
     {
-        auto wfcpw = this_->pw_wfc;
+        auto wfcpw = pw_wfc;
         // if (GlobalC::exx_info.info_lip.lambda == 0.0)
         // {
         //     return 0;
@@ -28,7 +27,7 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
         // here we follow the exx_divergence subroutine in q-e (PW/src/exx_base.f90)
         // double alpha = GlobalC::exx_info.info_lip.lambda;
         double alpha = 10.0 / wfcpw->gk_ecut;
-        double tpiba2 = this_->pw_rhod->tpiba2;
+        double tpiba2 = ucell->tpiba2;
         double div = 0;
 
         // this is the \sum_q F(q) part
@@ -96,7 +95,7 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
         aa *= 8 / ModuleBase::FOUR_PI;
         aa += 1.0 / std::sqrt(alpha * ModuleBase::PI);
 
-        double omega = this_->pelec->omega;
+        double omega = ucell->omega;
         div -= ModuleBase::e2 * omega * aa;
         return div * wfcpw->nks;
 
@@ -111,14 +110,14 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
     if (wf_wg == nullptr) return 0.0;
     // evaluate the Eexx
     // T Eexx_ik = 0.0;
-    Real Eexx_ik_real = 0.0;
-    for (int ik = 0; ik < this_->pw_wfc->nks; ik++)
+    double Eexx_ik_real = 0.0;
+    for (int ik = 0; ik < pw_wfc->nks; ik++)
     {
         //        auto k = this->pw_wfc->kvec_c[ik];
         //        std::cout << k << std::endl;
         for (int n_iband = 0; n_iband < psi.get_nbands(); n_iband++)
         {
-            setmem_complex_op()(h_psi_recip, 0, this_->pw_wfc->npwk_max);
+            setmem_complex_op()(h_psi_recip, 0, pw_wfc->npwk_max);
             setmem_complex_op()(h_psi_real, 0, rhopw->nrxx);
             setmem_complex_op()(density_real, 0, rhopw->nrxx);
             setmem_complex_op()(density_recip, 0, rhopw->npw);
@@ -137,16 +136,16 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
             psi.fix_kb(ik, n_iband);
             const T* psi_nk = psi.get_pointer();
             // retrieve \psi_nk in real space
-            this_->pw_wfc->recip_to_real(this_->ctx, psi_nk, psi_nk_real, ik);
+            pw_wfc->recip_to_real(ctx, psi_nk, psi_nk_real, ik);
 
             // for \psi_nk, get the pw of iq and band m
             // q_points is a vector of integers, 0 to nks-1
             std::vector<int> q_points;
-            for (int iq = 0; iq < this_->pw_wfc->nks; iq++)
+            for (int iq = 0; iq < pw_wfc->nks; iq++)
             {
                 q_points.push_back(iq);
             }
-            Real nqs = q_points.size();
+            double nqs = q_points.size();
 
             //            std::cout << "ik = " << ik << " ib = " << n_iband << " wg_kb = " << wg_ikb_real << " wk_ik = " << kv->wk[ik] << std::endl;
             for (int iq: q_points)
@@ -168,15 +167,15 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
                     psi.fix_kb(iq, m_iband);
                     const T* psi_mq = psi.get_pointer();
                     // const T* psi_mq = get_pw(m_iband, iq);
-                    this_->pw_wfc->recip_to_real(this_->ctx, psi_mq, psi_mq_real, iq);
+                    pw_wfc->recip_to_real(ctx, psi_mq, psi_mq_real, iq);
 
-                    Real omega_inv = 1.0 / this_->pelec->omega;
+                    T omega_inv = 1.0 / ucell->omega;
 
                     // direct multiplication in real space, \psi_nk(r) * \psi_mq(r)
                     #ifdef _OPENMP
                     #pragma omp parallel for
                     #endif
-                    for (int ir = 0; ir < this_->pw_wfc->nrxx; ir++)
+                    for (int ir = 0; ir < pw_wfc->nrxx; ir++)
                     {
                         // assert(is_finite(psi_nk_real[ir]));
                         // assert(is_finite(psi_mq_real[ir]));
@@ -187,17 +186,17 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
                     // bring the density to recip space
                     rhopw->real2recip(density_real, density_recip);
 
-                    Real tpiba2 = this_->pw_rho->tpiba2;
+                    double tpiba2 = pw_rho->tpiba2;
                     //                    std::cout << tpiba2 << std::endl;
-                    Real hse_omega2 = GlobalC::exx_info.info_global.hse_omega * GlobalC::exx_info.info_global.hse_omega;
+                    double hse_omega2 = GlobalC::exx_info.info_global.hse_omega * GlobalC::exx_info.info_global.hse_omega;
 
                     #ifdef _OPENMP
                     #pragma omp parallel for reduction(+:Eexx_ik_real) reduction(min:min_gg) reduction(max:max_gg)
                     #endif
                     for (int ig = 0; ig < rhopw->npw; ig++)
                     {
-                        auto k = this_->pw_wfc->kvec_c[ik];// * latvec;
-                        auto q = this_->pw_wfc->kvec_c[iq];// * latvec;
+                        auto k = pw_wfc->kvec_c[ik];// * latvec;
+                        auto q = pw_wfc->kvec_c[iq];// * latvec;
                         auto gcar = rhopw->gcar[ig];
                         double gg = (k - q + gcar).norm2() * tpiba2;
 
@@ -236,18 +235,18 @@ double ModuleESolver::ESolver_KS_PW<T, Device>::Exx_Helper::cal_exx_energy(psi::
         } // n_iband
 
     } // ik
-    Eexx_ik_real *= 0.5 * this_->pelec->omega;
+    Eexx_ik_real *= 0.5 * ucell->omega;
     Parallel_Reduce::reduce_pool(Eexx_ik_real);
 //    std::cout << "omega = " << this_->pelec->omega << " tpiba = " << this_->pw_rho->tpiba2 << " exx_div = " << exx_div << std::endl;
 
-    Real Eexx = Eexx_ik_real;
+    double Eexx = Eexx_ik_real;
     ModuleBase::timer::tick("ESolver_KS_PW", "cal_exx_energy");
     return Eexx;
 }
 
-template class ModuleESolver::ESolver_KS_PW<std::complex<float>, base_device::DEVICE_CPU>;
-template class ModuleESolver::ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>;
+template class Exx_Helper<std::complex<float>, base_device::DEVICE_CPU>;
+template class Exx_Helper<std::complex<double>, base_device::DEVICE_CPU>;
 #if ((defined __CUDA) || (defined __ROCM))
-template class ModuleESolver::ESolver_KS_PW<std::complex<float>, base_device::DEVICE_GPU>;
-template class ModuleESolver::ESolver_KS_PW<std::complex<double>, base_device::DEVICE_GPU>;
+template class Exx_Helper<std::complex<float>, base_device::DEVICE_GPU>;
+template class Exx_Helper<std::complex<double>, base_device::DEVICE_GPU>;
 #endif

source/module_hamilt_pw/hamilt_pwdft/module_exx_helper/exx_helper.h

@@ -0,0 +1,62 @@
+//
+// For EXX in PW.
+//
+#include "module_psi/psi.h"
+#include "module_base/matrix.h"
+#include "module_hamilt_pw/hamilt_pwdft/global.h"
+
+#ifndef EXX_HELPER_H
+#define EXX_HELPER_H
+template <typename T, typename Device>
+struct Exx_Helper
+{
+  public:
+    Exx_Helper() = default;
+    ModuleBase::matrix * wf_wg;
+    psi::Psi<T, Device> psi;
+    static constexpr double DIV_UNDEFINED = 0x0d000721;
+    double div = DIV_UNDEFINED;
+    bool construct_ace = false;
+
+    bool exx_after_converge(int &iter)
+    {
+        if (first_iter)
+        {
+            first_iter = false;
+        }
+        else if (!GlobalC::exx_info.info_global.separate_loop)
+        {
+            return true;
+        }
+        else if (iter == 1)
+        {
+            return true;
+        }
+        GlobalV::ofs_running << "Updating EXX and rerun SCF" << std::endl;
+        iter = 0;
+        return false;
+
+    }
+
+    void set_psi(psi::Psi<T, Device> &psi_)
+    {
+        this->psi = psi_;
+        construct_ace = true;
+    }
+
+    void reset_div()
+    {
+        this->div = DIV_UNDEFINED;
+    }
+
+    double cal_exx_energy(const Device *ctx,
+                          psi::Psi<T, Device>& psi,
+                          ModulePW::PW_Basis_K* pw_wfc,
+                          ModulePW::PW_Basis* pw_rho,
+                          UnitCell* ucell,
+                          K_Vectors *kv);
+
+
+    bool first_iter = false;
+};
+#endif // EXX_HELPER_H