Merge pull request #89 from PyLops/gradient

Add MPIGradient

docs/source/api/index.rst

@@ -59,6 +59,7 @@ Derivatives
     MPIFirstDerivative
     MPISecondDerivative
     MPILaplacian
+    MPIGradient
 
 Solvers
 -------

examples/plot_derivative.py

@@ -124,3 +124,37 @@
     plt.colorbar(im, ax=axs[2])
     plt.tight_layout()
     plt.subplots_adjust(top=0.8)
+
+###############################################################################
+# We now consider the :py:class:`pylops_mpi.basicoperators.MPIGradient` operator.
+# Given a 2-dimensional array, this operator applies first-order derivatives on both
+# dimensions and concatenates them.
+Gop = pylops_mpi.MPIGradient(dims=(nx, ny), dtype=np.float64)
+
+y_grad_dist = Gop @ x_dist
+# Reshaping to (ndims, nx, ny) for plotting
+y_grad = y_grad_dist.asarray().reshape((2, nx, ny))
+y_grad_adj_dist = Gop.H @ y_grad_dist
+# Reshaping to (nx, ny) for plotting
+y_grad_adj = y_grad_adj_dist.asarray().reshape((nx, ny))
+
+if rank == 0:
+    fig, axs = plt.subplots(2, 2, figsize=(10, 6), sharex=True, sharey=True)
+    fig.suptitle("Gradient", fontsize=12, fontweight="bold", y=0.95)
+    im = axs[0, 0].imshow(x, interpolation="nearest", cmap="rainbow")
+    axs[0, 0].axis("tight")
+    axs[0, 0].set_title("x")
+    plt.colorbar(im, ax=axs[0, 0])
+    im = axs[0, 1].imshow(y_grad[0, ...], interpolation="nearest", cmap="rainbow")
+    axs[0, 1].axis("tight")
+    axs[0, 1].set_title("y - 1st direction")
+    plt.colorbar(im, ax=axs[0, 1])
+    im = axs[1, 1].imshow(y_grad[1, ...], interpolation="nearest", cmap="rainbow")
+    axs[1, 1].axis("tight")
+    axs[1, 1].set_title("y - 2nd direction")
+    plt.colorbar(im, ax=axs[1, 1])
+    im = axs[1, 0].imshow(y_grad_adj, interpolation="nearest", cmap="rainbow")
+    axs[1, 0].axis("tight")
+    axs[1, 0].set_title("xadj")
+    plt.colorbar(im, ax=axs[1, 0])
+    plt.tight_layout()

pylops_mpi/StackedLinearOperator.py

@@ -15,12 +15,12 @@ class MPIStackedLinearOperator(ABC):
     for StackedLinearOperators.
 
     This class provides methods to perform matrix-vector product and adjoint matrix-vector
-    products on a stack of MPILinearOperator objects.
+    products on a stack of :class:`pylops_mpi.MPILinearOperator` objects.
 
     .. 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. This class is meant for
+      developers only, it has to be used as the parent class of any new operator
+      developed within pylops-mpi.
 
     Parameters
     ----------
@@ -45,6 +45,25 @@ def __init__(self, shape: Optional[ShapeLike] = None,
         self.rank = self.base_comm.Get_rank()
 
     def matvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[DistributedArray, StackedDistributedArray]:
+        """Matrix-vector multiplication.
+
+        Modified version of :class:`pylops_mpi.MPILinearOperator` matvec.
+        This method makes use of either :class:`pylops_mpi.DistributedArray` or
+        :class:`pylops_mpi.StackedDistributedArray` to calculate matrix vector multiplication
+        in a distributed fashion.
+
+        Parameters
+        ----------
+        x : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.DistributedArray`
+            A StackedDistributedArray or a DistributedArray of global shape (N, )
+
+        Returns
+        -------
+        y : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.DistributedArray`
+            A StackedDistributedArray or a DistributedArray of global shape (M, )
+
+        """
+
         M, N = self.shape
         if isinstance(x, StackedDistributedArray):
             stacked_shape = (np.sum([a.global_shape for a in x.distarrays]), )
@@ -59,6 +78,25 @@ def _matvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[
         pass
 
     def rmatvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[DistributedArray, StackedDistributedArray]:
+        """Adjoint Matrix-vector multiplication.
+
+        Modified version of :class:`pylops_mpi.MPILinearOperator` rmatvec
+        This method makes use of either :class:`pylops_mpi.DistributedArray` or
+        :class:`pylops_mpi.StackedDistributedArray` to calculate adjoint matrix vector multiplication
+        in a distributed fashion.
+
+        Parameters
+        ----------
+        x : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.DistributedArray`
+            A StackedDistributedArray or a DistributedArray of global shape (M, )
+
+        Returns
+        -------
+        y : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.DistributedArray`
+            A StackedDistributedArray or a DistributedArray of global shape (N, )
+
+        """
+
         M, N = self.shape
         if isinstance(x, StackedDistributedArray):
             stacked_shape = (np.sum([a.global_shape for a in x.distarrays]), )

pylops_mpi/basicoperators/Gradient.py

@@ -0,0 +1,119 @@
+from typing import Union
+from mpi4py import MPI
+import numpy as np
+
+from pylops import FirstDerivative
+from pylops.utils import InputDimsLike, DTypeLike
+from pylops.utils._internal import _value_or_sized_to_tuple
+
+from pylops_mpi import MPIStackedLinearOperator
+from pylops_mpi.basicoperators.VStack import MPIStackedVStack
+from pylops_mpi.basicoperators.BlockDiag import MPIBlockDiag
+from pylops_mpi.basicoperators.FirstDerivative import MPIFirstDerivative
+from pylops_mpi.DistributedArray import (
+    DistributedArray,
+    StackedDistributedArray,
+    Partition,
+    local_split
+)
+
+
+class MPIGradient(MPIStackedLinearOperator):
+    r"""MPI Gradient
+
+    Apply gradient operator to a multi-dimensional distributed
+    array.
+
+    .. note:: At least 2 dimensions are required, use
+        :py:func:`pylops_mpi.MPIFirstDerivative` for 1d arrays.
+
+    Parameters
+    ----------
+    dims : :obj:`int` or :obj:`tuple`
+        Number of samples for each dimension.
+    sampling : :obj:`int` or :obj:`tuple`, optional
+        Sampling steps for each direction.
+    edge : :obj:`bool`, optional
+        Use reduced order derivative at edges (``True``) or ignore them (``False``).
+    kind : :obj:`str`, optional
+        Derivative kind (``forward``, ``centered``, or ``backward``).
+    base_comm : : obj:`MPI.Comm`, optional
+        MPI Base Communicator. Defaults to ``mpi4py.MPI.COMM_WORLD``.
+    dtype : :obj:`str`, optional
+        Type of elements in input array.
+
+    Notes
+    -----
+    The MPIGradient operator applies a first-order derivative to each dimension of
+    a multi-dimensional distributed array in forward mode.
+
+    We utilize the :py:class:`pylops_mpi.basicoperators.MPIFirstDerivative` to
+    calculate the first derivative along the first direction (i.e., axis=0).
+    For all remaining axes, the :py:class:`pylops.FirstDerivative` operator is
+    pushed into the :py:class:`pylops_mpi.basicoperators.MPIBlockDiag` operator.
+
+    Finally, using the :py:class:`pylops_mpi.basicoperators.MPIStackedVStack` we vertically
+    stack the :py:class:`pylops_mpi.basicoperators.MPIFirstDerivative` and the
+    :py:class:`pylops_mpi.basicoperators.MPIBlockDiag` operators.
+
+    For the forward mode, the matrix vector product is performed between the
+    :py:class:`pylops_mpi.basicoperators.MPIStackedVStack` and the :py:class:`pylops_mpi.DistributedArray`.
+
+    For simplicity, given a three-dimensional array, the MPIGradient in forward mode using a
+    centered stencil can be expressed as:
+
+    .. math::
+        \mathbf{g}_{i, j, k} =
+            (f_{i+1, j, k} - f_{i-1, j, k}) / d_1 \mathbf{i_1} +
+            (f_{i, j+1, k} - f_{i, j-1, k}) / d_2 \mathbf{i_2} +
+            (f_{i, j, k+1} - f_{i, j, k-1}) / d_3 \mathbf{i_3}
+
+    In adjoint mode, the adjoint matrix vector product is performed between the
+    :py:class:`pylops_mpi.basicoperators.MPIStackedVStack` and the :py:class:`pylops_mpi.StackedDistributedArray`.
+
+    """
+
+    def __init__(self,
+                 dims: Union[int, InputDimsLike],
+                 sampling: Union[int, InputDimsLike] = 1,
+                 edge: bool = False,
+                 kind: str = "centered",
+                 base_comm: MPI.Comm = MPI.COMM_WORLD,
+                 dtype: DTypeLike = "float64",
+                 ):
+        self.dims = _value_or_sized_to_tuple(dims)
+        ndims = len(self.dims)
+        sampling = _value_or_sized_to_tuple(sampling, repeat=ndims)
+        self.sampling = sampling
+        self.edge = edge
+        self.kind = kind
+        self.base_comm = base_comm
+        self.dtype = np.dtype(dtype)
+        self.Op = self._calc_stack_op(ndims)
+        super().__init__(shape=self.Op.shape, dtype=dtype, base_comm=base_comm)
+
+    def _matvec(self, x: DistributedArray) -> StackedDistributedArray:
+        return self.Op._matvec(x)
+
+    def _rmatvec(self, x: StackedDistributedArray) -> DistributedArray:
+        return self.Op._rmatvec(x)
+
+    def _calc_stack_op(self, ndims):
+        local_dims = local_split(tuple(self.dims), self.base_comm, Partition.SCATTER, axis=0)
+        grad_ops = []
+        Op1 = MPIFirstDerivative(dims=self.dims, sampling=self.sampling[0],
+                                 kind=self.kind, edge=self.edge,
+                                 dtype=self.dtype)
+        grad_ops.append(Op1)
+        for iax in range(1, ndims):
+            diag = MPIBlockDiag([
+                FirstDerivative(
+                    dims=local_dims,
+                    axis=iax,
+                    sampling=self.sampling[iax],
+                    edge=self.edge,
+                    kind=self.kind,
+                    dtype=self.dtype)
+            ])
+            grad_ops.append(diag)
+        return MPIStackedVStack(ops=grad_ops)

pylops_mpi/basicoperators/__init__.py

@@ -12,9 +12,10 @@
     MPIVStack                         Vertical Stacking of PyLops operators
     MPIStackedVStack                  Vertical Stacking of PyLops-MPI operators
     MPIHStack                         Horizontal Stacking of PyLops operators
-    MPIFirstDerivative                First Derivative
-    MPISecondDerivative               Second Derivative
-    MPILaplacian                      Laplacian
+    MPIFirstDerivative                First Derivative operator
+    MPISecondDerivative               Second Derivative operator
+    MPILaplacian                      Laplacian operator
+    MPIGradient                       Gradient operator
 
 """
 
@@ -24,6 +25,7 @@
 from .FirstDerivative import *
 from .SecondDerivative import *
 from .Laplacian import *
+from .Gradient import *
 
 __all__ = [
     "MPIBlockDiag",
@@ -33,5 +35,6 @@
     "MPIHStack",
     "MPIFirstDerivative",
     "MPISecondDerivative",
-    "MPILaplacian"
+    "MPILaplacian",
+    "MPIGradient"
 ]

tests/test_derivative.py

@@ -396,3 +396,46 @@ def test_laplacian(par):
             y_adj_np = Lop.H @ x_global
             assert_allclose(y, y_np, rtol=1e-14)
             assert_allclose(y_adj, y_adj_np, rtol=1e-14)
+
+
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par", [(par5), (par5e), (par6), (par6e)])
+def test_gradient(par):
+    """MPIGradient Operator"""
+    for kind in ["forward", "backward", "centered"]:
+        Gop_MPI = pylops_mpi.basicoperators.MPIGradient(dims=par['n'], sampling=par['sampling'],
+                                                        kind=kind, edge=par['edge'],
+                                                        dtype=par['dtype'])
+        x_fwd = pylops_mpi.DistributedArray(global_shape=np.prod(par['n']), dtype=par['dtype'])
+        x_fwd[:] = np.random.normal(rank, 10, x_fwd.local_shape)
+        x_global = x_fwd.asarray()
+        # Forward
+        y_dist = Gop_MPI @ x_fwd
+        assert isinstance(y_dist, pylops_mpi.StackedDistributedArray)
+        y = y_dist.asarray()
+
+        # Adjoint
+        x_adj_dist1 = pylops_mpi.DistributedArray(global_shape=int(np.prod(par['n'])), dtype=par['dtype'])
+        x_adj_dist1[:] = np.random.normal(rank, 10, x_adj_dist1.local_shape)
+        x_adj_dist2 = pylops_mpi.DistributedArray(global_shape=int(np.prod(par['n'])), dtype=par['dtype'])
+        x_adj_dist2[:] = np.random.normal(rank, 20, x_adj_dist2.local_shape)
+        x_adj_dist3 = pylops_mpi.DistributedArray(global_shape=int(np.prod(par['n'])), dtype=par['dtype'])
+        x_adj_dist3[:] = np.random.normal(rank, 30, x_adj_dist3.local_shape)
+
+        x_adj = pylops_mpi.StackedDistributedArray(distarrays=[x_adj_dist1, x_adj_dist2, x_adj_dist3])
+
+        x_adj_global = x_adj.asarray()
+        y_adj_dist = Gop_MPI.H @ x_adj
+        assert isinstance(y_adj_dist, pylops_mpi.DistributedArray)
+
+        y_adj = y_adj_dist.asarray()
+
+        if rank == 0:
+            Gop = pylops.Gradient(dims=par['n'], sampling=par['sampling'],
+                                  kind=kind, edge=par['edge'],
+                                  dtype=par['dtype'])
+            assert Gop_MPI.shape == Gop.shape
+            y_np = Gop @ x_global
+            y_adj_np = Gop.H @ x_adj_global
+            assert_allclose(y, y_np, rtol=1e-14)
+            assert_allclose(y_adj, y_adj_np, rtol=1e-14)