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.

Changes:
- parallel_global: add base_comm parameter to divide_pools(),
  split_diag_world(), split_grid_world(), init_pools()
- parallel_reduce: add comm parameter to all reduce_all, reduce_min,
  reduce_max, reduce_double_allpool, gather_int_all functions
- parallel_common: add comm parameter to all bcast_* functions
- PyESolverLCAO: add mpi_comm_handle parameter to initialize()
- LCAOWorkflow: add mpi_comm parameter forwarded to ESolver

All parameters default to MPI_COMM_WORLD for full backward compatibility.
Existing code compiles and behaves identically without modification.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: dyzheng

docs/plans/2026-02-10-mpi-comm-injection-design.md

@@ -0,0 +1,128 @@
+# MPI Communicator Injection Design
+
+## Problem
+
+ABACUS hardcodes `MPI_COMM_WORLD` as the base communicator throughout its initialization
+and communication paths. This prevents workflows like parallel NEB, phonon calculations,
+or replica exchange MD from splitting MPI ranks into independent groups, each running its
+own ESolver instance in a separate subprocess.
+
+The pyabacus ESolver bindings have no mechanism to pass a custom MPI communicator from
+Python to C++. All sub-communicators (`POOL_WORLD`, `DIAG_WORLD`, `GRID_WORLD`, etc.)
+are derived from `MPI_COMM_WORLD`, and all global reduction/broadcast wrappers
+(`parallel_reduce`, `parallel_common`) hardcode `MPI_COMM_WORLD`.
+
+## Goal
+
+Enable general-purpose MPI communicator injection so that any workflow can pass a
+sub-communicator to the ESolver. One ESolver instance per process, but each process
+can use any communicator (not just `MPI_COMM_WORLD`). Default to `MPI_COMM_WORLD`
+when no communicator is provided, preserving full backward compatibility.
+
+## Architecture
+
+```
+Python (mpi4py)                  C++ pybind11 bindings            ABACUS core
+-------------------              ---------------------            -----------
+comm = MPI.COMM_WORLD            PyESolverLCAO::initialize(       Parallel_Global::divide_pools(
+sub = comm.Split(color, key)       input_dir,                       base_comm, ...)
+                                   mpi_comm_handle)             Parallel_Reduce::reduce_all(
+esolver.initialize(                                                ..., base_comm)
+  input_dir,                     Stores base_comm_ member        Parallel_Common::bcast_*(
+  comm_handle=sub.py2f())        Passes to Parallel_Global          ..., base_comm)
+```
+
+## Scope
+
+### In scope (this PR)
+
+1. Add `MPI_Comm base_comm` parameter to `Parallel_Global::divide_pools()`,
+   `split_diag_world()`, `split_grid_world()` — defaulting to `MPI_COMM_WORLD`
+2. Add `MPI_Comm comm` parameter to all `Parallel_Reduce::reduce_all` and
+   `Parallel_Common::bcast_*` functions — defaulting to `MPI_COMM_WORLD`
+3. Add `int mpi_comm_handle` parameter to `PyESolverLCAO::initialize()` and
+   `PyESolverPW::initialize()` — defaulting to `-1` (meaning `MPI_COMM_WORLD`)
+4. Update pybind11 bindings to expose the new parameter
+5. Unit tests for communicator injection
+
+### Out of scope (follow-up PRs)
+
+- Parallel NEB orchestration in Python
+- Multi-instance-per-process support (would require encapsulating GlobalV)
+- Changing the ~280 scattered `MPI_COMM_WORLD` references in computation/IO modules
+  (unnecessary for single-instance-per-process + subprocess model)
+
+## Detailed Changes
+
+### Layer 1: `source_base/parallel_global.h` and `.cpp`
+
+Add `MPI_Comm base_comm = MPI_COMM_WORLD` as the last parameter to:
+
+- `divide_pools()` — replace `MPI_COMM_WORLD` with `base_comm` in `MPICommGroup`
+  constructor (line 337), `MPI_Comm_split` for `INT_BGROUP` (line 365),
+  `MPI_Comm_dup` fallback (line 373), `MPI_Comm_split` fallback (line 374)
+- `split_diag_world()` — replace `MPI_COMM_WORLD` in `MPI_Comm_split` (line 60)
+- `split_grid_world()` — replace `MPI_COMM_WORLD` in `MPI_Comm_split` (line 80)
+
+~10 lines changed, ~3 signature changes.
+
+### Layer 2: `source_base/parallel_reduce.h` and `.cpp`
+
+Add `MPI_Comm comm = MPI_COMM_WORLD` to all `reduce_all` overloads (~13 functions),
+`reduce_double_allpool` (~2), `gather_int_all` (~1), `reduce_min` (~3),
+`reduce_max` (~2). Total: ~19 signature changes, ~19 substitutions.
+
+Functions using `POOL_WORLD`, `DIAG_WORLD`, `GRID_WORLD` are unchanged.
+
+### Layer 3: `source_base/parallel_common.h` and `.cpp`
+
+Add `MPI_Comm comm = MPI_COMM_WORLD` to all `bcast_*` functions (~10 functions).
+Total: ~10 signature changes, ~12 substitutions.
+
+### Layer 4: Pybind11 bindings
+
+**`py_esolver_lcao.hpp`**: Add `int mpi_comm_handle = -1` to `initialize()`.
+
+**`py_esolver_lcao_impl.cpp`**: Convert handle via `MPI_Comm_f2c()`, use resulting
+communicator for `MPI_Comm_size`, `MPI_Comm_rank`, and pass to `divide_pools()` etc.
+
+**`py_esolver_pw_impl.cpp`**: Same changes as LCAO.
+
+**`py_esolver_bindings.cpp`**: Update `.def("initialize", ...)` to include
+`"mpi_comm_handle"_a = -1`.
+
+~15 lines changed across 4 files.
+
+### Layer 5: Python interface
+
+**`esolver/workflow.py`**: Add optional `mpi_comm` parameter to `LCAOWorkflow` and
+`PWWorkflow` that forwards `comm.py2f()` to `esolver.initialize()`.
+
+**`ase/calculator.py`**: Add optional `mpi_comm` parameter to `AbacusCalculator`.
+
+### Scattered `MPI_COMM_WORLD` references (~280 in other modules)
+
+**No changes.** In the single-instance-per-process + subprocess model, each
+subprocess's `MPI_COMM_WORLD` is already the correct sub-communicator. These
+references are correct as-is.
+
+## Backward Compatibility
+
+1. **Standalone ABACUS binary**: `main.cpp` calls `read_pal_param()` which still
+   uses `MPI_COMM_WORLD` directly. Unaffected.
+2. **Existing pyabacus scripts**: `initialize(input_dir)` with no second argument
+   defaults to `MPI_COMM_WORLD`. Identical behavior.
+3. **Existing `parallel_reduce`/`parallel_common` callers**: Default parameter
+   means all call sites compile and behave identically without modification.
+
+## Testing
+
+1. **Unit test: `parallel_global` with sub-communicator** — Split `MPI_COMM_WORLD`,
+   call `divide_pools()` with sub-communicator, verify derived communicators are
+   correct relative to the sub-communicator.
+2. **Unit test: `parallel_reduce` with custom communicator** — Verify reduction
+   scoped to sub-communicator.
+3. **Integration test: pyabacus ESolver with sub-communicator** — `mpirun -np 4`,
+   split into 2 groups of 2, each initializes ESolver independently.
+4. **Regression test**: Run existing tests without `mpi_comm_handle` argument,
+   verify identical results.

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