Fix: Correct PW KPAR/BNDPAR occupations and distributed output

Author: AsTonyshment

docs/advanced/input_files/input-main.md

@@ -671,13 +671,22 @@
 ### kpar
 
 - **Type**: Integer
-- **Description**: Divide all processors into kpar groups, and k points will be distributed among each group. The value taken should be less than or equal to the number of k points as well as the number of MPI processes.
+- **Description**: Controls k-point parallelism. The value must be positive and should not exceed either the number of k-points or the number of MPI processes.
+  - For PW calculations, divide all MPI processes into persistent k-point pools. Each pool stores and processes a subset of the k-points.
+  - For LCAO calculations with lapack, genelpa, elpa, or scalapack_gvx, divide the diagonalization work into temporary k-point pools. After diagonalization, the eigenvalues and distributed wavefunctions are restored for all k-points before occupations, density matrices, and output are evaluated.
+  - Multi-process LCAO cusolver uses its own active-GPU distribution and does not use this value to define its k-point layout. Other LCAO eigensolvers do not use the temporary k-point-pool implementation.
 - **Default**: 1
 
 ### bndpar
 
 - **Type**: Integer
-- **Description**: Divide all processors into bndpar groups for SDFT or the BPCG solver. bndpar must be positive, no greater than the number of MPI processes, and kpar * bndpar must divide the number of MPI processes exactly.
+- **Availability**: *([`basis_type`](#basis_type)==pw and [`esolver_type`](#esolver_type)==sdft) or ([`basis_type`](#basis_type)==pw and [`esolver_type`](#esolver_type)==ksdft and [`ks_solver`](#ks_solver)==bpcg)*
+- **Description**: Controls band-group parallelism for PW SDFT and PW KSDFT calculations using the BPCG eigensolver.
+  - Within each k-point pool, divide the MPI processes into bndpar band groups. Each group contains NPROC / (kpar * bndpar) processes when bndpar is greater than 1.
+  - With BPCG, distribute contiguous ranges of global Kohn-Sham bands among the band groups. nbands does not need to be divisible by bndpar; groups with lower indices receive one additional band when necessary.
+  - In SDFT, distribute stochastic orbitals among the band groups. When the deterministic Kohn-Sham eigensolver is not BPCG, band group 0 calculates the deterministic orbitals and broadcasts them to the other groups.
+  - bndpar must be positive and no greater than the number of MPI processes. When bndpar is greater than 1, kpar * bndpar must divide the number of MPI processes exactly.
+  > Note: For PW calculations on GPU, if the input kpar * bndpar differs from the number of MPI processes, ABACUS automatically sets the effective kpar to NPROC / bndpar.
 - **Default**: 1
 
 ### latname
@@ -4492,9 +4501,9 @@
 ### ocp_set
 
 - **Type**: String
-- **Description**: Fixed occupation weights used when ocp is true. Values are assigned band by band for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x.
+- **Description**: Fixed occupation weights used when ocp is true. Values are assigned in band order for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x.
   - Example: 1 10*1 0 1 expands to 13 values, with the 12th value equal to 0 and all other values equal to 1.
-  - After expansion, the array length must equal nks * nbands.
+  - After expansion, provide one block of nbands values for each k-point. If nspin is 2, provide all k-point blocks for spin up followed by all k-point blocks for spin down; otherwise, provide one block per k-point.
   - The sum of all weights must equal nelec; otherwise the calculation terminates with an error.
 - **Default**: None
 

docs/parameters.yaml

@@ -100,18 +100,26 @@ parameters:
     category: System variables
     type: Integer
     description: |
-      Divide all processors into kpar groups, and k points will be distributed among each group. The value taken should be less than or equal to the number of k points as well as the number of MPI processes.
+      Controls k-point parallelism. The value must be positive and should not exceed either the number of k-points or the number of MPI processes.
+      * For PW calculations, divide all MPI processes into persistent k-point pools. Each pool stores and processes a subset of the k-points.
+      * For LCAO calculations with lapack, genelpa, elpa, or scalapack_gvx, divide the diagonalization work into temporary k-point pools. After diagonalization, the eigenvalues and distributed wavefunctions are restored for all k-points before occupations, density matrices, and output are evaluated.
+      * Multi-process LCAO cusolver uses its own active-GPU distribution and does not use this value to define its k-point layout. Other LCAO eigensolvers do not use the temporary k-point-pool implementation.
     default_value: "1"
     unit: ""
     availability: ""
   - name: bndpar
     category: System variables
     type: Integer
     description: |
-      Divide all processors into bndpar groups for SDFT or the BPCG solver. bndpar must be positive, no greater than the number of MPI processes, and kpar * bndpar must divide the number of MPI processes exactly.
+      Controls band-group parallelism for PW SDFT and PW KSDFT calculations using the BPCG eigensolver.
+      * Within each k-point pool, divide the MPI processes into bndpar band groups. Each group contains NPROC / (kpar * bndpar) processes when bndpar is greater than 1.
+      * With BPCG, distribute contiguous ranges of global Kohn-Sham bands among the band groups. nbands does not need to be divisible by bndpar; groups with lower indices receive one additional band when necessary.
+      * In SDFT, distribute stochastic orbitals among the band groups. When the deterministic Kohn-Sham eigensolver is not BPCG, band group 0 calculates the deterministic orbitals and broadcasts them to the other groups.
+      * bndpar must be positive and no greater than the number of MPI processes. When bndpar is greater than 1, kpar * bndpar must divide the number of MPI processes exactly.
+      [NOTE] For PW calculations on GPU, if the input kpar * bndpar differs from the number of MPI processes, ABACUS automatically sets the effective kpar to NPROC / bndpar.
     default_value: "1"
     unit: ""
-    availability: ""
+    availability: (basis_type==pw and esolver_type==sdft) or (basis_type==pw and esolver_type==ksdft and ks_solver==bpcg)
   - name: latname
     category: System variables
     type: String
@@ -2744,9 +2752,9 @@ parameters:
     category: "RT-TDDFT: Real-Time Time-Dependent Density Functional Theory"
     type: String
     description: |
-      Fixed occupation weights used when ocp is true. Values are assigned band by band for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x.
+      Fixed occupation weights used when ocp is true. Values are assigned in band order for each k-point, following k-point order. In LCAO RT-TDDFT, the initial ground-state SCF uses its normally determined occupations, and this array is applied only during subsequent real-time propagation steps. The repetition syntax N*x expands to N copies of x.
       * Example: 1 10*1 0 1 expands to 13 values, with the 12th value equal to 0 and all other values equal to 1.
-      * After expansion, the array length must equal nks * nbands.
+      * After expansion, provide one block of nbands values for each k-point. If nspin is 2, provide all k-point blocks for spin up followed by all k-point blocks for spin down; otherwise, provide one block per k-point.
       * The sum of all weights must equal nelec; otherwise the calculation terminates with an error.
     default_value: None
     unit: ""

source/Makefile.Objects

@@ -552,6 +552,7 @@ OBJS_IO=module_parameter/input_conv.o\
     cal_test.o\
     module_dos/write_dos_pw.o\
     nscf_fermi_surf.o\
+    band_parallel_output.o\
     write_bands.o\
     module_dos/cal_dos.o\
     module_dos/cal_pdos_gamma.o\

source/source_base/parallel_grid.cpp

@@ -331,60 +331,81 @@ void Parallel_Grid::zpiece_to_stogroup(double* zpiece, const int& iz, double* rh
     // ofs_running << "\n iz = " << iz << " Done.";
     return;
 }
+
+// Taoni modified on 2026-08-21, fixed BPCG out_chg MPI_ERR_RANK
 void Parallel_Grid::reduce(double* rhotot, const double* const rhoin, const bool reduce_all_pool) const
 {
-    // ModuleBase::TITLE("Parallel_Grid","reduce");
-
-    // if not the first pool, wait here until processpr 0
-    // send the Barrier command.
+    // POOL_WORLD communicators are disjoint, so inactive pools may return
+    // without skipping a collective required by an active pool.
     if (!reduce_all_pool && GlobalV::MY_POOL != 0)
     {
         return;
     }
 
-    double* zpiece = new double[this->ncxy];
+    assert(rhoin != nullptr);
+    assert(this->nrxx == this->ncxy * this->nczp);
 
-    for (int iz = 0; iz < this->ncz; iz++)
-    {
-        const int znow = iz - this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL];
-        const int proc = this->whichpro[GlobalV::MY_POOL][iz];
-        const int proc_loc = this->whichpro_loc[GlobalV::MY_POOL][iz]; // Obtain the local processor index in the pool
-        ModuleBase::GlobalFunc::ZEROS(zpiece, this->ncxy);
-        int tag = iz;
-        MPI_Status ierror;
-
-        // Local processor 0 collects data from all other processors in the pool
-        // proc = proc_loc if GlobalV::MY_POOL == 0
-        if (proc_loc == GlobalV::RANK_IN_POOL)
-        {
-            for (int ir = 0; ir < ncxy; ir++)
-            {
-                zpiece[ir] = rhoin[ir * this->nczp + znow];
-            }
-            // Send data to the root of the pool
-            if (GlobalV::RANK_IN_POOL != 0)
-            {
-                MPI_Send(zpiece, ncxy, MPI_DOUBLE, 0, tag, POOL_WORLD);
-            }
-        }
+    int pool_size = 0;
+    MPI_Comm_size(POOL_WORLD, &pool_size);
+
+    // Derive slab ownership from the communicator used below. Precomputed
+    // source ranks may belong to a larger BPCG layout and are not necessarily
+    // valid ranks in the current POOL_WORLD.
+    std::vector<int> local_z_counts(pool_size);
+    MPI_Allgather(&this->nczp, 1, MPI_INT, local_z_counts.data(), 1, MPI_INT, POOL_WORLD);
 
-        // The root of the pool receives data from other processors
-        if (GlobalV::RANK_IN_POOL == 0 && proc_loc != GlobalV::RANK_IN_POOL)
+    // MPI_Gatherv concatenates complete [xy][local_z] buffers by rank.
+    // These displacements therefore address rank blocks, not global z coordinates.
+    std::vector<int> receive_counts(pool_size);
+    std::vector<int> displacements(pool_size, 0);
+    int total_z = 0;
+    for (int rank = 0; rank < pool_size; ++rank)
+    {
+        receive_counts[rank] = local_z_counts[rank] * this->ncxy;
+        if (rank > 0)
         {
-            MPI_Recv(zpiece, ncxy, MPI_DOUBLE, proc_loc, tag, POOL_WORLD, &ierror);
+            displacements[rank] = displacements[rank - 1] + receive_counts[rank - 1];
         }
+        total_z += local_z_counts[rank];
+    }
+    assert(total_z == this->ncz);
 
-        if (GlobalV::RANK_IN_POOL == 0)
+    // Only the pool root needs storage for the gathered global grid.
+    // MPI ignores the receive buffer on all non-root ranks.
+    std::vector<double> gathered_data;
+    if (GlobalV::RANK_IN_POOL == 0)
+    {
+        assert(rhotot != nullptr);
+        gathered_data.resize(this->ncxyz);
+    }
+    MPI_Gatherv(rhoin,
+                this->nrxx,
+                MPI_DOUBLE,
+                gathered_data.data(),
+                receive_counts.data(),
+                displacements.data(),
+                MPI_DOUBLE,
+                0,
+                POOL_WORLD);
+
+    if (GlobalV::RANK_IN_POOL == 0)
+    {
+        // Rank blocks cannot be copied directly to rhotot: each block stores
+        // [xy][local_z], whereas Cube output expects [xy][global_z].
+        // The slab decomposition is contiguous and ordered by POOL_WORLD rank.
+        int global_z_start = 0;
+        for (int rank = 0; rank < pool_size; ++rank)
         {
+            const int local_z = local_z_counts[rank];
             for (int ixy = 0; ixy < this->ncxy; ++ixy)
             {
-                rhotot[ixy * ncz + iz] = zpiece[ixy];
+                for (int iz = 0; iz < local_z; ++iz)
+                {
+                    rhotot[ixy * this->ncz + global_z_start + iz] = gathered_data[displacements[rank] + ixy * local_z + iz];
+                }
             }
+            global_z_start += local_z;
         }
     }
-
-    delete[] zpiece;
-
-    return;
 }
 #endif

source/source_cell/parallel_kpoints.cpp

@@ -5,6 +5,7 @@
 #include "parallel_kpoints.h"
 
 #include "source_base/parallel_common.h"
+#include "source_base/parallel_comm.h"
 #include "source_base/parallel_global.h"
 
 /// @note the kpoints here are reduced after symmetry applied.
@@ -118,11 +119,13 @@ void Parallel_Kpoints::set_startpro_pool()
 void Parallel_Kpoints::gatherkvec(const std::vector<ModuleBase::Vector3<double>>& vec_local,
                                   std::vector<ModuleBase::Vector3<double>>& vec_global) const
 {
+    // Taoni fix bndpar on 2026-08-21
+    int band_rank = 0;
+    MPI_Comm_rank(BP_WORLD, &band_rank);
     vec_global.resize(this->nkstot_np, ModuleBase::Vector3<double>(0.0, 0.0, 0.0));
     for (int i = 0; i < this->nks_np; ++i)
     {
-
-        if (this->rank_in_pool == 0)
+        if (this->rank_in_pool == 0 && band_rank == 0)
         {
             vec_global[i + startk_pool[this->my_pool]] = vec_local[i];
         }

source/source_estate/elecstate_tools.cpp

@@ -1,6 +1,20 @@
 #include "elecstate_tools.h"
 #include "occupy.h"
+#include "source_base/parallel_comm.h"
 #include "source_base/parallel_reduce.h"
+
+namespace
+{
+int get_band_offset(const int local_nbands)
+{
+    int band_end = local_nbands;
+#ifdef __MPI
+    MPI_Scan(&local_nbands, &band_end, 1, MPI_INT, MPI_SUM, BP_WORLD);
+#endif
+    return band_end - local_nbands;
+}
+} // namespace
+
 namespace elecstate
 {
     void calEBand(const ModuleBase::matrix& ekb,const ModuleBase::matrix& wg,fenergy& f_en)
@@ -42,11 +56,14 @@ namespace elecstate
         const int nks = ekb.nr;
         if (!(Occupy::use_gaussian_broadening || Occupy::fixed_occupations))
         {
+            // Taoni fix smearing_method=fixed for BPCG on 2026-08-21
+            const int band_offset = get_band_offset(nbands);
             if (PARAM.globalv.two_fermi)
             {
                 Occupy::iweights(nks,
                                 klist->wk,
                                 nbands,
+                                band_offset,
                                 nelec_spin[0],
                                 ekb,
                                 eferm.ef_up,
@@ -56,6 +73,7 @@ namespace elecstate
                 Occupy::iweights(nks,
                                 klist->wk,
                                 nbands,
+                                band_offset,
                                 nelec_spin[1],
                                 ekb,
                                 eferm.ef_dw,
@@ -70,6 +88,7 @@ namespace elecstate
                 Occupy::iweights(nks,
                                 klist->wk,
                                 nbands,
+                                band_offset,
                                 PARAM.inp.nelec,
                                 ekb,
                                 eferm.ef,
@@ -151,15 +170,15 @@ namespace elecstate
 
         const double ne_thr = 1.0e-5;
 
-        const int num = klist->get_nks() * nbands;
-        if (num != ocp_kb.size())
+        const std::size_t expected_size = static_cast<std::size_t>(klist->get_nkstot()) * static_cast<std::size_t>(nbands);
+        if (expected_size != ocp_kb.size())
         {
             ModuleBase::WARNING_QUIT("ElecState::fixed_weights",
                                     "size of occupation array is wrong , please check ocp_set");
         }
 
         double num_elec = 0.0;
-        for (int i = 0; i < ocp_kb.size(); ++i)
+        for (std::size_t i = 0; i < ocp_kb.size(); ++i)
         {
             num_elec += ocp_kb[i];
         }
@@ -170,11 +189,28 @@ namespace elecstate
                                     "total number of occupations is wrong , please check ocp_set");
         }
 
+        const int band_offset = get_band_offset(wg.nc);
+        if (klist->ik2iktot.size() < static_cast<std::size_t>(wg.nr)
+            || band_offset < 0
+            || band_offset + wg.nc > nbands)
+        {
+            ModuleBase::WARNING_QUIT("ElecState::fixed_weights", "invalid distributed occupation layout");
+        }
+
         for (int ik = 0; ik < wg.nr; ++ik)
         {
+            const int global_k = klist->ik2iktot[ik];
+            if (global_k < 0 || global_k >= klist->get_nkstot())
+            {
+                ModuleBase::WARNING_QUIT("ElecState::fixed_weights", "invalid global k-point index");
+            }
             for (int ib = 0; ib < wg.nc; ++ib)
             {
-                wg(ik, ib) = ocp_kb[ik * wg.nc + ib];
+                // Taoni fix ocp_set under kpar and bndpar on 2026-08-21
+                const int global_band = band_offset + ib;
+                const std::size_t occupation_index
+                    = static_cast<std::size_t>(global_k) * static_cast<std::size_t>(nbands) + static_cast<std::size_t>(global_band);
+                wg(ik, ib) = ocp_kb[occupation_index];
             }
         }
         skip_weights = true;

source/source_estate/occupy.cpp

@@ -120,7 +120,8 @@ void Occupy::decision(const std::string& name, const std::string& smearing_metho
  *
  * @param nks number of k points.
  * @param wk weight of each k point (consider symmetry).
- * @param nbands number of bands.
+ * @param nbands number of locally owned bands.
+ * @param band_offset global index of the first locally owned band.
  * @param nelec number of electrons for this spin direction.
  * @param ekb the array save the band energy.
  * @param ef output: the highest occupied Kohn-Sham level.
@@ -132,6 +133,7 @@ void Occupy::iweights(
     const int nks,
     const std::vector<double>& wk,
     const int nbands,
+    const int band_offset,
     const double& nelec,
     const ModuleBase::matrix& ekb,
     double& ef,
@@ -168,7 +170,8 @@ void Occupy::iweights(
 
         for (int ib = 0; ib < nbands; ++ib)
         {
-            if (ib < ib_min)
+            const int global_band = band_offset + ib;
+            if (global_band < ib_min)
             {
                 wg(ik, ib) = wk[ik];
                 ef = std::max(ef, ekb(ik, ib));

source/source_estate/occupy.h

@@ -38,6 +38,7 @@ class Occupy
     static void iweights(const int nks,
                          const std::vector<double>& wk,
                          const int nband,
+                         const int band_offset,
                          const double& nelec,
                          const ModuleBase::matrix& ekb,
                          double& ef,

source/source_estate/test/elecstate_base_test.cpp

@@ -249,6 +249,8 @@ TEST_F(ElecStateTest, FixedWeights)
     PARAM.input.nelec = 30;
     K_Vectors* klist = new K_Vectors;
     klist->set_nks(5);
+    klist->set_nkstot(5);
+    klist->ik2iktot = {0, 1, 2, 3, 4};
     elecstate->klist = klist;
     elecstate->wg.create(klist->get_nks(), PARAM.input.nbands);
     std::vector<double> ocp_kb;
@@ -269,6 +271,8 @@ TEST_F(ElecStateDeathTest, FixedWeightsWarning1)
     PARAM.input.nelec = 30;
     K_Vectors* klist = new K_Vectors;
     klist->set_nks(5);
+    klist->set_nkstot(5);
+    klist->ik2iktot = {0, 1, 2, 3, 4};
     elecstate->klist = klist;
     elecstate->wg.create(klist->get_nks(), PARAM.input.nbands);
     std::vector<double> ocp_kb;
@@ -291,6 +295,8 @@ TEST_F(ElecStateDeathTest, FixedWeightsWarning2)
     PARAM.input.nelec = 29;
     K_Vectors* klist = new K_Vectors;
     klist->set_nks(5);
+    klist->set_nkstot(5);
+    klist->ik2iktot = {0, 1, 2, 3, 4};
     elecstate->klist = klist;
     elecstate->wg.create(klist->get_nks(), PARAM.input.nbands);
     std::vector<double> ocp_kb;