move calculate_weight() out of psi2rho()

YuLiu98 · YuLiu98 · commit bda3d3333b9c · 2024-11-12T16:00:53.000+08:00
diff --git a/source/module_elecstate/elecstate_lcao.cpp b/source/module_elecstate/elecstate_lcao.cpp
@@ -1,7 +1,5 @@
 #include "elecstate_lcao.h"
 
-#include <vector>
-
 #include "cal_dm.h"
 #include "module_base/timer.h"
 #include "module_elecstate/module_dm/cal_dm_psi.h"
@@ -11,6 +9,8 @@
 #include "module_hamilt_pw/hamilt_pwdft/global.h"
 #include "module_parameter/parameter.h"
 
+#include <vector>
+
 namespace elecstate
 {
 
@@ -21,34 +21,31 @@ void ElecStateLCAO<std::complex<double>>::psiToRho(const psi::Psi<std::complex<d
     ModuleBase::TITLE("ElecStateLCAO", "psiToRho");
     ModuleBase::timer::tick("ElecStateLCAO", "psiToRho");
 
-    this->calculate_weights();
-
-    // the calculations of dm, and dm -> rho are, technically, two separate
-    // functionalities, as we cannot rule out the possibility that we may have a
-    // dm from other sources, such as read from file. However, since we are not
-    // separating them now, I opt to add a flag to control how dm is obtained as
-    // of now
-    if (!PARAM.inp.dm_to_rho)
-    {
-        this->calEBand();
-
-        ModuleBase::GlobalFunc::NOTE("Calculate the density matrix.");
-
-        // this part for calculating DMK in 2d-block format, not used for charge
-        // now
-        //    psi::Psi<std::complex<double>> dm_k_2d();
-
-        if (PARAM.inp.ks_solver == "genelpa" || PARAM.inp.ks_solver == "elpa" || PARAM.inp.ks_solver == "scalapack_gvx" || PARAM.inp.ks_solver == "lapack"
-            || PARAM.inp.ks_solver == "cusolver" || PARAM.inp.ks_solver == "cusolvermp"
-            || PARAM.inp.ks_solver == "cg_in_lcao") // Peize Lin test 2019-05-15
-        {
-            elecstate::cal_dm_psi(this->DM->get_paraV_pointer(),
-                                  this->wg,
-                                  psi,
-                                  *(this->DM));
-            this->DM->cal_DMR();
-        }
-    }
+    // // the calculations of dm, and dm -> rho are, technically, two separate
+    // // functionalities, as we cannot rule out the possibility that we may have a
+    // // dm from other sources, such as read from file. However, since we are not
+    // // separating them now, I opt to add a flag to control how dm is obtained as
+    // // of now
+    // if (!PARAM.inp.dm_to_rho)
+    // {
+    //     ModuleBase::GlobalFunc::NOTE("Calculate the density matrix.");
+
+    //     // this part for calculating DMK in 2d-block format, not used for charge
+    //     // now
+    //     //    psi::Psi<std::complex<double>> dm_k_2d();
+
+    //     if (PARAM.inp.ks_solver == "genelpa" || PARAM.inp.ks_solver == "elpa" || PARAM.inp.ks_solver ==
+    //     "scalapack_gvx" || PARAM.inp.ks_solver == "lapack"
+    //         || PARAM.inp.ks_solver == "cusolver" || PARAM.inp.ks_solver == "cusolvermp"
+    //         || PARAM.inp.ks_solver == "cg_in_lcao") // Peize Lin test 2019-05-15
+    //     {
+    //         elecstate::cal_dm_psi(this->DM->get_paraV_pointer(),
+    //                               this->wg,
+    //                               psi,
+    //                               *(this->DM));
+    //         this->DM->cal_DMR();
+    //     }
+    // }
 
     for (int is = 0; is < PARAM.inp.nspin; is++)
     {
@@ -83,23 +80,6 @@ void ElecStateLCAO<double>::psiToRho(const psi::Psi<double>& psi)
     ModuleBase::TITLE("ElecStateLCAO", "psiToRho");
     ModuleBase::timer::tick("ElecStateLCAO", "psiToRho");
 
-    this->calculate_weights();
-    this->calEBand();
-
-    if (PARAM.inp.ks_solver == "genelpa" || PARAM.inp.ks_solver == "elpa" || PARAM.inp.ks_solver == "scalapack_gvx" || PARAM.inp.ks_solver == "lapack"
-        || PARAM.inp.ks_solver == "cusolver" || PARAM.inp.ks_solver == "cusolvermp" || PARAM.inp.ks_solver == "cg_in_lcao")
-    {
-        ModuleBase::timer::tick("ElecStateLCAO", "cal_dm_2d");
-
-        // get DMK in 2d-block format
-        elecstate::cal_dm_psi(this->DM->get_paraV_pointer(),
-                              this->wg,
-                              psi,
-                              *(this->DM));
-        this->DM->cal_DMR();
-        ModuleBase::timer::tick("ElecStateLCAO", "cal_dm_2d");
-    }
-
     for (int is = 0; is < PARAM.inp.nspin; is++)
     {
         ModuleBase::GlobalFunc::ZEROS(this->charge->rho[is],
diff --git a/source/module_elecstate/elecstate_lcao.h b/source/module_elecstate/elecstate_lcao.h
@@ -74,6 +74,8 @@ class ElecStateLCAO : public ElecState
     void dmToRho(std::vector<TK*> pexsi_DM, std::vector<TK*> pexsi_EDM);
 #endif
 
+    DensityMatrix<TK, double>* DM = nullptr;
+
   protected:
     // calculate electronic charge density on grid points or density matrix in real space
     // the consequence charge density rho saved into rho_out, preparing for charge mixing.
@@ -85,7 +87,6 @@ class ElecStateLCAO : public ElecState
 
     Gint_Gamma* gint_gamma = nullptr; // mohan add 2024-04-01
     Gint_k* gint_k = nullptr;         // mohan add 2024-04-01
-    DensityMatrix<TK, double>* DM = nullptr;
 };
 
 template <typename TK>
diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp
@@ -89,9 +89,6 @@ void ElecStatePW<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
     ModuleBase::timer::tick("ElecStatePW", "psiToRho");
 
     this->init_rho_data();
-    this->calculate_weights();
-
-    this->calEBand();
 
     for(int is=0; is<PARAM.inp.nspin; is++)
 	{
diff --git a/source/module_elecstate/elecstate_pw_sdft.cpp b/source/module_elecstate/elecstate_pw_sdft.cpp
@@ -17,8 +17,6 @@ void ElecStatePW_SDFT<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
 
     if (GlobalV::MY_STOGROUP == 0)
     {
-        this->calEBand();
-
         for (int is = 0; is < nspin; is++)
         {
             setmem_var_op()(this->ctx, this->rho[is], 0, this->charge->nrxx);
diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp
@@ -402,7 +402,7 @@ void ESolver_KS<T, Device>::hamilt2density(const int istep, const int iter, cons
 
         drho = p_chgmix->get_drho(pelec->charge, PARAM.inp.nelec);
         hsolver_error = 0.0;
-        if (iter == 1)
+        if (iter == 1 && PARAM.inp.calculation != "nscf")
         {
             hsolver_error
                 = hsolver::cal_hsolve_error(PARAM.inp.basis_type, PARAM.inp.esolver_type, diag_ethr, PARAM.inp.nelec);
@@ -628,7 +628,7 @@ void ESolver_KS<T, Device>::iter_finish(const int istep, int& iter)
 
         // If drho < hsolver_error in the first iter or drho < scf_thr, we
         // do not change rho.
-        if (drho < hsolver_error || this->conv_esolver)
+        if (drho < hsolver_error || this->conv_esolver || PARAM.inp.calculation == "nscf")
         {
             if (drho < hsolver_error)
             {
diff --git a/source/module_esolver/esolver_ks_lcao.cpp b/source/module_esolver/esolver_ks_lcao.cpp
@@ -3,6 +3,7 @@
 #include "module_base/formatter.h"
 #include "module_base/global_variable.h"
 #include "module_base/tool_title.h"
+#include "module_elecstate/module_dm/cal_dm_psi.h"
 #include "module_io/berryphase.h"
 #include "module_io/cube_io.h"
 #include "module_io/dos_nao.h"
@@ -51,7 +52,7 @@
 #include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
 #include "module_hsolver/hsolver_lcao.h"
 // function used by deepks
-#include "module_elecstate/cal_dm.h"
+// #include "module_elecstate/cal_dm.h"
 //---------------------------------------------------
 
 #include "module_hamilt_lcao/module_deltaspin/spin_constrain.h"
@@ -590,6 +591,14 @@ void ESolver_KS_LCAO<TK, TR>::iter_init(const int istep, const int iter)
             // and the ncalculate the charge density on grid.
 
             this->pelec->skip_weights = true;
+            this->pelec->calculate_weights();
+            if (!PARAM.inp.dm_to_rho)
+            {
+                auto _pelec = dynamic_cast<elecstate::ElecStateLCAO<TK>*>(this->pelec);
+                _pelec->calEBand();
+                elecstate::cal_dm_psi(_pelec->DM->get_paraV_pointer(), _pelec->wg, *this->psi, *(_pelec->DM));
+                _pelec->DM->cal_DMR();
+            }
             this->pelec->psiToRho(*this->psi);
             this->pelec->skip_weights = false;
 
@@ -718,9 +727,9 @@ void ESolver_KS_LCAO<TK, TR>::hamilt2density_single(int istep, int iter, double
     // reset energy
     this->pelec->f_en.eband = 0.0;
     this->pelec->f_en.demet = 0.0;
-
+    bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
     hsolver::HSolverLCAO<TK> hsolver_lcao_obj(&(this->pv), PARAM.inp.ks_solver);
-    hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec, false);
+    hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec, skip_charge);
 
     // 5) what's the exd used for?
 #ifdef __EXX
diff --git a/source/module_esolver/esolver_ks_lcao_tddft.cpp b/source/module_esolver/esolver_ks_lcao_tddft.cpp
@@ -165,8 +165,9 @@ void ESolver_KS_LCAO_TDDFT::hamilt2density_single(const int istep, const int ite
         this->pelec->f_en.demet = 0.0;
         if (this->psi != nullptr)
         {
+            bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
             hsolver::HSolverLCAO<std::complex<double>> hsolver_lcao_obj(&this->pv, PARAM.inp.ks_solver);
-            hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec_td, false);
+            hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec_td, skip_charge);
         }
     }
 
diff --git a/source/module_esolver/esolver_ks_lcaopw.cpp b/source/module_esolver/esolver_ks_lcaopw.cpp
@@ -126,6 +126,7 @@ namespace ModuleESolver
         hsolver::DiagoIterAssist<T>::SCF_ITER = iter;
         hsolver::DiagoIterAssist<T>::PW_DIAG_THR = ethr;
         hsolver::DiagoIterAssist<T>::PW_DIAG_NMAX = PARAM.inp.pw_diag_nmax;
+        bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
 
         // It is not a good choice to overload another solve function here, this will spoil the concept of
         // multiple inheritance and polymorphism. But for now, we just do it in this way.
@@ -138,7 +139,7 @@ namespace ModuleESolver
         }
 
         hsolver::HSolverLIP<T> hsolver_lip_obj(this->pw_wfc);
-        hsolver_lip_obj.solve(this->p_hamilt, this->kspw_psi[0], this->pelec, psig.lock().get()[0], false);
+        hsolver_lip_obj.solve(this->p_hamilt, this->kspw_psi[0], this->pelec, psig.lock().get()[0], skip_charge);
 
         // add exx
 #ifdef __EXX
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -354,6 +354,7 @@ void ESolver_KS_PW<T, Device>::hamilt2density_single(const int istep, const int
     hsolver::DiagoIterAssist<T, Device>::SCF_ITER = iter;
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX = PARAM.inp.pw_diag_nmax;
+    bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
 
     hsolver::HSolverPW<T, Device> hsolver_pw_obj(this->pw_wfc,
                                                  &this->wf,
@@ -375,7 +376,7 @@ void ESolver_KS_PW<T, Device>::hamilt2density_single(const int istep, const int
                          this->pelec->ekb.c,
                          GlobalV::RANK_IN_POOL,
                          GlobalV::NPROC_IN_POOL,
-                         false);
+                         skip_charge);
 
     this->init_psi = true;
 
diff --git a/source/module_esolver/esolver_sdft_pw.cpp b/source/module_esolver/esolver_sdft_pw.cpp
@@ -174,7 +174,7 @@ void ESolver_SDFT_PW<T, Device>::hamilt2density_single(int istep, int iter, doub
     this->pelec->f_en.demet = 0.0;
     // choose if psi should be diag in subspace
     // be careful that istep start from 0 and iter start from 1
-    if (istep == 0 && iter == 1)
+    if (istep == 0 && iter == 1 || PARAM.inp.calculation == "nscf")
     {
         hsolver::DiagoIterAssist<T, Device>::need_subspace = false;
     }
@@ -183,8 +183,8 @@ void ESolver_SDFT_PW<T, Device>::hamilt2density_single(int istep, int iter, doub
         hsolver::DiagoIterAssist<T, Device>::need_subspace = true;
     }
 
+    bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
-
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX = PARAM.inp.pw_diag_nmax;
 
     // hsolver only exists in this function
@@ -206,7 +206,15 @@ void ESolver_SDFT_PW<T, Device>::hamilt2density_single(int istep, int iter, doub
                                                            hsolver::DiagoIterAssist<T, Device>::need_subspace,
                                                            this->init_psi);
 
-    hsolver_pw_sdft_obj.solve(this->p_hamilt, this->kspw_psi[0], this->psi[0], this->pelec, this->pw_wfc, this->stowf, istep, iter, false);
+    hsolver_pw_sdft_obj.solve(this->p_hamilt,
+                              this->kspw_psi[0],
+                              this->psi[0],
+                              this->pelec,
+                              this->pw_wfc,
+                              this->stowf,
+                              istep,
+                              iter,
+                              skip_charge);
     this->init_psi = true;
 
     // set_diagethr need it
diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp
@@ -286,6 +286,7 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(const int istep)
             this->read_mat_npz(zipname, *(dm->get_DMR_pointer(2)));
         }
 
+        this->pelec->calculate_weights();
         this->pelec->psiToRho(*this->psi);
 
         int nspin0 = PARAM.inp.nspin == 2 ? 2 : 1;
diff --git a/source/module_hsolver/hsolver_lcao.cpp b/source/module_hsolver/hsolver_lcao.cpp
@@ -22,12 +22,14 @@
 
 #ifdef __PEXSI
 #include "diago_pexsi.h"
-#include "module_elecstate/elecstate_lcao.h"
 #endif
 
 #include "module_base/global_variable.h"
 #include "module_base/memory.h"
 #include "module_base/timer.h"
+#include "module_elecstate/elecstate_lcao.h"
+#include "module_elecstate/module_dm/cal_dm_psi.h"
+#include "module_elecstate/module_dm/density_matrix.h"
 #include "module_hsolver/parallel_k2d.h"
 #include "module_parameter/parameter.h"
 
@@ -73,6 +75,15 @@ void HSolverLCAO<T, Device>::solve(hamilt::Hamilt<T>* pHamilt,
                                      "This method and KPAR setting is not supported for lcao basis in ABACUS!");
         }
 
+        pes->calculate_weights();
+        if (!PARAM.inp.dm_to_rho)
+        {
+            auto _pes = dynamic_cast<elecstate::ElecStateLCAO<T>*>(pes);
+            _pes->calEBand();
+            elecstate::cal_dm_psi(_pes->DM->get_paraV_pointer(), _pes->wg, psi, *(_pes->DM));
+            _pes->DM->cal_DMR();
+        }
+
         if (!skip_charge)
         {
             // used in scf calculation
diff --git a/source/module_hsolver/hsolver_lcaopw.cpp b/source/module_hsolver/hsolver_lcaopw.cpp
@@ -272,6 +272,8 @@ void HSolverLIP<T>::solve(hamilt::Hamilt<T>* pHamilt, // ESolver_KS_PW::p_hamilt
         eigenvalues.data(),
         pes->ekb.nr * pes->ekb.nc);
 
+    reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calculate_weights();
+    reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();
     if (skip_charge)
     {
         ModuleBase::timer::tick("HSolverLIP", "solve");
diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp
@@ -327,6 +327,8 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
         // pes->ekb.nr * pes->ekb.nc
         this->wfc_basis->nks * psi.get_nbands());
 
+    reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calculate_weights();
+    reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();
     if (skip_charge)
     {
         ModuleBase::timer::tick("HSolverPW", "solve");
diff --git a/source/module_hsolver/hsolver_pw_sdft.cpp b/source/module_hsolver/hsolver_pw_sdft.cpp
@@ -80,14 +80,19 @@ void HSolverPW_SDFT<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
 
     // prepare sqrt{f(\hat{H})}|\chi> to calculate density, force and stress
     stoiter.calHsqrtchi(stowf);
+
+    elecstate::ElecStatePW<T, Device>* pes_pw = static_cast<elecstate::ElecStatePW<T, Device>*>(pes);
+    if (GlobalV::MY_STOGROUP == 0)
+    {
+        pes_pw->calEBand();
+    }
     if (skip_charge)
     {
         ModuleBase::timer::tick("HSolverPW_SDFT", "solve");
         return;
     }
     //(5) calculate new charge density
     // calculate KS rho.
-    elecstate::ElecStatePW<T, Device>* pes_pw = static_cast<elecstate::ElecStatePW<T, Device>*>(pes);
     pes_pw->init_rho_data();
     if (nbands > 0)
     {
diff --git a/source/module_io/read_input_item_system.cpp b/source/module_io/read_input_item_system.cpp
@@ -543,6 +543,10 @@ void ReadInput::item_system()
             {
                 ModuleBase::WARNING_QUIT("ReadInput", "dm_to_rho is not available for parallel calculations");
             }
+            if (para.input.dm_to_rho && para.inp.gamma_only)
+            {
+                ModuleBase::WARNING_QUIT("ReadInput", "dm_to_rho is not available for gamma_only calculations");
+            }
             if (para.input.dm_to_rho)
             {
 #ifndef __USECNPY
diff --git a/source/module_io/test_serial/read_input_item_test.cpp b/source/module_io/test_serial/read_input_item_test.cpp
@@ -928,6 +928,14 @@ TEST_F(InputTest, Item_test2)
         output = testing::internal::GetCapturedStdout();
         EXPECT_THAT(output, testing::HasSubstr("NOTICE"));
 
+        param.input.dm_to_rho = true;
+        param.input.gamma_only = true;
+        GlobalV::NPROC = 1;
+        testing::internal::CaptureStdout();
+        EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(0), "");
+        output = testing::internal::GetCapturedStdout();
+        EXPECT_THAT(output, testing::HasSubstr("NOTICE"));
+
 #ifndef __USECNPY
         param.input.dm_to_rho = true;
         GlobalV::NPROC = 1;

Original file line number	Diff line number	Diff line change
`@@ -89,9 +89,6 @@ void ElecStatePW<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)`
`89`	`89`	`ModuleBase::timer::tick("ElecStatePW", "psiToRho");`
`90`	`90`
`91`	`91`	`this->init_rho_data();`
`92`		`- this->calculate_weights();`
`93`		`-`
`94`		`- this->calEBand();`
`95`	`92`
`96`	`93`	`for(int is=0; is<PARAM.inp.nspin; is++)`
`97`	`94`	`{`
Original file line number	Diff line number	Diff line change
`@@ -17,8 +17,6 @@ void ElecStatePW_SDFT<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)`
`17`	`17`
`18`	`18`	`if (GlobalV::MY_STOGROUP == 0)`
`19`	`19`	`{`
`20`		`- this->calEBand();`
`21`		`-`
`22`	`20`	`for (int is = 0; is < nspin; is++)`
`23`	`21`	`{`
`24`	`22`	`setmem_var_op()(this->ctx, this->rho[is], 0, this->charge->nrxx);`
Original file line number	Diff line number	Diff line change
`@@ -402,7 +402,7 @@ void ESolver_KS<T, Device>::hamilt2density(const int istep, const int iter, cons`
`402`	`402`
`403`	`403`	`drho = p_chgmix->get_drho(pelec->charge, PARAM.inp.nelec);`
`404`	`404`	`hsolver_error = 0.0;`
`405`		`- if (iter == 1)`
	`405`	`+ if (iter == 1 && PARAM.inp.calculation != "nscf")`
`406`	`406`	`{`
`407`	`407`	`hsolver_error`
`408`	`408`	`= hsolver::cal_hsolve_error(PARAM.inp.basis_type, PARAM.inp.esolver_type, diag_ethr, PARAM.inp.nelec);`
`@@ -628,7 +628,7 @@ void ESolver_KS<T, Device>::iter_finish(const int istep, int& iter)`
`628`	`628`
`629`	`629`	`// If drho < hsolver_error in the first iter or drho < scf_thr, we`
`630`	`630`	`// do not change rho.`
`631`		`- if (drho < hsolver_error \|\| this->conv_esolver)`
	`631`	`+ if (drho < hsolver_error \|\| this->conv_esolver \|\| PARAM.inp.calculation == "nscf")`
`632`	`632`	`{`
`633`	`633`	`if (drho < hsolver_error)`
`634`	`634`	`{`
Original file line number	Diff line number	Diff line change
`@@ -165,8 +165,9 @@ void ESolver_KS_LCAO_TDDFT::hamilt2density_single(const int istep, const int ite`
`165`	`165`	`this->pelec->f_en.demet = 0.0;`
`166`	`166`	`if (this->psi != nullptr)`
`167`	`167`	`{`
	`168`	`+ bool skip_charge = PARAM.inp.calculation == "nscf" ? true : false;`
`168`	`169`	`hsolver::HSolverLCAO<std::complex<double>> hsolver_lcao_obj(&this->pv, PARAM.inp.ks_solver);`
`169`		`- hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec_td, false);`
	`170`	`+ hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec_td, skip_charge);`
`170`	`171`	`}`
`171`	`172`	`}`
`172`	`173`