Refactor: refactor the constructors of Psi class (#5761)

* remove Psi(const Psi& psi_in, const int nk_in, int nband_in);

* fix bug

* fix bug

* [pre-commit.ci lite] apply automatic fixes

* remove device value in psi

* update Psi(const Psi& psi_in, const int nk_in, int nband_in)

* update get_ngk usage

* fix bug about ngk

* [pre-commit.ci lite] apply automatic fixes

* fix bug

* format operator

* [pre-commit.ci lite] apply automatic fixes

* fix bug

* fix bug

* fix bug

* fix bug

* add get_cur_effective_basis func

* fix bug

* update get_cur_effective_basis

* check bugs

* update Constructor 8-1

* fix bug

* fix bug

* fix bug

* fix bug  maybe

* fix bug

* check correct

* check 1

* fix unit test

* fix unit bug

* update get_ngk func

* remove get-ngk in velocity-pw

* fix bug

* [pre-commit.ci lite] apply automatic fixes

* fix 186_PW_SKG_ALL bug

* format source/module_io/unk_overlap_pw.cpp

* update Constructor in  psi

* [pre-commit.ci lite] apply automatic fixes

* debug unit test

* fix ri test bug

* [pre-commit.ci lite] apply automatic fixes

* fix psi-ut bug

* remove Psi<T, Device>::Psi(T* psi_pointer, const Psi& psi_in, const int nk_in, int nband_in)

* remove useless code

* update Psi(const Psi& psi_in, const int nk_in, const int nband_in);

* remove Psi(const Psi& psi_in, const int nk_in, const int nband_in);

* refactor psi code

* fix sdft bug

* change to get_current_ngk

---------

Co-authored-by: pre-commit-ci-lite[bot] <117423508+pre-commit-ci-lite[bot]@users.noreply.github.com>

Author: haozhihan

source/module_elecstate/cal_dm.h

@@ -27,7 +27,12 @@ inline void cal_dm(const Parallel_Orbitals* ParaV, const ModuleBase::matrix& wg,
         //dm.fix_k(ik);
         dm[ik].create(ParaV->ncol, ParaV->nrow);
         // wg_wfc(ib,iw) = wg[ib] * wfc(ib,iw);
-        psi::Psi<double> wg_wfc(wfc, 1);
+        psi::Psi<double> wg_wfc(1, 
+                                wfc.get_nbands(), 
+                                wfc.get_nbasis(),
+                                wfc.get_nbasis(),
+                                true);
+        wg_wfc.set_all_psi(wfc.get_pointer(), wg_wfc.size());
 
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)
@@ -41,7 +46,8 @@ inline void cal_dm(const Parallel_Orbitals* ParaV, const ModuleBase::matrix& wg,
                     ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check global2local_col!");
                 }
             }
-            if (ib_global >= wg.nc) continue;
+            if (ib_global >= wg.nc) { continue;
+}
             const double wg_local = wg(ik, ib_global);
             double* wg_wfc_pointer = &(wg_wfc(0, ib_local, 0));
             BlasConnector::scal(nbasis_local, wg_local, wg_wfc_pointer, 1);
@@ -99,7 +105,8 @@ inline void cal_dm(const Parallel_Orbitals* ParaV, const ModuleBase::matrix& wg,
                     ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check global2local_col!");
                 }
             }
-            if (ib_global >= wg.nc) continue;
+            if (ib_global >= wg.nc) { continue;
+}
             const double wg_local = wg(ik, ib_global);
             std::complex<double>* wg_wfc_pointer = &(wg_wfc(0, ib_local, 0));
             BlasConnector::scal(nbasis_local, wg_local, wg_wfc_pointer, 1);

source/module_elecstate/elecstate_pw.cpp

@@ -183,7 +183,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)
 
     this->init_rho_data();
     int ik = psi.get_current_k();
-    int npw = psi.get_current_nbas();
+    int npw = psi.get_current_ngk();
     int current_spin = 0;
     if (PARAM.inp.nspin == 2)
     {
@@ -287,7 +287,7 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
         psi.fix_k(ik);
         const T* psi_now = psi.get_pointer();
         const int currect_spin = this->klist->isk[ik];
-        const int npw = psi.get_current_nbas();
+        const int npw = psi.get_current_ngk();
 
         // get |beta>
         if (this->ppcell->nkb > 0)

source/module_elecstate/elecstate_pw_cal_tau.cpp

@@ -15,7 +15,7 @@ void ElecStatePW<T, Device>::cal_tau(const psi::Psi<T, Device>& psi)
     for (int ik = 0; ik < psi.get_nk(); ++ik)
     {
         psi.fix_k(ik);
-        int npw = psi.get_current_nbas();
+        int npw = psi.get_current_ngk();
         int current_spin = 0;
         if (PARAM.inp.nspin == 2)
         {

source/module_elecstate/module_dm/cal_dm_psi.cpp

@@ -32,7 +32,14 @@ void cal_dm_psi(const Parallel_Orbitals* ParaV,
         // dm.fix_k(ik);
         // dm[ik].create(ParaV->ncol, ParaV->nrow);
         //  wg_wfc(ib,iw) = wg[ib] * wfc(ib,iw);
-        psi::Psi<double> wg_wfc(wfc, 1);
+
+        psi::Psi<double> wg_wfc(1, 
+                                wfc.get_nbands(), 
+                                wfc.get_nbasis(),
+                                wfc.get_nbasis(),
+                                true);
+        wg_wfc.set_all_psi(wfc.get_pointer(), wg_wfc.size());
+
 
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)
@@ -89,7 +96,12 @@ void cal_dm_psi(const Parallel_Orbitals* ParaV,
         // dm.fix_k(ik);
         //dm[ik].create(ParaV->ncol, ParaV->nrow);
         // wg_wfc(ib,iw) = wg[ib] * wfc(ib,iw);
-        psi::Psi<std::complex<double>> wg_wfc(1, wfc.get_nbands(), wfc.get_nbasis(), nullptr);
+        psi::Psi<std::complex<double>> wg_wfc(1, 
+                                              wfc.get_nbands(), 
+                                              wfc.get_nbasis(),
+                                              wfc.get_nbasis(),
+                                              true);
+        
         const std::complex<double>* pwfc = wfc.get_pointer();
         std::complex<double>* pwg_wfc = wg_wfc.get_pointer();
 #ifdef _OPENMP

source/module_esolver/esolver_ks_lcao.cpp

@@ -1083,7 +1083,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(UnitCell& ucell, const int istep)
         //! initialize the gradients of Etotal with respect to occupation numbers and wfc,
         //! and set all elements to 0.
         ModuleBase::matrix dE_dOccNum(this->pelec->wg.nr, this->pelec->wg.nc, true);
-        psi::Psi<TK> dE_dWfc(this->psi->get_nk(), this->psi->get_nbands(), this->psi->get_nbasis());
+        psi::Psi<TK> dE_dWfc(this->psi->get_nk(), this->psi->get_nbands(), this->psi->get_nbasis(), this->kv.ngk, true);
         dE_dWfc.zero_out();
 
         double Etotal_RDMFT = this->rdmft_solver.run(dE_dOccNum, dE_dWfc);

source/module_esolver/esolver_of.cpp

@@ -222,7 +222,11 @@ void ESolver_OF::before_opt(const int istep, UnitCell& ucell)
 
         // Refresh the arrays
         delete this->psi_;
-        this->psi_ = new psi::Psi<double>(1, PARAM.inp.nspin, this->pw_rho->nrxx);
+        this->psi_ = new psi::Psi<double>(1, 
+                                          PARAM.inp.nspin, 
+                                          this->pw_rho->nrxx,
+                                          this->pw_rho->nrxx,
+                                          true);
         for (int is = 0; is < PARAM.inp.nspin; ++is)
         {
             this->pphi_[is] = this->psi_->get_pointer(is);

source/module_esolver/esolver_of_tool.cpp

@@ -71,7 +71,11 @@ void ESolver_OF::init_elecstate(UnitCell& ucell)
 void ESolver_OF::allocate_array()
 {
     // Initialize the "wavefunction", which is sqrt(rho)
-    this->psi_ = new psi::Psi<double>(1, PARAM.inp.nspin, this->pw_rho->nrxx);
+    this->psi_ = new psi::Psi<double>(1, 
+                                      PARAM.inp.nspin, 
+                                      this->pw_rho->nrxx,
+                                      this->pw_rho->nrxx,
+                                      true);
     ModuleBase::Memory::record("OFDFT::Psi", sizeof(double) * PARAM.inp.nspin * this->pw_rho->nrxx);
     this->pphi_ = new double*[PARAM.inp.nspin];
     for (int is = 0; is < PARAM.inp.nspin; ++is)

source/module_hamilt_general/operator.cpp

@@ -4,28 +4,31 @@
 
 using namespace hamilt;
 
-
-template<typename T, typename Device>
-Operator<T, Device>::Operator(){}
-
-template<typename T, typename Device>
-Operator<T, Device>::~Operator() 
+template <typename T, typename Device>
+Operator<T, Device>::Operator()
 {
-    if(this->hpsi != nullptr) { delete this->hpsi;
 }
+
+template <typename T, typename Device>
+Operator<T, Device>::~Operator()
+{
+    if (this->hpsi != nullptr)
+    {
+        delete this->hpsi;
+    }
     Operator* last = this->next_op;
     Operator* last_sub = this->next_sub_op;
-    while(last != nullptr || last_sub != nullptr)
+    while (last != nullptr || last_sub != nullptr)
     {
-        if(last_sub != nullptr)
-        {//delete sub_chain first
+        if (last_sub != nullptr)
+        { // delete sub_chain first
             Operator* node_delete = last_sub;
             last_sub = last_sub->next_sub_op;
             node_delete->next_sub_op = nullptr;
             delete node_delete;
         }
         else
-        {//delete main chain if sub_chain is deleted
+        { // delete main chain if sub_chain is deleted
             Operator* node_delete = last;
             last_sub = last->next_sub_op;
             node_delete->next_sub_op = nullptr;
@@ -36,7 +39,7 @@ Operator<T, Device>::~Operator()
     }
 }
 
-template<typename T, typename Device>
+template <typename T, typename Device>
 typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& input) const
 {
     using syncmem_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
@@ -46,37 +49,51 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
 
     T* tmhpsi = this->get_hpsi(input);
     const T* tmpsi_in = std::get<0>(psi_info);
-    //if range in hpsi_info is illegal, the first return of to_range() would be nullptr
+    // if range in hpsi_info is illegal, the first return of to_range() would be nullptr
     if (tmpsi_in == nullptr)
     {
         ModuleBase::WARNING_QUIT("Operator", "please choose correct range of psi for hPsi()!");
     }
-    //if in_place, copy temporary hpsi to target hpsi_pointer, then delete hpsi and new a wrapper for return
+    // if in_place, copy temporary hpsi to target hpsi_pointer, then delete hpsi and new a wrapper for return
     T* hpsi_pointer = std::get<2>(input);
     if (this->in_place)
     {
         // ModuleBase::GlobalFunc::COPYARRAY(this->hpsi->get_pointer(), hpsi_pointer, this->hpsi->size());
         syncmem_op()(this->ctx, this->ctx, hpsi_pointer, this->hpsi->get_pointer(), this->hpsi->size());
         delete this->hpsi;
-        this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, *psi_input, 1, nbands / psi_input->npol);
+        this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, 
+                                             1, 
+                                             nbands / psi_input->npol,
+                                             psi_input->get_nbasis(),
+                                             psi_input->get_nbasis(),
+                                             true);
     }
 
     auto call_act = [&, this](const Operator* op, const bool& is_first_node) -> void {
-        
         // a "psi" with the bands of needed range
-        psi::Psi<T, Device> psi_wrapper(const_cast<T*>(tmpsi_in), 1, nbands, psi_input->get_nbasis(), true);
-        
-        
+        psi::Psi<T, Device> psi_wrapper(const_cast<T*>(tmpsi_in),
+                                        1,
+                                        nbands,
+                                        psi_input->get_nbasis(),
+                                        psi_input->get_nbasis(),
+                                        true);
+
         switch (op->get_act_type())
         {
         case 2:
             op->act(psi_wrapper, *this->hpsi, nbands);
             break;
         default:
-            op->act(nbands, psi_input->get_nbasis(), psi_input->npol, tmpsi_in, this->hpsi->get_pointer(), psi_input->get_ngk(op->ik), is_first_node);
+            op->act(nbands,
+                    psi_input->get_nbasis(),
+                    psi_input->npol,
+                    tmpsi_in,
+                    this->hpsi->get_pointer(),
+                    psi_input->get_current_nbas(),
+                    is_first_node);
             break;
         }
-        };
+    };
 
     ModuleBase::timer::tick("Operator", "hPsi");
     call_act(this, true); // first node
@@ -91,39 +108,43 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
     return hpsi_info(this->hpsi, psi::Range(1, 0, 0, nbands / psi_input->npol), hpsi_pointer);
 }
 
-
-template<typename T, typename Device>
-void Operator<T, Device>::init(const int ik_in) 
+template <typename T, typename Device>
+void Operator<T, Device>::init(const int ik_in)
 {
     this->ik = ik_in;
-    if(this->next_op != nullptr) {
+    if (this->next_op != nullptr)
+    {
         this->next_op->init(ik_in);
     }
 }
 
-template<typename T, typename Device>
-void Operator<T, Device>::add(Operator* next) 
+template <typename T, typename Device>
+void Operator<T, Device>::add(Operator* next)
 {
-    if(next==nullptr) { return;
-}
+    if (next == nullptr)
+    {
+        return;
+    }
     next->is_first_node = false;
-    if(next->next_op != nullptr) { this->add(next->next_op);
-}
+    if (next->next_op != nullptr)
+    {
+        this->add(next->next_op);
+    }
     Operator* last = this;
-    //loop to end of the chain
-    while(last->next_op != nullptr)
+    // loop to end of the chain
+    while (last->next_op != nullptr)
     {
-        if(next->cal_type==last->cal_type)
+        if (next->cal_type == last->cal_type)
         {
             break;
         }
         last = last->next_op;
     }
-    if(next->cal_type == last->cal_type)
+    if (next->cal_type == last->cal_type)
     {
-        //insert next to sub chain of current node
+        // insert next to sub chain of current node
         Operator* sub_last = last;
-        while(sub_last->next_sub_op != nullptr)
+        while (sub_last->next_sub_op != nullptr)
         {
             sub_last = sub_last->next_sub_op;
         }
@@ -136,34 +157,45 @@ void Operator<T, Device>::add(Operator* next)
     }
 }
 
-template<typename T, typename Device>
+template <typename T, typename Device>
 T* Operator<T, Device>::get_hpsi(const hpsi_info& info) const
 {
     const int nbands_range = (std::get<1>(info).range_2 - std::get<1>(info).range_1 + 1);
-    //in_place call of hPsi, hpsi inputs as new psi, 
-    //create a new hpsi and delete old hpsi later
+    // in_place call of hPsi, hpsi inputs as new psi,
+    // create a new hpsi and delete old hpsi later
     T* hpsi_pointer = std::get<2>(info);
     const T* psi_pointer = std::get<0>(info)->get_pointer();
-    if(this->hpsi != nullptr) 
+    if (this->hpsi != nullptr)
     {
         delete this->hpsi;
         this->hpsi = nullptr;
     }
-    if(!hpsi_pointer)
+    if (!hpsi_pointer)
     {
         ModuleBase::WARNING_QUIT("Operator::hPsi", "hpsi_pointer can not be nullptr");
     }
-    else if(hpsi_pointer == psi_pointer)
+    else if (hpsi_pointer == psi_pointer)
     {
         this->in_place = true;
-        this->hpsi = new psi::Psi<T, Device>(std::get<0>(info)[0], 1, nbands_range);
+        // this->hpsi = new psi::Psi<T, Device>(std::get<0>(info)[0], 1, nbands_range);
+        this->hpsi = new psi::Psi<T, Device>(1, 
+                                             nbands_range,
+                                             std::get<0>(info)->get_nbasis(),
+                                             std::get<0>(info)->get_nbasis(),
+                                             true);
     }
     else
     {
         this->in_place = false;
-        this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, std::get<0>(info)[0], 1, nbands_range);
+        
+        this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, 
+                                             1, 
+                                             nbands_range, 
+                                             std::get<0>(info)->get_nbasis(), 
+                                             std::get<0>(info)->get_nbasis(), 
+                                             true);
     }
-    
+
     hpsi_pointer = this->hpsi->get_pointer();
     size_t total_hpsi_size = nbands_range * this->hpsi->get_nbasis();
     // ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size);
@@ -172,7 +204,8 @@ T* Operator<T, Device>::get_hpsi(const hpsi_info& info) const
     return hpsi_pointer;
 }
 
-namespace hamilt {
+namespace hamilt
+{
 template class Operator<float, base_device::DEVICE_CPU>;
 template class Operator<std::complex<float>, base_device::DEVICE_CPU>;
 template class Operator<double, base_device::DEVICE_CPU>;
@@ -183,4 +216,4 @@ template class Operator<std::complex<float>, base_device::DEVICE_GPU>;
 template class Operator<double, base_device::DEVICE_GPU>;
 template class Operator<std::complex<double>, base_device::DEVICE_GPU>;
 #endif
-}
+} // namespace hamilt