minor: stylistic fixes

Author: mrava87

examples/plot_matrixmult.py

@@ -44,7 +44,7 @@
 X = np.random.rand(K * M).astype(dtype=np.float32).reshape(K, M)
 
 ################################################################################
-# The processes are now arranged in a :math:`\sqrt{P} \times \sqrt{P}` grid,
+# The processes are now arranged in a :math:`P' \times P'` grid,
 # where :math:`P` is the total number of processes.
 #
 # We define
@@ -78,14 +78,20 @@
 #   │ (r=1, c=0) │ (r=1, c=1) │
 #   └────────────┴────────────┘
 #    </div>
+#
+# This is obtained by invoking the
+# `:func:pylops_mpi.MPIMatrixMult.active_grid_comm` method, which is also
+# responsible to identify any rank that should be deactivated (if the number
+# of rows of the operator or columns of the input/output matrices are smaller
+# than the row or columm ranks.
 
 base_comm = MPI.COMM_WORLD
 comm, rank, row_id, col_id, is_active = MPIMatrixMult.active_grid_comm(base_comm, N, M)
 print(f"Process {base_comm.Get_rank()}  is {"active" if is_active else "inactive"}")
 if not is_active: exit(0)
-p_prime = math.isqrt(comm.Get_size())
 
 # Create sub‐communicators
+p_prime = math.isqrt(comm.Get_size())
 row_comm = comm.Split(color=row_id, key=col_id)  # all procs in same row
 col_comm = comm.Split(color=col_id, key=row_id)  # all procs in same col
 
@@ -127,11 +133,11 @@
 
 rs = col_id * blk_rows
 re = min(N, rs + blk_rows)
-my_own_rows = max(0,re - rs)
+my_own_rows = max(0, re - rs)
 
 cs = row_id * blk_cols
 ce = min(M, cs + blk_cols)
-my_own_cols = max(0,ce - cs)
+my_own_cols = max(0, ce - cs)
 
 A_p, X_p = A[rs:re, :].copy(), X[:, cs:ce].copy()
 
@@ -191,14 +197,14 @@
 for cnt in col_counts:
     block_size = K * cnt
     xadj_blk = xadj[offset: offset + block_size]
-    if len(xadj_blk)!= 0:
+    if len(xadj_blk) != 0:
         xadj_blocks.append(
             xadj_blk.reshape(K, cnt)
         )
     offset += block_size
 xadj = np.hstack(xadj_blocks)
 
-if comm.Get_rank() == 0:
+if rank == 0:
     y_loc = (A @ X).squeeze()
     xadj_loc = (A.T.dot(y_loc.conj())).conj().squeeze()
 

pylops_mpi/basicoperators/MatrixMult.py

@@ -190,30 +190,41 @@ def _matvec(self, x: DistributedArray) -> DistributedArray:
 
     @staticmethod
     def active_grid_comm(base_comm:MPI.Comm, N:int, M:int):
-        """
-        Configure a square process grid from a parent MPI communicator and select the subset of "active" processes.
+        r"""Configure active grid
 
-        Each process in base_comm is assigned to a logical 2D grid of size p_prime x p_prime,
-        where p_prime = floor(sqrt(total_ranks)). Only the first `active_dim x active_dim` processes
-        (by row-major order) are considered "active". Inactive ranks return immediately with no new communicator.
+        Configure a square process grid from a parent MPI communicator and
+        select the subset of "active" processes. Each process in ``base_comm``
+        is assigned to a logical 2D grid of size :math:`P' \times P'`,
+        where :math:`P' = \bigl \lceil \sqrt{P} \bigr \rceil`. Only the first
+        :math:`active_dim x active_dim` processes
+        (by row-major order) are considered "active". Inactive ranks return
+        immediately with no new communicator.
 
         Parameters:
         -----------
-        base_comm : MPI.Comm
-            The parent communicator (e.g., MPI.COMM_WORLD).
-        N : int
-            Number of rows of your global data domain.
-        M : int
-            Number of columns of your global data domain.
+        base_comm : :obj:`mpi4py.MPI.Comm`
+            MPI Parent Communicator. (e.g., ``mpi4py.MPI.COMM_WORLD``).
+        N : :obj:`int`
+            Number of rows of the global data domain.
+        M : :obj:`int`
+            Number of columns of the global data domain.
 
         Returns:
         --------
-        tuple:
-            comm (MPI.Comm or None) : Sub-communicator including only active ranks.
-            rank (int)              : Rank within the new sub-communicator (or original rank if inactive).
-            row (int)               : Grid row index of this process in the active grid (or original rank if inactive).
-            col (int)               : Grid column index of this process in the active grid (or original rank if inactive).
-            is_active (bool)        : Flag indicating whether this rank is in the active sub-grid.
+        comm : :obj:`mpi4py.MPI.Comm`
+            Sub-communicator including only active ranks.
+        rank : :obj:`int`
+            Rank within the new sub-communicator (or original rank
+            if inactive).
+        row : :obj:`int`
+            Grid row index of this process in the active grid (or original rank
+              if inactive).
+        col : :obj:`int`
+            Grid column index of this process in the active grid
+            (or original rank if inactive).
+        is_active : :obj:`bool`
+            Flag indicating whether this rank is in the active sub-grid.
+
         """
         rank = base_comm.Get_rank()
         size = base_comm.Get_size()
@@ -229,10 +240,10 @@ def active_grid_comm(base_comm:MPI.Comm, N:int, M:int):
                         if (r // p_prime) < active_dim and (r % p_prime) < active_dim]
         new_group = base_comm.Get_group().Incl(active_ranks)
         new_comm = base_comm.Create_group(new_group)
-
         p_prime_new = math.isqrt(len(active_ranks))
         new_rank = new_comm.Get_rank()
         new_row, new_col = divmod(new_rank, p_prime_new)
+
         return new_comm, new_rank, new_row, new_col, True
 
     def _rmatvec(self, x: DistributedArray) -> DistributedArray:

tests/test_matrixmult.py

@@ -13,6 +13,7 @@
 
 np.random.seed(42)
 base_comm = MPI.COMM_WORLD
+size = base_comm.Get_size()
 
 # Define test cases: (N, K, M, dtype_str)
 # M, K, N are matrix dimensions A(N,K), B(K,M)
@@ -29,14 +30,13 @@
 
 @pytest.mark.mpi(min_size=1)
 @pytest.mark.parametrize("M, K, N, dtype_str", test_params)
-def test_SUMMAMatrixMult(N, K, M, dtype_str):
+def test_MPIMatrixMult(N, K, M, dtype_str):
     dtype = np.dtype(dtype_str)
 
     cmplx = 1j if np.issubdtype(dtype, np.complexfloating) else 0
     base_float_dtype = np.float32 if dtype == np.complex64 else np.float64
 
     comm, rank, row_id, col_id, is_active = MPIMatrixMult.active_grid_comm(base_comm, N, M)
-    print(f"Process {base_comm.Get_rank()}  is {"active" if is_active else "inactive"}")
     if not is_active: return
 
     size = comm.Get_size()