add get_cur_effective_basis func

Author: haozhihan

source/module_elecstate/elecstate_pw.cpp

@@ -154,7 +154,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_cur_effective_basis();
     int current_spin = 0;
     if (PARAM.inp.nspin == 2)
     {
@@ -258,7 +258,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_cur_effective_basis();
 
         // 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_cur_effective_basis();
         int current_spin = 0;
         if (PARAM.inp.nspin == 2)
         {

source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp

@@ -60,7 +60,7 @@ void Stochastic_Iter<T, Device>::orthog(const int& ik, psi::Psi<T, Device>& psi,
     if (PARAM.inp.nbands > 0)
     {
         const int nchipk = stowf.nchip[ik];
-        const int npw = psi.get_current_nbas();
+        const int npw = psi.get_cur_effective_basis();
         const int npwx = psi.get_nbasis();
         stowf.chi0->fix_k(ik);
         stowf.chiortho->fix_k(ik);

source/module_psi/psi.cpp

@@ -295,6 +295,26 @@ const int& Psi<T, Device>::get_psi_bias() const
     return this->psi_bias;
 }
 
+template <typename T, typename Device>
+const int& Psi<T, Device>::get_cur_effective_basis() const
+{
+    if (this->npol == 1)
+    {
+        if (this->ngk != nullptr)
+        {
+            return this->ngk[this->current_k];
+        }
+        else
+        {
+            return this->current_nbasis;
+        }
+    }
+    else
+    {
+        return this->nbasis;
+    }
+}
+
 template <typename T, typename Device>
 const int& Psi<T, Device>::get_nk() const
 {

source/module_psi/psi.h

@@ -127,6 +127,8 @@ class Psi
     // return psi_bias
     const int& get_psi_bias() const;
 
+    const int& get_cur_effective_basis() const;
+
     // mark
     void zero_out();