Refactor: refactor npol, constructors and int *ngk of Psi class (#…

…5863) * change npol to private * fix cuda build bug * fix cuda build bug * fix build bug in cuda * remove npol value in psi * fix bug * fix bugs * update constructers * fix bug * fix test bug * remove Constructor 1-1 * fix bug * update psi * remove useless code
deepmodeling · Jan 20, 2025 · 51b4c88 · 51b4c88
1 parent e927eab
commit 51b4c88
Show file tree

Hide file tree

Showing 32 changed files with 170 additions and 182 deletions.
diff --git a/source/module_elecstate/cal_dm.h b/source/module_elecstate/cal_dm.h
@@ -82,7 +82,7 @@ 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<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

diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp
@@ -271,7 +271,7 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
 {
     const T one{1, 0};
     const T zero{0, 0};
-    const int npol = psi.npol;
+    const int npol = psi.get_npol();
     const int npwx = psi.get_nbasis() / npol;
     const int nbands = psi.get_nbands() * npol;
     const int nkb = this->ppcell->nkb;

diff --git a/source/module_esolver/esolver_ks_lcao_tddft.cpp b/source/module_esolver/esolver_ks_lcao_tddft.cpp
@@ -196,9 +196,9 @@ void ESolver_KS_LCAO_TDDFT::update_pot(UnitCell& ucell, const int istep, const i
         if (this->psi_laststep == nullptr)
         {
 #ifdef __MPI
-            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), ncol_nbands, nrow, nullptr);
+            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), ncol_nbands, nrow, kv.ngk, true);
 #else
-            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), nbands, nlocal, nullptr);
+            this->psi_laststep = new psi::Psi<std::complex<double>>(kv.get_nks(), nbands, nlocal, kv.ngk, true);
 #endif
         }
 

diff --git a/source/module_esolver/esolver_ks_lcaopw.cpp b/source/module_esolver/esolver_ks_lcaopw.cpp
@@ -93,9 +93,10 @@ namespace ModuleESolver
         ESolver_KS_PW<T>::before_all_runners(ucell, inp);
         delete this->psi_local;
         this->psi_local = new psi::Psi<T>(this->psi->get_nk(),
-                                                             this->p_psi_init->psi_initer->nbands_start(),
-                                                             this->psi->get_nbasis(),
-                                                             this->psi->get_ngk_pointer());
+                                          this->p_psi_init->psi_initer->nbands_start(),
+                                          this->psi->get_nbasis(),
+                                          this->kv.ngk,
+                                          true);
 #ifdef __EXX
         if (PARAM.inp.calculation == "scf" || PARAM.inp.calculation == "relax"
             || PARAM.inp.calculation == "cell-relax"

diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -212,7 +212,7 @@ void ESolver_KS_PW<T, Device>::before_all_runners(UnitCell& ucell, const Input_p
                                                    this->kv,
                                                    this->ppcell,
                                                    *this->pw_wfc);
-    allocate_psi(this->psi, this->kv.get_nks(), this->kv.ngk.data(), PARAM.inp.nbands, this->pw_wfc->npwk_max);
+    allocate_psi(this->psi, this->kv.get_nks(), this->kv.ngk, PARAM.inp.nbands, this->pw_wfc->npwk_max);
     this->p_psi_init->prepare_init(PARAM.inp.pw_seed);
 
     this->kspw_psi = PARAM.inp.device == "gpu" || PARAM.inp.precision == "single"

diff --git a/source/module_esolver/esolver_sdft_pw.cpp b/source/module_esolver/esolver_sdft_pw.cpp
@@ -78,13 +78,20 @@ void ESolver_SDFT_PW<T, Device>::before_all_runners(UnitCell& ucell, const Input
     // 4) allocate spaces for \sqrt(f(H))|chi> and |\tilde{chi}>
     size_t size = stowf.chi0->size();
     this->stowf.shchi
-        = new psi::Psi<T, Device>(this->kv.get_nks(), this->stowf.nchip_max, this->pw_wfc->npwk_max, this->kv.ngk.data());
+        = new psi::Psi<T, Device>(this->kv.get_nks(), 
+                                  this->stowf.nchip_max, 
+                                  this->pw_wfc->npwk_max, 
+                                  this->kv.ngk,
+                                  true);
     ModuleBase::Memory::record("SDFT::shchi", size * sizeof(T));
 
     if (PARAM.inp.nbands > 0)
     {
         this->stowf.chiortho
-            = new psi::Psi<T, Device>(this->kv.get_nks(), this->stowf.nchip_max, this->pw_wfc->npwk_max, this->kv.ngk.data());
+            = new psi::Psi<T, Device>(this->kv.get_nks(), 
+                                      this->stowf.nchip_max, 
+                                      this->pw_wfc->npwk_max, 
+                                      this->kv.ngk, true);
         ModuleBase::Memory::record("SDFT::chiortho", size * sizeof(T));
     }
 

diff --git a/source/module_esolver/lcao_before_scf.cpp b/source/module_esolver/lcao_before_scf.cpp
@@ -159,7 +159,7 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
             ncol = PARAM.inp.nbands;
 #endif
         }
-        this->psi = new psi::Psi<TK>(nsk, ncol, this->pv.nrow, nullptr);
+        this->psi = new psi::Psi<TK>(nsk, ncol, this->pv.nrow, this->kv.ngk, true);
     }
 
     // init wfc from file

diff --git a/source/module_esolver/lcao_others.cpp b/source/module_esolver/lcao_others.cpp
@@ -165,7 +165,7 @@ void ESolver_KS_LCAO<TK, TR>::others(UnitCell& ucell, const int istep)
             ncol = PARAM.inp.nbands;
 #endif
         }
-        this->psi = new psi::Psi<TK>(nsk, ncol, this->pv.nrow, nullptr);
+        this->psi = new psi::Psi<TK>(nsk, ncol, this->pv.nrow, this->kv.ngk, true);
     }
 
     // init wfc from file

diff --git a/source/module_hamilt_general/operator.cpp b/source/module_hamilt_general/operator.cpp
@@ -63,7 +63,7 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
         delete this->hpsi;
         this->hpsi = new psi::Psi<T, Device>(hpsi_pointer, 
                                              1, 
-                                             nbands / psi_input->npol,
+                                             nbands / psi_input->get_npol(),
                                              psi_input->get_nbasis(),
                                              psi_input->get_nbasis(),
                                              true);
@@ -86,7 +86,7 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
         default:
             op->act(nbands,
                     psi_input->get_nbasis(),
-                    psi_input->npol,
+                    psi_input->get_npol(),
                     tmpsi_in,
                     this->hpsi->get_pointer(),
                     psi_input->get_current_nbas(),
@@ -105,7 +105,7 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
     }
     ModuleBase::timer::tick("Operator", "hPsi");
 
-    return hpsi_info(this->hpsi, psi::Range(1, 0, 0, nbands / psi_input->npol), hpsi_pointer);
+    return hpsi_info(this->hpsi, psi::Range(1, 0, 0, nbands / psi_input->get_npol()), hpsi_pointer);
 }
 
 template <typename T, typename Device>

diff --git a/source/module_hamilt_lcao/module_deltaspin/cal_mw.cpp b/source/module_hamilt_lcao/module_deltaspin/cal_mw.cpp
@@ -66,7 +66,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mi_pw()
         psi::Psi<std::complex<double>, base_device::DEVICE_CPU>* psi_t = static_cast<psi::Psi<std::complex<double>, base_device::DEVICE_CPU>*>(this->psi);
         const int nbands = psi_t->get_nbands();
         const int nks = psi_t->get_nk();
-        const int npol = psi_t->npol;
+        const int npol = psi_t->get_npol();
         for(int ik = 0; ik < nks; ik++)
         {
             psi_t->fix_k(ik);
@@ -112,7 +112,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mi_pw()
         psi::Psi<std::complex<double>, base_device::DEVICE_GPU>* psi_t = static_cast<psi::Psi<std::complex<double>, base_device::DEVICE_GPU>*>(this->psi);
         const int nbands = psi_t->get_nbands();
         const int nks = psi_t->get_nk();
-        const int npol = psi_t->npol;
+        const int npol = psi_t->get_npol();
         for(int ik = 0; ik < nks; ik++)
         {
             psi_t->fix_k(ik);

diff --git a/source/module_hamilt_lcao/module_deltaspin/cal_mw_from_lambda.cpp b/source/module_hamilt_lcao/module_deltaspin/cal_mw_from_lambda.cpp
@@ -199,7 +199,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mw_from_lambda(int
                 hamilt::Hamilt<std::complex<double>, base_device::DEVICE_CPU>* hamilt_t = static_cast<hamilt::Hamilt<std::complex<double>, base_device::DEVICE_CPU>*>(this->p_hamilt);
                 auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_CPU>::get_instance();
                 nbands = psi_t->get_nbands();
-                npol = psi_t->npol;
+                npol = psi_t->get_npol();
                 nkb = onsite_p->get_tot_nproj();
                 nk = psi_t->get_nk();
                 nh_iat = &onsite_p->get_nh(0);
@@ -252,7 +252,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::cal_mw_from_lambda(int
                 hamilt::Hamilt<std::complex<double>, base_device::DEVICE_GPU>* hamilt_t = static_cast<hamilt::Hamilt<std::complex<double>, base_device::DEVICE_GPU>*>(this->p_hamilt);
                 auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_GPU>::get_instance();
                 nbands = psi_t->get_nbands();
-                npol = psi_t->npol;
+                npol = psi_t->get_npol();
                 nkb = onsite_p->get_tot_nproj();
                 nk = psi_t->get_nk();
                 nh_iat = &onsite_p->get_nh(0);
@@ -382,7 +382,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::update_psi_charge(const
             hamilt::Hamilt<std::complex<double>, base_device::DEVICE_CPU>* hamilt_t = static_cast<hamilt::Hamilt<std::complex<double>, base_device::DEVICE_CPU>*>(this->p_hamilt);
             auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_CPU>::get_instance();
             nbands = psi_t->get_nbands();
-            npol = psi_t->npol;
+            npol = psi_t->get_npol();
             nkb = onsite_p->get_tot_nproj();
             nk = psi_t->get_nk();
             nh_iat = &onsite_p->get_nh(0);
@@ -454,7 +454,7 @@ void spinconstrain::SpinConstrain<std::complex<double>>::update_psi_charge(const
             hamilt::Hamilt<std::complex<double>, base_device::DEVICE_GPU>* hamilt_t = static_cast<hamilt::Hamilt<std::complex<double>, base_device::DEVICE_GPU>*>(this->p_hamilt);
             auto* onsite_p = projectors::OnsiteProjector<double, base_device::DEVICE_GPU>::get_instance();
             nbands = psi_t->get_nbands();
-            npol = psi_t->npol;
+            npol = psi_t->get_npol();
             nkb = onsite_p->get_tot_nproj();
             nk = psi_t->get_nk();
             nh_iat = &onsite_p->get_nh(0);

diff --git a/source/module_hamilt_lcao/module_dftu/dftu_pw.cpp b/source/module_hamilt_lcao/module_dftu/dftu_pw.cpp
@@ -29,11 +29,11 @@ void DFTU::cal_occ_pw(const int iter, const void* psi_in, const ModuleBase::matr
             psi_p->fix_k(ik);
             onsite_p->tabulate_atomic(ik);
 
-            onsite_p->overlap_proj_psi(nbands*psi_p->npol, psi_p->get_pointer());
+            onsite_p->overlap_proj_psi(nbands*psi_p->get_npol(), psi_p->get_pointer());
             const std::complex<double>* becp = onsite_p->get_h_becp();
             // becp(nbands*npol , nkb)
             // mag = wg * \sum_{nh}becp * becp
-            int nkb = onsite_p->get_size_becp() / nbands / psi_p->npol;
+            int nkb = onsite_p->get_size_becp() / nbands / psi_p->get_npol();
             int begin_ih = 0;
             for(int iat = 0; iat < cell.nat; iat++)
             {
@@ -88,11 +88,11 @@ void DFTU::cal_occ_pw(const int iter, const void* psi_in, const ModuleBase::matr
             psi_p->fix_k(ik);
             onsite_p->tabulate_atomic(ik);
 
-            onsite_p->overlap_proj_psi(nbands*psi_p->npol, psi_p->get_pointer());
+            onsite_p->overlap_proj_psi(nbands*psi_p->get_npol(), psi_p->get_pointer());
             const std::complex<double>* becp = onsite_p->get_h_becp();
             // becp(nbands*npol , nkb)
             // mag = wg * \sum_{nh}becp * becp
-            int nkb = onsite_p->get_size_becp() / nbands / psi_p->npol;
+            int nkb = onsite_p->get_size_becp() / nbands / psi_p->get_npol();
             int begin_ih = 0;
             for(int iat = 0; iat < cell.nat; iat++)
             {

diff --git a/source/module_hamilt_pw/hamilt_pwdft/onsite_projector.cpp b/source/module_hamilt_pw/hamilt_pwdft/onsite_projector.cpp
@@ -165,7 +165,7 @@ void projectors::OnsiteProjector<T, Device>::init(const std::string& orbital_dir
         RadialProjection::RadialProjector::_build_backward_map(it2iproj, lproj, irow2it_, irow2iproj_, irow2m_);
         RadialProjection::RadialProjector::_build_forward_map(it2ia, it2iproj, lproj, itiaiprojm2irow_);
         //rp_._build_sbt_tab(rgrid, projs, lproj, nq, dq);
-        rp_._build_sbt_tab(nproj, rgrid, projs, lproj, nq, dq, ucell_in->omega, psi.npol, tab, nhtol);
+        rp_._build_sbt_tab(nproj, rgrid, projs, lproj, nq, dq, ucell_in->omega, psi.get_npol(), tab, nhtol);
         // For being compatible with present cal_force and cal_stress framework  
         // uncomment the following code block if you want to use the Onsite_Proj_tools
         if(this->tab_atomic_ == nullptr)
@@ -541,7 +541,7 @@ void projectors::OnsiteProjector<T, Device>::cal_occupations(const psi::Psi<std:
         }
         // std::cout << __FILE__ << ":" << __LINE__ << " nbands = " << nbands << std::endl;
         this->overlap_proj_psi(
-                        nbands * psi_in->npol,
+                        nbands * psi_in->get_npol(),
                         psi_in->get_pointer());
         const std::complex<double>* becp_p = this->get_h_becp();
         // becp(nbands*npol , nkb)

diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/velocity_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/velocity_pw.cpp
@@ -51,8 +51,8 @@ void Velocity::act
 
     const int npw = psi_in->get_current_nbas();
 
-    const int max_npw = psi_in->get_nbasis() / psi_in->npol;
-    const int npol = psi_in->npol;
+    const int max_npw = psi_in->get_nbasis() / psi_in->get_npol();
+    const int npol = psi_in->get_npol();
     const std::complex<double>* tmpsi_in = psi0;
     std::complex<double>* tmhpsi = vpsi;
     // -------------

diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_wf.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_wf.cpp
@@ -35,7 +35,7 @@ template <typename T, typename Device>
 void Stochastic_WF<T, Device>::init(K_Vectors* p_kv, const int npwx_in)
 {
     this->nks = p_kv->get_nks();
-    this->ngk = p_kv->ngk.data();
+    this->ngk = p_kv->ngk;
     this->npwx = npwx_in;
     nchip = new int[nks];
 
@@ -111,7 +111,7 @@ void Stochastic_WF<T, Device>::allocate_chi0()
 
     this->nchip_max = tmpnchip;
     size_t size = this->nchip_max * npwx * nks;
-    this->chi0_cpu = new psi::Psi<T>(nks, this->nchip_max, npwx, this->ngk);
+    this->chi0_cpu = new psi::Psi<T>(nks, this->nchip_max, npwx, this->ngk, true);
     ModuleBase::Memory::record("SDFT::chi0_cpu", size * sizeof(T));
 
     for (int ik = 0; ik < nks; ++ik)
@@ -123,7 +123,7 @@ void Stochastic_WF<T, Device>::allocate_chi0()
     Device* ctx = {};
     if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
     {
-        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk);
+        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk, true);
     }
     else
     {
@@ -207,7 +207,7 @@ void Stochastic_WF<T, Device>::init_com_orbitals()
         delete[] npwip;
     }
     size_t size = this->nchip_max * npwx * nks;
-    this->chi0_cpu = new psi::Psi<std::complex<double>>(nks, this->nchip_max, npwx, this->ngk);
+    this->chi0_cpu = new psi::Psi<std::complex<double>>(nks, this->nchip_max, npwx, this->ngk, true);
     this->chi0_cpu->zero_out();
     ModuleBase::Memory::record("SDFT::chi0_cpu", size * sizeof(std::complex<double>));
     for (int ik = 0; ik < nks; ++ik)
@@ -252,7 +252,7 @@ void Stochastic_WF<T, Device>::init_com_orbitals()
     Device* ctx = {};
     if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
     {
-        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk);
+        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk, true);
     }
     else
     {
@@ -266,7 +266,7 @@ void Stochastic_WF<T, Device>::init_com_orbitals()
     const int npwx = this->npwx;
     const int nks = this->nks;
     size_t size = this->nchip_max * npwx * nks;
-    this->chi0_cpu = new psi::Psi<std::complex<double>>(nks, npwx, npwx, this->ngk);
+    this->chi0_cpu = new psi::Psi<std::complex<double>>(nks, npwx, npwx, this->ngk, true);
     this->chi0_cpu->zero_out();
     ModuleBase::Memory::record("SDFT::chi0_cpu", size * sizeof(std::complex<double>));
     for (int ik = 0; ik < nks; ++ik)
@@ -284,7 +284,7 @@ void Stochastic_WF<T, Device>::init_com_orbitals()
     Device* ctx = {};
     if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
     {
-        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk);
+        this->chi0 = new psi::Psi<T, Device>(nks, this->nchip_max, npwx, this->ngk, true);
     }
     else
     {

diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_wf.h b/source/module_hamilt_pw/hamilt_stodft/sto_wf.h
@@ -30,10 +30,10 @@ class Stochastic_WF
     int* nchip = nullptr; ///< The number of stochatic orbitals in current process of each k point.
     int nchip_max = 0;    ///< Max number of stochastic orbitals among all k points.
     int nks = 0;          ///< number of k-points
-    int* ngk = nullptr;   ///< ngk in klist
     int npwx = 0;         ///< max ngk[ik] in all processors
     int nbands_diag = 0;  ///< number of bands obtained from diagonalization
     int nbands_total = 0; ///< number of bands in total, nbands_total=nchi+nbands_diag;
+    std::vector<int> ngk;   ///< ngk in klist
   public:
     // Tn(H)|chi>
     psi::Psi<T, Device>* chiallorder = nullptr;

diff --git a/source/module_hsolver/hsolver_lcao.cpp b/source/module_hsolver/hsolver_lcao.cpp
@@ -219,7 +219,7 @@ void HSolverLCAO<T, Device>::parakSolve(hamilt::Hamilt<T>* pHamilt,
         k2d.distribute_hsk(pHamilt, ik_kpar, nrow);
         /// global index of k point
         int ik_global = ik + k2d.get_pKpoints()->startk_pool[k2d.get_my_pool()];
-        auto psi_pool = psi::Psi<T>(1, ncol_bands_pool, k2d.get_p2D_pool()->nrow, nullptr);
+        auto psi_pool = psi::Psi<T>(1, ncol_bands_pool, k2d.get_p2D_pool()->nrow, k2d.get_p2D_pool()->nrow, true);
         ModuleBase::Memory::record("HSolverLCAO::psi_pool", nrow * ncol_bands_pool * sizeof(T));
         if (ik_global < psi.get_nk() && ik < k2d.get_pKpoints()->nks_pool[k2d.get_my_pool()])
         {

diff --git a/source/module_hsolver/test/diago_mock.h b/source/module_hsolver/test/diago_mock.h
@@ -214,7 +214,7 @@ class HPsi
     {
         Structure_Factor* sf;
         int* ngk = nullptr;
-        psi::Psi<T> psitmp(1, nband, npw, ngk);
+        psi::Psi<T> psitmp(1, nband, npw, npw, true);
         for(int i=0;i<nband;i++)
 	    {
 		    for(int j=0;j<npw;j++) { psitmp(0,i,j) = psimatrix[i * npw + j];

diff --git a/source/module_io/test/read_wfc_to_rho_test.cpp b/source/module_io/test/read_wfc_to_rho_test.cpp
@@ -181,7 +181,7 @@ TEST_F(ReadWfcRhoTest, ReadWfcRho)
     }
 
     // Init Psi
-    psi = new psi::Psi<std::complex<double>>(nks, nbands, wfcpw->npwk_max, wfcpw->npwk);
+    psi = new psi::Psi<std::complex<double>>(nks, nbands, wfcpw->npwk_max, kv->ngk, true);
     std::complex<double>* ptr = psi->get_pointer();
     for (int i = 0; i < nks * nbands * wfcpw->npwk_max; i++)
     {

diff --git a/source/module_io/to_wannier90_lcao_in_pw.cpp b/source/module_io/to_wannier90_lcao_in_pw.cpp
@@ -52,7 +52,11 @@ void toWannier90_LCAO_IN_PW::calculate(
     const int nks_psi = (PARAM.inp.calculation == "nscf" && PARAM.inp.mem_saver == 1)? 1 : wfcpw->nks;
     const int nks_psig = (PARAM.inp.basis_type == "pw")? 1 : nks_psi;
     const int nbands_actual = this->psi_initer_->nbands_start();
-    this->psi = new psi::Psi<std::complex<double>, base_device::DEVICE_CPU>(nks_psig, nbands_actual, wfcpw->npwk_max*PARAM.globalv.npol, wfcpw->npwk);
+    this->psi = new psi::Psi<std::complex<double>, base_device::DEVICE_CPU>(nks_psig, 
+                                                                            nbands_actual, 
+                                                                            wfcpw->npwk_max*PARAM.globalv.npol, 
+                                                                            kv.ngk,
+                                                                            true);
     read_nnkp(ucell,kv);
 
     if (PARAM.inp.nspin == 2)
@@ -117,7 +121,11 @@ psi::Psi<std::complex<double>>* toWannier90_LCAO_IN_PW::get_unk_from_lcao(
 {
     // init
     int npwx = wfcpw->npwk_max;
-    psi::Psi<std::complex<double>> *unk_inLcao = new psi::Psi<std::complex<double>>(num_kpts, num_bands, npwx*PARAM.globalv.npol, kv.ngk.data());
+    psi::Psi<std::complex<double>> *unk_inLcao = new psi::Psi<std::complex<double>>(num_kpts, 
+                                                                                    num_bands, 
+                                                                                    npwx*PARAM.globalv.npol, 
+                                                                                    kv.ngk,
+                                                                                    true);
     unk_inLcao->zero_out();
 
     // Orbital projection to plane wave