fix bug

Author: haozhihan

source/module_hamilt_general/operator.cpp

@@ -89,8 +89,8 @@ typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& inp
                     psi_input->npol,
                     tmpsi_in,
                     this->hpsi->get_pointer(),
-                    // psi_input->get_ngk(op->ik),
-                    psi_input->get_current_nbas(),
+                    psi_input->get_ngk(op->ik),
+                    // psi_input->get_current_nbas(),
                     is_first_node);
             break;
         }

source/module_hamilt_pw/hamilt_pwdft/operator_pw/velocity_pw.cpp

@@ -48,8 +48,20 @@ void Velocity::act
 ) const
 {
     ModuleBase::timer::tick("Operator", "Velocity");
-    // const int npw = psi_in->get_ngk(this->ik);
-    const int npw = psi_in->get_current_nbas();
+
+    // if (psi_in->get_ngk(this->ik) != psi_in->get_current_nbas()) 
+    // {
+    //     std::cout << "op->ik : " << this->ik << std::endl;
+    //     std::cout << "get_ngk(op->ik) : " << psi_in->get_ngk(this->ik) << std::endl;
+    //     std::cout << "get_current_nbas() : " << psi_in->get_current_nbas() << std::endl;
+
+    //     std::cout << "ik : " << this->ik << std::endl;
+    // }
+
+
+    const int npw = psi_in->get_ngk(this->ik);
+    // 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 std::complex<double>* tmpsi_in = psi0;

source/module_hsolver/diago_iter_assist.cpp

@@ -199,7 +199,7 @@ void DiagoIterAssist<T, Device>::diagH_subspace_init(hamilt::Hamilt<T, Device>*
 
     if (base_device::get_device_type(ctx) == base_device::GpuDevice)
     {
-        psi::Psi<T, Device> psi_temp(1, 1, psi_nc, dmin, true);
+        psi::Psi<T, Device> psi_temp(1, 1, psi_nc, &evc.get_ngk(0), dmin, true);
         T* ppsi = psi_temp.get_pointer();
         // hpsi and spsi share the temp space
         T* temp = nullptr;
@@ -246,7 +246,7 @@ void DiagoIterAssist<T, Device>::diagH_subspace_init(hamilt::Hamilt<T, Device>*
     }
     else if (base_device::get_device_type(ctx) == base_device::CpuDevice)
     {
-        psi::Psi<T, Device> psi_temp(1, nstart, psi_nc, dmin, true);
+        psi::Psi<T, Device> psi_temp(1, nstart, psi_nc, &evc.get_ngk(0), dmin, true);
         T* ppsi = psi_temp.get_pointer();
         syncmem_complex_op()(ctx, ctx, ppsi, psi, psi_temp.size());
         // hpsi and spsi share the temp space

source/module_psi/psi.cpp

@@ -118,6 +118,7 @@ template <typename T, typename Device>
 Psi<T, Device>::Psi(const int nk_in,
                     const int nbd_in,
                     const int nbs_in,
+                    const int* ngk_in, 
                     const int current_nbasis_in,
                     const bool k_first_in)
 {
@@ -129,7 +130,7 @@ Psi<T, Device>::Psi(const int nk_in,
     this->npol = PARAM.globalv.npol;
     this->allocate_inside = true;
 
-    this->ngk = nullptr;
+    this->ngk = ngk_in;
     assert(nk_in > 0 && nbd_in > 0 && nbs_in > 0);
     resize_memory_op()(this->ctx, this->psi, nk_in * static_cast<std::size_t>(nbd_in) * nbs_in, "no_record");
 

source/module_psi/psi.h

@@ -73,7 +73,12 @@ class Psi
 
 
     // Constructor 8-3: 2D Psi version 3
-    Psi(const int nk_in, const int nbd_in, const int nbs_in, const int current_nbasis_in, const bool k_first_in);
+    Psi(const int nk_in, 
+        const int nbd_in, 
+        const int nbs_in, 
+        const int* ngk_in, 
+        const int current_nbasis_in, 
+        const bool k_first_in);
 
 
     // Destructor for deleting the psi array manually