update Constructor in psi

Author: haozhihan

source/module_esolver/esolver_ks_lcao.cpp

@@ -1080,7 +1080,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

@@ -220,7 +220,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_io/get_pchg_lcao.cpp

@@ -541,8 +541,12 @@ void IState_Charge::idmatrix(const int& ib,
         }
 
         this->psi_k->fix_k(ik);
-        // psi::Psi<std::complex<double>> wg_wfc(*this->psi_k, 1);
-        psi::Psi<std::complex<double>> wg_wfc(1, this->psi_k->get_nbands(), this->psi_k->get_nbasis());
+        
+        psi::Psi<std::complex<double>> wg_wfc(1, 
+                                              this->psi_k->get_nbands(), 
+                                              this->psi_k->get_nbasis(),
+                                              this->psi_k->get_nbasis(),
+                                              true);
 
         for (int ir = 0; ir < wg_wfc.get_nbands(); ++ir)
         {

source/module_lr/esolver_lrtd_lcao.cpp

@@ -181,7 +181,11 @@ LR::ESolver_LR<T, TR>::ESolver_LR(ModuleESolver::ESolver_KS_LCAO<T, TR>&& ks_sol
     if (this->nbands == PARAM.inp.nbands) { move_gs(); }
     else    // copy the part of ground state info according to paraC_
     {
-        this->psi_ks = new psi::Psi<T>(this->kv.get_nks(), this->paraC_.get_col_size(), this->paraC_.get_row_size());
+        this->psi_ks = new psi::Psi<T>(this->kv.get_nks(), 
+                                       this->paraC_.get_col_size(), 
+                                       this->paraC_.get_row_size(),
+                                       this->kv.ngk,
+                                       true);
         this->eig_ks.create(this->kv.get_nks(), this->nbands);
         const int start_band = this->nocc_max - *std::max_element(nocc.begin(), nocc.end());
         for (int ik = 0;ik < this->kv.get_nks();++ik)
@@ -289,8 +293,10 @@ LR::ESolver_LR<T, TR>::ESolver_LR(const Input_para& inp, UnitCell& ucell) : inpu
     // now ModuleIO::read_wfc_nao needs `Parallel_Orbitals` and can only read all the bands
     // it need improvement to read only the bands needed
     this->psi_ks = new psi::Psi<T>(this->kv.get_nks(),
-        this->paraMat_.ncol_bands,
-        this->paraMat_.get_row_size());
+                                   this->paraMat_.ncol_bands,
+                                   this->paraMat_.get_row_size(), 
+                                   this->kv.ngk,
+                                   true);
     this->read_ks_wfc();
     if (nspin == 2)
     {

source/module_lr/hamilt_casida.cpp

@@ -18,13 +18,17 @@ namespace LR
                 {//calculate A^{ai} for each bj
                     int bj = j * nv + b;    //global
                     int kbj = ik * npairs + bj; //global
-                    psi::Psi<T> X_bj(1, 1, this->nk * px.get_local_size()); // k1-first, like in iterative solver
+                    psi::Psi<T> X_bj(1, 1, this->nk * px.get_local_size(), this->nk * px.get_local_size(), true); // k1-first, like in iterative solver
                     X_bj.zero_out();
                     // X_bj(0, 0, lj * px.get_row_size() + lb) = this->one();
                     int lj = px.global2local_col(j);
                     int lb = px.global2local_row(b);
                     if (px.in_this_processor(b, j)) { X_bj(0, 0, ik * px.get_local_size() + lj * px.get_row_size() + lb) = this->one(); }
-                    psi::Psi<T> A_aibj(1, 1, this->nk * px.get_local_size()); // k1-first
+                    psi::Psi<T> A_aibj(1, 
+                                       1, 
+                                       this->nk * px.get_local_size(),
+                                       this->nk * px.get_local_size(),
+                                       true); // k1-first
                     A_aibj.zero_out();
 
                     this->cal_dm_trans(0, X_bj.get_pointer());

source/module_psi/psi.cpp

@@ -93,7 +93,7 @@ Psi<T, Device>::Psi(T* psi_pointer,
 {
 
     // Currently this function only supports nk_in == 1 when called within diagH_subspace_init.
-    // assert(nk_in == 1);
+    // assert(nk_in == 1); // NOTE because lr/utils/lr_uril.hpp func
 
     this->k_first = k_first_in;
     this->npol = PARAM.globalv.npol;

source/module_psi/psi.h

@@ -40,7 +40,7 @@ class Psi
     Psi();
 
     // Constructor 3: specify nk, nbands, nbasis, ngk, and do not need to call resize() later
-    Psi(const int nk_in, const int nbd_in, const int nbs_in, const int* ngk_in = nullptr, const bool k_first_in = true);
+    Psi(const int nk_in, const int nbd_in, const int nbs_in, const int* ngk_in, const bool k_first_in = true);
 
     Psi(const int nk_in, const int nbd_in, const int nbs_in, const std::vector<int>& ngk_in, const bool k_first_in);