From 4b131dc5ab0d416546f959959a04767e1ff78807 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 2 Jun 2026 20:37:29 +0000 Subject: [PATCH 01/10] WIP: enable ND-Distributed arrays on operators --- pylops_mpi/DistributedArray.py | 27 ++++++++++++++++++++ pylops_mpi/basicoperators/FirstDerivative.py | 2 ++ 2 files changed, 29 insertions(+) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index a0647500..fdd3a707 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -863,7 +863,34 @@ def ravel(self, order: Optional[str] = "C"): x = local_array.copy() arr[:] = x return arr + + def reshape(self, local_shape, axis=0): + """Return a reshaped DistributedArray + Parameters + ---------- + + Returns + ------- + arr : :obj:`pylops_mpi.DistributedArray` + Reshaped N-D DistributedArray + """ + local_shapes = self.base_comm.allgather(local_shape) + global_shape = list(local_shapes[0]) + global_shape[axis] = np.sum([ls[axis] for ls in local_shapes]) + 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 = self.local_array.reshape(local_shapes[self.rank]) + x = local_array.copy() + arr[:] = x + return arr + def empty_like(self): """Creates an empty like DistributedArray with uninitialized values """ diff --git a/pylops_mpi/basicoperators/FirstDerivative.py b/pylops_mpi/basicoperators/FirstDerivative.py index 7cbd18a5..86b5d829 100644 --- a/pylops_mpi/basicoperators/FirstDerivative.py +++ b/pylops_mpi/basicoperators/FirstDerivative.py @@ -129,12 +129,14 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: # If Partition.BROADCAST, then convert to Partition.SCATTER if x.partition is Partition.BROADCAST: x = DistributedArray.to_dist(x=x.local_array, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl) + self._dims_dimsd_local = (self.dimsd[0] // x.size, *self.dimsd[1:]) return self._hmatvec(x) def _rmatvec(self, x: DistributedArray) -> DistributedArray: # If Partition.BROADCAST, then convert to Partition.SCATTER if x.partition is Partition.BROADCAST: x = DistributedArray.to_dist(x=x.local_array, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl) + self._dims_dimsd_local = (self.dims[0] // x.size, *self.dims[1:]) return self._hrmatvec(x) @reshaped From f80103d0069de889f83d1a481ae2d318f97069df Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Tue, 23 Jun 2026 14:35:09 +0530 Subject: [PATCH 02/10] Update reshaped decorator and dot method --- pylops_mpi/DistributedArray.py | 6 ++--- pylops_mpi/LinearOperator.py | 26 +++++++++++++++------ pylops_mpi/utils/decorators.py | 42 +++++++++++++++++++--------------- 3 files changed, 45 insertions(+), 29 deletions(-) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index fdd3a707..d66c4b81 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -863,13 +863,13 @@ def ravel(self, order: Optional[str] = "C"): x = local_array.copy() arr[:] = x return arr - + def reshape(self, local_shape, axis=0): """Return a reshaped DistributedArray Parameters ---------- - + Returns ------- arr : :obj:`pylops_mpi.DistributedArray` @@ -890,7 +890,7 @@ def reshape(self, local_shape, axis=0): x = local_array.copy() arr[:] = x return arr - + def empty_like(self): """Creates an empty like DistributedArray with uninitialized values """ diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 81db5ba8..1a537f51 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 @@ -186,7 +186,7 @@ def matvec(self, x: DistributedArray) -> DistributedArray: """ M, N = self.shape - if x.global_shape != (N,): + if np.prod(x.global_shape) != (N,): raise ValueError("dimension mismatch") return self._matvec(x) @@ -224,7 +224,7 @@ def rmatvec(self, x: DistributedArray) -> DistributedArray: M, N = self.shape - if x.global_shape != (M,): + if np.prod(x.global_shape) != (M,): raise ValueError("dimension mismatch") return self._rmatvec(x) @@ -270,11 +270,23 @@ 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 or x.ndim == 1: + y = self.matvec(x) + 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 diff --git a/pylops_mpi/utils/decorators.py b/pylops_mpi/utils/decorators.py index 5988e2d7..b40a1e21 100644 --- a/pylops_mpi/utils/decorators.py +++ b/pylops_mpi/utils/decorators.py @@ -57,26 +57,30 @@ def wrapper(self, x: DistributedArray): 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))) - 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) - # Calculate cells_front(0 means no ghost cells) - cells_front = abs(min(0, dif[self.rank - 1])) - # Calculate cells_back(0 means no ghost cells) - cells_back = max(0, dif[self.rank]) - 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) + global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") + # Check if global and local shapes match the desired array + if (local_shapes is None or x.local_shapes == local_shapes) and x.global_shape == global_shape: + arr = x + else: + 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))) + 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) + # Calculate cells_front(0 means no ghost cells) + cells_front = abs(min(0, dif[self.rank - 1])) + # Calculate cells_back(0 means no ghost cells) + cells_back = max(0, dif[self.rank]) + 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) y: DistributedArray = f(self, arr) - if len(y.global_shape) > 1: + if x.ndim == 1 and len(y.global_shape) > 1: # Make sure y is distributed along axis=0 before applying ravel y = y.redistribute(axis=0).ravel() return y From 228615229282ada6cbd9ca3bccaecf1fd9b8fb9e Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Tue, 23 Jun 2026 16:23:42 +0530 Subject: [PATCH 03/10] Update method to handle different axis --- pylops_mpi/utils/decorators.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/pylops_mpi/utils/decorators.py b/pylops_mpi/utils/decorators.py index b40a1e21..25d0b966 100644 --- a/pylops_mpi/utils/decorators.py +++ b/pylops_mpi/utils/decorators.py @@ -58,10 +58,16 @@ def wrapper(self, x: DistributedArray): ) local_shapes = None global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") - # Check if global and local shapes match the desired array - if (local_shapes is None or x.local_shapes == local_shapes) and x.global_shape == global_shape: + x_dim = x.ndim + # Reuse existing distribution if it already matches the target layout + if ( + (local_shapes is None or x.local_shapes == local_shapes) + and x.global_shape == global_shape + and x.axis == 0 + ): arr = x else: + x = x.redistribute(axis=0) arr = DistributedArray(global_shape=global_shape, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, @@ -80,7 +86,7 @@ def wrapper(self, x: DistributedArray): index = max(0, dif[self.rank - 1]) arr[:] = ghosted_array[index: arr_local_shapes[self.rank] + index].reshape(arr.local_shape) y: DistributedArray = f(self, arr) - if x.ndim == 1 and len(y.global_shape) > 1: + if x_dim == 1 and len(y.global_shape) > 1: # Make sure y is distributed along axis=0 before applying ravel y = y.redistribute(axis=0).ravel() return y From 6285e3020228e1e959a98e58beab8e5fe7a5495e Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sat, 4 Jul 2026 01:18:52 +0530 Subject: [PATCH 04/10] Update reshaped and matvec sections --- pylops_mpi/LinearOperator.py | 2 ++ pylops_mpi/basicoperators/BlockDiag.py | 44 +++++++++++++++++------- pylops_mpi/basicoperators/Laplacian.py | 1 + pylops_mpi/basicoperators/VStack.py | 24 ++++++++----- pylops_mpi/signalprocessing/Fredholm1.py | 8 ++--- pylops_mpi/utils/decorators.py | 14 ++++---- 6 files changed, 61 insertions(+), 32 deletions(-) diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 1a537f51..117015fc 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -279,6 +279,8 @@ def dot(self, x): is_dims_shaped = x.global_shape == self.dims if is_dims_shaped or x.ndim == 1: y = self.matvec(x) + if x.ndim == 1 and y.ndim > 1: + y = y.redistribute(axis=0).ravel() return y else: msg = ( diff --git a/pylops_mpi/basicoperators/BlockDiag.py b/pylops_mpi/basicoperators/BlockDiag.py index 85d5811e..2f3486c3 100644 --- a/pylops_mpi/basicoperators/BlockDiag.py +++ b/pylops_mpi/basicoperators/BlockDiag.py @@ -109,37 +109,57 @@ 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_shapes_n = [(len(ops), *dimsd[0]) for _ in range(base_comm.size)] + dimsd = (base_comm.allreduce(len(ops)), *dimsd[0]) + else: + self.local_shapes_n = base_comm.allgather((self.nops, )) + 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_shapes_m = [(len(ops), *dims[0]) for _ in range(base_comm.size)] + dims = (base_comm.allreduce(len(ops)), *dims[0]) + else: + self.local_shapes_m = base_comm.allgather((self.mops, )) + dims = (base_comm.allreduce(self.mops),) 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.dimsd, 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]])) + arr = x.local_array.flatten()[self.mmops[iop]:self.mmops[iop + 1]] + op_output = oper.matvec(arr) + if len(self.dimsd) > 1: + shape_rank = list(self.local_shapes_n[self.rank]) + shape_rank[0] //= len(self.ops) + op_output = op_output.reshape(shape_rank) + y1.append(op_output) 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.dims, 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]])) + arr = x.local_array.flatten()[self.nnops[iop]:self.nnops[iop + 1]] + op_output = oper.rmatvec(arr) + if len(self.dims) > 1: + shape_rank = list(self.local_shapes_m[self.rank]) + shape_rank[0] //= len(self.ops) + op_output = op_output.reshape(shape_rank) + y1.append(op_output) y[:] = ncp.concatenate(y1) return y diff --git a/pylops_mpi/basicoperators/Laplacian.py b/pylops_mpi/basicoperators/Laplacian.py index c3e92d74..322e8764 100644 --- a/pylops_mpi/basicoperators/Laplacian.py +++ b/pylops_mpi/basicoperators/Laplacian.py @@ -117,6 +117,7 @@ def _calc_l2op(self, dims: tuple, base_comm: MPI.Comm): edge=self.edge, dtype=self.dtype) else: + print("local dims", local_dims) l2op += weight * MPIBlockDiag(ops=[SecondDerivative(dims=local_dims, axis=ax, sampling=samp, diff --git a/pylops_mpi/basicoperators/VStack.py b/pylops_mpi/basicoperators/VStack.py index 7e3760dd..b1f1675d 100644 --- a/pylops_mpi/basicoperators/VStack.py +++ b/pylops_mpi/basicoperators/VStack.py @@ -106,7 +106,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 +113,12 @@ 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,) 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,18 +128,22 @@ 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.dimsd, + 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)) + y1.append(oper.matvec(x.local_array.flatten())) 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], + y = DistributedArray(global_shape=self.dims, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=Partition.BROADCAST, @@ -143,10 +151,10 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: 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.flatten()[self.nnops[iop]: self.nnops[iop + 1]])) y1 = ncp.sum(ncp.vstack(y1), axis=0) y[:] = self._allreduce(x.base_comm, x.base_comm_nccl, - y1, op=MPI.SUM, engine=x.engine) + y1, op=MPI.SUM, engine=x.engine).reshape(self.dims) return y diff --git a/pylops_mpi/signalprocessing/Fredholm1.py b/pylops_mpi/signalprocessing/Fredholm1.py index 964ae9d9..76ddd4c5 100644 --- a/pylops_mpi/signalprocessing/Fredholm1.py +++ b/pylops_mpi/signalprocessing/Fredholm1.py @@ -109,7 +109,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: if x.partition not in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST]: raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}" f"Got {x.partition} instead...") - y = DistributedArray(global_shape=self.shape[0], + y = DistributedArray(global_shape=self.dimsd, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=x.partition, @@ -127,7 +127,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: y1[isl] = ncp.dot(self.G[isl], x[isl]) # gather results y[:] = ncp.vstack(y._allgather(y.base_comm, y.base_comm_nccl, y1, - engine=y.engine)).ravel() + engine=y.engine)) return y def _rmatvec(self, x: DistributedArray) -> DistributedArray: @@ -135,7 +135,7 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: if x.partition not in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST]: raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}" f"Got {x.partition} instead...") - y = DistributedArray(global_shape=self.shape[1], + y = DistributedArray(global_shape=self.dims, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=x.partition, @@ -165,5 +165,5 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: # gather results y[:] = ncp.vstack(y._allgather(y.base_comm, y.base_comm_nccl, y1, - engine=y.engine)).ravel() + engine=y.engine)) return y diff --git a/pylops_mpi/utils/decorators.py b/pylops_mpi/utils/decorators.py index 25d0b966..1dc8fea2 100644 --- a/pylops_mpi/utils/decorators.py +++ b/pylops_mpi/utils/decorators.py @@ -46,19 +46,17 @@ def wrapper(self, x: DistributedArray): raise ValueError(f"x should have partition={Partition.SCATTER}, {x.partition} != {Partition.SCATTER}") if stacking and forward: local_shapes = getattr(self, "local_shapes_m") - global_shape = x.global_shape elif stacking and not forward: local_shapes = getattr(self, "local_shapes_n") - global_shape = x.global_shape else: - fwd = ( - "rmat" not in f.__name__ - and f.__name__ != "div" - and f.__name__ != "__truediv__" - ) local_shapes = None - global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") + fwd = ( + "rmat" not in f.__name__ + and f.__name__ != "div" + and f.__name__ != "__truediv__" + ) x_dim = x.ndim + global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") # Reuse existing distribution if it already matches the target layout if ( (local_shapes is None or x.local_shapes == local_shapes) From 5076b057d79d3b3f728efcf8fbdf9ce05136073f Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 5 Jul 2026 16:10:54 +0530 Subject: [PATCH 05/10] Add local_array, and update ravel, reshape and reshaped decorator --- pylops_mpi/DistributedArray.py | 34 ++++++++---- pylops_mpi/LinearOperator.py | 13 +++-- pylops_mpi/basicoperators/BlockDiag.py | 26 +++------- pylops_mpi/basicoperators/Laplacian.py | 1 - pylops_mpi/basicoperators/VStack.py | 12 ++--- pylops_mpi/signalprocessing/Fredholm1.py | 8 +-- pylops_mpi/utils/decorators.py | 66 ++++++++++++------------ 7 files changed, 76 insertions(+), 84 deletions(-) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index d66c4b81..eb5894ca 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,8 @@ 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, axis=0): @@ -878,6 +893,7 @@ def reshape(self, local_shape, axis=0): local_shapes = self.base_comm.allgather(local_shape) global_shape = list(local_shapes[0]) 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, @@ -885,10 +901,8 @@ def reshape(self, local_shape, axis=0): mask=self.mask, partition=self.partition, engine=self.engine, - dtype=self.dtype) - local_array = self.local_array.reshape(local_shapes[self.rank]) - x = local_array.copy() - arr[:] = x + 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 117015fc..398e1c67 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -186,7 +186,7 @@ def matvec(self, x: DistributedArray) -> DistributedArray: """ M, N = self.shape - if np.prod(x.global_shape) != (N,): + if x.global_shape != (N,): raise ValueError("dimension mismatch") return self._matvec(x) @@ -224,7 +224,7 @@ def rmatvec(self, x: DistributedArray) -> DistributedArray: M, N = self.shape - if np.prod(x.global_shape) != (M,): + if x.global_shape != (M,): raise ValueError("dimension mismatch") return self._rmatvec(x) @@ -277,11 +277,10 @@ def dot(self, x): ) raise ValueError(msg) is_dims_shaped = x.global_shape == self.dims - if is_dims_shaped or x.ndim == 1: - y = self.matvec(x) - if x.ndim == 1 and y.ndim > 1: - y = y.redistribute(axis=0).ravel() - return y + if is_dims_shaped: + x = x.redistribute(axis=0).ravel() + if x.ndim == 1: + return self.matvec(x) else: msg = ( "Wrong shape.\nExpects either a 1d array or, an ndarray of " diff --git a/pylops_mpi/basicoperators/BlockDiag.py b/pylops_mpi/basicoperators/BlockDiag.py index 2f3486c3..bd119cf5 100644 --- a/pylops_mpi/basicoperators/BlockDiag.py +++ b/pylops_mpi/basicoperators/BlockDiag.py @@ -112,18 +112,16 @@ def __init__(self, ops: Sequence[LinearOperator], self.nops = nops.sum() 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_shapes_n = [(len(ops), *dimsd[0]) for _ in range(base_comm.size)] dimsd = (base_comm.allreduce(len(ops)), *dimsd[0]) else: - self.local_shapes_n = base_comm.allgather((self.nops, )) 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_shapes_m = [(len(ops), *dims[0]) for _ in range(base_comm.size)] dims = (base_comm.allreduce(len(ops)), *dims[0]) else: - self.local_shapes_m = base_comm.allgather((self.mops, )) 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) dtype = _get_dtype(ops) if dtype is None else np.dtype(dtype) @@ -132,34 +130,22 @@ def __init__(self, ops: Sequence[LinearOperator], @reshaped(forward=True, stacking=True) def _matvec(self, x: DistributedArray) -> DistributedArray: ncp = get_module(x.engine) - y = DistributedArray(global_shape=self.dimsd, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, + 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): - arr = x.local_array.flatten()[self.mmops[iop]:self.mmops[iop + 1]] - op_output = oper.matvec(arr) - if len(self.dimsd) > 1: - shape_rank = list(self.local_shapes_n[self.rank]) - shape_rank[0] //= len(self.ops) - op_output = op_output.reshape(shape_rank) - y1.append(op_output) + 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.dims, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, + 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): - arr = x.local_array.flatten()[self.nnops[iop]:self.nnops[iop + 1]] - op_output = oper.rmatvec(arr) - if len(self.dims) > 1: - shape_rank = list(self.local_shapes_m[self.rank]) - shape_rank[0] //= len(self.ops) - op_output = op_output.reshape(shape_rank) - y1.append(op_output) + 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/Laplacian.py b/pylops_mpi/basicoperators/Laplacian.py index 322e8764..c3e92d74 100644 --- a/pylops_mpi/basicoperators/Laplacian.py +++ b/pylops_mpi/basicoperators/Laplacian.py @@ -117,7 +117,6 @@ def _calc_l2op(self, dims: tuple, base_comm: MPI.Comm): edge=self.edge, dtype=self.dtype) else: - print("local dims", local_dims) l2op += weight * MPIBlockDiag(ops=[SecondDerivative(dims=local_dims, axis=ax, sampling=samp, diff --git a/pylops_mpi/basicoperators/VStack.py b/pylops_mpi/basicoperators/VStack.py index b1f1675d..fca80f92 100644 --- a/pylops_mpi/basicoperators/VStack.py +++ b/pylops_mpi/basicoperators/VStack.py @@ -128,12 +128,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.dimsd, - 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.flatten())) @@ -143,7 +139,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: @reshaped(forward=False, stacking=True) def _rmatvec(self, x: DistributedArray) -> DistributedArray: ncp = get_module(x.engine) - y = DistributedArray(global_shape=self.dims, + y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=Partition.BROADCAST, @@ -154,7 +150,7 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: y1.append(oper.rmatvec(x.local_array.flatten()[self.nnops[iop]: self.nnops[iop + 1]])) y1 = ncp.sum(ncp.vstack(y1), axis=0) y[:] = self._allreduce(x.base_comm, x.base_comm_nccl, - y1, op=MPI.SUM, engine=x.engine).reshape(self.dims) + y1, op=MPI.SUM, engine=x.engine) return y diff --git a/pylops_mpi/signalprocessing/Fredholm1.py b/pylops_mpi/signalprocessing/Fredholm1.py index 76ddd4c5..964ae9d9 100644 --- a/pylops_mpi/signalprocessing/Fredholm1.py +++ b/pylops_mpi/signalprocessing/Fredholm1.py @@ -109,7 +109,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: if x.partition not in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST]: raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}" f"Got {x.partition} instead...") - y = DistributedArray(global_shape=self.dimsd, + y = DistributedArray(global_shape=self.shape[0], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=x.partition, @@ -127,7 +127,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: y1[isl] = ncp.dot(self.G[isl], x[isl]) # gather results y[:] = ncp.vstack(y._allgather(y.base_comm, y.base_comm_nccl, y1, - engine=y.engine)) + engine=y.engine)).ravel() return y def _rmatvec(self, x: DistributedArray) -> DistributedArray: @@ -135,7 +135,7 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: if x.partition not in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST]: raise ValueError(f"x should have partition={Partition.BROADCAST},{Partition.UNSAFE_BROADCAST}" f"Got {x.partition} instead...") - y = DistributedArray(global_shape=self.dims, + y = DistributedArray(global_shape=self.shape[1], base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl, partition=x.partition, @@ -165,5 +165,5 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: # gather results y[:] = ncp.vstack(y._allgather(y.base_comm, y.base_comm_nccl, y1, - engine=y.engine)) + engine=y.engine)).ravel() return y diff --git a/pylops_mpi/utils/decorators.py b/pylops_mpi/utils/decorators.py index 1dc8fea2..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( @@ -46,45 +47,42 @@ def wrapper(self, x: DistributedArray): raise ValueError(f"x should have partition={Partition.SCATTER}, {x.partition} != {Partition.SCATTER}") if stacking and forward: local_shapes = getattr(self, "local_shapes_m") + global_shape = x.global_shape elif stacking and not forward: local_shapes = getattr(self, "local_shapes_n") + global_shape = x.global_shape else: - local_shapes = None - fwd = ( - "rmat" not in f.__name__ - and f.__name__ != "div" - and f.__name__ != "__truediv__" + fwd = ( + "rmat" not in f.__name__ + and f.__name__ != "div" + and f.__name__ != "__truediv__" + ) + 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) + # Calculate cells_front(0 means no ghost cells) + cells_front = abs(min(0, dif[self.rank - 1])) + # Calculate cells_back(0 means no ghost cells) + cells_back = max(0, dif[self.rank]) + 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]) + 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 ) - x_dim = x.ndim - global_shape = getattr(self, "dims") if fwd else getattr(self, "dimsd") - # Reuse existing distribution if it already matches the target layout - if ( - (local_shapes is None or x.local_shapes == local_shapes) - and x.global_shape == global_shape - and x.axis == 0 - ): - arr = x - else: - x = x.redistribute(axis=0) - 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))) - 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) - # Calculate cells_front(0 means no ghost cells) - cells_front = abs(min(0, dif[self.rank - 1])) - # Calculate cells_back(0 means no ghost cells) - cells_back = max(0, dif[self.rank]) - 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) y: DistributedArray = f(self, arr) - if x_dim == 1 and len(y.global_shape) > 1: + if len(y.global_shape) > 1: # Make sure y is distributed along axis=0 before applying ravel y = y.redistribute(axis=0).ravel() return y From 90404958d0339e9ac17148bedcfa95f9d39b65f5 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Mon, 6 Jul 2026 21:15:15 +0530 Subject: [PATCH 06/10] Add _local_dims_dimsd_axis and Update First and second derivative --- pylops_mpi/LinearOperator.py | 9 ++++++++- pylops_mpi/basicoperators/FirstDerivative.py | 3 ++- pylops_mpi/basicoperators/SecondDerivative.py | 2 ++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 398e1c67..0e6c65c8 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -280,7 +280,12 @@ def dot(self, x): if is_dims_shaped: x = x.redistribute(axis=0).ravel() if x.ndim == 1: - return self.matvec(x) + y = self.matvec(x) + if is_dims_shaped and get_ndarray_multiplication(): + _local_dims_dimsd_axis = getattr(self, "_local_dims_dimsd_axis", None) + if _local_dims_dimsd_axis: + y = y.reshape(_local_dims_dimsd_axis[0], axis=_local_dims_dimsd_axis[1]) + return y else: msg = ( "Wrong shape.\nExpects either a 1d array or, an ndarray of " @@ -372,6 +377,8 @@ def _adjoint(self): ) Op.dims = self.dimsd Op.dimsd = self.dims + if getattr(self, "_local_dims_dimsd_axis", None) is not None: + Op._local_dims_dimsd_axis = self._local_dims_dimsd_axis return Op def _transpose(self): diff --git a/pylops_mpi/basicoperators/FirstDerivative.py b/pylops_mpi/basicoperators/FirstDerivative.py index 86b5d829..5f57faa8 100644 --- a/pylops_mpi/basicoperators/FirstDerivative.py +++ b/pylops_mpi/basicoperators/FirstDerivative.py @@ -11,7 +11,7 @@ MPILinearOperator, Partition ) - +from pylops_mpi.DistributedArray import local_split from pylops_mpi.utils.decorators import reshaped @@ -90,6 +90,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_dimsd_axis = (local_split(dims, base_comm, Partition.SCATTER, axis=0), 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..0fae576c 100644 --- a/pylops_mpi/basicoperators/SecondDerivative.py +++ b/pylops_mpi/basicoperators/SecondDerivative.py @@ -8,6 +8,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 +82,7 @@ def __init__( dtype: DTypeLike = np.float64, ) -> None: dims = _value_or_sized_to_tuple(dims) + self._local_dims_dimsd_axis = (local_split(dims, base_comm, Partition.SCATTER, axis=0), 0) super().__init__(dims=dims, dimsd=dims, dtype=np.dtype(dtype), base_comm=base_comm) self.sampling = sampling self.kind = kind From ce59583a46ffbf96f7e13891d56781b2daeeb115 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Wed, 8 Jul 2026 00:48:32 +0530 Subject: [PATCH 07/10] Add local_dims and local_dimsd, update First,SecondDerivate, and BlockDiag and VStack --- pylops_mpi/DistributedArray.py | 3 +- pylops_mpi/LinearOperator.py | 26 ++- pylops_mpi/basicoperators/BlockDiag.py | 7 +- pylops_mpi/basicoperators/FirstDerivative.py | 3 +- pylops_mpi/basicoperators/SecondDerivative.py | 3 +- pylops_mpi/basicoperators/VStack.py | 9 +- pylops_mpi/utils/dottest.py | 4 +- tests/test_blockdiag.py | 38 ++++ tests/test_derivative.py | 205 ++++++++++++++++++ 9 files changed, 280 insertions(+), 18 deletions(-) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index eb5894ca..1cce2c32 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -892,7 +892,8 @@ def reshape(self, local_shape, axis=0): """ local_shapes = self.base_comm.allgather(local_shape) global_shape = list(local_shapes[0]) - global_shape[axis] = np.sum([ls[axis] for ls in local_shapes]) + 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, diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 0e6c65c8..6cd394b0 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -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) @@ -282,9 +284,9 @@ def dot(self, x): if x.ndim == 1: y = self.matvec(x) if is_dims_shaped and get_ndarray_multiplication(): - _local_dims_dimsd_axis = getattr(self, "_local_dims_dimsd_axis", None) - if _local_dims_dimsd_axis: - y = y.reshape(_local_dims_dimsd_axis[0], axis=_local_dims_dimsd_axis[1]) + _local_dimsd = getattr(self, "_local_dimsd", None) + if _local_dimsd: + y = y.reshape(_local_dimsd.dim, axis=_local_dimsd.axis) return y else: msg = ( @@ -373,22 +375,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_dimsd_axis", None) is not None: - Op._local_dims_dimsd_axis = self._local_dims_dimsd_axis + 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): @@ -408,7 +416,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 bd119cf5..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 @@ -112,13 +113,17 @@ def __init__(self, ops: Sequence[LinearOperator], self.nops = nops.sum() 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,)) diff --git a/pylops_mpi/basicoperators/FirstDerivative.py b/pylops_mpi/basicoperators/FirstDerivative.py index 5f57faa8..ddf807f8 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 @@ -90,7 +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_dimsd_axis = (local_split(dims, base_comm, Partition.SCATTER, axis=0), 0) + 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 0fae576c..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 @@ -82,7 +83,7 @@ def __init__( dtype: DTypeLike = np.float64, ) -> None: dims = _value_or_sized_to_tuple(dims) - self._local_dims_dimsd_axis = (local_split(dims, base_comm, Partition.SCATTER, axis=0), 0) + 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 fca80f92..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 @@ -119,6 +120,8 @@ def __init__(self, ops: Sequence[LinearOperator], 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) @@ -132,7 +135,7 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: 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.flatten())) + y1.append(oper.matvec(x.local_array)) y[:] = ncp.concatenate(y1) return y @@ -147,7 +150,7 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray: dtype=self.dtype) y1 = [] for iop, oper in enumerate(self.ops): - y1.append(oper.rmatvec(x.local_array.flatten()[self.nnops[iop]: self.nnops[iop + 1]])) + y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: self.nnops[iop + 1]])) y1 = ncp.sum(ncp.vstack(y1), axis=0) y[:] = self._allreduce(x.base_comm, x.base_comm_nccl, y1, op=MPI.SUM, engine=x.engine) 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): From 95f8fc70522da7d1e90e67045553ae7ecf5560a2 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Wed, 8 Jul 2026 12:13:09 +0530 Subject: [PATCH 08/10] Update __add__ dunder --- pylops_mpi/LinearOperator.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pylops_mpi/LinearOperator.py b/pylops_mpi/LinearOperator.py index 6cd394b0..39fd7dbc 100644 --- a/pylops_mpi/LinearOperator.py +++ b/pylops_mpi/LinearOperator.py @@ -359,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): From 8452b9c594557fc0de7b5daf95cbfe6606dbcacc Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Thu, 9 Jul 2026 01:13:49 +0530 Subject: [PATCH 09/10] Remove self._local_dims_dimsd --- pylops_mpi/basicoperators/FirstDerivative.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pylops_mpi/basicoperators/FirstDerivative.py b/pylops_mpi/basicoperators/FirstDerivative.py index ddf807f8..336693ce 100644 --- a/pylops_mpi/basicoperators/FirstDerivative.py +++ b/pylops_mpi/basicoperators/FirstDerivative.py @@ -131,14 +131,12 @@ def _matvec(self, x: DistributedArray) -> DistributedArray: # If Partition.BROADCAST, then convert to Partition.SCATTER if x.partition is Partition.BROADCAST: x = DistributedArray.to_dist(x=x.local_array, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl) - self._dims_dimsd_local = (self.dimsd[0] // x.size, *self.dimsd[1:]) return self._hmatvec(x) def _rmatvec(self, x: DistributedArray) -> DistributedArray: # If Partition.BROADCAST, then convert to Partition.SCATTER if x.partition is Partition.BROADCAST: x = DistributedArray.to_dist(x=x.local_array, base_comm=x.base_comm, base_comm_nccl=x.base_comm_nccl) - self._dims_dimsd_local = (self.dims[0] // x.size, *self.dims[1:]) return self._hrmatvec(x) @reshaped From 06b3632d11e9159448b11eefbcdae3fa8e4e3d45 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Thu, 9 Jul 2026 01:21:30 +0530 Subject: [PATCH 10/10] Update docstrings of reshape method --- pylops_mpi/DistributedArray.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index 1cce2c32..9d243a23 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -879,16 +879,15 @@ def ravel(self, order: Optional[str] = "C"): local_array=local_array) return arr - def reshape(self, local_shape, axis=0): + def reshape(self, local_shape: Tuple, axis: Optional[int] = 0): """Return a reshaped DistributedArray Parameters ---------- - - Returns - ------- - arr : :obj:`pylops_mpi.DistributedArray` - Reshaped N-D DistributedArray + 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])