minor: fix flake8

Author: mrava87

examples/plot_matrixmult.py

@@ -1,21 +1,21 @@
-"""
+r"""
 Distributed Matrix Multiplication
 =================================
 This example shows how to use the :py:class:`pylops_mpi.basicoperators.MPIMatrixMult`
-operator to perform matrix-matrix multiplication between a matrix :math:`\mathbf{A}` 
+operator to perform matrix-matrix multiplication between a matrix :math:`\mathbf{A}`
 blocked over rows (i.e., blocks of rows are stored over different ranks) and a
-matrix :math:`\mathbf{X}` blocked over columns (i.e., blocks of columns are 
-stored over different ranks), with equal number of row and column blocks. 
-Similarly, the adjoint operation can be peformed with a matrix :math:`\mathbf{Y}` 
+matrix :math:`\mathbf{X}` blocked over columns (i.e., blocks of columns are
+stored over different ranks), with equal number of row and column blocks.
+Similarly, the adjoint operation can be peformed with a matrix :math:`\mathbf{Y}`
 blocked in the same fashion of matrix :math:`\mathbf{X}`.
 
-Note that whilst the different blocks of the matrix :math:`\mathbf{A}` are directly 
-stored in the operator on different ranks, the matrix :math:`\mathbf{X}` is 
-effectively represented by a 1-D :py:class:`pylops_mpi.DistributedArray` where 
+Note that whilst the different blocks of the matrix :math:`\mathbf{A}` are directly
+stored in the operator on different ranks, the matrix :math:`\mathbf{X}` is
+effectively represented by a 1-D :py:class:`pylops_mpi.DistributedArray` where
 the different blocks are flattened and stored on different ranks. Note that to
-optimize communications, the ranks are organized in a 2D grid and some of the 
-row blocks of :math:`\mathbf{A}` and column blocks of :math:`\mathbf{X}` are 
-replicated across different ranks - see below for details.  
+optimize communications, the ranks are organized in a 2D grid and some of the
+row blocks of :math:`\mathbf{A}` and column blocks of :math:`\mathbf{X}` are
+replicated across different ranks - see below for details.
 
 """
 
@@ -30,17 +30,17 @@
 plt.close("all")
 
 ###############################################################################
-# We set the seed such that all processes can create the input matrices filled 
-# with the same random number. In practical application, such matrices will be 
+# We set the seed such that all processes can create the input matrices filled
+# with the same random number. In practical application, such matrices will be
 # filled with data that is appropriate that is appropriate the use-case.
 np.random.seed(42)
 
 ###############################################################################
-# Next we obtain the MPI parameters for each rank and check that the number 
+# Next we obtain the MPI parameters for each rank and check that the number
 # of processes (``size``) is a square number
 comm = MPI.COMM_WORLD
-rank = comm.Get_rank() # rank of current process
-size = comm.Get_size() # number of processes
+rank = comm.Get_rank()  # rank of current process
+size = comm.Get_size()  # number of processes
 
 p_prime = math.isqrt(size)
 repl_factor = p_prime
@@ -58,7 +58,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:`\sqrt{P} \times \sqrt{P}` grid,
 # where :math:`P` is the total number of processes.
 #
 # We define
@@ -71,7 +71,7 @@
 # .. math::
 #    R = \bigl\lceil \tfrac{P}{P'} \bigr\rceil.
 #
-# Each process is therefore assigned a pair of coordinates 
+# Each process is therefore assigned a pair of coordinates
 # :math:`(r,c)` within this grid:
 #
 # .. math::
@@ -101,7 +101,7 @@
 col_comm = comm.Split(color=my_col, key=my_row)  # all procs in same col
 
 ################################################################################
-# At this point we divide the rows and columns of :math:`\mathbf{A}` and  
+# At this point we divide the rows and columns of :math:`\mathbf{A}` and
 # :math:`\mathbf{X}`, respectively, such that each rank ends up with:
 #
 #  - :math:`A_{p} \in \mathbb{R}^{\text{my_own_rows}\times K}`
@@ -147,12 +147,12 @@
 A_p, X_p = A[rs:re, :].copy(), X[:, cs:ce].copy()
 
 ################################################################################
-# We are now ready to create the :py:class:`pylops_mpi.basicoperators.MPIMatrixMult` 
+# We are now ready to create the :py:class:`pylops_mpi.basicoperators.MPIMatrixMult`
 # operator and the input matrix :math:`\mathbf{X}`
 Aop = MPIMatrixMult(A_p, M, dtype="float32")
 
 col_lens = comm.allgather(my_own_cols)
-total_cols =  np.sum(col_lens)
+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,
@@ -162,22 +162,22 @@
 x[:] = X_p.flatten()
 
 ################################################################################
-# We can now apply the forward pass :math:`\mathbf{y} = \mathbf{Ax}` (which effectively 
+# We can now apply the forward pass :math:`\mathbf{y} = \mathbf{Ax}` (which effectively
 # implements a distributed matrix-matrix multiplication :math:`Y = \mathbf{AX}`)
-# Note :math:`\mathbf{Y}` is distributed in the same way as the input 
+# Note :math:`\mathbf{Y}` is distributed in the same way as the input
 # :math:`\mathbf{X}`.
 y = Aop @ x
 
 ###############################################################################
-# Next we apply the adjoint pass :math:`\mathbf{x}_{adj} = \mathbf{A}^H \mathbf{x}` 
-# (which effectively implements a distributed matrix-matrix multiplication 
+# Next we apply the adjoint pass :math:`\mathbf{x}_{adj} = \mathbf{A}^H \mathbf{x}`
+# (which effectively implements a distributed matrix-matrix multiplication
 # :math:`\mathbf{X}_{adj} = \mathbf{A}^H \mathbf{X}`). Note that
-# :math:`\mathbf{X}_{adj}` is again distributed in the same way as the input 
+# :math:`\mathbf{X}_{adj}` is again distributed in the same way as the input
 # :math:`\mathbf{X}`.
 xadj = Aop.H @ y
 
 ###############################################################################
-# To conclude we verify our result against the equivalent serial version of 
+# To conclude we verify our result against the equivalent serial version of
 # the operation by gathering the resulting matrices in rank0 and reorganizing
 # the returned 1D-arrays into 2D-arrays.
 
@@ -210,15 +210,15 @@
     xadj_loc = (A.T.dot(y_loc.conj())).conj().squeeze()
 
     if not np.allclose(y, y_loc, rtol=1e-6):
-        print(f" FORWARD VERIFICATION FAILED")
-        print(f'distributed: {y}') 
+        print("FORWARD VERIFICATION FAILED")
+        print(f'distributed: {y}')
         print(f'expected: {y_loc}')
     else:
-        print(f"FORWARD VERIFICATION PASSED")
+        print("FORWARD VERIFICATION PASSED")
 
     if not np.allclose(xadj, xadj_loc, rtol=1e-6):
-        print(f" ADJOINT VERIFICATION FAILED")
-        print(f'distributed: {xadj}') 
+        print("ADJOINT VERIFICATION FAILED")
+        print(f'distributed: {xadj}')
         print(f'expected: {xadj_loc}')
     else:
-        print(f"ADJOINT VERIFICATION PASSED")
+        print("ADJOINT VERIFICATION PASSED")

pylops_mpi/basicoperators/MatrixMult.py

@@ -13,8 +13,8 @@
 
 class MPIMatrixMult(MPILinearOperator):
     r"""MPI Matrix multiplication
-    
-    Implement distributed matrix-matrix multiplication between a matrix 
+
+    Implement distributed matrix-matrix multiplication between a matrix
     :math:`\mathbf{A}` blocked over rows (i.e., blocks of rows are stored
     over different ranks) and the input model and data vector, which are both to
     be interpreted as matrices blocked over columns.
@@ -26,7 +26,7 @@ class MPIMatrixMult(MPILinearOperator):
         where :math:`N_{loc}` is the number of rows stored on this MPI rank and
         ``K`` is the global number of columns.
     M : :obj:`int`
-        Global leading dimension (i.e., number of columns) of the matrices 
+        Global leading dimension (i.e., number of columns) of the matrices
         representing the input model and data vectors.
     saveAt : :obj:`bool`, optional
         Save ``A`` and ``A.H`` to speed up the computation of adjoint
@@ -52,24 +52,24 @@ class MPIMatrixMult(MPILinearOperator):
 
     Notes
     -----
-    This operator performs a matrix-matrix multiplication, whose forward 
+    This operator performs a matrix-matrix multiplication, whose forward
     operation can be described as :math:`Y = A \cdot X` where:
 
     - :math:`\mathbf{A}` is the distributed matrix operator of shape :math:`[N \times K]`
     - :math:`\mathbf{X}` is the distributed operand matrix of shape :math:`[K \times M]`
     - :math:`\mathbf{Y}` is the resulting distributed matrix of shape :math:`[N \times M]`
 
-    whilst the adjoint operation is represented by 
-    :math:`\mathbf{X}_{adj} = \mathbf{A}^H \cdot \mathbf{Y}` where 
+    whilst the adjoint operation is represented by
+    :math:`\mathbf{X}_{adj} = \mathbf{A}^H \cdot \mathbf{Y}` where
     :math:`\mathbf{A}^H` is the complex conjugate and transpose of :math:`\mathbf{A}`.
-    
-    This implementation is based on a 1D block distribution of the operator 
+
+    This implementation is based on a 1D block distribution of the operator
     matrix and reshaped model and data vectors replicated across :math:`P`
-    processes by a factor equivalent to :math:`\sqrt{P}` across a square process 
+    processes by a factor equivalent to :math:`\sqrt{P}` across a square process
     grid (:math:`\sqrt{P}\times\sqrt{P}`). More specifically:
 
     - The matrix ``A`` is distributed across MPI processes in a block-row fashion
-      and each process holds a local block of ``A`` with shape 
+      and each process holds a local block of ``A`` with shape
       :math:`[N_{loc} \times K]`
     - The operand matrix ``X`` is distributed in a block-column fashion and
       each process holds a local block of ``X`` with shape
@@ -111,7 +111,7 @@ class MPIMatrixMult(MPILinearOperator):
        sum of the contributions from all column blocks of ``A^H``, processes in
        the same row perform an ``allreduce`` sum to combine their partial results.
        This gives the complete ``(K, M_local)`` result for their assigned column.
-    
+
     """
     def __init__(
             self,
@@ -125,7 +125,7 @@ def __init__(
         size = base_comm.Get_size()
 
         # Determine grid dimensions (P_prime × C) such that P_prime * C ≥ size
-        self._P_prime =  math.isqrt(size)
+        self._P_prime = math.isqrt(size)
         self._C = self._P_prime
         if self._P_prime * self._C != size:
             raise Exception(f"Number of processes must be a square number, provided {size} instead...")
@@ -138,7 +138,8 @@ def __init__(
         self._col_comm = base_comm.Split(color=self._col_id, key=self._row_id)
 
         self.A = A.astype(np.dtype(dtype))
-        if saveAt: self.At = A.T.conj()
+        if saveAt:
+            self.At = A.T.conj()
 
         self.N = self._row_comm.allreduce(self.A.shape[0], op=MPI.SUM)
         self.K = A.shape[1]

tests/test_matrixmult.py

@@ -23,9 +23,9 @@
     pytest.param(37, 37, 37, "float32", id="f32_37_37_37"),
     pytest.param(50, 30, 40, "float64", id="f64_50_30_40"),
     pytest.param(22, 20, 16, "complex64", id="c64_22_20_16"),
-    pytest.param( 3,  4,  5, "float32", id="f32_3_4_5"),
-    pytest.param( 1,  2,  1, "float64", id="f64_1_2_1",),
-    pytest.param( 2,  1,  3, "float32", id="f32_2_1_3",),
+    pytest.param(3, 4, 5, "float32", id="f32_3_4_5"),
+    pytest.param(1, 2, 1, "float64", id="f64_1_2_1",),
+    pytest.param(2, 1, 3, "float32", id="f32_2_1_3",),
 ]
 
 
@@ -34,8 +34,8 @@
 def test_SUMMAMatrixMult(N, K, M, dtype_str):
     p_prime = math.isqrt(size)
     C = p_prime
-    if  p_prime * C != size:
-        pytest.skip(f"Number of processes must be a square number, "
+    if p_prime * C != size:
+        pytest.skip("Number of processes must be a square number, "
                     "provided {size} instead...")
 
     dtype = np.dtype(dtype_str)
@@ -68,8 +68,8 @@ def test_SUMMAMatrixMult(N, K, M, dtype_str):
     X_glob_imag = np.arange(K * M, dtype=base_float_dtype).reshape(K, M) * 0.7
     X_glob = (X_glob_real + cmplx * X_glob_imag).astype(dtype)
 
-    A_p = A_glob[row_start_A:row_end_A,:]
-    X_p = X_glob[:,col_start_X:col_end_X]
+    A_p = A_glob[row_start_A:row_end_A, :]
+    X_p = X_glob[:, col_start_X:col_end_X]
 
     # Create MPIMatrixMult operator
     Aop = MPIMatrixMult(A_p, M, base_comm=comm, dtype=dtype_str)
@@ -137,4 +137,4 @@ def test_SUMMAMatrixMult(N, K, M, dtype_str):
         )
 
     col_comm.Free()
-    row_comm.Free()
+    row_comm.Free()