Feature: add MPI communicator injection for general-purpose parallel workflows

    Replace hardcoded MPI_COMM_WORLD with configurable base_comm parameter in
    the MPI initialization path, enabling workflows like parallel NEB, phonon
    calculations, or replica exchange MD to split MPI ranks into independent
    groups, each running its own ESolver instance in a separate subprocess.

Author: dyzheng

python/pyabacus/src/ModuleESolver/py_esolver_lcao.cpp

@@ -448,15 +448,20 @@ PyESolverLCAO<TK, TR>::~PyESolverLCAO()
 }
 
 template <typename TK, typename TR>
-void PyESolverLCAO<TK, TR>::initialize(const std::string& input_dir)
+void PyESolverLCAO<TK, TR>::initialize(const std::string& input_dir, int mpi_comm_handle)
 {
     // Placeholder: will be implemented in Phase 3
     // This will:
-    // 1. Read INPUT file from input_dir
-    // 2. Initialize UnitCell
-    // 3. Create ESolver_KS_LCAO instance
+    // 1. Convert mpi_comm_handle to MPI_Comm via MPI_Comm_f2c()
+    //    (if mpi_comm_handle >= 0, otherwise use MPI_COMM_WORLD)
+    // 2. Use the resulting communicator for MPI_Comm_size/rank
+    // 3. Pass base_comm to Parallel_Global::init_pools(), split_diag_world(), split_grid_world()
+    // 4. Read INPUT file from input_dir
+    // 5. Initialize UnitCell
+    // 6. Create ESolver_KS_LCAO instance
     initialized_ = true;
-    std::cout << "[PyESolverLCAO] Initialized with input directory: " << input_dir << std::endl;
+    std::cout << "[PyESolverLCAO] Initialized with input directory: " << input_dir
+              << ", mpi_comm_handle: " << mpi_comm_handle << std::endl;
 }
 
 template <typename TK, typename TR>
@@ -861,7 +866,10 @@ void bind_esolver_lcao(py::module& m, const std::string& suffix)
             ----------
             input_dir : str
                 Directory containing INPUT, STRU, and other input files
-            )pbdoc", "input_dir"_a)
+            mpi_comm_handle : int, optional
+                Fortran MPI communicator handle from mpi4py comm.py2f().
+                Use -1 (default) to use MPI_COMM_WORLD.
+            )pbdoc", "input_dir"_a, "mpi_comm_handle"_a = -1)
         .def("before_all_runners", &ESolver::before_all_runners,
             "Initialize calculation environment")
 

python/pyabacus/src/ModuleESolver/py_esolver_lcao.hpp

@@ -284,7 +284,10 @@ class PyESolverLCAO
     // ==================== Initialization ====================
 
     /// Initialize from INPUT file directory
-    void initialize(const std::string& input_dir);
+    /// @param input_dir Directory containing INPUT, STRU, and other input files
+    /// @param mpi_comm_handle Fortran MPI communicator handle from mpi4py comm.py2f().
+    ///        Use -1 (default) to use MPI_COMM_WORLD.
+    void initialize(const std::string& input_dir, int mpi_comm_handle = -1);
 
     /// Call before_all_runners
     void before_all_runners();

python/pyabacus/src/pyabacus/esolver/workflow.py

@@ -54,7 +54,7 @@ class LCAOWorkflow:
         'after_scf',         # Called after after_scf()
     ]
 
-    def __init__(self, input_dir: str, gamma_only: bool = True):
+    def __init__(self, input_dir: str, gamma_only: bool = True, mpi_comm=None):
         """
         Initialize LCAOWorkflow.
 
@@ -64,9 +64,13 @@ def __init__(self, input_dir: str, gamma_only: bool = True):
             Directory containing input files
         gamma_only : bool
             Use gamma-only calculation if True, multi-k if False
+        mpi_comm : mpi4py.MPI.Comm, optional
+            MPI communicator to use. If None, uses MPI_COMM_WORLD.
+            Pass a sub-communicator to run this ESolver on a subset of MPI ranks.
         """
         self._input_dir = input_dir
         self._gamma_only = gamma_only
+        self._mpi_comm = mpi_comm
         self._esolver = None
         self._initialized = False
         self._scf_running = False
@@ -96,10 +100,17 @@ def initialize(self) -> None:
                 "Make sure pyabacus is properly installed with ESolver support."
             ) from e
 
-        self._esolver.initialize(self._input_dir)
+        self._esolver.initialize(self._input_dir, mpi_comm_handle=self._mpi_comm_handle)
         self._esolver.before_all_runners()
         self._initialized = True
 
+    @property
+    def _mpi_comm_handle(self) -> int:
+        """Convert mpi4py communicator to Fortran handle, or -1 for default."""
+        if self._mpi_comm is None:
+            return -1
+        return self._mpi_comm.py2f()
+
     def register_callback(self, event: str, callback: Callable[['LCAOWorkflow'], None]) -> None:
         """
         Register a callback function for a specific event.

source/source_base/parallel_common.cpp

@@ -7,75 +7,75 @@
 #include <cstring>
 
 #ifdef __MPI
-void Parallel_Common::bcast_string(std::string& object) // Peize Lin fix bug 2019-03-18
+void Parallel_Common::bcast_string(std::string& object, MPI_Comm comm) // Peize Lin fix bug 2019-03-18
 {
     int size = object.size();
-    MPI_Bcast(&size, 1, MPI_INT, 0, MPI_COMM_WORLD);
-    
+    MPI_Bcast(&size, 1, MPI_INT, 0, comm);
+
     int my_rank;
-    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
-    
+    MPI_Comm_rank(comm, &my_rank);
+
     if (0 != my_rank)
     {
         object.resize(size);
     }
 
-    MPI_Bcast(&object[0], size, MPI_CHAR, 0, MPI_COMM_WORLD);
+    MPI_Bcast(&object[0], size, MPI_CHAR, 0, comm);
     return;
 }
 
-void Parallel_Common::bcast_string(std::string* object, const int n) // Peize Lin fix bug 2019-03-18
+void Parallel_Common::bcast_string(std::string* object, const int n, MPI_Comm comm) // Peize Lin fix bug 2019-03-18
 {
     for (int i = 0; i < n; i++)
-        bcast_string(object[i]);
+        bcast_string(object[i], comm);
     return;
 }
 
-void Parallel_Common::bcast_complex_double(std::complex<double>& object)
+void Parallel_Common::bcast_complex_double(std::complex<double>& object, MPI_Comm comm)
 {
-    MPI_Bcast(&object, 1, MPI_DOUBLE_COMPLEX, 0, MPI_COMM_WORLD);
+    MPI_Bcast(&object, 1, MPI_DOUBLE_COMPLEX, 0, comm);
 }
 
-void Parallel_Common::bcast_complex_double(std::complex<double>* object, const int n)
+void Parallel_Common::bcast_complex_double(std::complex<double>* object, const int n, MPI_Comm comm)
 {
-    MPI_Bcast(object, n, MPI_DOUBLE_COMPLEX, 0, MPI_COMM_WORLD);
+    MPI_Bcast(object, n, MPI_DOUBLE_COMPLEX, 0, comm);
 }
 
-void Parallel_Common::bcast_double(double& object)
+void Parallel_Common::bcast_double(double& object, MPI_Comm comm)
 {
-    MPI_Bcast(&object, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD);
+    MPI_Bcast(&object, 1, MPI_DOUBLE, 0, comm);
 }
 
-void Parallel_Common::bcast_double(double* object, const int n)
+void Parallel_Common::bcast_double(double* object, const int n, MPI_Comm comm)
 {
-    MPI_Bcast(object, n, MPI_DOUBLE, 0, MPI_COMM_WORLD);
+    MPI_Bcast(object, n, MPI_DOUBLE, 0, comm);
 }
 
-void Parallel_Common::bcast_int(int& object)
+void Parallel_Common::bcast_int(int& object, MPI_Comm comm)
 {
-    MPI_Bcast(&object, 1, MPI_INT, 0, MPI_COMM_WORLD);
+    MPI_Bcast(&object, 1, MPI_INT, 0, comm);
 }
 
-void Parallel_Common::bcast_int(int* object, const int n)
+void Parallel_Common::bcast_int(int* object, const int n, MPI_Comm comm)
 {
-    MPI_Bcast(object, n, MPI_INT, 0, MPI_COMM_WORLD);
+    MPI_Bcast(object, n, MPI_INT, 0, comm);
 }
 
-void Parallel_Common::bcast_bool(bool& object)
+void Parallel_Common::bcast_bool(bool& object, MPI_Comm comm)
 {
     int swap = object;
     int my_rank;
-    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
+    MPI_Comm_rank(comm, &my_rank);
     if (my_rank == 0)
         swap = object;
-    MPI_Bcast(&swap, 1, MPI_INT, 0, MPI_COMM_WORLD);
+    MPI_Bcast(&swap, 1, MPI_INT, 0, comm);
     if (my_rank != 0)
         object = static_cast<bool>(swap);
 }
 
-void Parallel_Common::bcast_char(char* object, const int n)
+void Parallel_Common::bcast_char(char* object, const int n, MPI_Comm comm)
 {
-    MPI_Bcast(object, n, MPI_CHAR, 0, MPI_COMM_WORLD);
+    MPI_Bcast(object, n, MPI_CHAR, 0, comm);
 }
 
 #endif

source/source_base/parallel_common.h

@@ -10,18 +10,58 @@
 namespace Parallel_Common
 {
 //(1) bcast array
-void bcast_complex_double(std::complex<double>* object, const int n);
-void bcast_string(std::string* object, const int n);
-void bcast_double(double* object, const int n);
-void bcast_int(int* object, const int n);
-void bcast_char(char* object, const int n);
+void bcast_complex_double(std::complex<double>* object, const int n
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_string(std::string* object, const int n
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_double(double* object, const int n
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_int(int* object, const int n
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_char(char* object, const int n
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
 
 //(2) bcast single
-void bcast_complex_double(std::complex<double>& object);
-void bcast_string(std::string& object);
-void bcast_double(double& object);
-void bcast_int(int& object);
-void bcast_bool(bool& object);
+void bcast_complex_double(std::complex<double>& object
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_string(std::string& object
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_double(double& object
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_int(int& object
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
+void bcast_bool(bool& object
+#ifdef __MPI
+    , MPI_Comm comm = MPI_COMM_WORLD
+#endif
+    );
 
 } // namespace Parallel_Common
 

source/source_base/parallel_global.cpp

@@ -49,15 +49,19 @@ void Parallel_Global::split_diag_world(const int& diag_np,
                                        const int& my_rank,
                                        int& drank,
                                        int& dsize,
-                                       int& dcolor)
+                                       int& dcolor
+#ifdef __MPI
+                                       , MPI_Comm base_comm
+#endif
+                                       )
 {
 #ifdef __MPI
     assert(diag_np > 0);
     int group_grid_np = -1;
     int color = -1;
     int key = -1;
     divide_mpi_groups(nproc, diag_np, my_rank, group_grid_np, key, color);
-    MPI_Comm_split(MPI_COMM_WORLD, color, key, &DIAG_WORLD);
+    MPI_Comm_split(base_comm, color, key, &DIAG_WORLD);
     MPI_Comm_rank(DIAG_WORLD, &drank);
     MPI_Comm_size(DIAG_WORLD, &dsize);
     dcolor = color;
@@ -69,15 +73,19 @@ void Parallel_Global::split_diag_world(const int& diag_np,
     return;
 }
 
-void Parallel_Global::split_grid_world(const int diag_np, const int& nproc, const int& my_rank, int& grank, int& gsize)
+void Parallel_Global::split_grid_world(const int diag_np, const int& nproc, const int& my_rank, int& grank, int& gsize
+#ifdef __MPI
+    , MPI_Comm base_comm
+#endif
+    )
 {
 #ifdef __MPI
     assert(diag_np > 0);
     int group_grid_np = -1;
     int color = -1;
     int key = -1;
     divide_mpi_groups(nproc, diag_np, my_rank, group_grid_np, color, key);
-    MPI_Comm_split(MPI_COMM_WORLD, color, key, &GRID_WORLD);
+    MPI_Comm_split(base_comm, color, key, &GRID_WORLD);
     MPI_Comm_rank(GRID_WORLD, &grank);
     MPI_Comm_size(GRID_WORLD, &gsize);
 #else
@@ -258,7 +266,11 @@ void Parallel_Global::init_pools(const int& NPROC,
                                  int& MY_BNDGROUP,
                                  int& NPROC_IN_POOL,
                                  int& RANK_IN_POOL,
-                                 int& MY_POOL)
+                                 int& MY_POOL
+#ifdef __MPI
+                                 , MPI_Comm base_comm
+#endif
+                                 )
 {
 #ifdef __MPI
     //----------------------------------------------------------
@@ -273,7 +285,8 @@ void Parallel_Global::init_pools(const int& NPROC,
                                   MY_BNDGROUP,
                                   NPROC_IN_POOL,
                                   RANK_IN_POOL,
-                                  MY_POOL);
+                                  MY_POOL,
+                                  base_comm);
 
     // for test
     // turn on when you want to check the index of pools.
@@ -321,7 +334,8 @@ void Parallel_Global::divide_pools(const int& NPROC,
                                    int& MY_BNDGROUP,
                                    int& NPROC_IN_POOL,
                                    int& RANK_IN_POOL,
-                                   int& MY_POOL)
+                                   int& MY_POOL,
+                                   MPI_Comm base_comm)
 {
     // note: the order of k-point parallelization and band parallelization is important
     //       The order will not change the behavior of KP_WORLD or BP_WORLD, and MY_POOL
@@ -334,7 +348,7 @@ void Parallel_Global::divide_pools(const int& NPROC,
             "When BNDPAR > 1, number of processes NPROC must be divisible by the number of groups BNDPAR * KPAR.");
     }
     // k-point parallelization
-    MPICommGroup kpar_group(MPI_COMM_WORLD);
+    MPICommGroup kpar_group(base_comm);
     kpar_group.divide_group_comm(KPAR, false);
 
     // band parallelization
@@ -362,16 +376,16 @@ void Parallel_Global::divide_pools(const int& NPROC,
         NPROC_IN_BNDGROUP = kpar_group.ngroups * bndpar_group.nprocs_in_group;
         RANK_IN_BPGROUP = kpar_group.my_group * bndpar_group.nprocs_in_group + bndpar_group.rank_in_group;
         MY_BNDGROUP = bndpar_group.my_group;
-        MPI_Comm_split(MPI_COMM_WORLD, MY_BNDGROUP, RANK_IN_BPGROUP, &INT_BGROUP);
+        MPI_Comm_split(base_comm, MY_BNDGROUP, RANK_IN_BPGROUP, &INT_BGROUP);
         MPI_Comm_dup(bndpar_group.inter_comm, &BP_WORLD);
     }
     else
     {
         NPROC_IN_BNDGROUP = NPROC;
         RANK_IN_BPGROUP = MY_RANK;
         MY_BNDGROUP = 0;
-        MPI_Comm_dup(MPI_COMM_WORLD, &INT_BGROUP);
-        MPI_Comm_split(MPI_COMM_WORLD, MY_RANK, 0, &BP_WORLD);
+        MPI_Comm_dup(base_comm, &INT_BGROUP);
+        MPI_Comm_split(base_comm, MY_RANK, 0, &BP_WORLD);
     }
     return;
 }