fix unit test

Author: haozhihan

source/module_io/test/write_wfc_nao_test.cpp

@@ -167,7 +167,7 @@ class WriteWfcLcaoTest : public testing::Test
 TEST_F(WriteWfcLcaoTest, WriteWfcLcao)
 {
     // create a psi object
-    psi::Psi<double> my_psi(psi_local_double.data(), nk, nbands_local, nbasis_local, true);
+    psi::Psi<double> my_psi(psi_local_double.data(), nk, nbands_local, nbasis_local, nbasis_local, true);
     PARAM.sys.global_out_dir = "./";
     ModuleIO::write_wfc_nao(2, my_psi, ekb, wg, kvec_c, pv, -1);
 

source/module_lr/utils/lr_util.hpp

@@ -104,14 +104,19 @@ namespace LR_Util
 
     /// psi(nk=1, nbands=nb, nk * nbasis) -> psi(nb, nk, nbasis) without memory copy
     template<typename T, typename Device>
-    psi::Psi<T, Device> k1_to_bfirst_wrapper(const psi::Psi<T, Device>& psi_kfirst, int nk_in, int nbasis_in)
+    psi::Psi<T, Device> c(const psi::Psi<T, Device>& psi_kfirst, int nk_in, int nbasis_in)
     {
         assert(psi_kfirst.get_nk() == 1);
         assert(nk_in * nbasis_in == psi_kfirst.get_nbasis());
 
         int ib_now = psi_kfirst.get_current_b();
         psi_kfirst.fix_b(0);    // for get_pointer() to get the head pointer
-        psi::Psi<T, Device> psi_bfirst(psi_kfirst.get_pointer(), nk_in, psi_kfirst.get_nbands(), nbasis_in, false);
+        psi::Psi<T, Device> psi_bfirst(psi_kfirst.get_pointer(), 
+                                       nk_in, 
+                                       psi_kfirst.get_nbands(), 
+                                       nbasis_in, 
+                                       nbasis_in, 
+                                       false);
         psi_kfirst.fix_b(ib_now);
         return psi_bfirst;
     }
@@ -124,7 +129,12 @@ namespace LR_Util
         int ik_now = psi_bfirst.get_current_k();
 
         psi_bfirst.fix_kb(0, 0);    // for get_pointer() to get the head pointer
-        psi::Psi<T, Device> psi_kfirst(psi_bfirst.get_pointer(), 1, psi_bfirst.get_nbands(), psi_bfirst.get_nk() * psi_bfirst.get_nbasis(), true);
+        psi::Psi<T, Device> psi_kfirst(psi_bfirst.get_pointer(), 
+                                       1, 
+                                       psi_bfirst.get_nbands(), 
+                                       psi_bfirst.get_nk() * psi_bfirst.get_nbasis(), 
+                                       psi_bfirst.get_nk() * psi_bfirst.get_nbasis(),
+                                       true);
         psi_bfirst.fix_kb(ik_now, ib_now);
         return psi_kfirst;
     }