Refactor occupy.h/cpp and update MPI kmesh_world

AGENTS.md

@@ -58,6 +58,10 @@ rules. Read the complete governance document before making or reviewing changes:
   pointers (e.g., `std::vector<bool>` has no `.data()`), and use
   `std::fill`/`std::copy` instead of `ZEROS`/`COPYARRAY` on vector buffers.
 
+## Token And Tool Budget
+
+- Before any non-trivial action, state the expected token cost; default to the cheapest path (e.g., one `Read`/`Grep` over sub-agents) and ask before deep exploration.
+
 ## Repository Map
 
 - Core C++ implementation lives under `source/`; source additions must be wired

docs/advanced/input_files/input-main.md

@@ -18,6 +18,7 @@
     - [cal\_force](#cal_force)
     - [kpar](#kpar)
     - [bndpar](#bndpar)
+    - [nimage](#nimage)
     - [latname](#latname)
     - [assume\_isolated](#assume_isolated)
     - [init\_wfc](#init_wfc)
@@ -725,6 +726,15 @@
   > 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
 
+### nimage
+
+- **Type**: Integer
+- **Description**: Number of independent calculation images that share the MPI processes.
+  - Each image runs its own esolver instance on a dedicated esolver_world communicator, split from MPI_COMM_WORLD by image id.
+  - The cross-image images_world communicator connects ranks with the same rank_in_esolver across images.
+  - Currently only nimage = 1 is supported; larger values are reserved for path-based methods such as NEB and will be rejected.
+- **Default**: 1
+
 ### latname
 
 - **Type**: String

docs/parameters.yaml

@@ -145,6 +145,17 @@ parameters:
     default_value: "1"
     unit: ""
     availability: (basis_type==pw and esolver_type==sdft) or (basis_type==pw and esolver_type==ksdft and ks_solver==bpcg)
+  - name: nimage
+    category: System variables
+    type: Integer
+    description: |
+      Number of independent calculation images that share the MPI processes.
+      * Each image runs its own esolver instance on a dedicated esolver_world communicator, split from MPI_COMM_WORLD by image id.
+      * The cross-image images_world communicator connects ranks with the same rank_in_esolver across images.
+      * Currently only nimage = 1 is supported; larger values are reserved for path-based methods such as NEB and will be rejected.
+    default_value: "1"
+    unit: ""
+    availability: ""
   - name: latname
     category: System variables
     type: String

source/CMakeLists.txt

@@ -551,7 +551,8 @@ add_library(
   driver
   OBJECT
   source_main/driver.cpp
-  source_main/driver_run.cpp)
+  source_main/driver_run.cpp
+  source_main/para_worlds_global.cpp)
 
 list(APPEND device_srcs
   source_pw/module_pwdft/kernels/nonlocal_op.cpp
@@ -754,6 +755,7 @@ abacus_apply_build_options(${ABACUS_BIN_NAME})
 # Register integration tests only after the final executable and its path are
 # available. Unit tests are added by source subdirectories through AddTest().
 if(BUILD_TESTING)
+  add_subdirectory(source_main/test)
   add_subdirectory("${ABACUS_TEST_DIR}" "${PROJECT_BINARY_DIR}/tests")
 endif()
 

source/Makefile.Objects

@@ -130,7 +130,8 @@ ${OBJS_RDMFT}
 OBJS_MAIN=main.o\
     driver.o\
     driver_run.o\
-    parameter.o
+    parameter.o\
+    para_worlds_global.o
 
 OBJS_BASE=assoc_laguerre.o\
     blas_connector_base.o\
@@ -231,6 +232,10 @@ OBJS_CELL=atom_pseudo.o\
     mdcell.o\
     cif_io.o\
     ucell_io.o\
+    read_cube.o\
+    write_cube.o\
+    write_pao.o\
+    output_log.o\
 
 OBJS_DEEPKS=lcao_deepks.o\
         deepks_basic.o\
@@ -630,7 +635,6 @@ OBJS_IO=module_parameter/input_conv.o\
     module_bessel/numerical_basis_output.o\
     output.o\
     module_output/print_info.o\
-    module_output/read_cube.o\
     module_wf/read_wfc_pw.o\
     module_wf/read_wf2rho_pw.o\
     module_restart/restart.o\
@@ -649,17 +653,14 @@ OBJS_IO=module_parameter/input_conv.o\
     module_wannier/to_w90_pw_setup.o\
     module_wannier/fr_overlap.o\
     module_unk/unk_overlap_pw.o\
-    module_output/write_pao.o\
     module_wf/write_wfc_pw.o\
-    module_output/write_cube.o\
     module_elf/write_elf.o\
     module_dipole/write_dipole.o\
     module_current/td_current_io.o\
     module_current/td_current_io_comm.o\
     td_efield_io.o\
     td_vector_pot_io.o\
     module_chgpot/write_libxc_r.o\
-    module_output/output_log.o\
     module_hs/output_mat_sparse.o\
     module_ctrl/ctrl_scf_lcao.o\
     module_ctrl/ctrl_runner_lcao.o\
@@ -696,7 +697,7 @@ OBJS_IO=module_parameter/input_conv.o\
     module_hs/cal_plpr.o\
 
 OBJS_IO_LCAO=module_hs/cal_r_overlap_r.o\
-      module_output/write_orb_info.o\
+      write_orb_info.o\
       module_dos/write_dos_lcao.o\
       module_energy/write_proj_band_lcao.o\
       module_energy/write_eig_occ.o\
@@ -801,7 +802,7 @@ OBJS_PARALLEL=parallel_common.o\
     para_pw_world.o\
     para_diag_world.o\
     para_rgrid_world.o\
-    para_bgroup_world.o\
+    para_bdiff_ksame_world.o\
     para_matrix_world.o\
     para_mpi_func.o\
     para_setup.o\

source/source_base/CMakeLists.txt

@@ -79,7 +79,7 @@ add_library(
     module_parallel/para_pw_world.cpp
     module_parallel/para_diag_world.cpp
     module_parallel/para_rgrid_world.cpp
-    module_parallel/para_bgroup_world.cpp
+    module_parallel/para_bdiff_ksame_world.cpp
     module_parallel/para_matrix_world.cpp
     module_parallel/para_mpi_func.cpp
     module_parallel/para_setup.cpp

source/source_base/module_parallel/para_bdiff_ksame_world.cpp

@@ -0,0 +1,33 @@
+#include "para_bdiff_ksame_world.h"
+
+namespace Parallel
+{
+
+ParaBdiffKsameWorld::ParaBdiffKsameWorld()
+    : ParaWorld("bdiff_ksame"), my_bndgroup_(0), nbndgroup_(1)
+{
+}
+
+#ifdef __MPI
+ParaBdiffKsameWorld::ParaBdiffKsameWorld(const MPI_Comm& intra_comm, const MPI_Comm& inter_comm, int nbndgroup)
+    : ParaWorld("bdiff_ksame", intra_comm), inter_comm_(inter_comm), nbndgroup_(nbndgroup)
+{
+    if (inter_comm != MPI_COMM_NULL)
+    {
+        MPI_Comm_rank(inter_comm, &my_bndgroup_);
+    }
+}
+#endif
+
+void ParaBdiffKsameWorld::reduce_across_bdiff_ksame(double& value) const
+{
+#ifdef __MPI
+    if (inter_comm_ == MPI_COMM_NULL || nbndgroup_ <= 1)
+    {
+        return;
+    }
+    MPI_Allreduce(MPI_IN_PLACE, &value, 1, MPI_DOUBLE, MPI_SUM, inter_comm_);
+#endif
+}
+
+} // namespace Parallel

source/source_base/module_parallel/para_bdiff_ksame_world.h

@@ -0,0 +1,87 @@
+#ifndef PARA_BDIFF_KSAME_WORLD_H
+#define PARA_BDIFF_KSAME_WORLD_H
+
+#include "para_world.h"
+
+namespace Parallel
+{
+
+/**
+ * @brief bdiff_ksame parallel domain: band-group communication topology
+ * inside one k-pool.
+ *
+ * Self-contained replacement for INT_BGROUP + BP_WORLD +
+ * GlobalV::MY_BNDGROUP/NPROC_IN_BNDGROUP/RANK_IN_BPGROUP.
+ *
+ * The domain has two communicators:
+ *   - intra: INT_BGROUP (bsame_kdiff; same band group, different k/pw)
+ *   - inter: BP_WORLD  (bdiff_ksame; different band groups, same k)
+ *
+ * Tests only need this header.
+ */
+class ParaBdiffKsameWorld : public ParaWorld
+{
+public:
+    /**
+     * @brief Construct a serial domain (single band group).
+     */
+    ParaBdiffKsameWorld();
+
+#ifdef __MPI
+    /**
+     * @brief Construct a domain from intra and inter communicators.
+     *
+     * @param[in] intra_comm  intra-group communicator (e.g. INT_BGROUP)
+     * @param[in] inter_comm  inter-group communicator (e.g. BP_WORLD)
+     * @param[in] nbndgroup   number of band groups
+     */
+    ParaBdiffKsameWorld(const MPI_Comm& intra_comm, const MPI_Comm& inter_comm, int nbndgroup);
+#endif
+
+    /// Band group index of this process.
+    int my_bndgroup() const { return my_bndgroup_; }
+
+    /// Number of band groups.
+    int nbndgroup() const { return nbndgroup_; }
+
+    /// Rank within the band group (alias for rank()).
+    int rank_in_bpgroup() const { return rank(); }
+
+    /// Number of processes in the band group (alias for size()).
+    int nproc_in_bndgroup() const { return size(); }
+
+#ifdef __MPI
+    /// Inter-group communicator (BP_WORLD / bdiff_ksame equivalent).
+    MPI_Comm inter_comm() const { return inter_comm_; }
+#endif
+
+    /**
+     * @brief Sum a scalar across the band groups of this k-pool.
+     *
+     * Band-parallel eigensolvers (bpcg) shard the band range across the
+     * BNDPAR band groups of a k-pool: every process only accumulates the
+     * partial sum over its own band window. This reduction combines those
+     * partial sums on the bdiff_ksame (BP_WORLD) communicator, which links
+     * the same rank position of every band group inside one k-pool, so each
+     * band window contributes exactly once.
+     *
+     * It must run BEFORE ParaKmeshWorld::reduce_across_pools so that the
+     * k-pool reduction receives one complete per-k-pool partial sum.
+     * No-op when there is only a single band group.
+     *
+     * @param[in,out] value  local partial sum, overwritten with the
+     *                       k-pool-wide total
+     */
+    void reduce_across_bdiff_ksame(double& value) const;
+
+private:
+    int my_bndgroup_ = 0;
+    int nbndgroup_ = 1;
+#ifdef __MPI
+    MPI_Comm inter_comm_ = MPI_COMM_NULL;
+#endif
+};
+
+} // namespace Parallel
+
+#endif // PARA_BDIFF_KSAME_WORLD_H