Feature: enable cal_force and cal_stress in nscf (#5752)

YuLiu98 · web-flow · commit dc7147dcbbb3 · 2024-12-24T20:47:59.000+08:00
* Feature: enable cal_force and cal_stress in nscf

* Fix: enable force and stress in uspp nscf

* update unit tests
diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp
@@ -238,7 +238,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)
 }
 
 template <typename T, typename Device>
-void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
+void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
 {
     const T one{1, 0};
     const T zero{0, 0};
@@ -392,6 +392,12 @@ void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
         }
     }
     delmem_complex_op()(this->ctx, becp);
+}
+
+template <typename T, typename Device>
+void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
+{
+    this->cal_becsum(psi);
 
     // transform soft charge to recip space using smooth grids
     T* rhog = nullptr;
diff --git a/source/module_elecstate/elecstate_pw.h b/source/module_elecstate/elecstate_pw.h
@@ -35,6 +35,9 @@ class ElecStatePW : public ElecState
 
     virtual void cal_tau(const psi::Psi<T, Device>& psi);
 
+    //! calculate becsum for uspp
+    void cal_becsum(const psi::Psi<T, Device>& psi);
+
     Real* becsum = nullptr;
 
     //! init rho_data and kin_r_data
@@ -61,7 +64,7 @@ class ElecStatePW : public ElecState
     
     //! calcualte rho for each k
     void rhoBandK(const psi::Psi<T, Device>& psi);
-    
+
     //! add to the charge density in reciprocal space the part which is due to the US augmentation.
     void add_usrho(const psi::Psi<T, Device>& psi);
 
diff --git a/source/module_hsolver/hsolver_lcaopw.cpp b/source/module_hsolver/hsolver_lcaopw.cpp
@@ -281,6 +281,10 @@ void HSolverLIP<T>::solve(hamilt::Hamilt<T>* pHamilt, // ESolver_KS_PW::p_hamilt
     reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();
     if (skip_charge)
     {
+        if (PARAM.globalv.use_uspp)
+        {
+            reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->cal_becsum(psi);
+        }
         ModuleBase::timer::tick("HSolverLIP", "solve");
         return;
     }
diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp
@@ -337,6 +337,10 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();
     if (skip_charge)
     {
+        if (PARAM.globalv.use_uspp)
+        {
+            reinterpret_cast<elecstate::ElecStatePW<T, Device>*>(pes)->cal_becsum(psi);
+        }
         ModuleBase::timer::tick("HSolverPW", "solve");
         return;
     }
diff --git a/source/module_hsolver/test/hsolver_supplementary_mock.h b/source/module_hsolver/test/hsolver_supplementary_mock.h
@@ -1,5 +1,5 @@
 #pragma once
-#include "module_elecstate/elecstate.h"
+#include "module_elecstate/elecstate_pw.h"
 #include "module_psi/wavefunc.h"
 
 namespace elecstate
@@ -62,6 +62,46 @@ void ElecState::init_ks(Charge* chg_in, // pointer for class Charge
     return;
 }
 
+template <typename T, typename Device>
+ElecStatePW<T, Device>::ElecStatePW(ModulePW::PW_Basis_K* wfc_basis_in,
+                                    Charge* chg_in,
+                                    K_Vectors* pkv_in,
+                                    UnitCell* ucell_in,
+                                    pseudopot_cell_vnl* ppcell_in,
+                                    ModulePW::PW_Basis* rhodpw_in,
+                                    ModulePW::PW_Basis* rhopw_in,
+                                    ModulePW::PW_Basis_Big* bigpw_in)
+    : basis(wfc_basis_in)
+{
+}
+
+template <typename T, typename Device>
+ElecStatePW<T, Device>::~ElecStatePW()
+{
+}
+
+template <typename T, typename Device>
+void ElecStatePW<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
+{
+}
+
+template <typename T, typename Device>
+void ElecStatePW<T, Device>::cal_tau(const psi::Psi<T, Device>& psi)
+{
+}
+
+template <typename T, typename Device>
+void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
+{
+}
+
+template class ElecStatePW<std::complex<float>, base_device::DEVICE_CPU>;
+template class ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>;
+#if ((defined __CUDA) || (defined __ROCM))
+template class ElecStatePW<std::complex<float>, base_device::DEVICE_GPU>;
+template class ElecStatePW<std::complex<double>, base_device::DEVICE_GPU>;
+#endif
+
 Potential::~Potential()
 {
 }
diff --git a/source/module_hsolver/test/test_hsolver_sdft.cpp b/source/module_hsolver/test/test_hsolver_sdft.cpp
@@ -23,40 +23,11 @@ Sto_Func<REAL>::Sto_Func()
 }
 template class Sto_Func<double>;
 
-
-template <>
-elecstate::ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>::ElecStatePW(ModulePW::PW_Basis_K* wfc_basis_in,
-                                    Charge* chg_in,
-                                    K_Vectors* pkv_in,
-                                    UnitCell* ucell_in,
-                                    pseudopot_cell_vnl* ppcell_in,
-                                    ModulePW::PW_Basis* rhodpw_in,
-                                    ModulePW::PW_Basis* rhopw_in,
-                                    ModulePW::PW_Basis_Big* bigpw_in)
-    : basis(wfc_basis_in)
-{
-}
-
-template<>
-elecstate::ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>::~ElecStatePW() 
-{
-}
-
 template<>
 void elecstate::ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>::init_rho_data() 
 {
 }
 
-template<>
-void elecstate::ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>::psiToRho(const psi::Psi<std::complex<double>, base_device::DEVICE_CPU>& psi)
-{
-}
-
-template<>
-void elecstate::ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>::cal_tau(const psi::Psi<std::complex<double>, base_device::DEVICE_CPU>& psi)
-{
-}
-
 template <typename REAL, typename Device>
 StoChe<REAL, Device>::StoChe(const int& nche, const int& method, const REAL& emax_sto, const REAL& emin_sto)
 {
diff --git a/source/module_io/read_input_item_system.cpp b/source/module_io/read_input_item_system.cpp
@@ -207,7 +207,7 @@ void ReadInput::item_system()
         item.annotation = "if calculate the force at the end of the electronic iteration";
         item.reset_value = [](const Input_Item& item, Parameter& para) {
             std::vector<std::string> use_force = {"cell-relax", "relax", "md"};
-            std::vector<std::string> not_use_force = {"get_wf", "get_pchg", "nscf", "get_S"};
+            std::vector<std::string> not_use_force = {"get_wf", "get_pchg", "get_S"};
             if (std::find(use_force.begin(), use_force.end(), para.input.calculation) != use_force.end())
             {
                 if (!para.input.cal_force)
diff --git a/source/module_relax/relax_driver.cpp b/source/module_relax/relax_driver.cpp
@@ -54,90 +54,86 @@ void Relax_Driver::relax_driver(ModuleESolver::ESolver* p_esolver, UnitCell& uce
         time_t fstart = time(nullptr);
         ModuleBase::matrix force;
         ModuleBase::matrix stress;
-        if (PARAM.inp.calculation == "scf" || PARAM.inp.calculation == "relax" || PARAM.inp.calculation == "cell-relax")
-        {
-            // I'm considering putting force and stress
-            // as part of ucell and use ucell to pass information
-            // back and forth between esolver and relaxation
-            // but I'll use force and stress explicitly here for now
 
-            // calculate the total energy
-            this->etot = p_esolver->cal_energy();
+        // I'm considering putting force and stress
+        // as part of ucell and use ucell to pass information
+        // back and forth between esolver and relaxation
+        // but I'll use force and stress explicitly here for now
+
+        // calculate the total energy
+        this->etot = p_esolver->cal_energy();
+
+        // calculate and gather all parts of total ionic forces
+        if (PARAM.inp.cal_force)
+        {
+            p_esolver->cal_force(ucell, force);
+        }
+        // calculate and gather all parts of stress
+        if (PARAM.inp.cal_stress)
+        {
+            p_esolver->cal_stress(ucell, stress);
+        }
 
-            // calculate and gather all parts of total ionic forces
-            if (PARAM.inp.cal_force)
+        if (PARAM.inp.calculation == "relax" || PARAM.inp.calculation == "cell-relax")
+        {
+            if (PARAM.inp.relax_new)
             {
-                p_esolver->cal_force(ucell, force);
+                stop = rl.relax_step(ucell, force, stress, this->etot);
             }
-            // calculate and gather all parts of stress
-            if (PARAM.inp.cal_stress)
+            else
             {
-                p_esolver->cal_stress(ucell, stress);
+                stop = rl_old.relax_step(istep,
+                                         this->etot,
+                                         ucell,
+                                         force,
+                                         stress,
+                                         force_step,
+                                         stress_step); // pengfei Li 2018-05-14
             }
-
-            if (PARAM.inp.calculation == "relax" || PARAM.inp.calculation == "cell-relax")
+            // print structure
+            // changelog 20240509
+            // because I move out the dependence on GlobalV from UnitCell::print_stru_file
+            // so its parameter is calculated here
+            bool need_orb = PARAM.inp.basis_type == "pw";
+            need_orb = need_orb && PARAM.inp.psi_initializer;
+            need_orb = need_orb && PARAM.inp.init_wfc.substr(0, 3) == "nao";
+            need_orb = need_orb || PARAM.inp.basis_type == "lcao";
+            need_orb = need_orb || PARAM.inp.basis_type == "lcao_in_pw";
+            std::stringstream ss, ss1;
+            ss << PARAM.globalv.global_out_dir << "STRU_ION_D";
+            ucell.print_stru_file(ss.str(),
+                                  PARAM.inp.nspin,
+                                  true,
+                                  PARAM.inp.calculation == "md",
+                                  PARAM.inp.out_mul,
+                                  need_orb,
+                                  PARAM.globalv.deepks_setorb,
+                                  GlobalV::MY_RANK);
+
+            if (Ions_Move_Basic::out_stru)
             {
-                if (PARAM.inp.relax_new)
-                {
-                    stop = rl.relax_step(ucell,force, stress, this->etot);
-                }
-                else
-                {
-                    stop = rl_old.relax_step(istep,
-                                             this->etot,
-                                             ucell,
-                                             force,
-                                             stress,
-                                             force_step,
-                                             stress_step); // pengfei Li 2018-05-14
-                }         
-                // print structure
-                // changelog 20240509
-                // because I move out the dependence on GlobalV from UnitCell::print_stru_file
-                // so its parameter is calculated here
-                bool need_orb = PARAM.inp.basis_type == "pw";
-                need_orb = need_orb && PARAM.inp.psi_initializer;
-                need_orb = need_orb && PARAM.inp.init_wfc.substr(0, 3) == "nao";
-                need_orb = need_orb || PARAM.inp.basis_type == "lcao";
-                need_orb = need_orb || PARAM.inp.basis_type == "lcao_in_pw";
-                std::stringstream ss, ss1;
-                ss << PARAM.globalv.global_out_dir << "STRU_ION_D";
-                ucell.print_stru_file(ss.str(),
+                ss1 << PARAM.globalv.global_out_dir << "STRU_ION";
+                ss1 << istep << "_D";
+                ucell.print_stru_file(ss1.str(),
                                       PARAM.inp.nspin,
                                       true,
                                       PARAM.inp.calculation == "md",
                                       PARAM.inp.out_mul,
                                       need_orb,
                                       PARAM.globalv.deepks_setorb,
                                       GlobalV::MY_RANK);
-
-                if (Ions_Move_Basic::out_stru)
-                {
-                    ss1 << PARAM.globalv.global_out_dir << "STRU_ION";
-                    ss1 << istep << "_D";
-                    ucell.print_stru_file(ss1.str(),
-                                          PARAM.inp.nspin,
-                                          true,
-                                          PARAM.inp.calculation == "md",
-                                          PARAM.inp.out_mul,
-                                          need_orb,
-                                          PARAM.globalv.deepks_setorb,
-                                          GlobalV::MY_RANK);
-                    ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU_NOW.cif",
-                                               ucell,
-                                               "# Generated by ABACUS ModuleIO::CifParser",
-                                               "data_?");
-                }
-
-                ModuleIO::output_after_relax(stop, p_esolver->conv_esolver, GlobalV::ofs_running);
+                ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU_NOW.cif",
+                                           ucell,
+                                           "# Generated by ABACUS ModuleIO::CifParser",
+                                           "data_?");
             }
 
-#ifdef __RAPIDJSON
-            // add the energy to outout
-            Json::add_output_energy(p_esolver->cal_energy() * ModuleBase::Ry_to_eV);
-#endif
+            ModuleIO::output_after_relax(stop, p_esolver->conv_esolver, GlobalV::ofs_running);
         }
+
 #ifdef __RAPIDJSON
+        // add the energy to outout
+        Json::add_output_energy(p_esolver->cal_energy() * ModuleBase::Ry_to_eV);
         // add Json of cell coo stress force
         double unit_transform = ModuleBase::RYDBERG_SI / pow(ModuleBase::BOHR_RADIUS_SI, 3) * 1.0e-8;
         double fac = ModuleBase::Ry_to_eV / 0.529177;

Original file line number	Diff line number	Diff line change
`@@ -238,7 +238,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)`
`238`	`238`	`}`
`239`	`239`
`240`	`240`	`template <typename T, typename Device>`
`241`		`-void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)`
	`241`	`+void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)`
`242`	`242`	`{`
`243`	`243`	`const T one{1, 0};`
`244`	`244`	`const T zero{0, 0};`
`@@ -392,6 +392,12 @@ void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)`
`392`	`392`	`}`
`393`	`393`	`}`
`394`	`394`	`delmem_complex_op()(this->ctx, becp);`
	`395`	`+}`
	`396`	`+`
	`397`	`+template <typename T, typename Device>`
	`398`	`+void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)`
	`399`	`+{`
	`400`	`+ this->cal_becsum(psi);`
`395`	`401`
`396`	`402`	`// transform soft charge to recip space using smooth grids`
`397`	`403`	`T* rhog = nullptr;`
Original file line number	Diff line number	Diff line change
`@@ -281,6 +281,10 @@ void HSolverLIP<T>::solve(hamilt::Hamilt<T>* pHamilt, // ESolver_KS_PW::p_hamilt`
`281`	`281`	`reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();`
`282`	`282`	`if (skip_charge)`
`283`	`283`	`{`
	`284`	`+ if (PARAM.globalv.use_uspp)`
	`285`	`+ {`
	`286`	`+ reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->cal_becsum(psi);`
	`287`	`+ }`
`284`	`288`	`ModuleBase::timer::tick("HSolverLIP", "solve");`
`285`	`289`	`return;`
`286`	`290`	`}`
Original file line number	Diff line number	Diff line change
`@@ -337,6 +337,10 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,`
`337`	`337`	`reinterpret_cast<elecstate::ElecStatePW<T>*>(pes)->calEBand();`
`338`	`338`	`if (skip_charge)`
`339`	`339`	`{`
	`340`	`+ if (PARAM.globalv.use_uspp)`
	`341`	`+ {`
	`342`	`+ reinterpret_cast<elecstate::ElecStatePW<T, Device>*>(pes)->cal_becsum(psi);`
	`343`	`+ }`
`340`	`344`	`ModuleBase::timer::tick("HSolverPW", "solve");`
`341`	`345`	`return;`
`342`	`346`	`}`
Original file line number	Diff line number	Diff line change
`@@ -207,7 +207,7 @@ void ReadInput::item_system()`
`207`	`207`	`item.annotation = "if calculate the force at the end of the electronic iteration";`
`208`	`208`	`item.reset_value = [](const Input_Item& item, Parameter& para) {`
`209`	`209`	`std::vector<std::string> use_force = {"cell-relax", "relax", "md"};`
`210`		`- std::vector<std::string> not_use_force = {"get_wf", "get_pchg", "nscf", "get_S"};`
	`210`	`+ std::vector<std::string> not_use_force = {"get_wf", "get_pchg", "get_S"};`
`211`	`211`	`if (std::find(use_force.begin(), use_force.end(), para.input.calculation) != use_force.end())`
`212`	`212`	`{`
`213`	`213`	`if (!para.input.cal_force)`