Feature: Support outputting partial charge and wave functions under PW basis when kpar > 1

docs/advanced/input_files/input-main.md

@@ -599,7 +599,7 @@ These variables are used to control general system parameters.
   - nao: from numerical atomic orbitals. If they are not enough, other wave functions are initialized with random numbers.
   - nao+random: add small random numbers on numerical atomic orbitals
 
-  > Only the `file` option is useful for the lcao basis set, which is mostly used when [calculation](#calculation) is set to `set_wf` and `get_pchg`. See more details in [out_wfc_lcao](#out_wfc_lcao).
+  > Only the `file` option is useful for the lcao basis set, which is mostly used when [calculation](#calculation) is set to `get_wf` and `get_pchg`. See more details in [out_wfc_lcao](#out_wfc_lcao).
 - **Default**: atomic
 
 ### init_chg
@@ -1921,7 +1921,7 @@ The band (KS orbital) energy for each (k-point, spin, band) will be printed in t
 ### if_separate_k
 
 - **Type**: Boolean
-- **Availability**: Only for LCAO, used only when `calculation = get_pchg` and `gamma_only` is turned off.
+- **Availability**: For both PW and LCAO. When `basis_type = pw`, used if `out_pchg` is set. When `basis_type = lcao`, used only when `calculation = get_pchg` and `gamma_only` is turned off.
 - **Description**: Specifies whether to write the partial charge densities for all k-points to individual files or merge them. **Warning**: Enabling symmetry may produce incorrect results due to incorrect k-point weights. Therefore, when calculating partial charge densities, it is strongly recommended to set `symmetry = -1`.
 - **Default**: false
 

source/module_elecstate/module_charge/symmetry_rho.cpp

@@ -97,7 +97,7 @@ void Symmetry_rho::psymm(double* rho_part,
         rhotot.resize(rho_basis->nxyz);
         ModuleBase::GlobalFunc::ZEROS(rhotot.data(), rho_basis->nxyz);
     }
-    Pgrid.reduce(rhotot.data(), rho_part);
+    Pgrid.reduce(rhotot.data(), rho_part, false);
 
     // (2)
     if (GlobalV::MY_RANK == 0)

source/module_esolver/esolver_ks_pw.cpp

@@ -645,14 +645,24 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
     //----------------------------------------------------------
     // 3) Print out electronic wavefunctions in pw basis
     //----------------------------------------------------------
-	if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
-	{
-		ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK, 
-				PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
-				GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, 
-				PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
-				this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
-	}
+    if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
+    {
+        ModuleIO::write_wfc_pw(GlobalV::KPAR,
+                               GlobalV::MY_POOL,
+                               GlobalV::MY_RANK,
+                               PARAM.inp.nbands,
+                               PARAM.inp.nspin,
+                               PARAM.globalv.npol,
+                               GlobalV::RANK_IN_POOL,
+                               GlobalV::NPROC_IN_POOL,
+                               PARAM.inp.out_wfc_pw,
+                               PARAM.inp.ecutwfc,
+                               PARAM.globalv.global_out_dir,
+                               this->psi[0],
+                               this->kv,
+                               this->pw_wfc,
+                               GlobalV::ofs_running);
+    }
 
     //----------------------------------------------------------
     // 4) check if oscillate for delta_spin method
@@ -718,11 +728,6 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
                               this->pw_rhod->ny,
                               this->pw_rhod->nz,
                               this->pw_rhod->nxyz,
-                              this->kv.get_nks(),
-                              this->kv.isk,
-                              this->kv.wk,
-                              this->pw_big->bz,
-                              this->pw_big->nbz,
                               this->chr.ngmc,
                               &ucell,
                               this->psi,
@@ -731,16 +736,29 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
                               this->ctx,
                               this->Pgrid,
                               PARAM.globalv.global_out_dir,
-                              PARAM.inp.if_separate_k);
+                              PARAM.inp.if_separate_k,
+                              this->kv,
+                              GlobalV::KPAR,
+                              GlobalV::MY_POOL,
+                              &this->chr);
     }
 
-
-	// tmp 2025-05-17, mohan note
-	ModuleIO::write_wfc_pw(GlobalV::KPAR, GlobalV::MY_POOL, GlobalV::MY_RANK, 
-			PARAM.inp.nbands, PARAM.inp.nspin, PARAM.globalv.npol,
-			GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, 
-			PARAM.inp.out_wfc_pw, PARAM.inp.ecutwfc, PARAM.globalv.global_out_dir,
-			this->psi[0], this->kv, this->pw_wfc, GlobalV::ofs_running);
+    // tmp 2025-05-17, mohan note
+    ModuleIO::write_wfc_pw(GlobalV::KPAR,
+                           GlobalV::MY_POOL,
+                           GlobalV::MY_RANK,
+                           PARAM.inp.nbands,
+                           PARAM.inp.nspin,
+                           PARAM.globalv.npol,
+                           GlobalV::RANK_IN_POOL,
+                           GlobalV::NPROC_IN_POOL,
+                           PARAM.inp.out_wfc_pw,
+                           PARAM.inp.ecutwfc,
+                           PARAM.globalv.global_out_dir,
+                           this->psi[0],
+                           this->kv,
+                           this->pw_wfc,
+                           GlobalV::ofs_running);
 
     //------------------------------------------------------------------
     //! 5) calculate Wannier functions in pw basis
@@ -931,7 +949,7 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
     //----------------------------------------------------------
 
     //----------------------------------------------------------
-    //! The write_psi_r_1 interface will be removed in the very 
+    //! The write_psi_r_1 interface will be removed in the very
     //! near future. Don't use it!
     //----------------------------------------------------------
     // if (PARAM.inp.out_wfc_r == 1) // Peize Lin add 2021.11.21
@@ -951,19 +969,15 @@ void ESolver_KS_PW<T, Device>::after_all_runners(UnitCell& ucell)
                             this->pw_rhod->ny,
                             this->pw_rhod->nz,
                             this->pw_rhod->nxyz,
-                            this->kv.get_nks(),
-                            this->kv.isk,
-                            this->kv.wk,
-                            this->pw_big->bz,
-                            this->pw_big->nbz,
-                            this->chr.ngmc,
                             &ucell,
                             this->psi,
-                            this->pw_rhod,
                             this->pw_wfc,
                             this->ctx,
                             this->Pgrid,
-                            PARAM.globalv.global_out_dir);
+                            PARAM.globalv.global_out_dir,
+                            this->kv,
+                            GlobalV::KPAR,
+                            GlobalV::MY_POOL);
     }
 
     //----------------------------------------------------------