More minor changes

Author: astroC86

pylops_mpi/basicoperators/MatrixMult.py

@@ -91,13 +91,15 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
         if x.partition != Partition.SCATTER:
             raise ValueError(f"x should have partition={Partition.SCATTER}. Got {x.partition} instead.")
+
         y = DistributedArray(
             global_shape=(self.K * self.dimsd[1]),
             local_shapes=[self.K * c for c in self._rank_col_lens],
             mask=x.mask,
             partition=Partition.SCATTER,
             dtype=self.dtype,
         )
+
         x_arr = x.local_array.reshape((self.M, self._local_ncols)).astype(self.dtype)
         X_tile = x_arr[self._row_start:self._row_end, :]
 

tests/test_matrixmult.py

@@ -35,27 +35,26 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     cmplx = 1j if np.issubdtype(dtype, np.complexfloating) else 0
     base_float_dtype = np.float32 if dtype == np.complex64 else np.float64
 
-    P_prime = int(math.ceil(math.sqrt(size)))
-    C = int(math.ceil(size / P_prime))
-    assert P_prime * C >= size  # Ensure process grid covers all processes
+    p_prime = int(math.ceil(math.sqrt(size)))
+    C = int(math.ceil(size / p_prime))
+    assert p_prime * C == size
 
-    my_group = rank % P_prime
-    my_layer = rank // P_prime
+    my_group = rank % p_prime
+    my_layer = rank // p_prime
 
     # Create sub-communicators
     layer_comm = comm.Split(color=my_layer, key=my_group)
     group_comm = comm.Split(color=my_group, key=my_layer)
 
     # Calculate local matrix dimensions
-    blk_rows_A = int(math.ceil(M / P_prime))
+    blk_rows_A = int(math.ceil(M / p_prime))
     row_start_A = my_group * blk_rows_A
     row_end_A = min(M, row_start_A + blk_rows_A)
-    my_own_rows_A = max(0, row_end_A - row_start_A)
 
-    blk_cols_BC = int(math.ceil(N / P_prime))
+    blk_cols_BC = int(math.ceil(N / p_prime))
     col_start_B = my_layer * blk_cols_BC
     col_end_B = min(N, col_start_B + blk_cols_BC)
-    my_own_cols_B = max(0, col_end_B - col_start_B)
+    local_col_B_len = max(0, col_end_B - col_start_B)
 
 
     A_glob_real = np.arange(M * K, dtype=base_float_dtype).reshape(M, K)
@@ -73,33 +72,28 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     Aop = MPIMatrixMult(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)
-    total_cols = np.sum(all_my_own_cols_B)
+    all_local_col_len = comm.allgather(local_col_B_len)
+    total_cols = np.sum(all_local_col_len)
 
     x_dist = DistributedArray(
         global_shape=(K * total_cols),
-        local_shapes=[K * cl_b for cl_b in all_my_own_cols_B],
+        local_shapes=[K * cl_b for cl_b in all_local_col_len],
         partition=Partition.SCATTER,
         base_comm=comm,
+        mask=[i // p_prime for i in range(size)],
         dtype=dtype
     )
 
-    if B_p.size > 0:
-        x_dist.local_array[:] = B_p.ravel()
-    else:
-        assert x_dist.local_array.size == 0, (
-            f"Rank {rank}: B_p empty but x_dist.local_array not empty "
-            f"(size {x_dist.local_array.size})"
-        )
+    x_dist.local_array[:] = B_p.ravel()
 
     # Forward operation: y = A @ B (distributed)
     y_dist = Aop @ x_dist
 
-    # Adjoint operation: z = A.H @ y (distributed y representing C)
-    z_dist = Aop.H @ y_dist
+    # Adjoint operation: xadj = A.H @ y (distributed)
+    xadj_dist = Aop.H @ y_dist
 
-    C_true = A_glob @ B_glob
-    Z_true = A_glob.conj().T @ C_true
+    y_loc = A_glob @ B_glob
+    xadj_loc = A_glob.conj().T @ y_loc
 
     col_start_C_dist   = my_layer * blk_cols_BC
     col_end_C_dist     = min(N, col_start_C_dist + blk_cols_BC)
@@ -110,27 +104,27 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
         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]
+    if y_dist.local_array.size > 0 and y_loc is not None and y_loc.size > 0:
+        expected_y_slice = y_loc[:, col_start_C_dist:col_end_C_dist]
         assert_allclose(
             y_dist.local_array,
             expected_y_slice.ravel(),
             rtol=np.finfo(np.dtype(dtype)).resolution,
             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 operation (xadj = A.H @ y)
+    expected_xadj_shape = (K * my_own_cols_C_dist,)
+    assert xadj_dist.local_array.shape == expected_xadj_shape, (
+        f"Rank {rank}: z_dist shape {xadj_dist.local_array.shape} != expected {expected_xadj_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]
+    if xadj_dist.local_array.size > 0 and xadj_loc  is not None and xadj_loc .size > 0:
+        expected_xadj_slice = xadj_loc [:, col_start_C_dist:col_end_C_dist]
         assert_allclose(
-            z_dist.local_array,
-            expected_z_slice.ravel(),
+            xadj_dist.local_array,
+            expected_xadj_slice.ravel(),
             rtol=np.finfo(np.dtype(dtype)).resolution,
             err_msg=f"Rank {rank}: Adjoint verification failed."
         )