Fix: nscf for pw code

Author: dyzheng

source/module_esolver/esolver_ks_pw.cpp

@@ -1241,6 +1241,12 @@ void ESolver_KS_PW<T, Device>::after_all_runners() {
 
 template <typename T, typename Device>
 void ESolver_KS_PW<T, Device>::hamilt2estates(const double ethr) {
+    if(PARAM.inp.dft_plus_u)
+    {
+        auto* dftu = ModuleDFTU::DFTU::get_instance();
+        dftu->cal_occ_pw(0, this->kspw_psi, this->pelec->wg, GlobalC::ucell, nullptr);
+        dftu->output();
+    }
     if (this->phsol != nullptr) {
         hsolver::DiagoIterAssist<T, Device>::need_subspace = false;
         hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;

source/module_hamilt_lcao/module_dftu/dftu_occup.cpp

@@ -56,6 +56,15 @@ void DFTU::copy_locale()
             {
                 locale_save[iat][target_l][0][0] = locale[iat][target_l][0][0];
                 locale_save[iat][target_l][0][1] = locale[iat][target_l][0][1];
+                if(this->uom_save.size() != 0)
+                {
+                    const int size = locale[iat][target_l][0][0].nr * locale[iat][target_l][0][0].nc;
+                    for(int mm=0; mm<size; mm++)
+                    {
+                        this->uom_save[eff_pot_pw_index[iat]+mm] = locale[iat][target_l][0][0].c[mm];
+                        this->uom_save[eff_pot_pw_index[iat]+mm+size] = locale[iat][target_l][0][1].c[mm];
+                    }
+                }
             }
         }
     }

source/module_hamilt_lcao/module_dftu/dftu_pw.cpp

@@ -212,6 +212,14 @@ void DFTU::cal_occ_pw(const int iter, const void* psi_in, const ModuleBase::matr
         if(GlobalV::NSPIN == 2)
         {
             Parallel_Reduce::reduce_double_allpool(GlobalV::KPAR, GlobalV::NPROC_IN_POOL, this->locale[iat][target_l][0][1].c, size);
+            //save locale matrix for spin=1 to uom_array
+            if(this->uom_array.size() != 0)
+            {
+                for(int mm=0;mm<size;mm++)
+                {
+                    this->uom_array[eff_pot_pw_index[iat]+mm+size] = this->locale[iat][target_l][0][1].c[mm];
+                }
+            }
         }
     }
     }//end if initialed_locale