doc: improved docstrings of MPILinearOperator and MPIStackedLinearOperator

Author: mrava87

pylops_mpi/LinearOperator.py

@@ -12,20 +12,27 @@
 
 
 class MPILinearOperator:
-    """Common interface for performing matrix-vector products in distributed fashion.
+    """MPI-enabled PyLops Linear Operator
 
-    This class provides methods to perform matrix-vector product and adjoint matrix-vector
-    products using MPI.
+    Common interface for performing matrix-vector products in distributed fashion.
 
-    .. note:: End users of pylops-mpi should not use this class directly but simply
-      use operators that are already implemented. This class is meant for
-      developers only, it has to be used as the parent class of any new operator
-      developed within pylops-mpi.
+    In practice, this class provides methods to perform matrix-vector and
+    adjoint matrix-vector products between any :obj:`pylops.LinearOperator`
+    (which must be the same across ranks) and a :class:`pylops_mpi.DistributedArray`
+    with ``Partition.BROADCAST`` and ``Partition.UNSAFE_BROADCAST`` partition. It
+    internally handles the extraction of the local array from the distributed array
+    and the creation of the output :class:`pylops_mpi.DistributedArray`.
+
+    Note that whilst this operator could also be used with different
+    :obj:`pylops.LinearOperator` across ranks, and with a
+    :class:`pylops_mpi.DistributedArray` with ``Partition.SCATTER``, it is however
+    reccomended to use the :class:`pylops_mpi.basicoperators.MPIBlockDiag` operator
+    instead as this can also handle distributed arrays with subcommunicators.
 
     Parameters
     ----------
-    Op : :obj:`pylops.LinearOperator`, optional
-        Linear Operator. Defaults to ``None``.
+    Op : :obj:`pylops.LinearOperator`
+        PyLops Linear Operator to wrap.
     shape : :obj:`tuple(int, int)`, optional
         Shape of the MPI Linear Operator. Defaults to ``None``.
     dtype : :obj:`str`, optional
@@ -35,7 +42,7 @@ class MPILinearOperator:
 
     """
 
-    def __init__(self, Op: Optional[LinearOperator] = None, shape: Optional[ShapeLike] = None,
+    def __init__(self, Op: LinearOperator, shape: Optional[ShapeLike] = None,
                  dtype: Optional[DTypeLike] = None, base_comm: MPI.Comm = MPI.COMM_WORLD):
         if Op:
             self.Op = Op

pylops_mpi/StackedLinearOperator.py

@@ -11,16 +11,20 @@
 
 
 class MPIStackedLinearOperator(ABC):
-    """Common interface for performing matrix-vector products in distributed fashion
-    for StackedLinearOperators.
+    """Stack of MPI-enabled PyLops Linear Operators
 
-    This class provides methods to perform matrix-vector product and adjoint matrix-vector
-    products on a stack of :class:`pylops_mpi.MPILinearOperator` objects.
+    Common interface for performing matrix-vector products in distributed fashion
+    for a stack of :class:`pylops_mpi.MPILinearOperator` operators.
+
+    In practice, this class provides methods to perform matrix-vector and adjoint
+    matrix-vector products on a stack of :class:`pylops_mpi.MPILinearOperator`
+    operators, allowing the actual execution of each operator to be distributed,
+    whilst dispatching the execution of the different operators in sequential order.
 
     .. note:: End users of pylops-mpi should not use this class directly but simply
-      use operators that are already implemented. This class is meant for
-      developers only, it has to be used as the parent class of any new operator
-      developed within pylops-mpi.
+      use operators that are already implemented as extensions of this class.
+      This class is meant for developers only, it has to be used as the parent
+      class of any new stacked operator developed within pylops-mpi.
 
     Parameters
     ----------