Addressed more issues

Author: astroC86

examples/plot_matrixmult.py

@@ -63,7 +63,7 @@
         re = min(M, rs + blk_rows)
         cs = pg * blk_cols;
         ce = min(N, cs + blk_cols)
-        a_block, b_block = A[rs:re, :].copy(), B[:, cs:ce].copy()
+        a_block, b_block = A[rs:re, :], B[:, cs:ce]
         if dest == 0:
             A_p, B_p = a_block, b_block
         else:
@@ -79,7 +79,7 @@
 
 Aop = MPISUMMAMatrixMult(A_p, N)
 col_lens = comm.allgather(my_own_cols)
-total_cols = np.add.reduce(col_lens, 0)
+total_cols =  np.sum(col_lens)
 x = DistributedArray(global_shape=K * total_cols,
                      local_shapes=[K * col_len for col_len in col_lens],
                      partition=Partition.SCATTER,
@@ -99,16 +99,17 @@
 my_own_cols = col_end - col_start
 expected_y = C_true[:, col_start:col_end].flatten()
 
+xadj = Aop.H @ y
+
 if not np.allclose(y.local_array, expected_y, atol=1e-6, rtol=1e-14):
     print(f"RANK {rank}: FORWARD VERIFICATION FAILED")
     print(f'{rank} local: {y.local_array}, expected: {C_true[:, col_start:col_end]}')
 else:
     print(f"RANK {rank}: FORWARD VERIFICATION PASSED")
 
-z = Aop.H @ y
 expected_z = Z_true[:, col_start:col_end].flatten()
-if not np.allclose(z.local_array, expected_z, atol=1e-6, rtol=1e-14):
+if not np.allclose(xadj.local_array, expected_z, atol=1e-6, rtol=1e-14):
     print(f"RANK {rank}: ADJOINT VERIFICATION FAILED")
-    print(f'{rank} local: {z.local_array}, expected: {Z_true[:, col_start:col_end]}')
+    print(f'{rank} local: {xadj.local_array}, expected: {Z_true[:, col_start:col_end]}')
 else:
     print(f"RANK {rank}: ADJOINT VERIFICATION PASSED")

pylops_mpi/basicoperators/MatrixMult.py

@@ -25,7 +25,8 @@ def __init__(
         # Determine grid dimensions (P_prime × C) such that P_prime * C ≥ size
         self._P_prime = int(math.ceil(math.sqrt(size)))
         self._C = int(math.ceil(size / self._P_prime))
-        assert self._P_prime * self._C >= size
+        if self._P_prime * self._C < size:
+            raise Exception("Number of Procs must be a square number")
 
         # Compute this process's group and layer indices
         self._group_id = rank % self._P_prime
@@ -36,7 +37,6 @@ def __init__(
         self._layer_comm = base_comm.Split(color=self._layer_id, key=self._group_id)
         self._group_comm = base_comm.Split(color=self._group_id, key=self._layer_id)
 
-        self.dtype = np.dtype(dtype)
         self.A    = np.array(A, dtype=self.dtype, copy=False)
 
         self.M = self._layer_comm.allreduce(self.A.shape[0], op=MPI.SUM)

tests/test_matrixmult.py

@@ -28,7 +28,7 @@
 
 @pytest.mark.mpi(min_size=1)  # SUMMA should also work for 1 process.
 @pytest.mark.parametrize("M, K, N, dtype_str", test_params)
-def test_SUMMAMatrixMult(M, K, N, dtype_str):
+def test_MPIMatrixMult(M, K, N, dtype_str):
     dtype = np.dtype(dtype_str)
 
     cmplx = 1j if np.issubdtype(dtype, np.complexfloating) else 0
@@ -133,56 +133,26 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
 
     # Forward operation: y = A @ B (distributed)
     y_dist = Aop @ x_dist
+    y = y_dist.asarray(),
 
     # Adjoint operation: z = A.H @ y (distributed y representing C)
-    z_dist = Aop.H @ y_dist
+    y_adj_dist = Aop.H @ y_dist
+    y_adj = y_adj_dist.asarray()
 
     if rank == 0:
-        if all(dim > 0 for dim in [M, K, N]):
-            C_true = A_glob @ B_glob
-            Z_true = A_glob.conj().T @ C_true
-        else:  # Handle cases with zero dimensions
-            C_true = np.zeros((M, N), dtype=dtype)
-            Z_true = np.zeros((K if K > 0 else 0, N), dtype=dtype) if K > 0 else np.zeros((0, N), dtype=dtype)
-    else:
-        C_true = Z_true = None
-
-    C_true = comm.bcast(C_true, root=0)
-    Z_true = comm.bcast(Z_true, root=0)
-
-    col_start_C_dist   = my_layer * blk_cols_BC
-    col_end_C_dist     = min(N, col_start_C_dist + blk_cols_BC)
-    my_own_cols_C_dist = max(0, col_end_C_dist - col_start_C_dist)
-    expected_y_shape   = (M * my_own_cols_C_dist,)
-
-    assert y_dist.local_array.shape == expected_y_shape, (
-        f"Rank {rank}: y_dist shape {y_dist.local_array.shape} != expected {expected_y_shape}"
-    )
-
-    if y_dist.local_array.size > 0 and C_true is not None and C_true.size > 0:
-        expected_y_slice = C_true[:, col_start_C_dist:col_end_C_dist]
+        y_np = A_glob @ B_glob
+        y_adj_np = A_glob.conj().T @ y_np
         assert_allclose(
-            y_dist.local_array,
-            expected_y_slice.ravel(),
+            y,
+            y_np.ravel(),
             rtol=1e-14,
-            atol=1e-14,
             err_msg=f"Rank {rank}: Forward verification failed."
         )
 
-    # Verify adjoint operation (z = A.H @ y)
-    expected_z_shape = (K * my_own_cols_C_dist,)
-    assert z_dist.local_array.shape == expected_z_shape, (
-        f"Rank {rank}: z_dist shape {z_dist.local_array.shape} != expected {expected_z_shape}"
-    )
-
-    # Verify adjoint result values
-    if z_dist.local_array.size > 0 and Z_true is not None and Z_true.size > 0:
-        expected_z_slice = Z_true[:, col_start_C_dist:col_end_C_dist]
         assert_allclose(
-            z_dist.local_array,
-            expected_z_slice.ravel(),
+            y_adj,
+            y_adj_np.ravel(),
             rtol=1e-14,
-            atol=1e-14,
             err_msg=f"Rank {rank}: Adjoint verification failed."
         )