diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index a0647500..9d243a23 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -140,6 +140,10 @@ class DistributedArray(DistributedMixIn): Engine used to store array (``numpy`` or ``cupy``) dtype : :obj:`str`, optional Type of elements in input array. Defaults to ``numpy.float64``. + local_array : :obj:`numpy.ndarray` or :obj:`cupy.ndarray`, optional + Existing local array to use as the underlying storage instead of + allocating a new one. The shape of the local_array must match the local + shape at that rank. """ def __init__(self, global_shape: Union[Tuple, Integral], @@ -149,7 +153,8 @@ def __init__(self, global_shape: Union[Tuple, Integral], local_shapes: Optional[List[Union[Tuple, Integral]]] = None, mask: Optional[List[Integral]] = None, engine: Optional[str] = "numpy", - dtype: Optional[DTypeLike] = np.float64): + dtype: Optional[DTypeLike] = np.float64, + local_array: Optional[NDArray] = None): if isinstance(global_shape, Integral): global_shape = (global_shape,) if len(global_shape) <= axis: @@ -179,7 +184,18 @@ def __init__(self, global_shape: Union[Tuple, Integral], partition, axis) self._engine = engine - self._local_array = get_module(engine).empty(shape=self.local_shape, dtype=self.dtype) + if local_array is None: + self._local_array = get_module(engine).empty( + shape=self.local_shape, + dtype=self.dtype, + ) + else: + if local_array.shape != self.local_shape: + raise ValueError( + f"local_array has shape {local_array.shape}, " + f"expected {self.local_shape}" + ) + self._local_array = local_array def __getitem__(self, index): return self.local_array[index] @@ -851,6 +867,7 @@ def ravel(self, order: Optional[str] = "C"): Flattened 1-D DistributedArray """ local_shapes = [(np.prod(local_shape, axis=-1), ) for local_shape in self.local_shapes] + local_array = np.ravel(self.local_array, order=order) arr = DistributedArray(global_shape=np.prod(self.global_shape), base_comm=self.base_comm, base_comm_nccl=self.base_comm_nccl, @@ -858,10 +875,34 @@ def ravel(self, order: Optional[str] = "C"): mask=self.mask, partition=self.partition, engine=self.engine, - dtype=self.dtype) - local_array = np.ravel(self.local_array, order=order) - x = local_array.copy() - arr[:] = x + dtype=self.dtype, + local_array=local_array) + return arr + + def reshape(self, local_shape: Tuple, axis: Optional[int] = 0): + """Return a reshaped DistributedArray + + Parameters + ---------- + local_shape : :obj:`tuple` + Shape of the local array on each MPI rank. + axis : :obj:`int`, optional + Distribution axis in the reshaped global array. Defaults to ``0``. + """ + local_shapes = self.base_comm.allgather(local_shape) + global_shape = list(local_shapes[0]) + if self.partition is Partition.SCATTER: + global_shape[axis] = np.sum([ls[axis] for ls in local_shapes]) + local_array = self.local_array.reshape(local_shapes[self.rank]) + arr = DistributedArray(global_shape=tuple(global_shape), + base_comm=self.base_comm, + base_comm_nccl=self.base_comm_nccl, + local_shapes=local_shapes, + mask=self.mask, + partition=self.partition, + engine=self.engine, + dtype=self.dtype, + local_array=local_array) return arr def empty_like(self): diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 81db5ba8..39fd7dbc 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -7,7 +7,7 @@ from scipy.sparse._sputils import isintlike, isshape from scipy.sparse.linalg._interface import _get_dtype -from pylops import LinearOperator +from pylops import LinearOperator, get_ndarray_multiplication from pylops.utils import DTypeLike, ShapeLike from pylops_mpi import DistributedArray @@ -259,9 +259,11 @@ def dot(self, x): Op = _ProductLinearOperator(self, x) self._copy_attributes( Op, - exclude=['dims'] + exclude=['dims', '_local_dims'] ) Op.dims = x.dims + if getattr(x, "_local_dims", None) is not None: + Op._local_dims = x._local_dims return Op elif np.isscalar(x): Op = _ScaledLinearOperator(self, x) @@ -270,11 +272,29 @@ def dot(self, x): ) return Op else: - if x is None or x.ndim == 1: - return self.matvec(x) + if not get_ndarray_multiplication() and x.ndim >= 2: + msg = ( + "Operator can only be applied to 1D vectors. " + "Enable ndarray multiplication with pylops.set_ndarray_multiplication(True)." + ) + raise ValueError(msg) + is_dims_shaped = x.global_shape == self.dims + if is_dims_shaped: + x = x.redistribute(axis=0).ravel() + if x.ndim == 1: + y = self.matvec(x) + if is_dims_shaped and get_ndarray_multiplication(): + _local_dimsd = getattr(self, "_local_dimsd", None) + if _local_dimsd: + y = y.reshape(_local_dimsd.dim, axis=_local_dimsd.axis) + return y else: - raise ValueError('expected 1-d DistributedArray, got %r' - % x.global_shape) + msg = ( + "Wrong shape.\nExpects either a 1d array or, an ndarray of " + f"size `dims` when `dims` and `dimsd` both are available.\n" + f"Instead, received an array of shape {x.global_shape}." + ) + raise ValueError(msg) def adjoint(self): """Adjoint MPI LinearOperator @@ -339,6 +359,10 @@ def __add__(self, x): self._copy_attributes( Op ) + if len(self.dims) == 1: + Op.dims = x.dims + if len(self.dimsd) == 1: + Op.dimsd = x.dimsd return Op def __neg__(self): @@ -355,20 +379,28 @@ def _adjoint(self): Op = _AdjointLinearOperator(self) self._copy_attributes( Op, - exclude=['dims', 'dimsd'] + exclude=['dims', 'dimsd', '_local_dims', '_local_dimsd'] ) Op.dims = self.dimsd Op.dimsd = self.dims + if getattr(self, "_local_dims", None) is not None: + Op._local_dimsd = self._local_dims + if getattr(self, "_local_dimsd", None) is not None: + Op._local_dims = self._local_dimsd return Op def _transpose(self): Op = _TransposedLinearOperator(self) self._copy_attributes( Op, - exclude=['dims', 'dimsd'] + exclude=['dims', 'dimsd', '_local_dims', '_local_dimsd'] ) Op.dims = self.dimsd Op.dimsd = self.dims + if getattr(self, "_local_dims", None) is not None: + Op._local_dimsd = self._local_dims + if getattr(self, "_local_dimsd", None) is not None: + Op._local_dims = self._local_dimsd return Op def conj(self): @@ -388,7 +420,7 @@ def _copy_attributes( exclude: list[str] | None = None, ) -> None: """Copy attributes from one MPILinearOperator to another""" - attrs = ["dims", "dimsd"] + attrs = ["dims", "dimsd", "_local_dims", "_local_dimsd"] if exclude is not None: for item in exclude: attrs.remove(item) diff --git a/pylops_mpi/basicoperators/BlockDiag.py b/pylops_mpi/basicoperators/BlockDiag.py index 85d5811e..b79cba6f 100644 --- a/pylops_mpi/basicoperators/BlockDiag.py +++ b/pylops_mpi/basicoperators/BlockDiag.py @@ -1,7 +1,8 @@ +from typing import Optional, Sequence, List +from types import SimpleNamespace import numpy as np from scipy.sparse.linalg._interface import _get_dtype from mpi4py import MPI -from typing import Optional, Sequence, List from numbers import Integral from pylops import LinearOperator @@ -109,37 +110,47 @@ def __init__(self, ops: Sequence[LinearOperator], nops[iop] = oper.shape[0] mops[iop] = oper.shape[1] self.mops = mops.sum() - self.local_shapes_m = base_comm.allgather((self.mops, )) self.nops = nops.sum() - self.local_shapes_n = base_comm.allgather((self.nops, )) + dimsd = [d for dimsd in base_comm.allgather([op.dimsd for op in self.ops]) for d in dimsd] + if len(set(dimsd)) == 1: + self._local_dimsd = SimpleNamespace(dim=(len(ops), *dimsd[0]), axis=0) + dimsd = (base_comm.allreduce(len(ops)), *dimsd[0]) + else: + self._local_dimsd = SimpleNamespace(dim=(self.nops, ), axis=0) + dimsd = (base_comm.allreduce(self.nops),) + dims = [d for dims in base_comm.allgather([op.dims for op in self.ops]) for d in dims] + if len(set(dims)) == 1: + self._local_dims = SimpleNamespace(dim=(len(ops), *dims[0]), axis=0) + dims = (base_comm.allreduce(len(ops)), *dims[0]) + else: + self._local_dims = SimpleNamespace(dim=(self.mops, ), axis=0) + dims = (base_comm.allreduce(self.mops),) + self.local_shapes_m = base_comm.allgather((self.mops, )) + self.local_shapes_n = base_comm.allgather((self.nops,)) self.nnops = np.insert(np.cumsum(nops), 0, 0) self.mmops = np.insert(np.cumsum(mops), 0, 0) - dimsd = (base_comm.allreduce(self.nops), ) - dims = (base_comm.allreduce(self.mops), ) dtype = _get_dtype(ops) if dtype is None else np.dtype(dtype) super().__init__(dims=dims, dimsd=dimsd, dtype=dtype, base_comm=base_comm) @reshaped(forward=True, stacking=True) def _matvec(self, x: DistributedArray) -> DistributedArray: ncp = get_module(x.engine) - y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_n, - mask=self.mask, engine=x.engine, dtype=self.dtype) + y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, + local_shapes=self.local_shapes_n, mask=self.mask, engine=x.engine, dtype=self.dtype) y1 = [] for iop, oper in enumerate(self.ops): - y1.append(oper.matvec(x.local_array[self.mmops[iop]: - self.mmops[iop + 1]])) + y1.append(oper.matvec(x.local_array[self.mmops[iop]: self.mmops[iop + 1]])) y[:] = ncp.concatenate(y1) return y @reshaped(forward=False, stacking=True) def _rmatvec(self, x: DistributedArray) -> DistributedArray: ncp = get_module(x.engine) - y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_m, - mask=self.mask, engine=x.engine, dtype=self.dtype) + y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, + local_shapes=self.local_shapes_m, mask=self.mask, engine=x.engine, dtype=self.dtype) y1 = [] for iop, oper in enumerate(self.ops): - y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: - self.nnops[iop + 1]])) + y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: self.nnops[iop + 1]])) y[:] = ncp.concatenate(y1) return y diff --git a/pylops_mpi/basicoperators/FirstDerivative.py b/pylops_mpi/basicoperators/FirstDerivative.py index 7cbd18a5..336693ce 100644 --- a/pylops_mpi/basicoperators/FirstDerivative.py +++ b/pylops_mpi/basicoperators/FirstDerivative.py @@ -1,4 +1,5 @@ from typing import Callable, Union +from types import SimpleNamespace import numpy as np from mpi4py import MPI @@ -11,7 +12,7 @@ MPILinearOperator, Partition ) - +from pylops_mpi.DistributedArray import local_split from pylops_mpi.utils.decorators import reshaped @@ -90,6 +91,7 @@ def __init__(self, base_comm: MPI.Comm = MPI.COMM_WORLD, dtype: DTypeLike = np.float64): dims = _value_or_sized_to_tuple(dims) + self._local_dims = self._local_dimsd = SimpleNamespace(dim=local_split(dims, base_comm, Partition.SCATTER, axis=0), axis=0) super().__init__(dims=dims, dimsd=dims, dtype=np.dtype(dtype), base_comm=base_comm) self.sampling = sampling self.kind = kind diff --git a/pylops_mpi/basicoperators/SecondDerivative.py b/pylops_mpi/basicoperators/SecondDerivative.py index f19b219b..f589b7ca 100644 --- a/pylops_mpi/basicoperators/SecondDerivative.py +++ b/pylops_mpi/basicoperators/SecondDerivative.py @@ -1,4 +1,5 @@ from typing import Callable, Union +from types import SimpleNamespace import numpy as np from mpi4py import MPI @@ -8,6 +9,7 @@ from pylops_mpi import DistributedArray, MPILinearOperator, Partition from pylops_mpi.utils.decorators import reshaped +from pylops_mpi.DistributedArray import local_split class MPISecondDerivative(MPILinearOperator): @@ -81,6 +83,7 @@ def __init__( dtype: DTypeLike = np.float64, ) -> None: dims = _value_or_sized_to_tuple(dims) + self._local_dims = self._local_dimsd = SimpleNamespace(dim=local_split(dims, base_comm, Partition.SCATTER, axis=0), axis=0) super().__init__(dims=dims, dimsd=dims, dtype=np.dtype(dtype), base_comm=base_comm) self.sampling = sampling self.kind = kind diff --git a/pylops_mpi/basicoperators/VStack.py b/pylops_mpi/basicoperators/VStack.py index 7e3760dd..966b032a 100644 --- a/pylops_mpi/basicoperators/VStack.py +++ b/pylops_mpi/basicoperators/VStack.py @@ -1,7 +1,8 @@ +from typing import Sequence, Optional +from types import SimpleNamespace import numpy as np from scipy.sparse.linalg._interface import _get_dtype from mpi4py import MPI -from typing import Sequence, Optional from pylops import LinearOperator from pylops.utils import DTypeLike @@ -106,7 +107,6 @@ def __init__(self, ops: Sequence[LinearOperator], for iop, oper in enumerate(self.ops): nops[iop] = oper.shape[0] self.nops = nops.sum() - self.local_shapes_n = base_comm.allgather((self.nops, )) mops = [oper.shape[1] for oper in self.ops] mops = np.concatenate(base_comm.allgather(mops), axis=0) if len(set(mops)) > 1: @@ -114,7 +114,14 @@ def __init__(self, ops: Sequence[LinearOperator], self.mops = int(mops[0]) self.nnops = np.insert(np.cumsum(nops), 0, 0) dimsd = (base_comm.allreduce(self.nops), ) - dims = (self.mops, ) + dims = [d for dims in base_comm.allgather([op.dims for op in self.ops]) for d in dims] + self.local_shapes_n = base_comm.allgather((self.nops,)) + if len(set(dims)) == 1: + dims = dims[0] + else: + dims = (self.mops,) + self._local_dimsd = SimpleNamespace(dim=(self.nops, ), axis=0) + self._local_dims = SimpleNamespace(dim=dims, axis=0) dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype) super().__init__(dims=dims, dimsd=dimsd, dtype=dtype, base_comm=base_comm) @@ -124,8 +131,8 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}" f"Got {x.partition} instead...") # the output y should use NCCL if the operand x uses it - y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, local_shapes=self.local_shapes_n, - engine=x.engine, dtype=self.dtype) + y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, + local_shapes=self.local_shapes_n, engine=x.engine, dtype=self.dtype) y1 = [] for iop, oper in enumerate(self.ops): y1.append(oper.matvec(x.local_array)) diff --git a/pylops_mpi/utils/decorators.py b/pylops_mpi/utils/decorators.py index 5988e2d7..8d752e7b 100644 --- a/pylops_mpi/utils/decorators.py +++ b/pylops_mpi/utils/decorators.py @@ -4,6 +4,7 @@ import numpy as np from pylops_mpi import DistributedArray, Partition +from pylops_mpi.DistributedArray import local_split def reshaped( @@ -56,14 +57,9 @@ def wrapper(self, x: DistributedArray): and f.__name__ != "div" and f.__name__ != "__truediv__" ) - local_shapes = None - global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd", getattr(self, "dims")) - arr = DistributedArray(global_shape=global_shape, - base_comm=x.base_comm, - base_comm_nccl=x.base_comm_nccl, - local_shapes=local_shapes, axis=0, - engine=x.engine, dtype=x.dtype) - arr_local_shapes = np.asarray(arr.base_comm.allgather(np.prod(arr.local_shape))) + global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") + local_shapes = self.base_comm.allgather(local_split(global_shape, self.base_comm, Partition.SCATTER, 0)) + arr_local_shapes = np.array([np.prod(shape) for shape in local_shapes]) x_local_shapes = np.asarray(x.base_comm.allgather(np.prod(x.local_shape))) # Calculate num_ghost_cells required for each rank dif = np.cumsum(arr_local_shapes - x_local_shapes) @@ -74,7 +70,17 @@ def wrapper(self, x: DistributedArray): ghosted_array = x.add_ghost_cells(cells_front=cells_front, cells_back=cells_back) # Fill the redistributed array index = max(0, dif[self.rank - 1]) - arr[:] = ghosted_array[index: arr_local_shapes[self.rank] + index].reshape(arr.local_shape) + local_array = ghosted_array[index:index + arr_local_shapes[self.rank]].reshape(local_shapes[self.rank]) + arr = DistributedArray( + global_shape=global_shape, + base_comm=x.base_comm, + base_comm_nccl=x.base_comm_nccl, + local_shapes=local_shapes, + engine=x.engine, + dtype=x.dtype, + local_array=local_array, + axis=0 + ) y: DistributedArray = f(self, arr) if len(y.global_shape) > 1: # Make sure y is distributed along axis=0 before applying ravel diff --git a/pylops_mpi/utils/dottest.py b/pylops_mpi/utils/dottest.py index e7c6c4cb..b64f07f8 100644 --- a/pylops_mpi/utils/dottest.py +++ b/pylops_mpi/utils/dottest.py @@ -81,8 +81,8 @@ def dottest( if (nr, nc) != Op.shape: raise AssertionError("Provided nr and nc do not match operator shape") - y = Op.matvec(u) # Op * u - x = Op.rmatvec(v) # Op'* v + y = Op * u # Op * u + x = Op.H * v # Op'* v yy = np.vdot(y.asarray(), v.asarray()) # (Op * u)' * v xx = np.vdot(u.asarray(), x.asarray()) # u' * (Op' * v) diff --git a/tests/test_blockdiag.py b/tests/test_blockdiag.py index e2688f42..4ca1f36d 100644 --- a/tests/test_blockdiag.py +++ b/tests/test_blockdiag.py @@ -71,6 +71,44 @@ def test_blockdiag(par): assert_allclose(y_rmat_mpi, y_rmat_np, rtol=1e-14) +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)]) +def test_blockdiag_(par): + size = MPI.COMM_WORLD.Get_size() + rank = MPI.COMM_WORLD.Get_rank() + Op = pylops.MatrixMult(A=((rank + 1) * np.ones(shape=(par['ny'], par['nx']))).astype(par['dtype'])) + BDiag_MPI = pylops_mpi.MPIBlockDiag(ops=[Op, ]) + + x = pylops_mpi.DistributedArray(global_shape=BDiag_MPI.dims, dtype=par['dtype'], engine=backend) + x[:] = np.ones(shape=x.local_shape, dtype=par['dtype']) + x_global = x.asarray() + + y = pylops_mpi.DistributedArray(global_shape=BDiag_MPI.dimsd, dtype=par['dtype'], engine=backend) + y[:] = np.ones(shape=y.local_shape, dtype=par['dtype']) + y_global = y.asarray() + + # Forward + x_mat = BDiag_MPI @ x + # Adjoint + y_rmat = BDiag_MPI.H @ y + assert isinstance(x_mat, pylops_mpi.DistributedArray) + assert isinstance(y_rmat, pylops_mpi.DistributedArray) + # Dot test + dottest(BDiag_MPI, x, y, size * par['ny'], size * par['nx']) + + x_mat_mpi = x_mat.asarray() + y_rmat_mpi = y_rmat.asarray() + + if rank == 0: + ops = [pylops.MatrixMult((i + 1) * np.ones(shape=(par['ny'], par['nx'])).astype(par['dtype'])) for i in range(size)] + BDiag = pylops.BlockDiag(ops=ops) + + x_mat_np = BDiag @ x_global + y_rmat_np = BDiag.H @ y_global + assert_allclose(x_mat_mpi, x_mat_np, rtol=1e-14) + assert_allclose(y_rmat_mpi, y_rmat_np, rtol=1e-14) + + @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)]) def test_stacked_blockdiag(par): diff --git a/tests/test_derivative.py b/tests/test_derivative.py index 5ef6eef6..205f3040 100644 --- a/tests/test_derivative.py +++ b/tests/test_derivative.py @@ -400,6 +400,211 @@ def test_second_derivative_centered(par): assert_allclose(y_adj, y_adj_np, rtol=1e-14) +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("par", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_first_derivative_forward_ND(par): + """MPIFirstDerivative operator (forward stencil)""" + Fop_MPI = pylops_mpi.MPIFirstDerivative(dims=par['nz'], sampling=par['dz'], + kind="forward", edge=par['edge'], + dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Fop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Fop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Fop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Fop = pylops.FirstDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="forward", edge=par['edge'], + dtype=par['dtype']) + assert Fop_MPI.shape == Fop.shape + y_np = Fop @ x_global + y_adj_np = Fop.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", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_first_derivative_backward_ND(par): + """MPIFirstDerivative operator (backward stencil)""" + Fop_MPI = pylops_mpi.MPIFirstDerivative(dims=par['nz'], sampling=par['dz'], + kind="backward", edge=par['edge'], + dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Fop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Fop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Fop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Fop = pylops.FirstDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="backward", edge=par['edge'], + dtype=par['dtype']) + assert Fop_MPI.shape == Fop.shape + y_np = Fop @ x_global + y_adj_np = Fop.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", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_first_derivative_centered_ND(par): + """MPIFirstDerivative operator (centered stencil)""" + for order in [3, 5]: + Fop_MPI = pylops_mpi.MPIFirstDerivative(dims=par['nz'], sampling=par['dz'], + kind="centered", edge=par['edge'], + order=order, dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Fop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Fop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Fop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Fop = pylops.FirstDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="centered", edge=par['edge'], + order=order, dtype=par['dtype']) + assert Fop_MPI.shape == Fop.shape + y_np = Fop @ x_global + y_adj_np = Fop.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", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_second_derivative_forward_ND(par): + """MPISecondDerivative operator (forward stencil)""" + Sop_MPI = pylops_mpi.basicoperators.MPISecondDerivative(dims=par['nz'], sampling=par['dz'], + kind="forward", edge=par['edge'], + dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Sop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Sop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Sop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Sop = pylops.SecondDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="forward", edge=par['edge'], + dtype=par['dtype']) + assert Sop_MPI.shape == Sop.shape + y_np = Sop @ x_global + y_adj_np = Sop.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", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_second_derivative_backward_ND(par): + """MPISecondDerivative operator (backward stencil)""" + Sop_MPI = pylops_mpi.basicoperators.MPISecondDerivative(dims=par['nz'], sampling=par['dz'], + kind="backward", edge=par['edge'], + dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Sop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Sop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Sop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Sop = pylops.SecondDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="backward", edge=par['edge'], + dtype=par['dtype']) + assert Sop_MPI.shape == Sop.shape + y_np = Sop @ x_global + y_adj_np = Sop.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", [(par1), (par1b), (par1j), (par1e), (par2), (par2b), + (par2j), (par2e), (par3), (par3b), (par3j), (par3e), + (par4), (par4b), (par4j), (par4e)]) +def test_second_derivative_centered_ND(par): + """MPISecondDerivative operator (centered stencil)""" + Sop_MPI = pylops_mpi.basicoperators.MPISecondDerivative(dims=par['nz'], sampling=par['dz'], + kind="centered", edge=par['edge'], + dtype=par['dtype']) + x = pylops_mpi.DistributedArray(global_shape=par['nz'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape) + x_global = x.asarray() + # Forward + y_dist = Sop_MPI @ x + y = y_dist.asarray() + # Adjoint + y_adj_dist = Sop_MPI.H @ x + y_adj = y_adj_dist.asarray() + # Dot test + dottest(Sop_MPI, x, y_dist, npp.prod(par['nz']), npp.prod(par['nz'])) + + if rank == 0: + Sop = pylops.SecondDerivative(dims=par['nz'], axis=0, + sampling=par['dz'], + kind="centered", edge=par['edge'], + dtype=par['dtype']) + assert Sop_MPI.shape == Sop.shape + y_np = Sop @ x_global + y_adj_np = Sop.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_laplacian(par):