Addressed some comments

Author: astroC86

examples/matrixmul.py renamed to examples/plot_matrixmult.py

@@ -1,10 +1,17 @@
-import sys
+"""
+Distributed Matrix Multiplication
+=========================
+This example shows how to use the :py:class:`pylops_mpi.basicoperators.MatrixMultiply.SUMMAMatrixMult`.
+This class provides a way to distribute arrays across multiple processes in
+a parallel computing environment.
+"""
+
 import math
 import numpy as np
 from mpi4py import MPI
 
 from pylops_mpi import DistributedArray, Partition
-from pylops_mpi.basicoperators.MatrixMultiply import SUMMAMatrixMult
+from pylops_mpi.basicoperators.MatrixMultiply import MPISUMMAMatrixMult
 
 np.random.seed(42)
 
@@ -15,12 +22,15 @@
 
 P_prime = int(math.ceil(math.sqrt(nProcs)))
 C       = int(math.ceil(nProcs / P_prime))
-assert P_prime * C >= nProcs
+
+if P_prime * C < nProcs:
+    print("No. of procs has to be a square number")
+    exit(-1)
 
 # matrix dims
-M = 32    # any M
-K = 32    # any K
-N = 35    # any N
+M = 32
+K = 32
+N = 35
 
 blk_rows = int(math.ceil(M / P_prime))
 blk_cols = int(math.ceil(N / P_prime))
@@ -66,7 +76,7 @@
 
 comm.Barrier()
 
-Aop        = SUMMAMatrixMult(A_p, N)
+Aop        = MPISUMMAMatrixMult(A_p, N)
 col_lens   = comm.allgather(my_own_cols)
 total_cols = np.add.reduce(col_lens, 0)
 x = DistributedArray(global_shape=K * total_cols,

pylops_mpi/basicoperators/MatrixMultiply.py

@@ -11,7 +11,7 @@
 )
 
 
-class SUMMAMatrixMult(MPILinearOperator):
+class MPISUMMAMatrixMult(MPILinearOperator):
     def __init__(
             self,
             A: NDArray,
@@ -90,7 +90,6 @@ def _matvec(self, x: DistributedArray) -> DistributedArray:
         y = DistributedArray(global_shape= (self.M * self.dimsd[1]),
                              local_shapes=[(self.M * c) for c in layer_col_lens],
                              mask=mask,
-                             #axis=1,
                              partition=Partition.SCATTER,
                              dtype=self.dtype)
         y[:] = C_local.flatten()
@@ -134,7 +133,6 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray:
             global_shape=(self.K * self.dimsd[1]),
             local_shapes=[self.K * c for c in layer_col_lens],
             mask=mask,
-            #axis=1
             partition=Partition.SCATTER,
             dtype=self.dtype,
         )

tests/test_matrixmult.py

@@ -3,10 +3,9 @@
 from numpy.testing import assert_allclose
 from mpi4py import MPI
 import math
-import sys
 
 from pylops_mpi import DistributedArray, Partition
-from pylops_mpi.basicoperators.MatrixMultiply import SUMMAMatrixMult
+from pylops_mpi.basicoperators.MatrixMultiply import MPISUMMAMatrixMult
 
 np.random.seed(42)
 
@@ -109,7 +108,7 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     comm.Barrier()
 
     # Create SUMMAMatrixMult operator
-    Aop = SUMMAMatrixMult(A_p, N, base_comm=comm, dtype=dtype_str)
+    Aop = MPISUMMAMatrixMult(A_p, N, base_comm=comm, dtype=dtype_str)
 
     # Create DistributedArray for input x (representing B flattened)
     all_my_own_cols_B = comm.allgather(my_own_cols_B)
@@ -166,7 +165,7 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
             y_dist.local_array,
             expected_y_slice.ravel(),
             rtol=1e-14,
-            atol=1e-7,
+            atol=1e-14,
             err_msg=f"Rank {rank}: Forward verification failed."
         )
 
@@ -183,7 +182,7 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
             z_dist.local_array,
             expected_z_slice.ravel(),
             rtol=1e-14,
-            atol=1e-7,
+            atol=1e-14,
             err_msg=f"Rank {rank}: Adjoint verification failed."
         )