fix bug

Author: haozhihan

source/module_elecstate/cal_dm.h

@@ -27,8 +27,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<double> wg_wfc(wfc, 1);
-        psi::Psi<double> wg_wfc(1, nbands_local, nbasis_local);
+        psi::Psi<double> wg_wfc(wfc, 1, nbands_local);
 
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)

source/module_elecstate/module_dm/cal_dm_psi.cpp

@@ -32,8 +32,8 @@ 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, nbands_local, nbasis_local);
+        
+        psi::Psi<double> wg_wfc(wfc, 1, nbands_local);
 
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)

source/module_hamilt_general/operator.cpp

@@ -156,8 +156,7 @@ T* Operator<T, Device>::get_hpsi(const hpsi_info& info) const
     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>(1, nbands_range, std::get<0>(info)->get_nbasis());
+        this->hpsi = new psi::Psi<T, Device>(std::get<0>(info)[0], 1, nbands_range);
     }
     else
     {

source/module_io/get_pchg_lcao.cpp

@@ -478,8 +478,8 @@ void IState_Charge::idmatrix(const int& ib,
 
         // wg_wfc(ib,iw) = wg[ib] * wfc(ib,iw);
         this->psi_gamma->fix_k(is);
-        // psi::Psi<double> wg_wfc(*this->psi_gamma, 1);
-        psi::Psi<double> wg_wfc(1, this->psi_gamma->get_nbands(), this->psi_gamma->get_nbasis());
+
+        psi::Psi<double> wg_wfc(*this->psi_gamma, 1, this->psi_gamma->get_nbands());
 
         for (int ir = 0; ir < wg_wfc.get_nbands(); ++ir)
         {

source/module_io/write_dos_lcao.cpp

@@ -461,8 +461,7 @@ void ModuleIO::write_dos_lcao(const UnitCell& ucell,
                     }
 
                     psi->fix_k(ik);
-                    // psi::Psi<std::complex<double>> Dwfc(psi[0], 1);
-                    psi::Psi<std::complex<double>> Dwfc(1, psi->get_nbands(), psi->get_nbasis());
+                    psi::Psi<std::complex<double>> Dwfc(*psi, 1, psi->get_nbands());
 
                     std::complex<double>* p_dwfc = Dwfc.get_pointer();
                     for (int index = 0; index < Dwfc.size(); ++index)

source/module_io/write_proj_band_lcao.cpp

@@ -221,8 +221,7 @@ void ModuleIO::write_proj_band_lcao(
 
                     // calculate Mulk
                     psi->fix_k(ik);
-                    // psi::Psi<std::complex<double>> Dwfc(psi[0], 1);
-                    psi::Psi<std::complex<double>> Dwfc(1, psi->get_nbands(), psi->get_nbasis());
+                    psi::Psi<std::complex<double>> Dwfc(psi[0], 1, psi->get_nbands());
 
                     std::complex<double>* p_dwfc = Dwfc.get_pointer();
                     for (int index = 0; index < Dwfc.size(); ++index)

source/module_psi/psi.cpp

@@ -107,6 +107,34 @@ Psi<T, Device>::Psi(T* psi_pointer, const int nk_in, const int nbd_in, const int
     base_device::information::print_device_info<Device>(this->ctx, GlobalV::ofs_device);
 }
 
+
+template <typename T, typename Device>
+Psi<T, Device>::Psi(const Psi& psi_in, const int nk_in, int nband_in)
+{
+    assert(nk_in <= psi_in.get_nk());
+    if (nband_in == 0)
+    {
+        nband_in = psi_in.get_nbands();
+    }
+    this->k_first = psi_in.get_k_first();
+    this->device = psi_in.device;
+    this->resize(nk_in, nband_in, psi_in.get_nbasis());
+    this->ngk = psi_in.ngk;
+    this->npol = psi_in.npol;
+    if (nband_in <= psi_in.get_nbands())
+    {
+        // copy from Psi from psi_in(current_k, 0, 0),
+        // if size of k is 1, current_k in new Psi is psi_in.current_k
+        if (nk_in == 1)
+        {
+            // current_k for this Psi only keep the spin index same as the copied Psi
+            this->current_k = psi_in.get_current_k();
+        }
+        synchronize_memory_op()(this->ctx, psi_in.get_device(), this->psi, psi_in.get_pointer(), this->size());
+    }
+}
+
+
 template <typename T, typename Device>
 Psi<T, Device>::Psi(T* psi_pointer, const Psi& psi_in, const int nk_in, int nband_in)
 {
@@ -208,11 +236,11 @@ template <typename T, typename Device>
 void Psi<T, Device>::resize(const int nks_in, const int nbands_in, const int nbasis_in)
 {
     assert(nks_in > 0 && nbands_in >= 0 && nbasis_in > 0);
-    
+
     // This function will delete the psi array first(if psi exist), then malloc a new memory for it.
     resize_memory_op()(this->ctx, this->psi, nks_in * static_cast<std::size_t>(nbands_in) * nbasis_in, "no_record");
 
-    this->zero_out();
+    // this->zero_out();
 
     this->nk = nks_in;
     this->nbands = nbands_in;

source/module_psi/psi.h

@@ -42,6 +42,10 @@ class 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);
 
+    // Constructor 4: copy a new Psi which have several k-points and several bands from inputted psi_in
+    Psi(const Psi& psi_in, const int nk_in, int nband_in = 0);
+    
+    
     // Constructor 5: a wrapper of a data pointer, used for Operator::hPsi()
     // in this case, fix_k can not be used
     Psi(T* psi_pointer, const Psi& psi_in, const int nk_in, int nband_in = 0);