From e887ccadc32a106cf340a27c01cdc5422adfcb0f Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Sun, 14 Jun 2026 20:33:39 +0000 Subject: [PATCH 01/47] feat: WIP proximal module --- examples/plot_prox.py | 132 ++++++++++++++++ pylops_mpi/__init__.py | 1 + pylops_mpi/proximal/ProxOperator.py | 112 +++++++++++++ pylops_mpi/proximal/__init__.py | 28 ++++ pylops_mpi/proximal/proximal/L2.py | 190 +++++++++++++++++++++++ pylops_mpi/proximal/proximal/__init__.py | 21 +++ 6 files changed, 484 insertions(+) create mode 100644 examples/plot_prox.py create mode 100644 pylops_mpi/proximal/ProxOperator.py create mode 100644 pylops_mpi/proximal/__init__.py create mode 100644 pylops_mpi/proximal/proximal/L2.py create mode 100644 pylops_mpi/proximal/proximal/__init__.py diff --git a/examples/plot_prox.py b/examples/plot_prox.py new file mode 100644 index 00000000..9f101c31 --- /dev/null +++ b/examples/plot_prox.py @@ -0,0 +1,132 @@ +r""" +Proximal operators +================== + +""" +import numpy as np +from mpi4py import MPI +from matplotlib import pyplot as plt + +import pylops +import pyproximal + +import pylops_mpi + +np.random.seed(42) +plt.close("all") +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() + +n = 10 + +# L1 norm +arr = pylops_mpi.DistributedArray(global_shape=n * size, + partition=pylops_mpi.Partition.SCATTER) + +arr[:] = rank * np.arange(n) + +l1 = pyproximal.L1(sigma=2.0) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) +f = l1d(arr) +prox = l1d.prox(arr, .1) +proxdlocal = prox.asarray() + +dprox = l1d.proxdual(arr, .1) +dproxdlocal = dprox.asarray() + +arrlocal = arr.asarray() +if rank == 0: + flocal = l1(arrlocal) + proxlocal = l1.prox(arrlocal, .1) + dproxlocal = l1.proxdual(arrlocal, .1) + print("||x||_1: ", f, flocal) + print("prox_||x||_1: ", all(proxdlocal == proxlocal)) + print("proxd_||x||_1: ", all(dproxdlocal == dproxlocal)) + +# Box norm +arr = pylops_mpi.DistributedArray(global_shape=n * size, + partition=pylops_mpi.Partition.SCATTER) + +arr[:] = 3 * np.ones(n) +if rank == 0: + arr[n//2] = 20 + +box = pyproximal.Box(lower=1., upper=5.) +boxd = pylops_mpi.proximal.MPIProxOperator(box) +f = boxd(arr) +prox = boxd.prox(arr, .1) +proxdlocal = prox.asarray() + +dprox = boxd.proxdual(arr, .1) +dproxdlocal = dprox.asarray() + +arrlocal = arr.asarray() +if rank == 0: + flocal = box(arrlocal) + proxlocal = box.prox(arrlocal, .1) + dproxlocal = box.proxdual(arrlocal, .1) + print("Box(x): ", f, flocal) + print("prox_Box ", all(proxdlocal == proxlocal)) + print("proxd_Box ", all(dproxdlocal == dproxlocal)) + + +# L2 norm ||x||_2^2 +arr = pylops_mpi.DistributedArray(global_shape=n * size, + partition=pylops_mpi.Partition.SCATTER) + +arr[:] = rank * np.arange(n) + +l2 = pyproximal.L2(sigma=2.0) +l2d = pylops_mpi.proximal.MPIL2(sigma=2.0) +f = l2d(arr) +prox = l2d.prox(arr, .1) +proxdlocal = prox.asarray() +grad = l2d.grad(arr) +graddlocal = grad.asarray() + +arrlocal = arr.asarray() +if rank == 0: + flocal = l2(arrlocal) + proxlocal = l2.prox(arrlocal, .1) + gradlocal = l2.grad(arrlocal) + print("||x||_2^2: ", f, flocal) + print("prox_||x||_2^2: ", all(proxdlocal == proxlocal)) + print("grad_||x||_2^2: ", all(graddlocal == gradlocal)) + +# L2 norm ||Op * x - d||_2^2 +solver="cgls" +Op = pylops.FirstDerivative(n * size, sampling=0.001) +Opd = pylops_mpi.MPIFirstDerivative(n * size, sampling=0.001) +# Op = pylops.Diagonal(np.ones(n * size)) +# Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(n)),]) + +b = pylops_mpi.DistributedArray(global_shape=n * size, + partition=pylops_mpi.Partition.SCATTER) + +b[:] = rank * np.ones(n) +blocal = b.asarray() + +x0 = arr.zeros_like() +x0local = x0.asarray() + +l2 = pyproximal.L2( + Op=Op, b=blocal, sigma=2.0, + solver=solver, x0=x0local, kwargs_solver=dict(show=True)) +l2d = pylops_mpi.proximal.MPIL2( + Op=Opd, b=b, sigma=2.0, + solver=solver, x0=x0, kwargs_solver=dict(show=True)) +f = l2d(arr) +prox = l2d.prox(arr, .1) +proxdlocal = prox.asarray() +grad = l2d.grad(arr) +graddlocal = grad.asarray() + +arrlocal = arr.asarray() +if rank == 0: + flocal = l2(arrlocal) + proxlocal = l2.prox(arrlocal, .1) + gradlocal = l2.grad(arrlocal) + print("||x||_2^2: ", f, flocal) + print("prox_||x||_2^2: ", all(proxdlocal == proxlocal), np.linalg.norm(proxdlocal - proxlocal)) + print("grad_||x||_2^2: ", all(graddlocal == gradlocal)) diff --git a/pylops_mpi/__init__.py b/pylops_mpi/__init__.py index f3ad5a2a..80d1b524 100644 --- a/pylops_mpi/__init__.py +++ b/pylops_mpi/__init__.py @@ -12,6 +12,7 @@ from .plotting.plotting import * from .optimization.basic import * from .optimization.sparsity import * +from .proximal.ProxOperator import * try: from .version import version as __version__ diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py new file mode 100644 index 00000000..a51a80b2 --- /dev/null +++ b/pylops_mpi/proximal/ProxOperator.py @@ -0,0 +1,112 @@ +from mpi4py import MPI +from typing import Any, Callable + +from pyproximal import ProxOperator +from pylops.utils.backend import get_module + +from pylops_mpi import DistributedArray + + +_call_reduce_op = dict( + Box=MPI.LAND, + L1=MPI.SUM, +) + + +class MPIProxOperator: + """MPI-enabled PyProximal Proximal Operator + + Common interface for applying (separable) proximal operators in a + distributed fashion. + + In practice, this class provides methods to compute the norm, proximal + operator and gradient between any :obj:`pyproximal.ProxOperator` + (which must be the same across ranks) and a :class:`pylops_mpi.DistributedArray`. + It internally handles the extraction of the local array from the distributed + array and the creation of the output :class:`pylops_mpi.DistributedArray`. + + Parameters + ---------- + prox : :obj:`pyproximal.ProxOperator` + PyProximal Proximal Operator to wrap. + + """ + + def __init__( + self, + prox: ProxOperator, + ) -> None: + # Check if prox is separable (by looking if is listed in + # the mapping dictionary) + prox_name = str(type(prox).__name__) + if prox_name not in _call_reduce_op: + raise NotImplementedError( + f"{prox_name} is not a separable proximal " + "operator, must be implemented directly...") + self.proxop = prox + self.hasgrad = prox.hasgrad + + def __call__(self, x: DistributedArray) -> DistributedArray: + """Functional evaluation of the oprator. + + Modified version of pyproximal `__call__`. This method makes use + of :class:`pylops_mpi.DistributedArray` to evaluate + the functional of the operator in a distributed fashion. + + Parameters + ---------- + x : :obj:`pylops_mpi.DistributedArray` + A DistributedArray of global shape (N, ). + + Returns + ------- + f : :obj:`bool` or :obj:`float` or :obj:`int` + Function evaluation + + """ + # Compute local function evaluation + f = self.proxop(x.local_array) + + # Create receiver buffer + ncp = get_module(x.engine) + recv_buf = ncp.empty(shape=1, dtype=ncp.float64) + + # Reduce local function evaluations into final evaluation + reduce_op = _call_reduce_op[str(type(self.proxop).__name__)] + recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, + ncp.asarray(f), recv_buf, + reduce_op, + engine=x.engine) + return recv_buf[0] + + def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: + """Proximal operator applied to a vector + """ + y = DistributedArray(global_shape=x.global_shape, + base_comm=x.base_comm, + base_comm_nccl=x.base_comm_nccl, + partition=x.partition, + axis=x.axis, + local_shapes=x.local_shapes, + mask=x.mask, + engine=x.engine, + dtype=x.dtype) + y[:] = self.proxop.prox(x.local_array, tau) + + return y + + def proxdual(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: + """Dual Proximal operator applied to a vector + """ + y = DistributedArray(global_shape=x.global_shape, + base_comm=x.base_comm, + base_comm_nccl=x.base_comm_nccl, + partition=x.partition, + axis=x.axis, + local_shapes=x.local_shapes, + mask=x.mask, + engine=x.engine, + dtype=x.dtype) + y[:] = self.proxop.proxdual(x.local_array, tau) + + return y diff --git a/pylops_mpi/proximal/__init__.py b/pylops_mpi/proximal/__init__.py new file mode 100644 index 00000000..c933daed --- /dev/null +++ b/pylops_mpi/proximal/__init__.py @@ -0,0 +1,28 @@ +""" +Proximal Operators and Solvers using MPI +======================================== + +The subpackage proximal extends the pyproximal library providing +proximal operators and solvers using MPI. + +A common interface for applying (separable) proximal operators in a +distributed fashion is provided by the MPIProxOperator operator. + +A list of proximal operators present in pylops_mpi.proximal.proximal: + MPIXX XX + +A list of proximal solvers present in pylops_mpi.proximal.optimization.primal: + MPIXX XX + +and in pylops_mpi.proximal.optimization.primaldual: + MPIXX XX + +""" + +from .ProxOperator import * +from .proximal import * + + +__all__ = [ + "MPIProxOperator", +] \ No newline at end of file diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py new file mode 100644 index 00000000..4726fc8d --- /dev/null +++ b/pylops_mpi/proximal/proximal/L2.py @@ -0,0 +1,190 @@ +from math import sqrt +from mpi4py import MPI +from typing import TYPE_CHECKING, Any, Callable + +from pylops.basicoperators import Identity +from pylops.utils.backend import get_module +from pyproximal.ProxOperator import _check_tau + +from pylops_mpi import DistributedArray, StackedDistributedArray +from pylops_mpi.basicoperators import MPIBlockDiag, MPIStackedVStack +from pylops_mpi.optimization.basic import cg, cgls +from pylops_mpi.proximal import MPIProxOperator + +if TYPE_CHECKING: + from pylops_mpi import MPILinearOperator + + +class MPIL2(MPIProxOperator): + """L2 Norm proximal operator. + + Implement a distributed version of the L2 norm proximal operator. + + Parameters + ---------- + Op : :obj:`pylops_mpi.MPILinearOperator`, optional + MPI-enabled PyLops Linear Operator + b : :obj:`pylops_mpi.DistributedArray`, optional + Data vector + q : :obj:`pylops_mpi.DistributedArray`, optional + Dot vector + sigma : :obj:`int`, optional + Multiplicative coefficient of L2 norm + alpha : :obj:`float`, optional + Multiplicative coefficient of dot product + qgrad : :obj:`bool`, optional + Add q term to gradient (``True``) or not (``False``) + niter : :obj:`int` or :obj:`func`, optional + Number of iterations of iterative scheme used to compute the proximal. + This can be a constant number or a function that is called passing a + counter which keeps track of how many times the ``prox`` method has + been invoked before and returns the ``niter`` to be used. + x0 : :obj:`pylops_mpi.DistributedArray`, optional + Initial vector. If ``Op`` is not None, this must be passed. + warm : :obj:`bool`, optional + Warm start (``True``) or not (``False``). Uses estimate from previous + call of ``prox`` method. + solver : :obj:`str`, optional + .. versionadded:: 0.11.0 + + Name of solver to use with non-explicit operators: + + - ``cg`` to use :py:func:`pylops_mpi.optimization.basic.cg` on the + normal equations; + - ``cgls`` to use :py:func:`pylops.optimization.basic.cgls` on the + regularized system of equations; + **kwargs_solver : :obj:`dict`, optional + Dictionary containing extra arguments for the solver selected + via the ``solver`` parameter. + + """ + + def __init__( + self, + Op: "MPILinearOperator" = None, + b: DistributedArray | None = None, + q: DistributedArray | None = None, + sigma: float = 1.0, + alpha: float = 1.0, + qgrad: bool = True, + niter: int | Callable[[int], int] = 10, + x0: DistributedArray | None = None, + warm: bool = True, + solver: str | None = "cgls", + kwargs_solver: dict[str, Any] | None = None, + ) -> None: + if Op is not None and x0 is None: + raise ValueError("x0 must be passed when Op is not None") + self.Op = Op + self.hasgrad = True + + self.b = b + self.q = q + self.sigma = sigma + self.alpha = alpha + self.qgrad = qgrad + self.niter = niter + self.x0 = x0 + self.warm = warm + self.solver = solver + self.count = 0 + self.kwargs_solver = {} if kwargs_solver is None else kwargs_solver + + # define whether the normal equations or the regularized system + # of equations are solved + if self.solver == "cg": + self.normaleqs = True + elif self.solver == "cgls": + self.normaleqs = False + else: + msg = ( + f"Provided solver={self.solver}. " + "Available options are 'cg' or 'cgls'." + ) + raise ValueError(msg) + + # create data term + if ( + self.Op is not None + and self.b is not None + and self.normaleqs + ): + self.OpTb = self.sigma * self.Op.H @ self.b + + def __call__(self, x: DistributedArray) -> DistributedArray: + if self.Op is not None and self.b is not None: + f = (self.sigma / 2.0) * ((self.Op * x - self.b).norm() ** 2) + elif self.b is not None: + f = (self.sigma / 2.0) * ((x - self.b).norm() ** 2) + else: + f = (self.sigma / 2.0) * (x.norm() ** 2) + if self.q is not None: + f += self.alpha * self.q.dot(x) + return float(f) + + + def _increment_count(func: Callable[..., Any]) -> Callable[..., Any]: + """Increment counter""" + + def wrapped(self, *args: Any, **kwargs: Any) -> Any: + self.count += 1 + return func(self, *args, **kwargs) + + return wrapped + + @_increment_count + @_check_tau + def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: + """Proximal operator applied to a vector + """ + # define current number of iterations + if isinstance(self.niter, int): + niter = self.niter + else: + niter = self.niter(self.count) + + # solve proximal optimization + if self.Op is not None and self.b is not None: + if self.normaleqs: + y = x + tau * self.OpTb + if self.q is not None: + y -= tau * self.alpha * self.q + if self.normaleqs: + Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( + tau * self.sigma + ) * (self.Op.H * self.Op) + x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] + else: + y = x + if self.q is not None: + y -= tau * self.alpha * self.q + + Opreg = MPIStackedVStack([ + sqrt(tau * self.sigma) * self.Op, + MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, ),])]) + breg = StackedDistributedArray([sqrt(tau * self.sigma) * self.b, y]) + x = cgls(Opreg, breg, x0=self.x0, niter=niter, **self.kwargs_solver)[0] + if self.warm: + self.x0 = x + elif self.b is not None: + num = x + tau * self.sigma * self.b + if self.q is not None: + num -= tau * self.alpha * self.q + x = (1. / (1.0 + tau * self.sigma)) * num + else: + num = x + if self.q is not None: + num -= tau * self.alpha * self.q + x = (1.0 / (1.0 + tau * self.sigma)) * num + return x + + def grad(self, x: DistributedArray) -> DistributedArray: + if self.Op is not None and self.b is not None: + g = self.sigma * self.Op.H @ (self.Op @ x - self.b) + elif self.b is not None: + g = self.sigma * (x - self.b) + else: + g = self.sigma * x + if self.q is not None and self.qgrad: + g += self.alpha * self.q + return g diff --git a/pylops_mpi/proximal/proximal/__init__.py b/pylops_mpi/proximal/proximal/__init__.py new file mode 100644 index 00000000..53ce760b --- /dev/null +++ b/pylops_mpi/proximal/proximal/__init__.py @@ -0,0 +1,21 @@ +""" +Proximal Operators using MPI +======================================== + +The subpackage proximal extends the pyproximal.proximal module providing +custom (non-separable) proximal operator using MPI, which cannot be +directly implemented by wrapping a PyProximal operator in +pylops_mpi.proximal.MPIProxOperator. + + +A list of proximal operators: + MPIL2 L2 Norm + +""" + +from .L2 import * + + +__all__ = [ + "MPIL2", +] \ No newline at end of file From 668ba69be7a5649286fcf0091f3beba01a728f25 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 15 Jun 2026 20:34:26 +0000 Subject: [PATCH 02/47] feat: added proximal optimization module --- examples/plot_prox.py | 43 ++++- pylops_mpi/optimization/cls_sparsity.py | 15 +- pylops_mpi/proximal/optimization/__init__.py | 19 ++ pylops_mpi/proximal/optimization/primal.py | 176 +++++++++++++++++++ tutorials/reflectivity.py | 72 ++++++-- 5 files changed, 302 insertions(+), 23 deletions(-) create mode 100644 pylops_mpi/proximal/optimization/__init__.py create mode 100644 pylops_mpi/proximal/optimization/primal.py diff --git a/examples/plot_prox.py b/examples/plot_prox.py index 9f101c31..8f1eb2d3 100644 --- a/examples/plot_prox.py +++ b/examples/plot_prox.py @@ -115,7 +115,7 @@ solver=solver, x0=x0local, kwargs_solver=dict(show=True)) l2d = pylops_mpi.proximal.MPIL2( Op=Opd, b=b, sigma=2.0, - solver=solver, x0=x0, kwargs_solver=dict(show=True)) + solver=solver, x0=x0, kwargs_solver=dict(show=True if rank==0 else False)) f = l2d(arr) prox = l2d.prox(arr, .1) proxdlocal = prox.asarray() @@ -130,3 +130,44 @@ print("||x||_2^2: ", f, flocal) print("prox_||x||_2^2: ", all(proxdlocal == proxlocal), np.linalg.norm(proxdlocal - proxlocal)) print("grad_||x||_2^2: ", all(graddlocal == gradlocal)) + + + +# Proximal gradient +arr = pylops_mpi.DistributedArray(global_shape=n, + partition=pylops_mpi.Partition.BROADCAST) +arr[:] = 0.0 +arr[n//4] = 1.0 +arr[n//2] = -0.5 + +Op = pylops.MatrixMult(np.random.normal(0, 1, (n-2, n,))) +Opd = pylops_mpi.MPILinearOperator(Op) + +b = Opd @ arr +blocal = b.asarray() + +l2d = pylops_mpi.proximal.MPIL2( + Op=Opd, b=b, solver=solver, x0=arr.zeros_like()) +l1 = pyproximal.L1(sigma=8e-1) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +arrpg = pylops_mpi.proximal.optimization.primal.ProximalGradient( + l2d, l1d, x0=arr.zeros_like(), tau=1e-2, niter=400, + show=True + ) +arrpgdlocal = arrpg.asarray() + +arrlocal = arr.asarray() +if rank == 0: + l2local = pyproximal.L2( + Op=Op, b=blocal, + solver=solver) + l1local = pyproximal.L1(sigma=8e-1) + + arrpglocal = pyproximal.optimization.primal.ProximalGradient( + l2local, l1local, x0=np.zeros(n), tau=1e-2, niter=400, show=True + ) + + print('PG true', arrlocal) + print('PG distr', arrpgdlocal) + print('PG local', arrpglocal) \ No newline at end of file diff --git a/pylops_mpi/optimization/cls_sparsity.py b/pylops_mpi/optimization/cls_sparsity.py index 8dcc0116..91dc04c9 100644 --- a/pylops_mpi/optimization/cls_sparsity.py +++ b/pylops_mpi/optimization/cls_sparsity.py @@ -1,7 +1,8 @@ +from typing import Any, Callable, Dict, Optional, Tuple, Union +import sys import time import logging from math import sqrt -from typing import Any, Callable, Dict, Optional, Tuple, Union import numpy as np @@ -110,6 +111,7 @@ def _print_setup(self) -> None: print(strpar1) print("-" * 80) print(head1) + sys.stdout.flush() def _print_step( self, @@ -129,6 +131,7 @@ def _print_step( + f"{costdata:10.3e} {costdata + costreg:9.3e} {xupdate:10.3e}" ) print(msg) + sys.stdout.flush() def memory_usage( self, @@ -250,10 +253,10 @@ def setup( **self.eigsdict )[0] ) - print("maxeieur", maxeig) self.alpha = float(1.0 / maxeig) self.thresh = eps * self.alpha * 0.5 x = x0.copy() + self.rank = x.rank # create variable to track residual if monitorres: @@ -263,7 +266,7 @@ def setup( # create variables to track the residual norm and iterations self.cost = [] self.iiter = 0 - if show: + if show and self.rank == 0: self._print_setup() return x @@ -338,7 +341,7 @@ def step( costreg = self.eps * x.norm(ord=1).item() self.cost.append(float(costdata + costreg)) self.iiter += 1 - if show: + if show and self.rank == 0: self._print_step(x, costdata, costreg, xupdate) return x, xupdate @@ -404,7 +407,7 @@ def finalize(self, show: bool = False) -> None: self.tend = time.time() self.telapsed = self.tend - self.tstart self.cost = np.array(self.cost) - if show: + if show and self.rank == 0: self._print_finalize() def solve( @@ -657,7 +660,7 @@ def step( self.cost.append(float(costdata + costreg)) self.iiter += 1 - if show: + if show and self.rank == 0: self._print_step(x, costdata, costreg, xupdate) return x, z, xupdate diff --git a/pylops_mpi/proximal/optimization/__init__.py b/pylops_mpi/proximal/optimization/__init__.py new file mode 100644 index 00000000..6899fd92 --- /dev/null +++ b/pylops_mpi/proximal/optimization/__init__.py @@ -0,0 +1,19 @@ +""" +Proximal Solvers using MPI +========================== + +The subpackage proximal extends the pyproximal.optimization module +providing proximal solvers using MPI. + + +A list of proximal solvers: + ProximalGradient Proximal Gradient + +""" + +from .primal import * + + +__all__ = [ + "ProximalGradient", +] \ No newline at end of file diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py new file mode 100644 index 00000000..8ad9a8f4 --- /dev/null +++ b/pylops_mpi/proximal/optimization/primal.py @@ -0,0 +1,176 @@ +import sys +import time +from collections.abc import Callable +from math import sqrt +from typing import TYPE_CHECKING, Any, Optional + +import numpy as np +from pylops.optimization.leastsquares import regularized_inversion +from pylops.utils.backend import get_array_module, to_numpy +from pylops.utils.typing import NDArray + +from pyproximal.proximal import L2 + +from pylops_mpi import DistributedArray +from pylops_mpi.proximal.ProxOperator import MPIProxOperator + + +if TYPE_CHECKING: + from pylops_mpi.linearoperator import MPILinearOperator + + +def ProximalGradient( + proxf: MPIProxOperator, + proxg: MPIProxOperator, + x0: DistributedArray, + epsg: float | NDArray = 1.0, + tau: float | None = None, + # backtracking: bool = False, + beta: float = 0.5, + eta: float = 1.0, + niter: int = 10, + niterback: int = 100, + acceleration: str | None = None, + tol: float | None = None, + callback: Callable[[NDArray], None] | None = None, + show: bool = False, +) -> NDArray: + r"""Proximal gradient (optionally accelerated) + + """ + rank = x0.rank + + # TODO: implement backtracking + backtracking = False + + # check if epgs is a vector + epsg = np.asarray(epsg, dtype=float) + if epsg.size == 1: + epsg = epsg * np.ones(niter) + epsg_print = str(epsg[0]) + else: + epsg_print = "Multi" + + if acceleration not in [None, "None", "vandenberghe", "fista"]: + msg = "Acceleration should be None, vandenberghe or fista" + raise NotImplementedError(msg) + if show and rank == 0: + tstart = time.time() + print( + "Accelerated Proximal Gradient\n" + "---------------------------------------------------------\n" + "Proximal operator (f): %s\n" + "Proximal operator (g): %s\n" + "tau = %s\tbacktrack = %s\tbeta = %10e\n" + "epsg = %s\tniter = %d\ttol = %s\n" + "" + "niterback = %d\tacceleration = %s\n" + % ( + type(proxf), + type(proxg), + str(tau), + backtracking, + beta, + epsg_print, + niter, + str(tol), + niterback, + acceleration, + ) + ) + head = " Itn x[0] f g J=f+eps*g tau" + print(head) + sys.stdout.flush() + + # if tau is None: + # backtracking = True + # tau = 1.0 + + # initialize model + t = 1.0 + x = x0.copy() + y = x.copy() + pfg = np.inf + tolbreak = False + + # iterate + for iiter in range(niter): + xold = x.copy() + + # proximal step + if not backtracking: + if eta == 1.0: + x = proxg.prox(y - tau * proxf.grad(y), epsg[iiter] * tau) + else: + x = x + eta * ( + proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x + ) + else: + pass + # x, tau = _backtracking( + # y, tau, proxf, proxg, epsg[iiter], beta=beta, niterback=niterback + # ) + # if eta != 1.0: + # x = x + eta * ( + # proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x + # ) + + # update internal parameters for bilinear operator + # if isinstance(proxf, BilinearOperator): + # proxf.updatexy(x) + + # update y + if acceleration == "vandenberghe": + omega = iiter / (iiter + 3) + elif acceleration == "fista": + told = t + t = (1.0 + np.sqrt(1.0 + 4.0 * t**2)) / 2.0 + omega = (told - 1.0) / t + else: + omega = 0 + y = x + omega * (x - xold) + + # run callback + if callback is not None: + callback(x) + + # tolerance check: break iterations if overall + # objective does not decrease below tolerance + if tol is not None: + pfgold = pfg + pf, pg = proxf(x), proxg(x) + pfg = pf + np.sum(epsg[iiter] * pg) + if np.abs(1.0 - pfg / pfgold) < tol: + tolbreak = True + + # show iteration logger + if show: + if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: + if tol is None: + pf, pg = proxf(x), proxg(x) + pfg = pf + np.sum(epsg[iiter] * pg) + if rank == 0: + msg = "%6g %12.5e %10.3e %10.3e %10.3e %10.3e" % ( + iiter + 1, + ( + np.real(to_numpy(x[0])) + if x.ndim == 1 + else np.real(to_numpy(x[0, 0])) + ), + pf, + pg, + pfg, + tau, + ) + print(msg) + sys.stdout.flush() + + # break if tolerance condition is met + if tolbreak: + break + + if show and rank == 0: + print("\nTotal time (s) = %.2f" % (time.time() - tstart)) + print("---------------------------------------------------------\n") + return x + diff --git a/tutorials/reflectivity.py b/tutorials/reflectivity.py index 6e6f27b1..77584b3b 100644 --- a/tutorials/reflectivity.py +++ b/tutorials/reflectivity.py @@ -16,8 +16,12 @@ the y-dimension) is distributed across ranks and each of them is in charge of performing modelling for a subvolume of the entire domain. -However, since a reflectivity model is sparse, the :py:class:`pylops.optimization.ISTA` -solver is used here. +However, since a reflectivity model is sparse, a sparsity-promoting solver is used here. +We will consider two options: +- :py:class:`pylops_mpi.optimization.ISTA`: ad-hoc distributed ISTA solver + acting like PyLops' ISTA; +- :py:class:`pylops_mpi.proximal.optimization.ProximalGradient`: general purpose + distributed Proximal Gradient solver acting like PyProximal's ProximalGradient; """ @@ -25,6 +29,7 @@ from matplotlib import pyplot as plt from mpi4py import MPI +import pyproximal from pylops.utils.wavelets import ricker from pylops.basicoperators import FirstDerivative from pylops.signalprocessing import Convolve1D @@ -87,15 +92,37 @@ d = d_dist.asarray().reshape((ny, nx, nz)) ############################################################################### -# We now perform sparsity-promotion inversion +# We now perform sparsity-promotion inversion with the ISTA solver r0_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz) r0_dist[:] = 0. -rinv3d_dist = pylops_mpi.optimization.sparsity.ista( +rfista3d_dist = pylops_mpi.optimization.sparsity.fista( CDiag, d_dist, x0=r0_dist, - niter=200, eps=1e-2, tol=1e-8, show=True)[0] -rinv3d = rinv3d_dist.asarray().reshape((ny, nx, nz)) + niter=200, eps=2e-2, tol=1e-8, show=True)[0] +rfista3d = rfista3d_dist.asarray().reshape((ny, nx, nz)) + +############################################################################### +# And now with the Proximal Gradient solver + +l2d = pylops_mpi.proximal.MPIL2(Op=CDiag, b=d_dist, x0=r0_dist) +l1 = pyproximal.L1(sigma=1e-2) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +CDiag1 = CDiag.H @ CDiag +maxeig = np.abs( + pylops_mpi.optimization.eigs.power_iteration( + CDiag1, + b_k=r0_dist.empty_like(), + dtype=CDiag1.dtype + )[0] +) + +rpg3d_dist = pylops_mpi.proximal.optimization.primal.ProximalGradient( + l2d, l1d, x0=r0_dist, tau=.99 / maxeig, niter=200, acceleration="fista", + show=True + ) +rpg3d = rpg3d_dist.asarray().reshape((ny, nx, nz)) ############################################################################### # Finally, we display the modeling and inversion results @@ -110,11 +137,11 @@ d0 = Cop0 @ r0 # Check the two distributed implementations give the same modelling results - print('Reflectivity Distr == Local', np.allclose(d, d0)) - print('Data Distr == Local', np.allclose(r, r0)) + print('Reflectivity Distr == Local', np.allclose(r, r0)) + print('Data Distr == Local', np.allclose(d, d0)) # Visualize - fig, axs = plt.subplots(nrows=4, ncols=3, figsize=(9, 14), constrained_layout=True) + fig, axs = plt.subplots(nrows=5, ncols=3, figsize=(9, 14), constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") @@ -145,16 +172,29 @@ axs[2][2].set_title('Data x-y') axs[2][2].axis('tight') - axs[3][0].imshow(rinv3d[5, :, :].T, cmap='gray', vmin=-.1, vmax=.1) - axs[3][0].set_title("Inverted Reflectivity iter x-z") + axs[3][0].imshow(rfista3d[5, :, :].T, cmap='gray', vmin=-.1, vmax=.1) + axs[3][0].set_title("FISTA Reflectivity iter x-z") axs[3][0].axis("tight") - axs[3][1].imshow(rinv3d[:, 200, :].T, cmap='gray', vmin=-.1, vmax=.1) - axs[3][1].set_title('Inverted Reflectivity iter y-z') + axs[3][1].imshow(rfista3d[:, 200, :].T, cmap='gray', vmin=-.1, vmax=.1) + axs[3][1].set_title('FISTA Reflectivity iter y-z') axs[3][1].axis('tight') - axs[3][2].imshow(rinv3d[:, :, 220].T, cmap='gray', vmin=-.1, vmax=.1) - axs[3][2].set_title('Inverted Reflectivity iter x-y') + axs[3][2].imshow(rfista3d[:, :, 220].T, cmap='gray', vmin=-.1, vmax=.1) + axs[3][2].set_title('FISTA Reflectivity iter x-y') axs[3][2].axis('tight') + axs[4][0].imshow(rpg3d[5, :, :].T, cmap='gray', vmin=-.1, vmax=.1) + axs[4][0].set_title("PG Reflectivity iter x-z") + axs[4][0].axis("tight") + axs[4][1].imshow(rpg3d[:, 200, :].T, cmap='gray', vmin=-.1, vmax=.1) + axs[4][1].set_title('PG Reflectivity iter y-z') + axs[4][1].axis('tight') + axs[4][2].imshow(rpg3d[:, :, 220].T, cmap='gray', vmin=-.1, vmax=.1) + axs[4][2].set_title('PG Reflectivity iter x-y') + axs[4][2].axis('tight') + + plt.savefig('Reflectivity') + ############################################################################### # To run this tutorial with our NCCL backend, refer to -# `Reflectivity Inversion with NCCL tutorial `_ in the repository. +# `Reflectivity Inversion with NCCL tutorial `_ +# in the repository. From fd3fc3ba6f5745e6d16b3c74cc4d1e6b3bdcbc82 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 15 Jun 2026 20:34:47 +0000 Subject: [PATCH 03/47] feat: improvements to ProxOperator --- pylops_mpi/proximal/ProxOperator.py | 28 ++++++++++++++---------- pylops_mpi/proximal/__init__.py | 1 + pylops_mpi/proximal/proximal/__init__.py | 2 +- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index a51a80b2..49c6c9bc 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -4,7 +4,7 @@ from pyproximal import ProxOperator from pylops.utils.backend import get_module -from pylops_mpi import DistributedArray +from pylops_mpi import DistributedArray, Partition _call_reduce_op = dict( @@ -67,17 +67,21 @@ def __call__(self, x: DistributedArray) -> DistributedArray: # Compute local function evaluation f = self.proxop(x.local_array) - # Create receiver buffer - ncp = get_module(x.engine) - recv_buf = ncp.empty(shape=1, dtype=ncp.float64) - - # Reduce local function evaluations into final evaluation - reduce_op = _call_reduce_op[str(type(self.proxop).__name__)] - recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, - ncp.asarray(f), recv_buf, - reduce_op, - engine=x.engine) - return recv_buf[0] + if Partition.SCATTER: + # Create receiver buffer + ncp = get_module(x.engine) + recv_buf = ncp.empty(shape=1, dtype=ncp.float64) + + # Reduce local function evaluations into final evaluation + reduce_op = _call_reduce_op[str(type(self.proxop).__name__)] + recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, + ncp.asarray(f), recv_buf, + reduce_op, + engine=x.engine) + return recv_buf[0] + else: + # For broadcasted arrays, simply return the local f + return f def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: """Proximal operator applied to a vector diff --git a/pylops_mpi/proximal/__init__.py b/pylops_mpi/proximal/__init__.py index c933daed..c97dce78 100644 --- a/pylops_mpi/proximal/__init__.py +++ b/pylops_mpi/proximal/__init__.py @@ -21,6 +21,7 @@ from .ProxOperator import * from .proximal import * +from .optimization import * __all__ = [ diff --git a/pylops_mpi/proximal/proximal/__init__.py b/pylops_mpi/proximal/proximal/__init__.py index 53ce760b..d5902cdd 100644 --- a/pylops_mpi/proximal/proximal/__init__.py +++ b/pylops_mpi/proximal/proximal/__init__.py @@ -1,6 +1,6 @@ """ Proximal Operators using MPI -======================================== +============================ The subpackage proximal extends the pyproximal.proximal module providing custom (non-separable) proximal operator using MPI, which cannot be From fa1badf4670a2aed341fe4fb465e5c0bbf07db64 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Thu, 18 Jun 2026 21:32:01 +0000 Subject: [PATCH 04/47] feat: enabled TV-reg with ADMML2 --- examples/plot_prox.py | 95 ++++++++++++++++++++- pylops_mpi/DistributedArray.py | 61 ++++++++++++-- pylops_mpi/proximal/ProxOperator.py | 67 ++++++++------- pylops_mpi/proximal/optimization/primal.py | 97 ++++++++++++++++++++-- pylops_mpi/proximal/proximal/L2.py | 1 - tutorials/poststack.py | 40 ++++++++- 6 files changed, 311 insertions(+), 50 deletions(-) diff --git a/examples/plot_prox.py b/examples/plot_prox.py index 8f1eb2d3..eff2026a 100644 --- a/examples/plot_prox.py +++ b/examples/plot_prox.py @@ -170,4 +170,97 @@ print('PG true', arrlocal) print('PG distr', arrpgdlocal) - print('PG local', arrpglocal) \ No newline at end of file + print('PG local', arrpglocal) + + +# ADMML2 with stacked operator +ny, nx = 40, 40 +arrlocal = np.ones((ny, nx)) +arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 +arr = pylops_mpi.DistributedArray(global_shape=ny * nx, + partition=pylops_mpi.Partition.SCATTER) +arr[:] = arrlocal[ny//4 * rank: ny//4 * (rank +1)].flatten() + +Op = pylops.Diagonal(np.ones(ny*nx)) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//4)),]) + +b = Opd @ arr +blocal = b.asarray() + +Iop = pylops.Identity(ny*nx) +Iopd = pylops_mpi.MPIBlockDiag([pylops.Identity(ny*nx//4),]) + +L = 8.0 # maxeig(Gop^H Gop) + +l1 = pyproximal.L1(sigma=8e-1) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +x0distr = arr.zeros_like() +arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, Opd, b, Iopd, x0=x0distr, tau=.99/L, niter=5, + show=True, kwargs_solver=dict(niter=20), + )[0] +arradmmdlocal = arradmm.asarray() + +arrlocal = arr.asarray() +if rank == 0: + + l1local = pyproximal.L1(sigma=8e-1) + + arradmmlocal = pyproximal.optimization.primal.ADMML2( + l1local, Op, blocal, Iop, x0=np.zeros(ny*nx), + tau=.99/L, niter=5, show=True, iter_lim=20, + )[0] + + print('ADMML2 true', arrlocal) + print('ADMML2 distr', arradmmdlocal) + print('ADMML2 local', arradmmlocal) + + +# ADMML2 with stacked operator for A +ny, nx = 40, 40 +arrlocal = np.ones((ny, nx)) +arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 +arr = pylops_mpi.DistributedArray(global_shape=ny * nx, + partition=pylops_mpi.Partition.SCATTER) +arr[:] = arrlocal[ny//4 * rank: ny//4 * (rank +1)].flatten() + +Op = pylops.Diagonal(np.ones(ny*nx)) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//4)),]) + +b = Opd @ arr +blocal = b.asarray() + +Gopd = pylops_mpi.MPIGradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward") + +L = 8.0 # maxeig(Gop^H Gop) + +l1 = pyproximal.L1(sigma=8e-1) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +x0distr = arr.zeros_like() +arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, + show=True, kwargs_solver=dict(niter=5), + )[0] +arradmmdlocal = arradmm.asarray() + +arrlocal = arr.asarray() +if rank == 0: + + Gop = pylops.Gradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward", + ) + l1local = pyproximal.L1(sigma=8e-1) + + arradmmlocal = pyproximal.optimization.primal.ADMML2( + l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), + tau=.99/L, niter=5, show=True, iter_lim=5, + )[0] + + print('ADMML2 true', arrlocal) + print('ADMML2 distr', arradmmdlocal) + print('ADMML2 local', arradmmlocal) + print(arradmmdlocal - arradmmlocal) + diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index a0647500..7a5085e9 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -982,11 +982,27 @@ def __init__(self, distarrays: List, base_comm: MPI.Comm = MPI.COMM_WORLD): self.rank = base_comm.Get_rank() self.size = base_comm.Get_size() + # Define global shape as sum as shapes + self.global_shape = distarrays[0].global_shape + for iarr in range(1, self.narrays): + self.global_shape = tuple([g1 + g2 for g1, g2 in \ + zip(self.global_shape, distarrays[iarr].global_shape)]) + def __getitem__(self, index): return self.distarrays[index] def __setitem__(self, index, value): - self.distarrays[index][:] = value + target = self.distarrays[index] + if isinstance(target, StackedDistributedArray): + # nested StackedDistributedArray: assign each sub-array in turn + for iarr in range(target.narrays): + target[iarr] = value[iarr] + elif isinstance(value, DistributedArray): + # plain DistributedArray: assign the local array + target[:] = value[:] + else: + # plain DistributedArray assigned from an array-like / scalar + target[:] = value def asarray(self): """Global view of the array @@ -1017,7 +1033,10 @@ def _check_stacked_size(self, stacked_array): def __neg__(self): arr = self.copy() for iarr in range(self.narrays): - arr[iarr][:] = -arr[iarr][:] + if isinstance(arr[iarr], StackedDistributedArray): + arr[iarr] = -arr[iarr] + else: + arr[iarr][:] = -arr[iarr][:] return arr def __add__(self, x): @@ -1044,7 +1063,10 @@ def add(self, stacked_array): self._check_stacked_size(stacked_array) SumArray = self.copy() for iarr in range(self.narrays): - SumArray[iarr][:] = (self[iarr] + stacked_array[iarr])[:] + if isinstance(stacked_array[iarr], StackedDistributedArray): + SumArray[iarr] = (self[iarr] + stacked_array[iarr]) + else: + SumArray[iarr][:] = (self[iarr] + stacked_array[iarr])[:] return SumArray def iadd(self, stacked_array): @@ -1052,7 +1074,10 @@ def iadd(self, stacked_array): """ self._check_stacked_size(stacked_array) for iarr in range(self.narrays): - self[iarr][:] = (self[iarr] + stacked_array[iarr])[:] + if isinstance(self[iarr], StackedDistributedArray): + self[iarr] = (self[iarr] + stacked_array[iarr]) + else: + self[iarr][:] = (self[iarr] + stacked_array[iarr])[:] return self def multiply(self, stacked_array): @@ -1063,13 +1088,19 @@ def multiply(self, stacked_array): ProductArray = self.copy() if isinstance(stacked_array, StackedDistributedArray): - # multiply two DistributedArray + # multiply two StackedDistributedArray for iarr in range(self.narrays): - ProductArray[iarr][:] = (self[iarr] * stacked_array[iarr])[:] + if isinstance(self[iarr], StackedDistributedArray): + ProductArray[iarr] = (self[iarr] * stacked_array[iarr]) + else: + ProductArray[iarr][:] = (self[iarr] * stacked_array[iarr])[:] else: # multiply with scalar for iarr in range(self.narrays): - ProductArray[iarr][:] = (self[iarr] * stacked_array)[:] + if isinstance(self[iarr], StackedDistributedArray): + ProductArray[iarr] = (self[iarr] * stacked_array) + else: + ProductArray[iarr][:] = (self[iarr] * stacked_array)[:] return ProductArray def dot(self, stacked_array, vdot: bool = False): @@ -1107,7 +1138,7 @@ def norm(self, ord: Optional[int] = None): Order of the norm. """ ncp = get_module(self.distarrays[0].engine) - norms = ncp.array([distarray.norm(ord) for distarray in self.distarrays]) + norms = ncp.hstack([distarray.norm(ord) for distarray in self.distarrays]) ord = 2 if ord is None else ord if ord in ['fro', 'nuc']: raise ValueError(f"norm-{ord} not possible for vectors") @@ -1136,6 +1167,20 @@ def copy(self): arr = StackedDistributedArray([distarray.copy() for distarray in self.distarrays]) return arr + def zeros_like(self): + """Creates a zero like StackedDistributedArray + """ + dists = [] + for iarr in range(self.narrays): + distarray = self.distarrays[iarr] + dist = DistributedArray(global_shape=distarray.global_shape, base_comm=distarray.base_comm, + base_comm_nccl=distarray.base_comm_nccl, partition=distarray.partition, + axis=distarray.axis, local_shapes=distarray.local_shapes, mask=distarray.mask, + engine=distarray.engine, dtype=distarray.dtype) + dist[:] = 0. + dists.append(dist) + return StackedDistributedArray(distarrays=dists) + def empty_like(self): """Creates an empty like StackedDistributedArray with uninitialized values """ diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 49c6c9bc..39896ee7 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -8,8 +8,9 @@ _call_reduce_op = dict( - Box=MPI.LAND, - L1=MPI.SUM, + Box=(MPI.LAND, all), + L0=(MPI.SUM, sum), + L1=(MPI.SUM, sum), ) @@ -46,6 +47,9 @@ def __init__( self.proxop = prox self.hasgrad = prox.hasgrad + def __repr__(self) -> str: + return f"<{type(self).__name__} ({type(self.proxop).__name__})>" + def __call__(self, x: DistributedArray) -> DistributedArray: """Functional evaluation of the oprator. @@ -64,39 +68,40 @@ def __call__(self, x: DistributedArray) -> DistributedArray: Function evaluation """ - # Compute local function evaluation - f = self.proxop(x.local_array) - - if Partition.SCATTER: - # Create receiver buffer - ncp = get_module(x.engine) - recv_buf = ncp.empty(shape=1, dtype=ncp.float64) - - # Reduce local function evaluations into final evaluation - reduce_op = _call_reduce_op[str(type(self.proxop).__name__)] - recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, - ncp.asarray(f), recv_buf, - reduce_op, - engine=x.engine) - return recv_buf[0] - else: - # For broadcasted arrays, simply return the local f + if isinstance(x, DistributedArray): + # Compute local function evaluation + f = self.proxop(x.local_array) + + if x.partition == Partition.SCATTER: + # Create receiver buffer + ncp = get_module(x.engine) + + # Reduce local function evaluations into final evaluation + reduce_op = _call_reduce_op[str(type(self.proxop).__name__)][0] + recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, + ncp.asarray(f), + op=reduce_op, + engine=x.engine) + return recv_buf + else: + # For broadcasted arrays, simply return the local f + return f + else: # StackedDistributedArray + reduce_op = _call_reduce_op[str(type(self.proxop).__name__)][1] + fs = [self(x[iarr]) for iarr in range(x.narrays)] + f = reduce_op(fs) return f - + def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: """Proximal operator applied to a vector """ - y = DistributedArray(global_shape=x.global_shape, - base_comm=x.base_comm, - base_comm_nccl=x.base_comm_nccl, - partition=x.partition, - axis=x.axis, - local_shapes=x.local_shapes, - mask=x.mask, - engine=x.engine, - dtype=x.dtype) - y[:] = self.proxop.prox(x.local_array, tau) - + if isinstance(x, DistributedArray): + y = x.empty_like() + y[:] = self.proxop.prox(x.local_array, tau) + else: # StackedDistributedArray + y = x.empty_like() + for iarr in range(x.narrays): + y[iarr][:] = self.proxop.prox(x[iarr].local_array, tau) return y def proxdual(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index 8ad9a8f4..b5309e34 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -5,16 +5,19 @@ from typing import TYPE_CHECKING, Any, Optional import numpy as np +from pylops.basicoperators import Identity from pylops.optimization.leastsquares import regularized_inversion from pylops.utils.backend import get_array_module, to_numpy from pylops.utils.typing import NDArray from pyproximal.proximal import L2 +from pyproximal.optimization.primal import _x0z0_init -from pylops_mpi import DistributedArray +from pylops_mpi import DistributedArray, StackedDistributedArray +from pylops_mpi.basicoperators import MPIBlockDiag, MPIStackedVStack +from pylops_mpi.optimization.basic import cg, cgls from pylops_mpi.proximal.ProxOperator import MPIProxOperator - if TYPE_CHECKING: from pylops_mpi.linearoperator import MPILinearOperator @@ -32,9 +35,9 @@ def ProximalGradient( niterback: int = 100, acceleration: str | None = None, tol: float | None = None, - callback: Callable[[NDArray], None] | None = None, + callback: Callable[[DistributedArray], None] | None = None, show: bool = False, -) -> NDArray: +) -> DistributedArray: r"""Proximal gradient (optionally accelerated) """ @@ -66,8 +69,8 @@ def ProximalGradient( "" "niterback = %d\tacceleration = %s\n" % ( - type(proxf), - type(proxg), + proxf, + proxg, str(tau), backtracking, beta, @@ -172,5 +175,87 @@ def ProximalGradient( if show and rank == 0: print("\nTotal time (s) = %.2f" % (time.time() - tstart)) print("---------------------------------------------------------\n") + sys.stdout.flush() return x + + +def ADMML2( + proxg: MPIProxOperator, + Op: "MPILinearOperator", + b: DistributedArray, + A: "MPILinearOperator", + x0: DistributedArray, + tau: float, + niter: int = 10, + z0: DistributedArray | None = None, + gfirst: bool = False, + callback: Callable[[DistributedArray], None] | None = None, + show: bool = False, + kwargs_solver: dict[str, Any] = {}, +) -> tuple[DistributedArray, DistributedArray]: + r"""Alternating Direction Method of Multipliers for L2 misfit term + + """ + rank = x0.rank + + # initialize variables + x, z = _x0z0_init(x0, z0, A, Opname="A") + u = z.zeros_like() + + if show and rank == 0: + tstart = time.time() + print( + "ADMM\n" + "---------------------------------------------------------\n" + "Proximal operator (g): %s\n" + "tau = %10e\tniter = %d\n" % (proxg, tau, niter) + ) + head = " Itn x[0] f g J = f + g" + print(head) + sys.stdout.flush() + + # run iterations + sqrttau = 1.0 / sqrt(tau) + for iiter in range(niter): + if gfirst: + Ax = A @ x + z = proxg.prox(Ax + u, tau) + + # solve augumented system + Opreg = MPIStackedVStack([Op, sqrttau * A]) + breg = StackedDistributedArray([b, sqrttau * (z - u)]) + x = cgls(Opreg, breg, x0=x, **kwargs_solver)[0] + else: + # solve augumented system + Opreg = MPIStackedVStack([Op, sqrttau * A]) + breg = StackedDistributedArray([b, sqrttau * (z - u)]) + x = cgls(Opreg, breg, x0=x, **kwargs_solver)[0] + + Ax = A @ x + z = proxg.prox(Ax + u, tau) + u = u + Ax - z + + # run callback + if callback is not None: + callback(x) + + if show: + if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: + pf, pg = 0.5 * (Op @ x - b).norm() ** 2, proxg(Ax) + if rank == 0: + msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( + iiter + 1, + np.real(to_numpy(x[0])), + pf, + pg, + pf + pg, + ) + print(msg) + sys.stdout.flush() + if show and rank == 0: + print("\nTotal time (s) = %.2f" % (time.time() - tstart)) + print("---------------------------------------------------------\n") + sys.stdout.flush() + return x, z + diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 4726fc8d..4e5e20b6 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any, Callable from pylops.basicoperators import Identity -from pylops.utils.backend import get_module from pyproximal.ProxOperator import _check_tau from pylops_mpi import DistributedArray, StackedDistributedArray diff --git a/tutorials/poststack.py b/tutorials/poststack.py index cdccce8f..41541a55 100644 --- a/tutorials/poststack.py +++ b/tutorials/poststack.py @@ -38,7 +38,7 @@ \mathbf{ai}_{2} \\ \vdots \\ \mathbf{ai}_{N} - \end{bmatrix} + \end{bmatrix} \rightarrow \mathbf{d} = \mathbf{G} \mathbf{ai} where :math:`\mathbf{G}_i` is a post-stack modelling operator, :math:`\mathbf{d}_i` is the data, and :math:`\mathbf{ai}_i` is the input model for the i-th portion of the model. @@ -56,6 +56,7 @@ from pylops.utils.wavelets import ricker from pylops.basicoperators import Transpose from pylops.avo.poststack import PoststackLinearModelling +from pyproximal.proximal import L1 import pylops_mpi @@ -129,7 +130,7 @@ d_0 = d_dist.asarray().reshape((ny, nx, nz)) ############################################################################### -# We perform 2 different kinds of inversions: +# We perform 3 different kinds of inversions: # # * Inversion calculated iteratively using the :py:class:`pylops_mpi.optimization.cls_basic.CGLS` solver. # @@ -174,6 +175,13 @@ # where :math:`\mathbf{L}` is the :py:class:`pylops_mpi.basicoperators.MPILaplacian` operator # which is used to apply second derivative along all three axes, :math:`\mathbf{N}` is an operator computing the # normal equations, and :math:`\mathbf{d}^{Norm}` is the data of the normal equation operator used for inversion. +# +# * Inversion with anisotropic Total Variation (TV) +# +# .. math:: +# \| \mathbf{d} + \mathbf{G} \mathbf{ai} \|_2^2 + \epsilon \| \boldsymbol \nabla \mathbf{ai} \|_1 +# +# where :math:`\boldsymbol \nabla` is the :py:class:`pylops_mpi.basicoperators.MPIGradient` operator. # Inversion using CGLS solver minv3d_iter_dist = pylops_mpi.optimization.basic.cgls(BDiag, d_dist, x0=mback3d_dist, niter=100, show=True)[0] @@ -203,6 +211,22 @@ StackOp, dstack_dist, x0=mback3d_dist, niter=100, show=False)[0] minv3d_reg = minv3d_reg_dist.asarray().reshape((ny, nx, nz)) +############################################################################### + +# Inversion with TV +Gopd = pylops_mpi.MPIGradient( + dims=(ny, nx, nz), sampling=1., edge=False, kind="forward") + +l1 = L1(sigma=1e-2) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +L = 12.0 # maxeig(Gopd^H Gopd) +minv3d_tv_dist = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99/L, niter=40, + show=True, kwargs_solver=dict(niter=20), + )[0] +minv3d_tv = minv3d_tv_dist.asarray().reshape((ny, nx, nz)) + ############################################################################### # Finally, we display the modeling and inversion results @@ -216,7 +240,7 @@ print('Distr == Local', np.allclose(d, d0)) # Visualize - fig, axs = plt.subplots(nrows=6, ncols=3, figsize=(9, 14), constrained_layout=True) + fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(9, 14), constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") @@ -277,6 +301,16 @@ axs[5][2].set_title('Regularized Inverted Model iter x-y') axs[5][2].axis('tight') + axs[5][0].imshow(minv3d_tv[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) + axs[5][0].set_title("TV-Regularized Inverted Model iter x-z") + axs[5][0].axis("tight") + axs[5][1].imshow(minv3d_tv[:, 200, :].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[5][1].set_title('TV-Regularized Inverted Model iter y-z') + axs[5][1].axis('tight') + axs[5][2].imshow(minv3d_tv[:, :, 220].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[5][2].set_title('TV-Regularized Inverted Model iter x-y') + axs[5][2].axis('tight') + ############################################################################### # To run this tutorial with our NCCL backend, refer to `Post Stack Inversion with NCCL # tutorial `_ From 9da416160edc5adee8a7b9016a74e3ba27ae067d Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Thu, 18 Jun 2026 21:53:53 +0000 Subject: [PATCH 05/47] test: added tests for nested StackedDistributedArray --- examples/plot_prox.py | 10 +-- pylops_mpi/DistributedArray.py | 16 +++- pylops_mpi/proximal/ProxOperator.py | 7 +- tests/test_stackedarray.py | 114 +++++++++++++++++++++++++++- 4 files changed, 137 insertions(+), 10 deletions(-) diff --git a/examples/plot_prox.py b/examples/plot_prox.py index eff2026a..107abce3 100644 --- a/examples/plot_prox.py +++ b/examples/plot_prox.py @@ -179,16 +179,16 @@ arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 arr = pylops_mpi.DistributedArray(global_shape=ny * nx, partition=pylops_mpi.Partition.SCATTER) -arr[:] = arrlocal[ny//4 * rank: ny//4 * (rank +1)].flatten() +arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() Op = pylops.Diagonal(np.ones(ny*nx)) -Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//4)),]) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) b = Opd @ arr blocal = b.asarray() Iop = pylops.Identity(ny*nx) -Iopd = pylops_mpi.MPIBlockDiag([pylops.Identity(ny*nx//4),]) +Iopd = pylops_mpi.MPIBlockDiag([pylops.Identity(ny*nx//size),]) L = 8.0 # maxeig(Gop^H Gop) @@ -223,10 +223,10 @@ arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 arr = pylops_mpi.DistributedArray(global_shape=ny * nx, partition=pylops_mpi.Partition.SCATTER) -arr[:] = arrlocal[ny//4 * rank: ny//4 * (rank +1)].flatten() +arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() Op = pylops.Diagonal(np.ones(ny*nx)) -Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//4)),]) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) b = Opd @ arr blocal = b.asarray() diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index 7a5085e9..7b66a847 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -988,6 +988,18 @@ def __init__(self, distarrays: List, base_comm: MPI.Comm = MPI.COMM_WORLD): self.global_shape = tuple([g1 + g2 for g1, g2 in \ zip(self.global_shape, distarrays[iarr].global_shape)]) + @property + def engine(self): + """Engine + + Find engine by inspecting the first :class:`pylops_mpi.DistributedArray` + among ``distarrays`` (some may be nested + :class:`pylops_mpi.StackedDistributedArray`, which expose this same + property). + """ + return next(distarr.engine for distarr in self.distarrays + if isinstance(distarr, (DistributedArray, StackedDistributedArray))) + def __getitem__(self, index): return self.distarrays[index] @@ -1015,7 +1027,7 @@ def asarray(self): Global Array gathered at all ranks """ - ncp = get_module(self.distarrays[0].engine) + ncp = get_module(self.engine) return ncp.hstack([distarr.asarray().ravel() for distarr in self.distarrays]) def _check_stacked_size(self, stacked_array): @@ -1137,7 +1149,7 @@ def norm(self, ord: Optional[int] = None): ord : :obj:`int`, optional Order of the norm. """ - ncp = get_module(self.distarrays[0].engine) + ncp = get_module(self.engine) norms = ncp.hstack([distarray.norm(ord) for distarray in self.distarrays]) ord = 2 if ord is None else ord if ord in ['fro', 'nuc']: diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 39896ee7..6404d177 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -48,8 +48,11 @@ def __init__( self.hasgrad = prox.hasgrad def __repr__(self) -> str: - return f"<{type(self).__name__} ({type(self.proxop).__name__})>" - + if hasattr(self, "proxop"): + return f"<{type(self).__name__} ({type(self.proxop).__name__})>" + else: + return f"<{type(self).__name__}>" + def __call__(self, x: DistributedArray) -> DistributedArray: """Functional evaluation of the oprator. diff --git a/tests/test_stackedarray.py b/tests/test_stackedarray.py index df9fcdfd..7060c79a 100644 --- a/tests/test_stackedarray.py +++ b/tests/test_stackedarray.py @@ -86,7 +86,8 @@ def test_creation(par): @pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j), (par3), (par3j)]) def test_stacked_math(par): - """Test the Element-Wise Addition, Subtraction and Multiplication, Dot-product, Norm""" + """Test Element-Wise Addition, Subtraction and Multiplication, Dot-product, Norm + for stacked distributed arrays""" distributed_array0 = DistributedArray(global_shape=par['global_shape'], partition=par['partition'], dtype=par['dtype'], axis=par['axis'], @@ -141,3 +142,114 @@ def test_stacked_math(par): # TODO (tharitt): FAIL at inf norm - see above # assert_allclose(linfnorm, np.linalg.norm(stacked_array1.asarray().flatten(), np.inf), # rtol=1e-14) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("par", [(par1), (par1j), (par2), + (par2j), (par3), (par3j)]) +def test_creation_nested(par): + """Test creation of nested stacked distributed arrays""" + # Create stacked array + distributed_array0 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array1 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array2 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array0[:] = 0 + distributed_array1[:] = 1 + distributed_array2[:] = 2 + + stacked_arrays_01 = StackedDistributedArray([distributed_array0, distributed_array1]) + stacked_arrays = StackedDistributedArray([stacked_arrays_01, distributed_array2]) + assert isinstance(stacked_arrays, StackedDistributedArray) + assert_allclose(stacked_arrays[0][0].local_array, + np.zeros(shape=distributed_array0.local_shape, + dtype=par['dtype']), rtol=1e-14) + assert_allclose(stacked_arrays[0][1].local_array, + np.ones(shape=distributed_array0.local_shape, + dtype=par['dtype']), rtol=1e-14) + assert_allclose(stacked_arrays[1].local_array, + 2 * np.ones(shape=distributed_array1.local_shape, + dtype=par['dtype']), rtol=1e-14) + + # Modify array in place + distributed_array0[:] = 2 + assert_allclose(stacked_arrays[0][0].local_array, + 2 * np.ones(shape=distributed_array0.local_shape, + dtype=par['dtype']), rtol=1e-14) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize("par", [(par1), (par1j), (par2), + (par2j), (par3), (par3j)]) +def test_stacked_nested_math(par): + """Test Element-Wise Addition, Subtraction and Multiplication, Dot-product, Norm + for nested stacked distributed arrays""" + distributed_array0 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array1 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array2 = DistributedArray(global_shape=par['global_shape'], + partition=par['partition'], + dtype=par['dtype'], axis=par['axis'], + engine=backend) + distributed_array0[:] = 0 + distributed_array1[:] = np.arange(npp.prod(distributed_array1.local_shape)).reshape(distributed_array1.local_shape) + distributed_array2[:] = np.ones(npp.prod(distributed_array2.local_shape)).reshape(distributed_array1.local_shape) + + stacked_array_01 = StackedDistributedArray([distributed_array0, distributed_array1]) + stacked_array_12 = StackedDistributedArray([distributed_array1, distributed_array2]) + stacked_array1 = StackedDistributedArray([stacked_array_01, distributed_array2]) + stacked_array2 = StackedDistributedArray([stacked_array_12, distributed_array0]) + + # Addition + sum_array = stacked_array1 + stacked_array2 + assert isinstance(sum_array, StackedDistributedArray) + assert_allclose(sum_array.asarray(), np.add(stacked_array1.asarray(), + stacked_array2.asarray()), + rtol=1e-14) + # Subtraction + sub_array = stacked_array1 - stacked_array2 + assert isinstance(sub_array, StackedDistributedArray) + assert_allclose(sub_array.asarray(), np.subtract(stacked_array1.asarray(), + stacked_array2.asarray()), + rtol=1e-14) + # Multiplication + mult_array = stacked_array1 * stacked_array2 + assert isinstance(mult_array, StackedDistributedArray) + assert_allclose(mult_array.asarray(), np.multiply(stacked_array1.asarray(), + stacked_array2.asarray()), + rtol=1e-14) + # Dot-product + dot_prod = stacked_array1.dot(stacked_array2) + assert_allclose(dot_prod, np.dot(stacked_array1.asarray().flatten(), + stacked_array2.asarray().flatten()), + rtol=1e-14) + # Norm + l0norm = stacked_array1.norm(0) + l1norm = stacked_array1.norm(1) + l2norm = stacked_array1.norm(2) + + # TODO (tharitt): FAIL with CuPy + MPI for inf norm - see test_distributedarray.py + # test_distributed_nrom(par) as well +# linfnorm = stacked_array1.norm(np.inf) + assert_allclose(l0norm, np.linalg.norm(stacked_array1.asarray().flatten(), 0), + rtol=1e-14) + assert_allclose(l1norm, np.linalg.norm(stacked_array1.asarray().flatten(), 1), + rtol=1e-14) + assert_allclose(l2norm, np.linalg.norm(stacked_array1.asarray(), 2), + rtol=1e-10) # needed to raise it due to how partial norms are combined (with power applied) + # TODO (tharitt): FAIL at inf norm - see above +# assert_allclose(linfnorm, np.linalg.norm(stacked_array1.asarray().flatten(), np.inf), +# rtol=1e-14) From 3b45e9a864934efa64fd7f932211eef32b5f00d9 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 22 Jun 2026 08:06:54 +0000 Subject: [PATCH 06/47] doc: fix figure numbering in poststack tutorial --- tutorials/poststack.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tutorials/poststack.py b/tutorials/poststack.py index 41541a55..7c0094e0 100644 --- a/tutorials/poststack.py +++ b/tutorials/poststack.py @@ -240,7 +240,7 @@ print('Distr == Local', np.allclose(d, d0)) # Visualize - fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(9, 14), constrained_layout=True) + fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(9, 17), constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") @@ -301,15 +301,15 @@ axs[5][2].set_title('Regularized Inverted Model iter x-y') axs[5][2].axis('tight') - axs[5][0].imshow(minv3d_tv[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) - axs[5][0].set_title("TV-Regularized Inverted Model iter x-z") - axs[5][0].axis("tight") - axs[5][1].imshow(minv3d_tv[:, 200, :].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][1].set_title('TV-Regularized Inverted Model iter y-z') - axs[5][1].axis('tight') - axs[5][2].imshow(minv3d_tv[:, :, 220].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][2].set_title('TV-Regularized Inverted Model iter x-y') - axs[5][2].axis('tight') + axs[6][0].imshow(minv3d_tv[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) + axs[6][0].set_title("TV-Regularized Inverted Model iter x-z") + axs[6][0].axis("tight") + axs[6][1].imshow(minv3d_tv[:, 200, :].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][1].set_title('TV-Regularized Inverted Model iter y-z') + axs[6][1].axis('tight') + axs[6][2].imshow(minv3d_tv[:, :, 220].T, cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][2].set_title('TV-Regularized Inverted Model iter x-y') + axs[6][2].axis('tight') ############################################################################### # To run this tutorial with our NCCL backend, refer to `Post Stack Inversion with NCCL From 609602710402aacad5d4a8116070c85bac014596 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 23 Jun 2026 19:24:03 +0000 Subject: [PATCH 07/47] test: added test_prox --- tests/test_prox.py | 164 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 tests/test_prox.py diff --git a/tests/test_prox.py b/tests/test_prox.py new file mode 100644 index 00000000..a4ea067c --- /dev/null +++ b/tests/test_prox.py @@ -0,0 +1,164 @@ +"""Test proximal operators + Designed to run with n processes + $ mpiexec -n 10 pytest test_prox.py --with-mpi +""" +import os + +if int(os.environ.get("TEST_CUPY_PYLOPS", 0)): + import cupy as np + from cupy.testing import assert_allclose + + backend = "cupy" +else: + import numpy as np + from numpy.testing import assert_allclose + + backend = "numpy" +from mpi4py import MPI +import pytest + +import pylops_mpi +from pylops.basicoperators import FirstDerivative +from pyproximal.proximal import ( + Box, + L0, + L1, + L2 +) +from pylops_mpi.proximal import MPIL2 + + +size = MPI.COMM_WORLD.Get_size() +rank = MPI.COMM_WORLD.Get_rank() +if backend == "cupy": + device_id = rank % np.cuda.runtime.getDeviceCount() + np.cuda.Device(device_id).use() + + +par1 = { + "n": 101, + "imag": 0, + "dtype": np.float64, + "partition": pylops_mpi.Partition.SCATTER +} + +par1b = { + "n": 101, + "imag": 0, + "dtype": np.float64, + "partition": pylops_mpi.Partition.BROADCAST +} + +par1j = { + "n": 101, + "imag": 1j, + "dtype": np.complex128, + "partition": pylops_mpi.Partition.SCATTER +} + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1b), (par1j),] +) +def test_separable_prox(par): + """Separable proximal operators""" + np.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + + # Box (does not support complex numbers) + if par['imag'] == 0: + box = Box(lower=0.0, upper=1.0) + boxd = pylops_mpi.proximal.MPIProxOperator(box) + + f = boxd(x) + prox = boxd.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = box(x_global) + prox_np = box.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-14) + + # L0 + l0 = L0(sigma=2.0) + l0d = pylops_mpi.proximal.MPIProxOperator(l0) + + f = l0d(x) + prox = l0d.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = l0(x_global) + prox_np = l0.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-14) + + # L1 + l1 = L1(sigma=2.0) + l1d = pylops_mpi.proximal.MPIProxOperator(l1) + + f = l1d(x) + prox = l1d.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = l1(x_global) + prox_np = l1.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-14) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1b), (par1j),] +) +def test_L2(par): + """L2 proximal operator""" + np.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + + b = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + b[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + b_global = b.asarray() + + Op_global = FirstDerivative( + par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), + sampling=0.001) + Opd = pylops_mpi.MPIFirstDerivative( + par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), + sampling=0.001) + + l2x = L2(sigma=2.0) + l2xd = MPIL2(sigma=2.0) + + l2b = L2(b=b_global, sigma=2.0) + l2bd = MPIL2(b=b, sigma=2.0) + + # l2Op = L2(Op=Op_global, b=b_global, sigma=2.0) + # l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0) + + # for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): + for l2, l2d in zip([l2x, l2b,], [l2xd, l2bd,]): + f = l2d(x) + prox = l2d.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = l2(x_global) + prox_np = l2.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-14) From 70eb6c06a50782a937803cffd03fb1c69d8a9f01 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 23 Jun 2026 22:06:18 +0000 Subject: [PATCH 08/47] build: added pyproximal to dependencies --- docs/source/installation.rst | 1 + environment-dev.yml | 1 + environment.yml | 1 + pyproject.toml | 1 + requirements-dev.txt | 1 + 5 files changed, 5 insertions(+) diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 5f7e7211..2c3b3044 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -14,6 +14,7 @@ The minimal set of dependencies for the PyLops-MPI project is: * `Matplotlib `_ * `MPI4py `_ * `PyLops `_ +* `PyProximal `_ Additionally, to use the CUDA-aware MPI engine, the following additional dependencies are required: diff --git a/environment-dev.yml b/environment-dev.yml index 101151e8..1aca6838 100644 --- a/environment-dev.yml +++ b/environment-dev.yml @@ -9,6 +9,7 @@ dependencies: - scipy>=1.8.0 - mpi4py - pylops>=2.0.0 + - pyproximal - matplotlib - ipython - pytest diff --git a/environment.yml b/environment.yml index 9b080fbb..02c17141 100644 --- a/environment.yml +++ b/environment.yml @@ -7,5 +7,6 @@ dependencies: - numpy>=1.15.0 - scipy>=1.8.0 - pylops>=2.0.0 + - pyproximal - matplotlib - mpi4py diff --git a/pyproject.toml b/pyproject.toml index e2df251d..bad83069 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "numpy>=1.15.0", "scipy >= 1.4.0", "pylops >= 2.0", + "pyproximal", "mpi4py", "matplotlib", ] diff --git a/requirements-dev.txt b/requirements-dev.txt index 24d12f34..1ef40437 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,6 +1,7 @@ numpy>=1.15.0 scipy>=1.8.0 pylops>=2.0.0 +pyproximal mpi4py matplotlib pytest From 74e8e35037c0f5f9b4124fc284ae80f44ef1b6e5 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Sun, 5 Jul 2026 22:13:05 +0000 Subject: [PATCH 09/47] feat: align DistributedArray with main branch --- pylops_mpi/DistributedArray.py | 103 +++++++++++++++++++++------------ 1 file changed, 67 insertions(+), 36 deletions(-) diff --git a/pylops_mpi/DistributedArray.py b/pylops_mpi/DistributedArray.py index 7b66a847..3058fa4a 100644 --- a/pylops_mpi/DistributedArray.py +++ b/pylops_mpi/DistributedArray.py @@ -1,6 +1,6 @@ from enum import Enum from numbers import Integral -from typing import Any, List, Optional, Tuple, Union, NewType +from typing import Any, List, NewType, Optional, Self, Tuple, Union import numpy as np from mpi4py import MPI @@ -273,11 +273,12 @@ def mask(self): @property def engine(self): - """Engine of the Distributed array + """Engine of the Distributed Array Returns ------- engine : :obj:`str` + Engine """ return self._engine @@ -963,8 +964,11 @@ class StackedDistributedArray: r"""Stacked DistributedArrays Stack DistributedArray objects and power them with basic mathematical operations. - This class allows one to work with a series of distributed arrays to avoid having to create - a single distributed array with some special internal sorting. + This class allows one to work with a series of distributed arrays to avoid + having to create a single distributed array with some special internal sorting. + + .. note:: All :class:`pylops_mpi.DistributedArray` objects passed to + ``distarrays`` must share the same engine (``numpy`` or ``cupy``). Parameters ---------- @@ -973,6 +977,11 @@ class StackedDistributedArray: base_comm : :obj:`mpi4py.MPI.Comm`, optional Base MPI Communicator. Defaults to ``mpi4py.MPI.COMM_WORLD``. + + Raises + ------ + ValueError + Stacked distributed arrays have different engine """ def __init__(self, distarrays: List, base_comm: MPI.Comm = MPI.COMM_WORLD): @@ -982,23 +991,21 @@ def __init__(self, distarrays: List, base_comm: MPI.Comm = MPI.COMM_WORLD): self.rank = base_comm.Get_rank() self.size = base_comm.Get_size() + # Ensure all stacked arrays share the same engine + engines = [distarr.engine for distarr in distarrays] + if any(engine != engines[0] for engine in engines[1:]): + raise ValueError(f"Stacked arrays have mismatching engines: {engines}") + # Define global shape as sum as shapes - self.global_shape = distarrays[0].global_shape + self._global_shape = distarrays[0].global_shape for iarr in range(1, self.narrays): - self.global_shape = tuple([g1 + g2 for g1, g2 in \ - zip(self.global_shape, distarrays[iarr].global_shape)]) - - @property - def engine(self): - """Engine + self._global_shape = \ + tuple([g1 + g2 for g1, g2 in + zip(self._global_shape, distarrays[iarr].global_shape)]) - Find engine by inspecting the first :class:`pylops_mpi.DistributedArray` - among ``distarrays`` (some may be nested - :class:`pylops_mpi.StackedDistributedArray`, which expose this same - property). - """ - return next(distarr.engine for distarr in self.distarrays - if isinstance(distarr, (DistributedArray, StackedDistributedArray))) + def __repr__(self) -> str: + repr_dist = "\n".join([distarray.__repr__() for distarray in self.distarrays]) + return f" NDArray: """Global view of the array - Gather all the distributed arrays + Gather all the distributed arrays and return + a flattened version Returns ------- final_array : :obj:`numpy.ndarray` - Global Array gathered at all ranks + Global Array gathered at all ranks and flattened """ ncp = get_module(self.engine) return ncp.hstack([distarr.asarray().ravel() for distarr in self.distarrays]) - def _check_stacked_size(self, stacked_array): + def _check_stacked_size(self, stacked_array: Self) -> None: """Check that arrays have consistent size """ @@ -1069,7 +1104,7 @@ def __mul__(self, x): def __rmul__(self, x): return self.multiply(x) - def add(self, stacked_array): + def add(self, stacked_array: Self) -> Self: """Stacked Distributed Addition of arrays """ self._check_stacked_size(stacked_array) @@ -1081,7 +1116,7 @@ def add(self, stacked_array): SumArray[iarr][:] = (self[iarr] + stacked_array[iarr])[:] return SumArray - def iadd(self, stacked_array): + def iadd(self, stacked_array: Self) -> Self: """Stacked Distributed In-Place Addition of arrays """ self._check_stacked_size(stacked_array) @@ -1092,7 +1127,7 @@ def iadd(self, stacked_array): self[iarr][:] = (self[iarr] + stacked_array[iarr])[:] return self - def multiply(self, stacked_array): + def multiply(self, stacked_array: float | int | Self) -> Self: """Stacked Distributed Multiplication of arrays """ if isinstance(stacked_array, StackedDistributedArray): @@ -1115,7 +1150,7 @@ def multiply(self, stacked_array): ProductArray[iarr][:] = (self[iarr] * stacked_array)[:] return ProductArray - def dot(self, stacked_array, vdot: bool = False): + def dot(self, stacked_array: Self, vdot: bool = False) -> Self: """ Compute the distributed dot product between this array and another distributed array. @@ -1141,7 +1176,7 @@ def dot(self, stacked_array, vdot: bool = False): dotprod += self[iarr].dot(stacked_array[iarr], vdot=vdot) return dotprod - def norm(self, ord: Optional[int] = None): + def norm(self, ord: Optional[int] = None) -> bool | float: """numpy.linalg.norm method on stacked Distributed arrays Parameters @@ -1167,19 +1202,19 @@ def norm(self, ord: Optional[int] = None): norm = ncp.power(ncp.sum(ncp.power(norms, ord)), 1. / ord) return norm - def conj(self): + def conj(self) -> Self: """Distributed conj() method """ ConjArray = StackedDistributedArray([distarray.conj() for distarray in self.distarrays]) return ConjArray - def copy(self): + def copy(self) -> Self: """Creates a copy of the DistributedArray """ arr = StackedDistributedArray([distarray.copy() for distarray in self.distarrays]) return arr - def zeros_like(self): + def zeros_like(self) -> Self: """Creates a zero like StackedDistributedArray """ dists = [] @@ -1193,7 +1228,7 @@ def zeros_like(self): dists.append(dist) return StackedDistributedArray(distarrays=dists) - def empty_like(self): + def empty_like(self) -> Self: """Creates an empty like StackedDistributedArray with uninitialized values """ dists = [] @@ -1204,8 +1239,4 @@ def empty_like(self): axis=distarray.axis, local_shapes=distarray.local_shapes, mask=distarray.mask, engine=distarray.engine, dtype=distarray.dtype) dists.append(dist) - return StackedDistributedArray(distarrays=dists) - - def __repr__(self): - repr_dist = "\n".join([distarray.__repr__() for distarray in self.distarrays]) - return f" Date: Sun, 5 Jul 2026 22:14:00 +0000 Subject: [PATCH 10/47] feat: align test_distributedarray with main beanch --- tests/test_distributedarray.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_distributedarray.py b/tests/test_distributedarray.py index a7bd350b..7ea59226 100644 --- a/tests/test_distributedarray.py +++ b/tests/test_distributedarray.py @@ -343,7 +343,7 @@ def test_distributed_maskednorm(par): if par['axis'] != 0: x = np.swapaxes(x, 0, par['axis']) arr = DistributedArray.to_dist(x=x, mask=mask, axis=par['axis']) - # TODO (tharitt): Fail with CuPy + MPI + assert_allclose( arr.norm(ord=1, axis=par['norm_axis']), np.linalg.norm(par['x'], ord=1, axis=par['norm_axis']) / (nsub if par['axis'] == par['norm_axis'] else 1), @@ -358,4 +358,4 @@ def test_distributed_maskednorm(par): arr.norm(ord=np.inf, axis=par['norm_axis']), np.linalg.norm(par['x'], ord=np.inf, axis=par['norm_axis']), rtol=1e-14 - ) + ) \ No newline at end of file From 5a55c19bdd1caf949cbc1a40aa5b3dbcca101043 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Sun, 5 Jul 2026 22:23:59 +0000 Subject: [PATCH 11/47] doc: improved example and tests for prox --- examples/plot_prox.py | 22 ++++++++++++++---- pylops_mpi/proximal/proximal/L2.py | 16 ++++++++----- tests/test_prox.py | 36 ++++++++++++++++-------------- 3 files changed, 48 insertions(+), 26 deletions(-) diff --git a/examples/plot_prox.py b/examples/plot_prox.py index 107abce3..82b84d4e 100644 --- a/examples/plot_prox.py +++ b/examples/plot_prox.py @@ -2,6 +2,10 @@ Proximal operators ================== +This example demonstrates the use of the :py:module:`pylops_mpi.proximal` +module, and more specifically how to create and apply PyProximal operators +to distributed array. + """ import numpy as np from mpi4py import MPI @@ -18,17 +22,27 @@ rank = comm.Get_rank() size = comm.Get_size() -n = 10 +############################################################################### +# Let's start with so-called separable proximal operators. These are functionals +# whose proximal operator can be computed in a element-wise fashion. As such, +# no special implementation is required for the distributed counterpart of +# those operators. Instead, we can simply wrap the PyProximal operator into +# a :py:class:`pylops_mpi.proximal.MPIProxOperator`. +# +# We take the :py:class:`pyproximal.proximal.L1` norm as an example. -# L1 norm +n = 10 arr = pylops_mpi.DistributedArray(global_shape=n * size, partition=pylops_mpi.Partition.SCATTER) - arr[:] = rank * np.arange(n) -l1 = pyproximal.L1(sigma=2.0) +l1 = pyproximal.proximal.L1(sigma=2.0) l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +# Call f = l1d(arr) + +# Proximal prox = l1d.prox(arr, .1) proxdlocal = prox.asarray() diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 4e5e20b6..b3b5a148 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -5,7 +5,7 @@ from pylops.basicoperators import Identity from pyproximal.ProxOperator import _check_tau -from pylops_mpi import DistributedArray, StackedDistributedArray +from pylops_mpi import DistributedArray, StackedDistributedArray, Partition from pylops_mpi.basicoperators import MPIBlockDiag, MPIStackedVStack from pylops_mpi.optimization.basic import cg, cgls from pylops_mpi.proximal import MPIProxOperator @@ -149,10 +149,16 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr if self.q is not None: y -= tau * self.alpha * self.q if self.normaleqs: - Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( - tau * self.sigma - ) * (self.Op.H * self.Op) - x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] + if x.partition == Partition.SCATTER: + Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( + tau * self.sigma + ) * (self.Op.H * self.Op) + x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] + else: + # TODO: handle case of x BROADCAST + raise NotImplementedError( + "L2 proximal operator currently " + f"not supporter for {x.partition} partition") else: y = x if self.q is not None: diff --git a/tests/test_prox.py b/tests/test_prox.py index a4ea067c..ce9dc5c9 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -18,7 +18,7 @@ import pytest import pylops_mpi -from pylops.basicoperators import FirstDerivative +from pylops.basicoperators import Diagonal from pyproximal.proximal import ( Box, L0, @@ -117,30 +117,33 @@ def test_separable_prox(par): @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( - "par", [(par1), (par1b), (par1j),] + "par", [(par1), (par1j),] # (par1b), ) def test_L2(par): """L2 proximal operator""" np.random.seed(42) - x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], - partition=par['partition'], engine=backend) + x = pylops_mpi.DistributedArray(global_shape=par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), + dtype=par['dtype'], partition=par['partition'], engine=backend) x[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) x_global = x.asarray() - b = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], - partition=par['partition'], engine=backend) + b = pylops_mpi.DistributedArray(global_shape=par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), + dtype=par['dtype'], partition=par['partition'], engine=backend) b[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) b_global = b.asarray() - Op_global = FirstDerivative( - par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), - sampling=0.001) - Opd = pylops_mpi.MPIFirstDerivative( - par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), - sampling=0.001) + Op_global = Diagonal( + np.ones(par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), dtype=par['dtype']), + dtype=par['dtype']) + if pylops_mpi.Partition.SCATTER: + Op_local = Diagonal(np.ones(par['n'], dtype=par['dtype']), dtype=par['dtype']) + Opd = pylops_mpi.MPIBlockDiag([Op_local, ]) + else: + Op_local = Diagonal(np.ones(par['n'] * size, dtype=par['dtype']), dtype=par['dtype']) + Opd = pylops_mpi.MPILinearOperator(Op_local) l2x = L2(sigma=2.0) l2xd = MPIL2(sigma=2.0) @@ -148,11 +151,10 @@ def test_L2(par): l2b = L2(b=b_global, sigma=2.0) l2bd = MPIL2(b=b, sigma=2.0) - # l2Op = L2(Op=Op_global, b=b_global, sigma=2.0) - # l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0) + l2Op = L2(Op=Op_global, b=b_global, sigma=2.0, solver="cgls") + l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") - # for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): - for l2, l2d in zip([l2x, l2b,], [l2xd, l2bd,]): + for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): f = l2d(x) prox = l2d.prox(x, .1) prox = prox.asarray() @@ -161,4 +163,4 @@ def test_L2(par): f_np = l2(x_global) prox_np = l2.prox(x_global, .1) assert_allclose(f, f_np, rtol=1e-14) - assert_allclose(prox, prox_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-12) From e61ece9ef0cbf05b6e11c688a3cf20804ca76ba2 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Sun, 5 Jul 2026 22:38:23 +0000 Subject: [PATCH 12/47] fix: extract item from f --- pylops_mpi/proximal/proximal/L2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index b3b5a148..6a1a303b 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -119,7 +119,7 @@ def __call__(self, x: DistributedArray) -> DistributedArray: f = (self.sigma / 2.0) * (x.norm() ** 2) if self.q is not None: f += self.alpha * self.q.dot(x) - return float(f) + return float(f.item()) def _increment_count(func: Callable[..., Any]) -> Callable[..., Any]: From 62353b5a6117d4ec9db59f74a185aba7119c08df Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Sun, 5 Jul 2026 22:43:27 +0000 Subject: [PATCH 13/47] minor: fix flake8 --- pylops_mpi/proximal/ProxOperator.py | 10 +++++----- pylops_mpi/proximal/__init__.py | 2 +- pylops_mpi/proximal/optimization/primal.py | 17 ++++++----------- pylops_mpi/proximal/proximal/L2.py | 4 +--- 4 files changed, 13 insertions(+), 20 deletions(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 6404d177..07f93e15 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -1,5 +1,5 @@ from mpi4py import MPI -from typing import Any, Callable +from typing import Any from pyproximal import ProxOperator from pylops.utils.backend import get_module @@ -30,7 +30,7 @@ class MPIProxOperator: ---------- prox : :obj:`pyproximal.ProxOperator` PyProximal Proximal Operator to wrap. - + """ def __init__( @@ -52,7 +52,7 @@ def __repr__(self) -> str: return f"<{type(self).__name__} ({type(self.proxop).__name__})>" else: return f"<{type(self).__name__}>" - + def __call__(self, x: DistributedArray) -> DistributedArray: """Functional evaluation of the oprator. @@ -92,9 +92,9 @@ def __call__(self, x: DistributedArray) -> DistributedArray: else: # StackedDistributedArray reduce_op = _call_reduce_op[str(type(self.proxop).__name__)][1] fs = [self(x[iarr]) for iarr in range(x.narrays)] - f = reduce_op(fs) + f = reduce_op(fs) return f - + def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: """Proximal operator applied to a vector """ diff --git a/pylops_mpi/proximal/__init__.py b/pylops_mpi/proximal/__init__.py index c97dce78..e1a0ef48 100644 --- a/pylops_mpi/proximal/__init__.py +++ b/pylops_mpi/proximal/__init__.py @@ -26,4 +26,4 @@ __all__ = [ "MPIProxOperator", -] \ No newline at end of file +] diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index b5309e34..a79f67b2 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -2,20 +2,17 @@ import time from collections.abc import Callable from math import sqrt -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any import numpy as np -from pylops.basicoperators import Identity -from pylops.optimization.leastsquares import regularized_inversion -from pylops.utils.backend import get_array_module, to_numpy +from pylops.utils.backend import to_numpy from pylops.utils.typing import NDArray -from pyproximal.proximal import L2 from pyproximal.optimization.primal import _x0z0_init from pylops_mpi import DistributedArray, StackedDistributedArray -from pylops_mpi.basicoperators import MPIBlockDiag, MPIStackedVStack -from pylops_mpi.optimization.basic import cg, cgls +from pylops_mpi.basicoperators import MPIStackedVStack +from pylops_mpi.optimization.basic import cgls from pylops_mpi.proximal.ProxOperator import MPIProxOperator if TYPE_CHECKING: @@ -45,7 +42,7 @@ def ProximalGradient( # TODO: implement backtracking backtracking = False - + # check if epgs is a vector epsg = np.asarray(epsg, dtype=float) if epsg.size == 1: @@ -179,7 +176,6 @@ def ProximalGradient( return x - def ADMML2( proxg: MPIProxOperator, Op: "MPILinearOperator", @@ -198,7 +194,7 @@ def ADMML2( """ rank = x0.rank - + # initialize variables x, z = _x0z0_init(x0, z0, A, Opname="A") u = z.zeros_like() @@ -258,4 +254,3 @@ def ADMML2( print("---------------------------------------------------------\n") sys.stdout.flush() return x, z - diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 6a1a303b..c408bb5b 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -1,5 +1,4 @@ from math import sqrt -from mpi4py import MPI from typing import TYPE_CHECKING, Any, Callable from pylops.basicoperators import Identity @@ -101,7 +100,7 @@ def __init__( "Available options are 'cg' or 'cgls'." ) raise ValueError(msg) - + # create data term if ( self.Op is not None @@ -120,7 +119,6 @@ def __call__(self, x: DistributedArray) -> DistributedArray: if self.q is not None: f += self.alpha * self.q.dot(x) return float(f.item()) - def _increment_count(func: Callable[..., Any]) -> Callable[..., Any]: """Increment counter""" From b151dc88234d1871b86541d153a049d1b7bfa07e Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 20:43:22 +0000 Subject: [PATCH 14/47] feat: add error for L2 with broadcast --- pylops_mpi/proximal/proximal/L2.py | 21 ++++++++------- tests/test_prox.py | 42 +++++++++++++++++------------- 2 files changed, 35 insertions(+), 28 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index c408bb5b..3cb5b19c 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -134,6 +134,13 @@ def wrapped(self, *args: Any, **kwargs: Any) -> Any: def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: """Proximal operator applied to a vector """ + # check partition + if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: + raise NotImplementedError( + "L2 proximal operator not " + f"supported for {x.partition} partition" + ) + # define current number of iterations if isinstance(self.niter, int): niter = self.niter @@ -147,16 +154,10 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr if self.q is not None: y -= tau * self.alpha * self.q if self.normaleqs: - if x.partition == Partition.SCATTER: - Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( - tau * self.sigma - ) * (self.Op.H * self.Op) - x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] - else: - # TODO: handle case of x BROADCAST - raise NotImplementedError( - "L2 proximal operator currently " - f"not supporter for {x.partition} partition") + Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( + tau * self.sigma + ) * (self.Op.H * self.Op) + x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] else: y = x if self.q is not None: diff --git a/tests/test_prox.py b/tests/test_prox.py index ce9dc5c9..bcb5253d 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -40,26 +40,26 @@ "imag": 0, "dtype": np.float64, "partition": pylops_mpi.Partition.SCATTER -} - -par1b = { - "n": 101, - "imag": 0, - "dtype": np.float64, - "partition": pylops_mpi.Partition.BROADCAST -} +} # scatter, real par1j = { "n": 101, "imag": 1j, "dtype": np.complex128, "partition": pylops_mpi.Partition.SCATTER -} +} # scatter, complex + +par1b = { + "n": 101, + "imag": 0, + "dtype": np.float64, + "partition": pylops_mpi.Partition.BROADCAST +} # broadcast, real @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( - "par", [(par1), (par1b), (par1j),] + "par", [(par1), (par1j), (par1b)] ) def test_separable_prox(par): """Separable proximal operators""" @@ -117,10 +117,14 @@ def test_separable_prox(par): @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( - "par", [(par1), (par1j),] # (par1b), + "par", [(par1), (par1j), (par1b)] ) def test_L2(par): - """L2 proximal operator""" + """L2 proximal operator + + Test call/prox for L2 without Op/b (scatter and broadcast), + with b (scatter and broadcast), and with Op/b (only scatter) + """ np.random.seed(42) x = pylops_mpi.DistributedArray(global_shape=par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), @@ -138,7 +142,7 @@ def test_L2(par): Op_global = Diagonal( np.ones(par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), dtype=par['dtype']), dtype=par['dtype']) - if pylops_mpi.Partition.SCATTER: + if par["partition"] == pylops_mpi.Partition.SCATTER: Op_local = Diagonal(np.ones(par['n'], dtype=par['dtype']), dtype=par['dtype']) Opd = pylops_mpi.MPIBlockDiag([Op_local, ]) else: @@ -155,12 +159,14 @@ def test_L2(par): l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): - f = l2d(x) - prox = l2d.prox(x, .1) - prox = prox.asarray() + if par["partition"] == pylops_mpi.Partition.SCATTER: + f = l2d(x) + prox = l2d.prox(x, .1) + prox = prox.asarray() if rank == 0: f_np = l2(x_global) prox_np = l2.prox(x_global, .1) - assert_allclose(f, f_np, rtol=1e-14) - assert_allclose(prox, prox_np, rtol=1e-12) + if par["partition"] == pylops_mpi.Partition.SCATTER: + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-12) From 301a637e467b0b3bab4ffb72a0b3eca13898e37e Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 21:04:49 +0000 Subject: [PATCH 15/47] feat: more on skipping broadcast --- pylops_mpi/proximal/proximal/L2.py | 30 ++++++++++++++- tests/test_prox.py | 61 +++++++++++++++++++++++++----- 2 files changed, 80 insertions(+), 11 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 3cb5b19c..3cb6ea40 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -18,6 +18,16 @@ class MPIL2(MPIProxOperator): Implement a distributed version of the L2 norm proximal operator. + .. note:: + + This operator does not support input arrays with + ``pylops_mpi.Partition.BROADCAST`` or + ``pylops_mpi.Partition.UNSAFE_BROADCAST`` partition when + ``Op`` and ``b`` are not ``None`` because it does not make + sense to have both the model and data replicated across + ranks. + + Parameters ---------- Op : :obj:`pylops_mpi.MPILinearOperator`, optional @@ -38,7 +48,7 @@ class MPIL2(MPIProxOperator): counter which keeps track of how many times the ``prox`` method has been invoked before and returns the ``niter`` to be used. x0 : :obj:`pylops_mpi.DistributedArray`, optional - Initial vector. If ``Op`` is not None, this must be passed. + Initial vector. If ``Op`` is not ``None``, this must be passed. warm : :obj:`bool`, optional Warm start (``True``) or not (``False``). Uses estimate from previous call of ``prox`` method. @@ -73,6 +83,13 @@ def __init__( ) -> None: if Op is not None and x0 is None: raise ValueError("x0 must be passed when Op is not None") + # check partition + if Op is not None and x0.partition != Partition.SCATTER: + raise NotImplementedError( + "L2 proximal operator not " + f"supported for inputs with {x0.partition} " + "partition" + ) self.Op = Op self.hasgrad = True @@ -110,6 +127,14 @@ def __init__( self.OpTb = self.sigma * self.Op.H @ self.b def __call__(self, x: DistributedArray) -> DistributedArray: + # check partition + if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: + raise NotImplementedError( + "L2 proximal operator not " + f"supported for inputs with {x.partition} " + "partition" + ) + if self.Op is not None and self.b is not None: f = (self.sigma / 2.0) * ((self.Op * x - self.b).norm() ** 2) elif self.b is not None: @@ -138,7 +163,8 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: raise NotImplementedError( "L2 proximal operator not " - f"supported for {x.partition} partition" + f"supported for inputs with {x.partition} " + "partition" ) # define current number of iterations diff --git a/tests/test_prox.py b/tests/test_prox.py index bcb5253d..e0d8df55 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -115,6 +115,45 @@ def test_separable_prox(par): assert_allclose(prox, prox_np, rtol=1e-14) +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1b)] +) +def test_L2_broadcast(par): + """Test L2 proximal operator raises error with broadcast partition + + """ + x = pylops_mpi.DistributedArray(global_shape=par['n'], + dtype=par['dtype'], + partition=pylops_mpi.Partition.BROADCAST, + engine=backend) + b = pylops_mpi.DistributedArray(global_shape=par['n'], + dtype=par['dtype'], + partition=pylops_mpi.Partition.BROADCAST, + engine=backend) + + Opd = pylops_mpi.MPILinearOperator( + Diagonal(np.ones(par['n'] * size, dtype=par['dtype']), dtype=par['dtype']) + ) + + # creation + with pytest.raises(NotImplementedError, match="not supported for"): + _ = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") + + # call/prox + x0 = pylops_mpi.DistributedArray(global_shape=par['n'] * size, + dtype=par['dtype'], + partition=pylops_mpi.Partition.SCATTER, + engine=backend) + l2d = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x0, solver="cgls") + + with pytest.raises(NotImplementedError, match="not supported for"): + _ = l2d(x) + + with pytest.raises(NotImplementedError, match="not supported for"): + _ = l2d.prox(x, .1) + + @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( "par", [(par1), (par1j), (par1b)] @@ -155,18 +194,22 @@ def test_L2(par): l2b = L2(b=b_global, sigma=2.0) l2bd = MPIL2(b=b, sigma=2.0) - l2Op = L2(Op=Op_global, b=b_global, sigma=2.0, solver="cgls") - l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") + if par["partition"] == pylops_mpi.Partition.SCATTER: + l2Op = L2(Op=Op_global, b=b_global, sigma=2.0, solver="cgls") + l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") + else: + l2Op = l2Opd = None # to skip tests with broadcast for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): - if par["partition"] == pylops_mpi.Partition.SCATTER: - f = l2d(x) - prox = l2d.prox(x, .1) - prox = prox.asarray() + if l2 is None: + continue + + f = l2d(x) + prox = l2d.prox(x, .1) + prox = prox.asarray() if rank == 0: f_np = l2(x_global) prox_np = l2.prox(x_global, .1) - if par["partition"] == pylops_mpi.Partition.SCATTER: - assert_allclose(f, f_np, rtol=1e-14) - assert_allclose(prox, prox_np, rtol=1e-12) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox, prox_np, rtol=1e-12) From 48c444c06e6674765d847cf981d8c87791b34854 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 21:31:27 +0000 Subject: [PATCH 16/47] doc: finalized plot_prox --- examples/plot_prox.py | 192 ++++++++++-------------------------------- 1 file changed, 44 insertions(+), 148 deletions(-) diff --git a/examples/plot_prox.py b/examples/plot_prox.py index 82b84d4e..c42dd3a9 100644 --- a/examples/plot_prox.py +++ b/examples/plot_prox.py @@ -58,20 +58,27 @@ print("prox_||x||_1: ", all(proxdlocal == proxlocal)) print("proxd_||x||_1: ", all(dproxdlocal == dproxlocal)) -# Box norm +############################################################################### +# We repeat now the same with the :py:class:`pyproximal.proximal.Box` operator. + arr = pylops_mpi.DistributedArray(global_shape=n * size, partition=pylops_mpi.Partition.SCATTER) arr[:] = 3 * np.ones(n) if rank == 0: - arr[n//2] = 20 + arr[n//2] = 20 # outside of the box box = pyproximal.Box(lower=1., upper=5.) boxd = pylops_mpi.proximal.MPIProxOperator(box) + +# Call f = boxd(arr) + +# Proximal prox = boxd.prox(arr, .1) proxdlocal = prox.asarray() +# Dual-Proximal dprox = boxd.proxdual(arr, .1) dproxdlocal = dprox.asarray() @@ -84,8 +91,16 @@ print("prox_Box ", all(proxdlocal == proxlocal)) print("proxd_Box ", all(dproxdlocal == dproxlocal)) +############################################################################### +# We move on now to a operator that is not separable and must be fully +# re-implemented in a distributed fashion, namely the +# :py:class:`pylops_mpi.proximal.MPIL2` norm. +# +# More precisely, when ``Op`` and ``b`` are passed to this operator, +# its proximal does call for the solution of a distributed inverse problem. +# +# However, let's start with the simplest case: :math:`||\mathbf{x}||_2^2` -# L2 norm ||x||_2^2 arr = pylops_mpi.DistributedArray(global_shape=n * size, partition=pylops_mpi.Partition.SCATTER) @@ -93,9 +108,15 @@ l2 = pyproximal.L2(sigma=2.0) l2d = pylops_mpi.proximal.MPIL2(sigma=2.0) + +# Call f = l2d(arr) + +# Proximal prox = l2d.prox(arr, .1) proxdlocal = prox.asarray() + +# Gradient grad = l2d.grad(arr) graddlocal = grad.asarray() @@ -108,16 +129,16 @@ print("prox_||x||_2^2: ", all(proxdlocal == proxlocal)) print("grad_||x||_2^2: ", all(graddlocal == gradlocal)) -# L2 norm ||Op * x - d||_2^2 +############################################################################### +# Next we move onto the more general case +# :math:`||\mathbf{Op} \mathbf{x} - \mathbf{b}||_2^2` + solver="cgls" -Op = pylops.FirstDerivative(n * size, sampling=0.001) -Opd = pylops_mpi.MPIFirstDerivative(n * size, sampling=0.001) -# Op = pylops.Diagonal(np.ones(n * size)) -# Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(n)),]) +Op = pylops.Diagonal(np.ones(n * size)) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(n)),]) b = pylops_mpi.DistributedArray(global_shape=n * size, partition=pylops_mpi.Partition.SCATTER) - b[:] = rank * np.ones(n) blocal = b.asarray() @@ -126,13 +147,21 @@ l2 = pyproximal.L2( Op=Op, b=blocal, sigma=2.0, - solver=solver, x0=x0local, kwargs_solver=dict(show=True)) + solver=solver, x0=x0local, + kwargs_solver=dict(show=True)) l2d = pylops_mpi.proximal.MPIL2( Op=Opd, b=b, sigma=2.0, - solver=solver, x0=x0, kwargs_solver=dict(show=True if rank==0 else False)) + solver=solver, x0=x0, + kwargs_solver=dict(show=True if rank==0 else False)) + +# Call f = l2d(arr) + +# Proximal prox = l2d.prox(arr, .1) proxdlocal = prox.asarray() + +# Gradient grad = l2d.grad(arr) graddlocal = grad.asarray() @@ -141,140 +170,7 @@ flocal = l2(arrlocal) proxlocal = l2.prox(arrlocal, .1) gradlocal = l2.grad(arrlocal) - print("||x||_2^2: ", f, flocal) - print("prox_||x||_2^2: ", all(proxdlocal == proxlocal), np.linalg.norm(proxdlocal - proxlocal)) - print("grad_||x||_2^2: ", all(graddlocal == gradlocal)) - - - -# Proximal gradient -arr = pylops_mpi.DistributedArray(global_shape=n, - partition=pylops_mpi.Partition.BROADCAST) -arr[:] = 0.0 -arr[n//4] = 1.0 -arr[n//2] = -0.5 - -Op = pylops.MatrixMult(np.random.normal(0, 1, (n-2, n,))) -Opd = pylops_mpi.MPILinearOperator(Op) - -b = Opd @ arr -blocal = b.asarray() - -l2d = pylops_mpi.proximal.MPIL2( - Op=Opd, b=b, solver=solver, x0=arr.zeros_like()) -l1 = pyproximal.L1(sigma=8e-1) -l1d = pylops_mpi.proximal.MPIProxOperator(l1) - -arrpg = pylops_mpi.proximal.optimization.primal.ProximalGradient( - l2d, l1d, x0=arr.zeros_like(), tau=1e-2, niter=400, - show=True - ) -arrpgdlocal = arrpg.asarray() - -arrlocal = arr.asarray() -if rank == 0: - l2local = pyproximal.L2( - Op=Op, b=blocal, - solver=solver) - l1local = pyproximal.L1(sigma=8e-1) - - arrpglocal = pyproximal.optimization.primal.ProximalGradient( - l2local, l1local, x0=np.zeros(n), tau=1e-2, niter=400, show=True - ) - - print('PG true', arrlocal) - print('PG distr', arrpgdlocal) - print('PG local', arrpglocal) - - -# ADMML2 with stacked operator -ny, nx = 40, 40 -arrlocal = np.ones((ny, nx)) -arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 -arr = pylops_mpi.DistributedArray(global_shape=ny * nx, - partition=pylops_mpi.Partition.SCATTER) -arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() - -Op = pylops.Diagonal(np.ones(ny*nx)) -Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) - -b = Opd @ arr -blocal = b.asarray() - -Iop = pylops.Identity(ny*nx) -Iopd = pylops_mpi.MPIBlockDiag([pylops.Identity(ny*nx//size),]) - -L = 8.0 # maxeig(Gop^H Gop) - -l1 = pyproximal.L1(sigma=8e-1) -l1d = pylops_mpi.proximal.MPIProxOperator(l1) - -x0distr = arr.zeros_like() -arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, Opd, b, Iopd, x0=x0distr, tau=.99/L, niter=5, - show=True, kwargs_solver=dict(niter=20), - )[0] -arradmmdlocal = arradmm.asarray() - -arrlocal = arr.asarray() -if rank == 0: - - l1local = pyproximal.L1(sigma=8e-1) - - arradmmlocal = pyproximal.optimization.primal.ADMML2( - l1local, Op, blocal, Iop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=True, iter_lim=20, - )[0] - - print('ADMML2 true', arrlocal) - print('ADMML2 distr', arradmmdlocal) - print('ADMML2 local', arradmmlocal) - - -# ADMML2 with stacked operator for A -ny, nx = 40, 40 -arrlocal = np.ones((ny, nx)) -arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 -arr = pylops_mpi.DistributedArray(global_shape=ny * nx, - partition=pylops_mpi.Partition.SCATTER) -arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() - -Op = pylops.Diagonal(np.ones(ny*nx)) -Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) - -b = Opd @ arr -blocal = b.asarray() - -Gopd = pylops_mpi.MPIGradient( - dims=(ny, nx), sampling=1., edge=False, kind="forward") - -L = 8.0 # maxeig(Gop^H Gop) - -l1 = pyproximal.L1(sigma=8e-1) -l1d = pylops_mpi.proximal.MPIProxOperator(l1) - -x0distr = arr.zeros_like() -arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, - show=True, kwargs_solver=dict(niter=5), - )[0] -arradmmdlocal = arradmm.asarray() - -arrlocal = arr.asarray() -if rank == 0: - - Gop = pylops.Gradient( - dims=(ny, nx), sampling=1., edge=False, kind="forward", - ) - l1local = pyproximal.L1(sigma=8e-1) - - arradmmlocal = pyproximal.optimization.primal.ADMML2( - l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=True, iter_lim=5, - )[0] - - print('ADMML2 true', arrlocal) - print('ADMML2 distr', arradmmdlocal) - print('ADMML2 local', arradmmlocal) - print(arradmmdlocal - arradmmlocal) - + print("||Op . x - b||_2^2: ", f, flocal) + print("prox_||Op . x - b||_2^2 - norm diff=", + np.linalg.norm(proxdlocal - proxlocal)) + print("grad_||Op . x - b||_2^2: ", all(graddlocal == gradlocal)) From a043afc512949e227e351138764cb463384e237c Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 21:33:20 +0000 Subject: [PATCH 17/47] fix: changed output of MPIProxOperator.__call__ to be scalar --- pylops_mpi/proximal/ProxOperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 07f93e15..f7321bfd 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -85,7 +85,7 @@ def __call__(self, x: DistributedArray) -> DistributedArray: ncp.asarray(f), op=reduce_op, engine=x.engine) - return recv_buf + return recv_buf[0] else: # For broadcasted arrays, simply return the local f return f From bd485513985e1529d1673953a59c7f553fe608c1 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 21:34:04 +0000 Subject: [PATCH 18/47] doc: added TV inversion to poststack_cupy --- tutorials/poststack.py | 7 ++--- tutorials_cupy/poststack_cupy.py | 44 +++++++++++++++++++++++++++----- 2 files changed, 41 insertions(+), 10 deletions(-) diff --git a/tutorials/poststack.py b/tutorials/poststack.py index 7c0094e0..97016768 100644 --- a/tutorials/poststack.py +++ b/tutorials/poststack.py @@ -199,8 +199,8 @@ minv3d_ne = minv3d_ne_dist.asarray().reshape((ny, nx, nz)) ############################################################################### - # Regularized inversion with regularized equations + StackOp = pylops_mpi.MPIStackedVStack([BDiag, np.sqrt(epsR) * LapOp]) d0_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz) d0_dist[:] = 0. @@ -212,8 +212,8 @@ minv3d_reg = minv3d_reg_dist.asarray().reshape((ny, nx, nz)) ############################################################################### +# TV-Regularized inversion -# Inversion with TV Gopd = pylops_mpi.MPIGradient( dims=(ny, nx, nz), sampling=1., edge=False, kind="forward") @@ -240,7 +240,8 @@ print('Distr == Local', np.allclose(d, d0)) # Visualize - fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(9, 17), constrained_layout=True) + fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(12, 18), + constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") diff --git a/tutorials_cupy/poststack_cupy.py b/tutorials_cupy/poststack_cupy.py index 16f53a4f..0eef4f51 100644 --- a/tutorials_cupy/poststack_cupy.py +++ b/tutorials_cupy/poststack_cupy.py @@ -15,6 +15,7 @@ from pylops.utils.wavelets import ricker from pylops.basicoperators import Transpose from pylops.avo.poststack import PoststackLinearModelling +from pyproximal.proximal import L1 import pylops_mpi @@ -116,8 +117,8 @@ minv3d_iter = minv3d_iter_dist.asarray().reshape((ny, nx, nz)) ############################################################################### - # Regularized inversion with normal equations + epsR = 1e2 LapOp = pylops_mpi.MPILaplacian(dims=(ny, nx, nz), axes=(0, 1, 2), weights=(1, 1, 1), @@ -131,8 +132,8 @@ minv3d_ne = minv3d_ne_dist.asarray().reshape((ny, nx, nz)) ############################################################################### - # Regularized inversion with regularized equations + StackOp = pylops_mpi.MPIStackedVStack([BDiag, np.sqrt(epsR) * LapOp]) d0_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz, engine="cupy") d0_dist[:] = 0. @@ -144,6 +145,24 @@ niter=100, show=True)[0] minv3d_reg = minv3d_reg_dist.asarray().reshape((ny, nx, nz)) + +############################################################################### +# TV-Regularized inversion + +Gopd = pylops_mpi.MPIGradient( + dims=(ny, nx, nz), sampling=1., edge=False, kind="forward") + +L = 12.0 # maxeig(Gop^H Gop) + +l1 = L1(sigma=1e-2) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +minv3d_admm_dist = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99/L, niter=40, + show=True, kwargs_solver=dict(niter=20), + )[0] +minv3d_admm = minv3d_admm_dist.asarray().reshape((ny, nx, nz)) + ############################################################################### # Finally we visualize the results. Note that the array must be copied back # to the CPU by calling the :code:`get()` method on the CuPy arrays. @@ -158,7 +177,8 @@ print('Distr == Local', np.allclose(cp.asnumpy(d), d0, atol=1e-6)) # Visualize - fig, axs = plt.subplots(nrows=6, ncols=3, figsize=(9, 14), constrained_layout=True) + fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(12, 18), + constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") @@ -203,18 +223,28 @@ axs[4][0].set_title("Normal Equations Inverted Model iter x-z") axs[4][0].axis("tight") axs[4][1].imshow(minv3d_ne[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[4][1].set_title('Normal Equations Inverted Model iter y-z') + axs[4][1].set_title("Normal Equations Inverted Model iter y-z") axs[4][1].axis('tight') axs[4][2].imshow(minv3d_ne[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[4][2].set_title('Normal Equations Inverted Model iter x-y') + axs[4][2].set_title("Normal Equations Inverted Model iter x-y") axs[4][2].axis('tight') axs[5][0].imshow(minv3d_reg[5, :, :].T.get(), cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[5][0].set_title("Regularized Inverted Model iter x-z") axs[5][0].axis("tight") axs[5][1].imshow(minv3d_reg[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][1].set_title('Regularized Inverted Model iter y-z') + axs[5][1].set_title("Regularized Inverted Model iter y-z") axs[5][1].axis('tight') axs[5][2].imshow(minv3d_reg[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][2].set_title('Regularized Inverted Model iter x-y') + axs[5][2].set_title("Regularized Inverted Model iter x-y") axs[5][2].axis('tight') + + axs[6][0].imshow(minv3d_admm[5, :, :].T.get(), cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) + axs[6][0].set_title("TV-Regularized Inverted Model iter x-z") + axs[6][0].axis("tight") + axs[6][1].imshow(minv3d_admm[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][1].set_title("TV-Regularized Inverted Model iter y-z") + axs[6][1].axis('tight') + axs[6][2].imshow(minv3d_admm[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][2].set_title("TV-Regularized Inverted Model iter x-y") + axs[6][2].axis('tight') From 417d864f4e9adeef0d0b4039940535d6789cc336 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 21:38:39 +0000 Subject: [PATCH 19/47] minor: fix flake8 --- pylops_mpi/proximal/optimization/__init__.py | 3 +- pylops_mpi/proximal/optimization/primal.py | 30 ++++---------------- pylops_mpi/proximal/proximal/__init__.py | 2 +- 3 files changed, 8 insertions(+), 27 deletions(-) diff --git a/pylops_mpi/proximal/optimization/__init__.py b/pylops_mpi/proximal/optimization/__init__.py index 6899fd92..199fbfba 100644 --- a/pylops_mpi/proximal/optimization/__init__.py +++ b/pylops_mpi/proximal/optimization/__init__.py @@ -16,4 +16,5 @@ __all__ = [ "ProximalGradient", -] \ No newline at end of file + "ADMML2", +] diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index a79f67b2..3481046d 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -25,7 +25,6 @@ def ProximalGradient( x0: DistributedArray, epsg: float | NDArray = 1.0, tau: float | None = None, - # backtracking: bool = False, beta: float = 0.5, eta: float = 1.0, niter: int = 10, @@ -40,7 +39,6 @@ def ProximalGradient( """ rank = x0.rank - # TODO: implement backtracking backtracking = False # check if epgs is a vector @@ -82,10 +80,6 @@ def ProximalGradient( print(head) sys.stdout.flush() - # if tau is None: - # backtracking = True - # tau = 1.0 - # initialize model t = 1.0 x = x0.copy() @@ -98,26 +92,12 @@ def ProximalGradient( xold = x.copy() # proximal step - if not backtracking: - if eta == 1.0: - x = proxg.prox(y - tau * proxf.grad(y), epsg[iiter] * tau) - else: - x = x + eta * ( - proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x - ) + if eta == 1.0: + x = proxg.prox(y - tau * proxf.grad(y), epsg[iiter] * tau) else: - pass - # x, tau = _backtracking( - # y, tau, proxf, proxg, epsg[iiter], beta=beta, niterback=niterback - # ) - # if eta != 1.0: - # x = x + eta * ( - # proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x - # ) - - # update internal parameters for bilinear operator - # if isinstance(proxf, BilinearOperator): - # proxf.updatexy(x) + x = x + eta * ( + proxg.prox(x - tau * proxf.grad(x), epsg[iiter] * tau) - x + ) # update y if acceleration == "vandenberghe": diff --git a/pylops_mpi/proximal/proximal/__init__.py b/pylops_mpi/proximal/proximal/__init__.py index d5902cdd..57d990fb 100644 --- a/pylops_mpi/proximal/proximal/__init__.py +++ b/pylops_mpi/proximal/proximal/__init__.py @@ -18,4 +18,4 @@ __all__ = [ "MPIL2", -] \ No newline at end of file +] From 7826d50095c6623b35ee4d1c676e7aabdabbf1e6 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:16:46 +0000 Subject: [PATCH 20/47] fix: revert change --- pylops_mpi/proximal/ProxOperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index f7321bfd..07f93e15 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -85,7 +85,7 @@ def __call__(self, x: DistributedArray) -> DistributedArray: ncp.asarray(f), op=reduce_op, engine=x.engine) - return recv_buf[0] + return recv_buf else: # For broadcasted arrays, simply return the local f return f From f78e01e9408cf4ecd2ee0759a177b5e49e3d273f Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:18:07 +0000 Subject: [PATCH 21/47] minor: removed backtracking from ProximalGradient --- pylops_mpi/proximal/optimization/primal.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index 3481046d..f8d9c3ed 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -39,8 +39,6 @@ def ProximalGradient( """ rank = x0.rank - backtracking = False - # check if epgs is a vector epsg = np.asarray(epsg, dtype=float) if epsg.size == 1: @@ -59,7 +57,7 @@ def ProximalGradient( "---------------------------------------------------------\n" "Proximal operator (f): %s\n" "Proximal operator (g): %s\n" - "tau = %s\tbacktrack = %s\tbeta = %10e\n" + "tau = %s\tbeta = %10e\n" "epsg = %s\tniter = %d\ttol = %s\n" "" "niterback = %d\tacceleration = %s\n" @@ -67,7 +65,6 @@ def ProximalGradient( proxf, proxg, str(tau), - backtracking, beta, epsg_print, niter, From 7eafb19d07ece1fcdc95f98493ddda5e597cd25f Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:18:33 +0000 Subject: [PATCH 22/47] minor: reverted raises on broadcast in L2 --- pylops_mpi/proximal/proximal/L2.py | 30 ++---------------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 3cb6ea40..fa1d52f1 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -18,16 +18,6 @@ class MPIL2(MPIProxOperator): Implement a distributed version of the L2 norm proximal operator. - .. note:: - - This operator does not support input arrays with - ``pylops_mpi.Partition.BROADCAST`` or - ``pylops_mpi.Partition.UNSAFE_BROADCAST`` partition when - ``Op`` and ``b`` are not ``None`` because it does not make - sense to have both the model and data replicated across - ranks. - - Parameters ---------- Op : :obj:`pylops_mpi.MPILinearOperator`, optional @@ -83,13 +73,6 @@ def __init__( ) -> None: if Op is not None and x0 is None: raise ValueError("x0 must be passed when Op is not None") - # check partition - if Op is not None and x0.partition != Partition.SCATTER: - raise NotImplementedError( - "L2 proximal operator not " - f"supported for inputs with {x0.partition} " - "partition" - ) self.Op = Op self.hasgrad = True @@ -127,14 +110,6 @@ def __init__( self.OpTb = self.sigma * self.Op.H @ self.b def __call__(self, x: DistributedArray) -> DistributedArray: - # check partition - if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: - raise NotImplementedError( - "L2 proximal operator not " - f"supported for inputs with {x.partition} " - "partition" - ) - if self.Op is not None and self.b is not None: f = (self.sigma / 2.0) * ((self.Op * x - self.b).norm() ** 2) elif self.b is not None: @@ -162,9 +137,8 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr # check partition if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: raise NotImplementedError( - "L2 proximal operator not " - f"supported for inputs with {x.partition} " - "partition" + "L2 proximal operator currently not implemented " + f"for inputs with {x.partition} partition" ) # define current number of iterations From e4344a1f61aadeba8f2e4c07ee216f4b8ff59fe4 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:19:10 +0000 Subject: [PATCH 23/47] doc: added proximal methods to doc --- docs/source/api/index.rst | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index d3da2779..2eab6e90 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -49,7 +49,6 @@ Basic Operators MPIStackedVStack MPIHStack MPIHalo - Derivatives ~~~~~~~~~~~ @@ -88,6 +87,30 @@ Wave-Equation processing MPIMDC +Proximal operators +------------------ + +Templates +~~~~~~~~~ + +.. currentmodule:: pylops_mpi.proximal + +.. autosummary:: + :toctree: generated/ + + MPIProxOperator + +Basic Operators +~~~~~~~~~~~~~~~ + +.. currentmodule:: pylops_mpi.proximal.proximal + +.. autosummary:: + :toctree: generated/ + + MPIL2 + + Solvers ------- @@ -129,6 +152,19 @@ Sparsity ista fista + +Proximal +~~~~~~~~ + +.. currentmodule:: pylops_mpi.proximal.optimization.primal + +.. autosummary:: + :toctree: generated/ + + ProximalGradient + ADMML2 + + Utils ----- From 7ef61ec30bea0dea893cd0cbd64a5bc0a4c8699a Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:21:45 +0000 Subject: [PATCH 24/47] minor: fix flake8 --- pylops_mpi/proximal/optimization/primal.py | 2 +- pylops_mpi/proximal/proximal/L2.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index f8d9c3ed..256fbfbe 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -38,7 +38,7 @@ def ProximalGradient( """ rank = x0.rank - + # check if epgs is a vector epsg = np.asarray(epsg, dtype=float) if epsg.size == 1: diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index fa1d52f1..2b125db4 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -157,12 +157,12 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( tau * self.sigma ) * (self.Op.H * self.Op) - x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] + x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] else: y = x if self.q is not None: y -= tau * self.alpha * self.q - + Opreg = MPIStackedVStack([ sqrt(tau * self.sigma) * self.Op, MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, ),])]) From 0cdcf6d2cbe098d73b5ad5b7d9110d2dd37ebc1e Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 13 Jul 2026 22:29:59 +0000 Subject: [PATCH 25/47] fix: removed test for L2 broadcast --- tests/test_prox.py | 39 --------------------------------------- 1 file changed, 39 deletions(-) diff --git a/tests/test_prox.py b/tests/test_prox.py index e0d8df55..553bf33d 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -115,45 +115,6 @@ def test_separable_prox(par): assert_allclose(prox, prox_np, rtol=1e-14) -@pytest.mark.mpi(min_size=2) -@pytest.mark.parametrize( - "par", [(par1b)] -) -def test_L2_broadcast(par): - """Test L2 proximal operator raises error with broadcast partition - - """ - x = pylops_mpi.DistributedArray(global_shape=par['n'], - dtype=par['dtype'], - partition=pylops_mpi.Partition.BROADCAST, - engine=backend) - b = pylops_mpi.DistributedArray(global_shape=par['n'], - dtype=par['dtype'], - partition=pylops_mpi.Partition.BROADCAST, - engine=backend) - - Opd = pylops_mpi.MPILinearOperator( - Diagonal(np.ones(par['n'] * size, dtype=par['dtype']), dtype=par['dtype']) - ) - - # creation - with pytest.raises(NotImplementedError, match="not supported for"): - _ = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") - - # call/prox - x0 = pylops_mpi.DistributedArray(global_shape=par['n'] * size, - dtype=par['dtype'], - partition=pylops_mpi.Partition.SCATTER, - engine=backend) - l2d = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x0, solver="cgls") - - with pytest.raises(NotImplementedError, match="not supported for"): - _ = l2d(x) - - with pytest.raises(NotImplementedError, match="not supported for"): - _ = l2d.prox(x, .1) - - @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( "par", [(par1), (par1j), (par1b)] From 195929c3844d613f4d636b75c05e2d8e084dc373 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:21:30 +0000 Subject: [PATCH 26/47] feat: added support for broadcast partition in MPIL2 --- pylops_mpi/proximal/proximal/L2.py | 24 ++++++++++++------------ tests/test_prox.py | 19 ++++++++----------- 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 2b125db4..d7912f22 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -5,6 +5,7 @@ from pyproximal.ProxOperator import _check_tau from pylops_mpi import DistributedArray, StackedDistributedArray, Partition +from pylops_mpi import MPILinearOperator from pylops_mpi.basicoperators import MPIBlockDiag, MPIStackedVStack from pylops_mpi.optimization.basic import cg, cgls from pylops_mpi.proximal import MPIProxOperator @@ -134,13 +135,6 @@ def wrapped(self, *args: Any, **kwargs: Any) -> Any: def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray: """Proximal operator applied to a vector """ - # check partition - if self.Op is not None and self.b is not None and x.partition != Partition.SCATTER: - raise NotImplementedError( - "L2 proximal operator currently not implemented " - f"for inputs with {x.partition} partition" - ) - # define current number of iterations if isinstance(self.niter, int): niter = self.niter @@ -154,18 +148,24 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr if self.q is not None: y -= tau * self.alpha * self.q if self.normaleqs: - Op1 = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + float( - tau * self.sigma - ) * (self.Op.H * self.Op) + if x.partition == Partition.SCATTER: + Iop = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, )]) + else: + Iop = MPILinearOperator(Identity(x.local_shape, dtype=self.Op.dtype, )) + Op1 = Iop + float(tau * self.sigma) * (self.Op.H * self.Op) x = cg(Op1, y, niter=niter, x0=self.x0, **self.kwargs_solver)[0] else: y = x if self.q is not None: y -= tau * self.alpha * self.q - + if x.partition == Partition.SCATTER: + Iop = MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, ),]) + else: + Iop = MPILinearOperator(Identity(x.local_shape, dtype=self.Op.dtype, )) Opreg = MPIStackedVStack([ sqrt(tau * self.sigma) * self.Op, - MPIBlockDiag([Identity(x.local_shape, dtype=self.Op.dtype, ),])]) + Iop, + ]) breg = StackedDistributedArray([sqrt(tau * self.sigma) * self.b, y]) x = cgls(Opreg, breg, x0=self.x0, niter=niter, **self.kwargs_solver)[0] if self.warm: diff --git a/tests/test_prox.py b/tests/test_prox.py index 553bf33d..34df73fa 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -119,7 +119,10 @@ def test_separable_prox(par): @pytest.mark.parametrize( "par", [(par1), (par1j), (par1b)] ) -def test_L2(par): +@pytest.mark.parametrize( + "solver", ["cg", "cgls"] +) +def test_L2(par, solver): """L2 proximal operator Test call/prox for L2 without Op/b (scatter and broadcast), @@ -146,7 +149,7 @@ def test_L2(par): Op_local = Diagonal(np.ones(par['n'], dtype=par['dtype']), dtype=par['dtype']) Opd = pylops_mpi.MPIBlockDiag([Op_local, ]) else: - Op_local = Diagonal(np.ones(par['n'] * size, dtype=par['dtype']), dtype=par['dtype']) + Op_local = Diagonal(np.ones(par['n'], dtype=par['dtype']), dtype=par['dtype']) Opd = pylops_mpi.MPILinearOperator(Op_local) l2x = L2(sigma=2.0) @@ -155,16 +158,10 @@ def test_L2(par): l2b = L2(b=b_global, sigma=2.0) l2bd = MPIL2(b=b, sigma=2.0) - if par["partition"] == pylops_mpi.Partition.SCATTER: - l2Op = L2(Op=Op_global, b=b_global, sigma=2.0, solver="cgls") - l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver="cgls") - else: - l2Op = l2Opd = None # to skip tests with broadcast - + l2Op = L2(Op=Op_global, b=b_global, sigma=2.0, solver=solver) + l2Opd = MPIL2(Op=Opd, b=b, sigma=2.0, x0=x.zeros_like(), solver=solver) + for l2, l2d in zip([l2x, l2b, l2Op], [l2xd, l2bd, l2Opd]): - if l2 is None: - continue - f = l2d(x) prox = l2d.prox(x, .1) prox = prox.asarray() From ab0dda0ceff86aeeb6b21c88570e68167a3a2daf Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:21:47 +0000 Subject: [PATCH 27/47] doc: added example with proximal solvers --- examples/plot_proxsolver.py | 196 ++++++++++++++++++++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 examples/plot_proxsolver.py diff --git a/examples/plot_proxsolver.py b/examples/plot_proxsolver.py new file mode 100644 index 00000000..f2d7bddb --- /dev/null +++ b/examples/plot_proxsolver.py @@ -0,0 +1,196 @@ +r""" +Proximal solvers +================ + +This example demonstrates the use of the solvers in the +:py:module:`pylops_mpi.proximal.optimization` module. + +""" +import numpy as np +from mpi4py import MPI +from matplotlib import pyplot as plt + +import pylops +import pyproximal + +import pylops_mpi + +np.random.seed(42) +plt.close("all") +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() + +############################################################################### +# Let's start with an example of sparsity promoting inversion using the +# :py:class:`pylops_mpi.proximal.optimization.primal.ProximalGradient` solver. +# Here for illustrative purposes, we consider a case where the model is +# broadcasted whilst the data is scattered across ranks. + +# Sparse input +n = 16 +arr = pylops_mpi.DistributedArray(global_shape=n, + partition=pylops_mpi.Partition.BROADCAST) +arr[:] = 0.0 +arr[n//4] = 1.0 +arr[n//2] = -0.5 + +# Operator and data +A = np.random.normal(0, 1, (n // size, n,)) +As = np.vstack(comm.allgather(A)) + +Op = pylops.MatrixMult(As) +Opd = pylops_mpi.MPIVStack([pylops.MatrixMult(A), ]) + +b = Opd @ arr +blocal = b.asarray() + +# L2 prox +l2d = pylops_mpi.proximal.MPIL2( + Op=Opd, b=b, x0=arr.zeros_like()) + +# L1 prox +l1 = pyproximal.L1(sigma=1e-1) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +# Distributed inversion +arrpg = pylops_mpi.proximal.optimization.primal.ProximalGradient( + l2d, l1d, x0=arr.zeros_like(), tau=1e-2, niter=400, + show=True, + ) +arrpgdlocal = arrpg.asarray() + +# Benchmark serial inversion +arrlocal = arr.asarray() +if rank == 0: + l2local = pyproximal.L2( + Op=Op, b=blocal) + l1local = pyproximal.L1(sigma=1e-1) + + arrpglocal = pyproximal.optimization.primal.ProximalGradient( + l2local, l1local, x0=np.zeros(n), tau=1e-2, niter=400, show=False + ) + + plt.figure(figsize=(12, 3)) + plt.plot(arrlocal, "k", label="True") + plt.plot(arrpgdlocal, "b", label="Distr") + plt.plot(arrpglocal, "--r", label="Local") + plt.legend() + plt.savefig('pg.png') + +############################################################################### +# Next we use the :py:class:`pylops_mpi.proximal.optimization.primal.ADMML2` +# solver for a similar problem. However we consider here a 2d array and impose +# blockiness in the solution. Once again, the model is broadcasted, +# whilst the data is scattered. + +# Input +ny, nx = 10 * size, 40 +arrlocal = np.zeros((ny, nx)) +arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 +arr = pylops_mpi.DistributedArray(global_shape=ny * nx, + partition=pylops_mpi.Partition.BROADCAST) +arr[:] = arrlocal.flatten() + +# Operator and data +Op = pylops.VStack([pylops.Diagonal(np.ones(ny*nx)) for _ in range(size)]) +Opd = pylops_mpi.MPIVStack([pylops.Diagonal(np.ones(ny*nx)),]) + +b = Opd @ arr +blocal = b.asarray() + +# Regularizer +Gopd = pylops_mpi.MPILinearOperator(pylops.Gradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward")) + +l1 = pyproximal.L1(sigma=2e0) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +# Distributed inversion +L = 8.0 # max eig of Gopd.H @ Gop +x0distr = arr.zeros_like() +arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, + show=True, kwargs_solver=dict(niter=5), + )[0] +arradmmdlocal = arradmm.asarray() + +# Benchmark serial inversion +arrlocal = arr.asarray() +if rank == 0: + + Gop = pylops.Gradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward", + ) + l1local = pyproximal.L1(sigma=2e0) + + arradmmlocal = pyproximal.optimization.primal.ADMML2( + l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), + tau=.99/L, niter=5, show=True, iter_lim=5, + )[0] + + fig, axs = plt.subplots(1, 3, figsize=(12, 6)) + axs[0].imshow(arrlocal.reshape(10 * size, 40)) + axs[0].set_title("True") + axs[1].imshow(arradmmdlocal.reshape(10 * size, 40)) + axs[1].set_title("ADMML2 distr") + axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) + axs[2].set_title("ADMML2 local") + fig.savefig('solver.png') + +############################################################################### +# And finally we repeat the same with a scattered model. + +# Input +ny, nx = 10 * size, 40 +arrlocal = np.zeros((ny, nx)) +arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 +arr = pylops_mpi.DistributedArray(global_shape=ny * nx, + partition=pylops_mpi.Partition.SCATTER) +arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() + +# Operator and data +Op = pylops.Diagonal(np.ones(ny*nx)) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) + +b = Opd @ arr +blocal = b.asarray() + +# Regularizer +Gopd = pylops_mpi.MPIGradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward") + +l1 = pyproximal.L1(sigma=2e0) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +# Distributed inversion +L = 8.0 # max eig of Gopd.H @ Gop +x0distr = arr.zeros_like() +arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, + show=True, kwargs_solver=dict(niter=5), + )[0] +arradmmdlocal = arradmm.asarray() + +# Benchmark serial inversion +arrlocal = arr.asarray() +if rank == 0: + + Gop = pylops.Gradient( + dims=(ny, nx), sampling=1., edge=False, kind="forward", + ) + l1local = pyproximal.L1(sigma=2e0) + + arradmmlocal = pyproximal.optimization.primal.ADMML2( + l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), + tau=.99/L, niter=5, show=True, iter_lim=5, + )[0] + + fig, axs = plt.subplots(1, 3, figsize=(12, 6)) + axs[0].imshow(arrlocal.reshape(10 * size, 40)) + axs[0].set_title("True") + axs[1].imshow(arradmmdlocal.reshape(10 * size, 40)) + axs[1].set_title("ADMML2 distr") + axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) + axs[2].set_title("ADMML2 local ") + fig.savefig('solver1.png') \ No newline at end of file From 6c475d116b177cd12397e972e6f7465a0d6afaa8 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:35:17 +0000 Subject: [PATCH 28/47] doc: finalized proxsolver example --- examples/plot_proxsolver.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/examples/plot_proxsolver.py b/examples/plot_proxsolver.py index f2d7bddb..a50fb247 100644 --- a/examples/plot_proxsolver.py +++ b/examples/plot_proxsolver.py @@ -3,7 +3,7 @@ ================ This example demonstrates the use of the solvers in the -:py:module:`pylops_mpi.proximal.optimization` module. +``pylops_mpi.proximal.optimization`` module. """ import numpy as np @@ -15,12 +15,14 @@ import pylops_mpi -np.random.seed(42) plt.close("all") comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() +np.random.seed(rank) + + ############################################################################### # Let's start with an example of sparsity promoting inversion using the # :py:class:`pylops_mpi.proximal.optimization.primal.ProximalGradient` solver. @@ -37,9 +39,6 @@ # Operator and data A = np.random.normal(0, 1, (n // size, n,)) -As = np.vstack(comm.allgather(A)) - -Op = pylops.MatrixMult(As) Opd = pylops_mpi.MPIVStack([pylops.MatrixMult(A), ]) b = Opd @ arr @@ -61,8 +60,10 @@ arrpgdlocal = arrpg.asarray() # Benchmark serial inversion +As = np.vstack(comm.allgather(A)) arrlocal = arr.asarray() if rank == 0: + Op = pylops.MatrixMult(As) l2local = pyproximal.L2( Op=Op, b=blocal) l1local = pyproximal.L1(sigma=1e-1) @@ -76,7 +77,8 @@ plt.plot(arrpgdlocal, "b", label="Distr") plt.plot(arrpglocal, "--r", label="Local") plt.legend() - plt.savefig('pg.png') + plt.tight_layout() + ############################################################################### # Next we use the :py:class:`pylops_mpi.proximal.optimization.primal.ADMML2` @@ -126,7 +128,7 @@ arradmmlocal = pyproximal.optimization.primal.ADMML2( l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=True, iter_lim=5, + tau=.99/L, niter=5, show=False, iter_lim=5, )[0] fig, axs = plt.subplots(1, 3, figsize=(12, 6)) @@ -136,7 +138,8 @@ axs[1].set_title("ADMML2 distr") axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) axs[2].set_title("ADMML2 local") - fig.savefig('solver.png') + fig.tight_layout() + ############################################################################### # And finally we repeat the same with a scattered model. @@ -183,7 +186,7 @@ arradmmlocal = pyproximal.optimization.primal.ADMML2( l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=True, iter_lim=5, + tau=.99/L, niter=5, show=False, iter_lim=5, )[0] fig, axs = plt.subplots(1, 3, figsize=(12, 6)) @@ -193,4 +196,4 @@ axs[1].set_title("ADMML2 distr") axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) axs[2].set_title("ADMML2 local ") - fig.savefig('solver1.png') \ No newline at end of file + fig.tight_layout() From 2fb45603ee7e478da7ca91ae26cda10472a97b36 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:49:09 +0000 Subject: [PATCH 29/47] minor: fix flake8 --- pylops_mpi/proximal/proximal/L2.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index d7912f22..80ae2ff3 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -10,9 +10,6 @@ from pylops_mpi.optimization.basic import cg, cgls from pylops_mpi.proximal import MPIProxOperator -if TYPE_CHECKING: - from pylops_mpi import MPILinearOperator - class MPIL2(MPIProxOperator): """L2 Norm proximal operator. @@ -60,7 +57,7 @@ class MPIL2(MPIProxOperator): def __init__( self, - Op: "MPILinearOperator" = None, + Op: MPILinearOperator = None, b: DistributedArray | None = None, q: DistributedArray | None = None, sigma: float = 1.0, @@ -163,8 +160,8 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr else: Iop = MPILinearOperator(Identity(x.local_shape, dtype=self.Op.dtype, )) Opreg = MPIStackedVStack([ - sqrt(tau * self.sigma) * self.Op, - Iop, + sqrt(tau * self.sigma) * self.Op, + Iop, ]) breg = StackedDistributedArray([sqrt(tau * self.sigma) * self.b, y]) x = cgls(Opreg, breg, x0=self.x0, niter=niter, **self.kwargs_solver)[0] From 5a50a4bf927aa19d39cd604b1da089ef820a4cb2 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:51:27 +0000 Subject: [PATCH 30/47] doc: added docstrings to proximal solvers --- pylops_mpi/proximal/optimization/primal.py | 122 ++++++++++++++++++++- 1 file changed, 118 insertions(+), 4 deletions(-) diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index 256fbfbe..cfa946dd 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -25,7 +25,6 @@ def ProximalGradient( x0: DistributedArray, epsg: float | NDArray = 1.0, tau: float | None = None, - beta: float = 0.5, eta: float = 1.0, niter: int = 10, niterback: int = 100, @@ -36,6 +35,61 @@ def ProximalGradient( ) -> DistributedArray: r"""Proximal gradient (optionally accelerated) + Solves the following minimization problem using (Accelerated) Proximal + gradient algorithm: + + .. math:: + + \mathbf{x} = \arg\,min_\mathbf{x} f(\mathbf{x}) + \epsilon g(\mathbf{x}) + + where :math:`f(\mathbf{x})` is a smooth convex function with a uniquely + defined gradient and :math:`g(\mathbf{x})` is any convex function that + has a known proximal operator. Both ``f`` and ``g`` must be of + :class:`pylops_mpi.proximal.MPIProxOperator` kind. + + Parameters + ---------- + proxf : :obj:`pylops_mpi.proximal.MPIProxOperator` + Proximal operator of f function (must have ``grad`` implemented) + proxg : :obj:`pylops_mpi.proximal.MPIProxOperator` + Proximal operator of g function + x0 : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Initial vector + epsg : :obj:`float` or :obj:`numpy.ndarray`, optional + Scaling factor of g function. Can be a scalar + for iteration-independent scaling or a a 1d vector for + iteration-dependent scaling + tau : :obj:`float`, optional + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/L]` where ``L`` is + the Lipschitz constant of :math:`\nabla f`. + eta : :obj:`float`, optional + Relaxation parameter (must be between 0 and 1, 0 excluded). + niter : :obj:`int`, optional + Number of iterations of iterative scheme + niterback : :obj:`int`, optional + Max number of iterations of backtracking + acceleration : :obj:`str`, optional + Acceleration (``None``, ``vandenberghe`` or ``fista``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached or the other tolerance + criterion is met + callback : :obj:`callable`, optional + Function with signature (``callback(x)``) to call after each iteration + where ``x`` is the current model vector + show : :obj:`bool`, optional + Display iterations log + + Returns + ------- + x : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Inverted model + + Notes + ----- + See :class:`pyproximal.optimization.primal.ProximalGradient` + """ rank = x0.rank @@ -57,15 +111,14 @@ def ProximalGradient( "---------------------------------------------------------\n" "Proximal operator (f): %s\n" "Proximal operator (g): %s\n" - "tau = %s\tbeta = %10e\n" - "epsg = %s\tniter = %d\ttol = %s\n" + "tau = %s\tepsg = %s\n" + "niter = %d\ttol = %s\n" "" "niterback = %d\tacceleration = %s\n" % ( proxf, proxg, str(tau), - beta, epsg_print, niter, str(tol), @@ -169,6 +222,67 @@ def ADMML2( ) -> tuple[DistributedArray, DistributedArray]: r"""Alternating Direction Method of Multipliers for L2 misfit term + Solves the following minimization problem using Alternating Direction + Method of Multipliers: + + .. math:: + + \mathbf{x},\mathbf{z} = \arg\,min_{\mathbf{x},\mathbf{z}} + \frac{1}{2}||\mathbf{Op}\mathbf{x} - \mathbf{b}||_2^2 + g(\mathbf{z}) \\ + s.t. \; \mathbf{Ax}=\mathbf{z} + + where :math:`g(\mathbf{z})` is any convex function that has a known proximal operator. + + Parameters + ---------- + proxg : :obj:`pylops_mpi.proximal.MPIProxOperator` + Proximal operator of g function + Op : :obj:`pylops_mpi.MPILinearOperator` or :obj:`pylops_mpi.MPIStackedLinearOperator` + Linear operator of data misfit term + b : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Data + A : :obj:`pylops_mpi.MPILinearOperator` or :obj:`pylops_mpi.MPIStackedLinearOperator` + Linear operator of regularization term + x0 : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Initial vector + tau : :obj:`float` + Positive scalar weight, which should satisfy the following condition + to guarantees convergence: :math:`\tau \in (0, 1/\lambda_{max}(\mathbf{A}^H\mathbf{A})]`. + niter : :obj:`int`, optional + Number of iterations of iterative scheme + z0 : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Initial auxiliary vector. If ``None``, initialized to ``A @ x0``. + gfirst : :obj:`bool`, optional + Apply Proximal of operator ``g`` first (``True``) or Proximal of + operator ``f`` first (``False``) + tol : :obj:`float`, optional + Tolerance on change of objective function (used as stopping criterion). If + ``tol=None``, run until ``niter`` is reached + callback : :obj:`callable`, optional + Function with signature (``callback(x)``) to call after each iteration + where ``x`` is the current model vector + show : :obj:`bool`, optional + Display iterations log + **kwargs_solver + Arbitrary keyword arguments for :py:func:`pylops_mpi.optimization.basic.cgls` used + to solve the x-update + + Returns + ------- + x : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Inverted model + z : :obj:`pylops_mpi.DistributedArray` or :obj:`pylops_mpi.StackedDistributedArray` + Inverted second model + + Raises + ------ + ValueError + If both ``x0`` and ``z0`` are set to ``None`` or ``x0`` is set to None + + Notes + ----- + See :class:`pyproximal.optimization.primal.ADMML2` + """ rank = x0.rank From 6a152b06de2d1f9b04f6533d03025d761e44f08e Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Tue, 14 Jul 2026 22:53:27 +0000 Subject: [PATCH 31/47] minor: fix flake8 --- pylops_mpi/proximal/proximal/L2.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index 80ae2ff3..e62cb5f6 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -1,5 +1,5 @@ from math import sqrt -from typing import TYPE_CHECKING, Any, Callable +from typing import Any, Callable from pylops.basicoperators import Identity from pyproximal.ProxOperator import _check_tau @@ -160,9 +160,9 @@ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArr else: Iop = MPILinearOperator(Identity(x.local_shape, dtype=self.Op.dtype, )) Opreg = MPIStackedVStack([ - sqrt(tau * self.sigma) * self.Op, - Iop, - ]) + sqrt(tau * self.sigma) * self.Op, + Iop, + ]) breg = StackedDistributedArray([sqrt(tau * self.sigma) * self.b, y]) x = cgls(Opreg, breg, x0=self.x0, niter=niter, **self.kwargs_solver)[0] if self.warm: From 40bcf3e994032fffe866d132c6548422a7baa0f2 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Fri, 17 Jul 2026 21:59:48 +0000 Subject: [PATCH 32/47] feat: force __call__ output out of GPU --- pylops_mpi/proximal/ProxOperator.py | 9 ++++-- tests/test_prox.py | 44 +++++++++++++++++++++++------ 2 files changed, 43 insertions(+), 10 deletions(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 07f93e15..ac4abc42 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -2,7 +2,7 @@ from typing import Any from pyproximal import ProxOperator -from pylops.utils.backend import get_module +from pylops.utils.backend import get_module, to_numpy from pylops_mpi import DistributedArray, Partition @@ -81,10 +81,15 @@ def __call__(self, x: DistributedArray) -> DistributedArray: # Reduce local function evaluations into final evaluation reduce_op = _call_reduce_op[str(type(self.proxop).__name__)][0] - recv_buf = x._allreduce_subcomm(x.sub_comm, x.base_comm_nccl, + recv_buf = x._allreduce_subcomm(x.sub_comm, + x.base_comm_nccl, ncp.asarray(f), op=reduce_op, engine=x.engine) + + # Ensure that a bool/float/int is returned + if not isinstance(recv_buf, bool): + recv_buf = to_numpy(recv_buf) return recv_buf else: # For broadcasted arrays, simply return the local f diff --git a/tests/test_prox.py b/tests/test_prox.py index 34df73fa..7d28a79f 100644 --- a/tests/test_prox.py +++ b/tests/test_prox.py @@ -40,29 +40,30 @@ "imag": 0, "dtype": np.float64, "partition": pylops_mpi.Partition.SCATTER -} # scatter, real +} # scatter, real par1j = { "n": 101, "imag": 1j, "dtype": np.complex128, "partition": pylops_mpi.Partition.SCATTER -} # scatter, complex +} # scatter, complex par1b = { "n": 101, "imag": 0, "dtype": np.float64, "partition": pylops_mpi.Partition.BROADCAST -} # broadcast, real +} # broadcast, real @pytest.mark.mpi(min_size=2) @pytest.mark.parametrize( "par", [(par1), (par1j), (par1b)] ) -def test_separable_prox(par): - """Separable proximal operators""" +def test_box(par): + """Box call/prox vs pyproximal + (does not support complex numbers)""" np.random.seed(42) x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], @@ -71,7 +72,6 @@ def test_separable_prox(par): par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) x_global = x.asarray() - # Box (does not support complex numbers) if par['imag'] == 0: box = Box(lower=0.0, upper=1.0) boxd = pylops_mpi.proximal.MPIProxOperator(box) @@ -86,7 +86,21 @@ def test_separable_prox(par): assert_allclose(f, f_np, rtol=1e-14) assert_allclose(prox, prox_np, rtol=1e-14) - # L0 + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par1b)] +) +def test_l0(par): + """L0 call/prox vs pyproximal""" + np.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + l0 = L0(sigma=2.0) l0d = pylops_mpi.proximal.MPIProxOperator(l0) @@ -100,7 +114,21 @@ def test_separable_prox(par): assert_allclose(f, f_np, rtol=1e-14) assert_allclose(prox, prox_np, rtol=1e-14) - # L1 + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par1b)] +) +def test_l1(par): + """L1 call/prox vs pyproximal""" + np.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine=backend) + x[:] = np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * np.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + l1 = L1(sigma=2.0) l1d = pylops_mpi.proximal.MPIProxOperator(l1) From fa78c63c3bce911a2d041a11116da3d8810577ff Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 20 Jul 2026 21:40:11 +0000 Subject: [PATCH 33/47] test: added tests for ProximalGradient --- tests/test_proxsolver.py | 222 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 222 insertions(+) create mode 100644 tests/test_proxsolver.py diff --git a/tests/test_proxsolver.py b/tests/test_proxsolver.py new file mode 100644 index 00000000..58eb1bf8 --- /dev/null +++ b/tests/test_proxsolver.py @@ -0,0 +1,222 @@ +"""Test proximal solvers + Designed to run with n processes + $ mpiexec -n 10 pytest test_proxsolver.py --with-mpi +""" +import os + +if int(os.environ.get("TEST_CUPY_PYLOPS", 0)): + import cupy as np + from cupy.testing import assert_allclose + + backend = "cupy" +else: + import numpy as np + from numpy.testing import assert_allclose + + backend = "numpy" +from mpi4py import MPI +import pytest +import pylops +from pylops import ( + BlockDiag, + MatrixMult, +) +from pyproximal import L1, L2 +from pyproximal.optimization.primal import ProximalGradient + +from pylops_mpi import DistributedArray, Partition +from pylops_mpi.basicoperators import MPIBlockDiag, MPIVStack +from pylops_mpi.proximal import MPIProxOperator, MPIL2 +from pylops_mpi.proximal.optimization.primal import ProximalGradient as MPIProximalGradient +from pylops_mpi.proximal.optimization.primal import ADMML2 as MPIADMML2 + + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() +size = comm.Get_size() +if backend == "cupy": + device_id = rank % np.cuda.runtime.getDeviceCount() + np.cuda.Device(device_id).use() + +par1 = { + "ny": 11, + "nx": 11, + "imag": 0, + "x0": False, + "dtype": "float64", +} # square real, zero initial guess +par2 = { + "ny": 11, + "nx": 11, + "imag": 0, + "x0": True, + "dtype": "float64", +} # square real, non-zero initial guess +par3 = { + "ny": 31, + "nx": 11, + "imag": 0, + "x0": False, + "dtype": "float64", +} # overdetermined real, zero initial guess +par4 = { + "ny": 31, + "nx": 11, + "imag": 0, + "x0": True, + "dtype": "float64", +} # overdetermined real, non-zero initial guess +par1j = { + "ny": 11, + "nx": 11, + "imag": 1j, + "x0": False, + "dtype": "complex128", +} # square complex, zero initial guess +par2j = { + "ny": 11, + "nx": 11, + "imag": 1j, + "x0": True, + "dtype": "complex128", +} # square complex, non-zero initial guess +par3j = { + "ny": 31, + "nx": 11, + "imag": 1j, + "x0": False, + "dtype": "complex128", +} # overdetermined complex, zero initial guess +par4j = { + "ny": 31, + "nx": 11, + "imag": 1j, + "x0": True, + "dtype": "complex128", +} # overdetermined complex, non-zero initial guess + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par2), (par2j), (par3), (par3j), (par4), (par4j)] +) +def test_proximalgradient_broadcast(par): + """ProximalGradient with broabcasted model""" + np.random.seed(rank) + + A = np.random.normal(0, 1, (par["ny"], par["nx"])) + par[ + "imag"] * np.random.normal(0, 1, (par["ny"], par["nx"])) + AVStack_MPI = MPIVStack(ops=[pylops.MatrixMult(A), ]) + + x = DistributedArray(global_shape=par['nx'], dtype=par['dtype'], + partition=Partition.BROADCAST, engine=backend) + x[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x_global = x.asarray() + if par["x0"]: + x0 = DistributedArray(global_shape=par['nx'], dtype=par['dtype'], + partition=Partition.BROADCAST, engine=backend) + x0[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x0_global = x0.asarray() + else: + # Set TO 0s if x0 = False + x0 = DistributedArray(global_shape=par['nx'], dtype=par['dtype'], + partition=Partition.BROADCAST, engine=backend) + x0[:] = 0 + x0_global = x0.asarray() + + y = AVStack_MPI * x + + # L2 prox + l2d = MPIL2( + Op=AVStack_MPI, b=y, x0=x0) + + # L1 prox + l1 = L1(sigma=1e-1) + l1d = MPIProxOperator(l1) + + xinv = MPIProximalGradient( + l2d, l1d, x0=x0, tau=1e-3, niter=400, show=True) + assert isinstance(xinv, DistributedArray) + xinv_array = xinv.asarray() + + As = np.vstack(comm.allgather(A)) + if rank == 0: + AVStack = MatrixMult(As) + if par["x0"]: + x0 = x0_global + else: + x0 = np.zeros(par['nx'], dtype=par['dtype']) + y1 = AVStack * x_global + + l2local = L2(Op=AVStack, b=y1, x0=x0) + l1local = L1(sigma=1e-1) + + xinv1 = ProximalGradient( + l2local, l1local, x0=x0, tau=1e-3, + niter=400, show=False) + assert_allclose(xinv_array, xinv1, rtol=1e-12) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par2), (par2j), (par3), (par3j), (par4), (par4j)] +) +def test_proximalgradient_scatter(par): + """ProximalGradient with broabcasted model""" + np.random.seed(rank) + + A = np.random.normal(0, 1, (par["ny"], par["nx"])) + par[ + "imag"] * np.random.normal(0, 1, (par["ny"], par["nx"])) + ABDiag_MPI = MPIBlockDiag(ops=[pylops.MatrixMult(A), ]) + + x = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x_global = x.asarray() + if par["x0"]: + x0 = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x0[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x0_global = x0.asarray() + else: + # Set TO 0s if x0 = False + x0 = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x0[:] = 0 + x0_global = x0.asarray() + + y = ABDiag_MPI * x + + # L2 prox + l2d = MPIL2( + Op=ABDiag_MPI, b=y, x0=x0) + + # L1 prox + l1 = L1(sigma=1e-1) + l1d = MPIProxOperator(l1) + + xinv = MPIProximalGradient( + l2d, l1d, x0=x0, tau=1e-3, niter=400, show=True) + assert isinstance(xinv, DistributedArray) + xinv_array = xinv.asarray() + + As = comm.allgather(A) + if rank == 0: + ABDiag = BlockDiag([MatrixMult(A) for A in As]) + if par["x0"]: + x0 = x0_global + else: + x0 = np.zeros(par['nx'] * size, dtype=par['dtype']) + y1 = ABDiag * x_global + + l2local = L2(Op=ABDiag, b=y1, x0=x0) + l1local = L1(sigma=1e-1) + + xinv1 = ProximalGradient( + l2local, l1local, x0=x0, tau=1e-3, + niter=400, show=False) + assert_allclose(xinv_array, xinv1, rtol=1e-12) From 5751995477b201e00f03b4aa3dbef1018683100f Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 20 Jul 2026 21:40:39 +0000 Subject: [PATCH 34/47] fix: force outputs of call to be always numpy --- pylops_mpi/proximal/optimization/primal.py | 2 +- pylops_mpi/proximal/proximal/L2.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index cfa946dd..bb84b817 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -329,7 +329,7 @@ def ADMML2( if show: if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = 0.5 * (Op @ x - b).norm() ** 2, proxg(Ax) + pf, pg = to_numpy(0.5 * (Op @ x - b).norm() ** 2), proxg(Ax) if rank == 0: msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( iiter + 1, diff --git a/pylops_mpi/proximal/proximal/L2.py b/pylops_mpi/proximal/proximal/L2.py index e62cb5f6..1bc6e231 100644 --- a/pylops_mpi/proximal/proximal/L2.py +++ b/pylops_mpi/proximal/proximal/L2.py @@ -2,6 +2,7 @@ from typing import Any, Callable from pylops.basicoperators import Identity +from pylops.utils.backend import to_numpy from pyproximal.ProxOperator import _check_tau from pylops_mpi import DistributedArray, StackedDistributedArray, Partition @@ -116,7 +117,7 @@ def __call__(self, x: DistributedArray) -> DistributedArray: f = (self.sigma / 2.0) * (x.norm() ** 2) if self.q is not None: f += self.alpha * self.q.dot(x) - return float(f.item()) + return float(to_numpy(f.item())) def _increment_count(func: Callable[..., Any]) -> Callable[..., Any]: """Increment counter""" From b5a6724acade7ca02509401d51fc90c5faea799b Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 20 Jul 2026 21:41:16 +0000 Subject: [PATCH 35/47] doc: more examples of proximal in tutorials --- examples/plot_proxsolver.py | 75 ++++++++++++++++++-------------- tutorials_cupy/poststack_cupy.py | 17 ++++---- tutorials_nccl/poststack_nccl.py | 50 ++++++++++++++++----- 3 files changed, 90 insertions(+), 52 deletions(-) diff --git a/examples/plot_proxsolver.py b/examples/plot_proxsolver.py index a50fb247..4575f3f0 100644 --- a/examples/plot_proxsolver.py +++ b/examples/plot_proxsolver.py @@ -31,11 +31,12 @@ # Sparse input n = 16 -arr = pylops_mpi.DistributedArray(global_shape=n, - partition=pylops_mpi.Partition.BROADCAST) +arr = pylops_mpi.DistributedArray( + global_shape=n, + partition=pylops_mpi.Partition.BROADCAST) arr[:] = 0.0 -arr[n//4] = 1.0 -arr[n//2] = -0.5 +arr[n // 4] = 1.0 +arr[n // 2] = -0.5 # Operator and data A = np.random.normal(0, 1, (n // size, n,)) @@ -54,9 +55,9 @@ # Distributed inversion arrpg = pylops_mpi.proximal.optimization.primal.ProximalGradient( - l2d, l1d, x0=arr.zeros_like(), tau=1e-2, niter=400, - show=True, - ) + l2d, l1d, x0=arr.zeros_like(), tau=1e-2, niter=400, + show=True, +) arrpgdlocal = arrpg.asarray() # Benchmark serial inversion @@ -70,7 +71,7 @@ arrpglocal = pyproximal.optimization.primal.ProximalGradient( l2local, l1local, x0=np.zeros(n), tau=1e-2, niter=400, show=False - ) + ) plt.figure(figsize=(12, 3)) plt.plot(arrlocal, "k", label="True") @@ -89,14 +90,16 @@ # Input ny, nx = 10 * size, 40 arrlocal = np.zeros((ny, nx)) -arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 -arr = pylops_mpi.DistributedArray(global_shape=ny * nx, - partition=pylops_mpi.Partition.BROADCAST) +arrlocal[ny // 2 - 5:ny // 2 + 5, nx // 2 - 5:nx // 2 + 5] = 2 +arr = pylops_mpi.DistributedArray( + global_shape=ny * nx, + partition=pylops_mpi.Partition.BROADCAST +) arr[:] = arrlocal.flatten() # Operator and data -Op = pylops.VStack([pylops.Diagonal(np.ones(ny*nx)) for _ in range(size)]) -Opd = pylops_mpi.MPIVStack([pylops.Diagonal(np.ones(ny*nx)),]) +Op = pylops.VStack([pylops.Diagonal(np.ones(ny * nx)) for _ in range(size)]) +Opd = pylops_mpi.MPIVStack([pylops.Diagonal(np.ones(ny * nx)),]) b = Opd @ arr blocal = b.asarray() @@ -112,9 +115,9 @@ L = 8.0 # max eig of Gopd.H @ Gop x0distr = arr.zeros_like() arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, - show=True, kwargs_solver=dict(niter=5), - )[0] + l1d, Opd, b, Gopd, x0=x0distr, tau=.99 / L, niter=5, + show=True, kwargs_solver=dict(niter=5), +)[0] arradmmdlocal = arradmm.asarray() # Benchmark serial inversion @@ -123,21 +126,24 @@ Gop = pylops.Gradient( dims=(ny, nx), sampling=1., edge=False, kind="forward", - ) + ) l1local = pyproximal.L1(sigma=2e0) arradmmlocal = pyproximal.optimization.primal.ADMML2( - l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=False, iter_lim=5, + l1local, Op, blocal, Gop, x0=np.zeros(ny * nx), + tau=.99 / L, niter=5, show=False, iter_lim=5, )[0] - fig, axs = plt.subplots(1, 3, figsize=(12, 6)) + fig, axs = plt.subplots(1, 3, figsize=(12, 3)) axs[0].imshow(arrlocal.reshape(10 * size, 40)) axs[0].set_title("True") + axs[0].axis("tight") axs[1].imshow(arradmmdlocal.reshape(10 * size, 40)) axs[1].set_title("ADMML2 distr") + axs[1].axis("tight") axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) axs[2].set_title("ADMML2 local") + axs[2].axis("tight") fig.tight_layout() @@ -147,14 +153,14 @@ # Input ny, nx = 10 * size, 40 arrlocal = np.zeros((ny, nx)) -arrlocal[ny//2-5:ny//2+5, nx//2-5:nx//2+5] = 2 +arrlocal[ny // 2 - 5:ny // 2 + 5, nx // 2 - 5:nx // 2 + 5] = 2 arr = pylops_mpi.DistributedArray(global_shape=ny * nx, partition=pylops_mpi.Partition.SCATTER) -arr[:] = arrlocal[ny//size * rank: ny//size * (rank +1)].flatten() +arr[:] = arrlocal[ny // size * rank: ny // size * (rank + 1)].flatten() # Operator and data -Op = pylops.Diagonal(np.ones(ny*nx)) -Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones(ny*nx//size)),]) +Op = pylops.Diagonal(np.ones(ny * nx)) +Opd = pylops_mpi.MPIBlockDiag([pylops.Diagonal(np.ones((ny * nx) // size)),]) b = Opd @ arr blocal = b.asarray() @@ -170,9 +176,9 @@ L = 8.0 # max eig of Gopd.H @ Gop x0distr = arr.zeros_like() arradmm = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, Opd, b, Gopd, x0=x0distr, tau=.99/L, niter=5, - show=True, kwargs_solver=dict(niter=5), - )[0] + l1d, Opd, b, Gopd, x0=x0distr, tau=.99 / L, niter=5, + show=True, kwargs_solver=dict(niter=5), +)[0] arradmmdlocal = arradmm.asarray() # Benchmark serial inversion @@ -180,20 +186,23 @@ if rank == 0: Gop = pylops.Gradient( - dims=(ny, nx), sampling=1., edge=False, kind="forward", - ) + dims=(ny, nx), sampling=1., edge=False, kind="forward" + ) l1local = pyproximal.L1(sigma=2e0) arradmmlocal = pyproximal.optimization.primal.ADMML2( - l1local, Op, blocal, Gop, x0=np.zeros(ny*nx), - tau=.99/L, niter=5, show=False, iter_lim=5, + l1local, Op, blocal, Gop, x0=np.zeros(ny * nx), + tau=.99 / L, niter=5, show=False, iter_lim=5, )[0] - fig, axs = plt.subplots(1, 3, figsize=(12, 6)) + fig, axs = plt.subplots(1, 3, figsize=(12, 3)) axs[0].imshow(arrlocal.reshape(10 * size, 40)) axs[0].set_title("True") + axs[0].axis("tight") axs[1].imshow(arradmmdlocal.reshape(10 * size, 40)) axs[1].set_title("ADMML2 distr") + axs[1].axis("tight") axs[2].imshow(arradmmlocal.reshape(10 * size, 40)) - axs[2].set_title("ADMML2 local ") + axs[2].set_title("ADMML2 local") + axs[2].axis("tight") fig.tight_layout() diff --git a/tutorials_cupy/poststack_cupy.py b/tutorials_cupy/poststack_cupy.py index 0eef4f51..5521173e 100644 --- a/tutorials_cupy/poststack_cupy.py +++ b/tutorials_cupy/poststack_cupy.py @@ -126,8 +126,8 @@ dtype=BDiag.dtype) NormEqOp = BDiag.H @ BDiag + epsR * LapOp.H @ LapOp dnorm_dist = BDiag.H @ d_dist -minv3d_ne_dist = pylops_mpi.optimization.basic.cg(NormEqOp, dnorm_dist, - x0=mback3d_dist, +minv3d_ne_dist = pylops_mpi.optimization.basic.cg(NormEqOp, dnorm_dist, + x0=mback3d_dist, niter=100, show=True)[0] minv3d_ne = minv3d_ne_dist.asarray().reshape((ny, nx, nz)) @@ -140,12 +140,11 @@ dstack_dist = pylops_mpi.StackedDistributedArray([d_dist, d0_dist]) dnorm_dist = BDiag.H @ d_dist -minv3d_reg_dist = pylops_mpi.optimization.basic.cgls(StackOp, dstack_dist, - x0=mback3d_dist, +minv3d_reg_dist = pylops_mpi.optimization.basic.cgls(StackOp, dstack_dist, + x0=mback3d_dist, niter=100, show=True)[0] minv3d_reg = minv3d_reg_dist.asarray().reshape((ny, nx, nz)) - ############################################################################### # TV-Regularized inversion @@ -158,9 +157,9 @@ l1d = pylops_mpi.proximal.MPIProxOperator(l1) minv3d_admm_dist = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99/L, niter=40, - show=True, kwargs_solver=dict(niter=20), - )[0] + l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99 / L, niter=40, + show=True, kwargs_solver=dict(niter=20), +)[0] minv3d_admm = minv3d_admm_dist.asarray().reshape((ny, nx, nz)) ############################################################################### @@ -175,7 +174,7 @@ # Check the two distributed implementations give the same modelling results print('Distr == Local', np.allclose(cp.asnumpy(d), d0, atol=1e-6)) - + # Visualize fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(12, 18), constrained_layout=True) diff --git a/tutorials_nccl/poststack_nccl.py b/tutorials_nccl/poststack_nccl.py index e81d9323..66a6176a 100644 --- a/tutorials_nccl/poststack_nccl.py +++ b/tutorials_nccl/poststack_nccl.py @@ -15,6 +15,7 @@ from pylops.utils.wavelets import ricker from pylops.basicoperators import Transpose from pylops.avo.poststack import PoststackLinearModelling +from pyproximal.proximal import L1 import pylops_mpi @@ -114,8 +115,8 @@ minv3d_iter = minv3d_iter_dist.asarray().reshape((ny, nx, nz)) ############################################################################### - # Regularized inversion with normal equations + epsR = 1e2 LapOp = pylops_mpi.MPILaplacian(dims=(ny, nx, nz), axes=(0, 1, 2), weights=(1, 1, 1), @@ -129,20 +130,37 @@ minv3d_ne = minv3d_ne_dist.asarray().reshape((ny, nx, nz)) ############################################################################### - # Regularized inversion with regularized equations + StackOp = pylops_mpi.MPIStackedVStack([BDiag, np.sqrt(epsR) * LapOp]) -d0_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz, +d0_dist = pylops_mpi.DistributedArray(global_shape=ny * nx * nz, base_comm_nccl=nccl_comm, engine="cupy") d0_dist[:] = 0. dstack_dist = pylops_mpi.StackedDistributedArray([d_dist, d0_dist]) dnorm_dist = BDiag.H @ d_dist -minv3d_reg_dist = pylops_mpi.optimization.basic.cgls(StackOp, dstack_dist, - x0=mback3d_dist, +minv3d_reg_dist = pylops_mpi.optimization.basic.cgls(StackOp, dstack_dist, + x0=mback3d_dist, niter=100, show=True)[0] minv3d_reg = minv3d_reg_dist.asarray().reshape((ny, nx, nz)) +############################################################################### +# TV-Regularized inversion + +Gopd = pylops_mpi.MPIGradient( + dims=(ny, nx, nz), sampling=1., edge=False, kind="forward") + +L = 12.0 # maxeig(Gop^H Gop) + +l1 = L1(sigma=1e-2) +l1d = pylops_mpi.proximal.MPIProxOperator(l1) + +minv3d_admm_dist = pylops_mpi.proximal.optimization.primal.ADMML2( + l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99 / L, niter=40, + show=True, kwargs_solver=dict(niter=20), +)[0] +minv3d_admm = minv3d_admm_dist.asarray().reshape((ny, nx, nz)) + ############################################################################### # Finally we visualize the results. Note that the array must be copied back # to the CPU by calling the :code:`get()` method on the CuPy arrays. @@ -159,7 +177,9 @@ print('Smooth Distr == Local', np.allclose(d_0, d0_0)) # Visualize - fig, axs = plt.subplots(nrows=6, ncols=3, figsize=(9, 14), constrained_layout=True) + # Visualize + fig, axs = plt.subplots(nrows=7, ncols=3, figsize=(12, 18), + constrained_layout=True) axs[0][0].imshow(m3d[5, :, :].T, cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[0][0].set_title("Model x-z") axs[0][0].axis("tight") @@ -204,18 +224,28 @@ axs[4][0].set_title("Normal Equations Inverted Model iter x-z") axs[4][0].axis("tight") axs[4][1].imshow(minv3d_ne[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[4][1].set_title('Normal Equations Inverted Model iter y-z') + axs[4][1].set_title("Normal Equations Inverted Model iter y-z") axs[4][1].axis('tight') axs[4][2].imshow(minv3d_ne[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[4][2].set_title('Normal Equations Inverted Model iter x-y') + axs[4][2].set_title("Normal Equations Inverted Model iter x-y") axs[4][2].axis('tight') axs[5][0].imshow(minv3d_reg[5, :, :].T.get(), cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) axs[5][0].set_title("Regularized Inverted Model iter x-z") axs[5][0].axis("tight") axs[5][1].imshow(minv3d_reg[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][1].set_title('Regularized Inverted Model iter y-z') + axs[5][1].set_title("Regularized Inverted Model iter y-z") axs[5][1].axis('tight') axs[5][2].imshow(minv3d_reg[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) - axs[5][2].set_title('Regularized Inverted Model iter x-y') + axs[5][2].set_title("Regularized Inverted Model iter x-y") axs[5][2].axis('tight') + + axs[6][0].imshow(minv3d_admm[5, :, :].T.get(), cmap="gist_rainbow", vmin=m.min(), vmax=m.max()) + axs[6][0].set_title("TV-Regularized Inverted Model iter x-z") + axs[6][0].axis("tight") + axs[6][1].imshow(minv3d_admm[:, 200, :].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][1].set_title("TV-Regularized Inverted Model iter y-z") + axs[6][1].axis('tight') + axs[6][2].imshow(minv3d_admm[:, :, 220].T.get(), cmap='gist_rainbow', vmin=m.min(), vmax=m.max()) + axs[6][2].set_title("TV-Regularized Inverted Model iter x-y") + axs[6][2].axis('tight') From 02e4a5b0e14b9f0bdef4b0b342dbdbf2b37b59cc Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Mon, 20 Jul 2026 22:30:54 +0000 Subject: [PATCH 36/47] test: more tests for proximal --- pylops_mpi/proximal/__init__.py | 9 +- pylops_mpi/proximal/optimization/__init__.py | 1 + tests/test_proxsolver.py | 80 ++++++++++- tests_nccl/test_prox_nccl.py | 135 +++++++++++++++++++ tutorials/poststack.py | 6 +- 5 files changed, 215 insertions(+), 16 deletions(-) create mode 100644 tests_nccl/test_prox_nccl.py diff --git a/pylops_mpi/proximal/__init__.py b/pylops_mpi/proximal/__init__.py index e1a0ef48..3a4e9c82 100644 --- a/pylops_mpi/proximal/__init__.py +++ b/pylops_mpi/proximal/__init__.py @@ -8,14 +8,9 @@ A common interface for applying (separable) proximal operators in a distributed fashion is provided by the MPIProxOperator operator. -A list of proximal operators present in pylops_mpi.proximal.proximal: - MPIXX XX +A list of proximal operators present in pylops_mpi.proximal.proximal. -A list of proximal solvers present in pylops_mpi.proximal.optimization.primal: - MPIXX XX - -and in pylops_mpi.proximal.optimization.primaldual: - MPIXX XX +A list of proximal solvers present in pylops_mpi.proximal.optimization.primal. """ diff --git a/pylops_mpi/proximal/optimization/__init__.py b/pylops_mpi/proximal/optimization/__init__.py index 199fbfba..cf7c2bb7 100644 --- a/pylops_mpi/proximal/optimization/__init__.py +++ b/pylops_mpi/proximal/optimization/__init__.py @@ -8,6 +8,7 @@ A list of proximal solvers: ProximalGradient Proximal Gradient + ADMML2 ADMM with L2 misfit term """ diff --git a/tests/test_proxsolver.py b/tests/test_proxsolver.py index 58eb1bf8..55f6c64f 100644 --- a/tests/test_proxsolver.py +++ b/tests/test_proxsolver.py @@ -22,7 +22,7 @@ MatrixMult, ) from pyproximal import L1, L2 -from pyproximal.optimization.primal import ProximalGradient +from pyproximal.optimization.primal import ProximalGradient, ADMML2 from pylops_mpi import DistributedArray, Partition from pylops_mpi.basicoperators import MPIBlockDiag, MPIVStack @@ -106,7 +106,7 @@ def test_proximalgradient_broadcast(par): A = np.random.normal(0, 1, (par["ny"], par["nx"])) + par[ "imag"] * np.random.normal(0, 1, (par["ny"], par["nx"])) - AVStack_MPI = MPIVStack(ops=[pylops.MatrixMult(A), ]) + AVStack_MPI = MPIVStack(ops=[pylops.MatrixMult(A, dtype=par['dtype']), ]) x = DistributedArray(global_shape=par['nx'], dtype=par['dtype'], partition=Partition.BROADCAST, engine=backend) @@ -143,7 +143,7 @@ def test_proximalgradient_broadcast(par): As = np.vstack(comm.allgather(A)) if rank == 0: - AVStack = MatrixMult(As) + AVStack = MatrixMult(As, dtype=par['dtype']) if par["x0"]: x0 = x0_global else: @@ -164,12 +164,12 @@ def test_proximalgradient_broadcast(par): "par", [(par1), (par1j), (par2), (par2j), (par3), (par3j), (par4), (par4j)] ) def test_proximalgradient_scatter(par): - """ProximalGradient with broabcasted model""" + """ProximalGradient with scattered model""" np.random.seed(rank) A = np.random.normal(0, 1, (par["ny"], par["nx"])) + par[ "imag"] * np.random.normal(0, 1, (par["ny"], par["nx"])) - ABDiag_MPI = MPIBlockDiag(ops=[pylops.MatrixMult(A), ]) + ABDiag_MPI = MPIBlockDiag(ops=[pylops.MatrixMult(A, dtype=par['dtype']), ]) x = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], partition=Partition.SCATTER, engine=backend) @@ -206,7 +206,7 @@ def test_proximalgradient_scatter(par): As = comm.allgather(A) if rank == 0: - ABDiag = BlockDiag([MatrixMult(A) for A in As]) + ABDiag = BlockDiag([MatrixMult(A, dtype=par['dtype']) for A in As]) if par["x0"]: x0 = x0_global else: @@ -220,3 +220,71 @@ def test_proximalgradient_scatter(par): l2local, l1local, x0=x0, tau=1e-3, niter=400, show=False) assert_allclose(xinv_array, xinv1, rtol=1e-12) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par2), (par2j), (par3), (par3j), (par4), (par4j)] +) +def test_admml2_scatter(par): + """ADMML2 with scattered model""" + np.random.seed(rank) + + A = np.random.normal(0, 1, (par["ny"], par["nx"])) + par[ + "imag"] * np.random.normal(0, 1, (par["ny"], par["nx"])) + ABDiag_MPI = MPIBlockDiag(ops=[pylops.MatrixMult(A, dtype=par['dtype']), ]) + + x = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x_global = x.asarray() + if par["x0"]: + x0 = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x0[:] = np.random.normal(1, 10, par["nx"]) + \ + par["imag"] * np.random.normal(10, 10, par["nx"]) + x0_global = x0.asarray() + else: + # Set TO 0s if x0 = False + x0 = DistributedArray(global_shape=par['nx'] * size, dtype=par['dtype'], + partition=Partition.SCATTER, engine=backend) + x0[:] = 0 + x0_global = x0.asarray() + + y = ABDiag_MPI * x + + # Regularizer (just make identity to solve the same problem + # as ProximalGradient) + Iopd = MPIBlockDiag(ops=[pylops.Identity(par['nx'], dtype=par['dtype']), ]) + + # L1 prox + l1 = L1(sigma=1e-1) + l1d = MPIProxOperator(l1) + + xinv = MPIADMML2( + l1d, ABDiag_MPI, y, Iopd, x0=x0, tau=1e-3, + niter=400, show=True)[0] + assert isinstance(xinv, DistributedArray) + xinv_array = xinv.asarray() + + As = comm.allgather(A) + if rank == 0: + ABDiag = BlockDiag([MatrixMult(A, dtype=par['dtype']) for A in As]) + if par["x0"]: + x0 = x0_global + else: + x0 = np.zeros(par['nx'] * size, dtype=par['dtype']) + y1 = ABDiag * x_global + + Iop = pylops.Identity(par['nx'] * size, dtype=par['dtype']) + l1local = L1(sigma=1e-1) + + xinv1 = ADMML2( + l1local, ABDiag, y1, Iop, x0=x0, tau=1e-3, + niter=400, show=False)[0] + + # Pretty high tolerance because a different + # linear solver is used internally in the + # serial vs distributed versions of ADMML2 + assert_allclose(xinv_array, xinv1, rtol=1e-4) diff --git a/tests_nccl/test_prox_nccl.py b/tests_nccl/test_prox_nccl.py new file mode 100644 index 00000000..fac00693 --- /dev/null +++ b/tests_nccl/test_prox_nccl.py @@ -0,0 +1,135 @@ +"""Test proximal operators + Designed to run with n GPUs (with 1 MPI process per GPU) + $ mpiexec -n 10 pytest test_prox_nccl.py --with-mpi +""" +import os + +import numpy as np +import cupy as cp +from numpy.testing import assert_allclose +from mpi4py import MPI +import pytest + +import pylops_mpi +from pylops.basicoperators import Diagonal +from pyproximal.proximal import ( + Box, + L0, + L1, + L2 +) +from pylops_mpi.proximal import MPIL2 +from pylops_mpi.utils._nccl import initialize_nccl_comm + + +nccl_comm = initialize_nccl_comm() +base_comm = MPI.COMM_WORLD +size = base_comm.Get_size() +rank = base_comm.Get_rank() + + +par1 = { + "n": 101, + "imag": 0, + "dtype": np.float64, + "partition": pylops_mpi.Partition.SCATTER +} # scatter, real + +par1j = { + "n": 101, + "imag": 1j, + "dtype": np.complex128, + "partition": pylops_mpi.Partition.SCATTER +} # scatter, complex + +par1b = { + "n": 101, + "imag": 0, + "dtype": np.float64, + "partition": pylops_mpi.Partition.BROADCAST +} # broadcast, real + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par1b)] +) +def test_box(par): + """Box call/prox vs pyproximal""" + cp.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine="cupy") + x[:] = cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + + if par['imag'] == 0: + box = Box(lower=0.0, upper=1.0) + boxd = pylops_mpi.proximal.MPIProxOperator(box) + + f = boxd(x) + prox = boxd.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = box(x_global) + prox_np = box.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox.get(), prox_np.get(), rtol=1e-14) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par1b)] +) +def test_l0(par): + """L0 call/prox vs pyproximal""" + cp.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine="cupy") + x[:] = cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + + l0 = L0(sigma=2.0) + l0d = pylops_mpi.proximal.MPIProxOperator(l0) + + f = l0d(x) + prox = l0d.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = l0(x_global) + prox_np = l0.prox(x_global, .1) + assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox.get(), prox_np.get(), rtol=1e-14) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.parametrize( + "par", [(par1), (par1j), (par1b)] +) +def test_l1(par): + """L1 call/prox vs pyproximal""" + cp.random.seed(42) + + x = pylops_mpi.DistributedArray(global_shape=par['n'], dtype=par['dtype'], + partition=par['partition'], engine="cupy") + x[:] = cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + \ + par['imag'] * cp.random.normal(rank, 10, x.local_shape).astype(par['dtype']) + x_global = x.asarray() + + l1 = L1(sigma=2.0) + l1d = pylops_mpi.proximal.MPIProxOperator(l1) + + f = l1d(x) + prox = l1d.prox(x, .1) + prox = prox.asarray() + + if rank == 0: + f_np = l1(x_global) + prox_np = l1.prox(x_global, .1) + # assert_allclose(f, f_np, rtol=1e-14) + assert_allclose(prox.get(), prox_np.get(), rtol=1e-14) diff --git a/tutorials/poststack.py b/tutorials/poststack.py index 97016768..a794f512 100644 --- a/tutorials/poststack.py +++ b/tutorials/poststack.py @@ -222,9 +222,9 @@ L = 12.0 # maxeig(Gopd^H Gopd) minv3d_tv_dist = pylops_mpi.proximal.optimization.primal.ADMML2( - l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99/L, niter=40, - show=True, kwargs_solver=dict(niter=20), - )[0] + l1d, BDiag, d_dist, Gopd, x0=mback3d_dist, tau=.99 / L, niter=40, + show=True, kwargs_solver=dict(niter=20), +)[0] minv3d_tv = minv3d_tv_dist.asarray().reshape((ny, nx, nz)) ############################################################################### From 8ee5c26cac7267a3aa5a3fbd66cc4e9454c30cb6 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 14:52:16 +0530 Subject: [PATCH 37/47] Update prox __call__ to check for scalar --- pylops_mpi/proximal/ProxOperator.py | 25 ++++++++++++++++------ pylops_mpi/proximal/optimization/primal.py | 2 +- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index ac4abc42..67026d6c 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -2,7 +2,7 @@ from typing import Any from pyproximal import ProxOperator -from pylops.utils.backend import get_module, to_numpy +from pylops.utils.backend import get_module from pylops_mpi import DistributedArray, Partition @@ -71,6 +71,20 @@ def __call__(self, x: DistributedArray) -> DistributedArray: Function evaluation """ + + def _as_scalar(value): + """Convert NumPy/CuPy/Python scalar-like objects to a Python scalar.""" + # Ensure that a bool/int/float is returned + if isinstance(value, (bool, int, float)): + return value + + ncp = get_module(x.engine) + if ncp.size(value) != 1: + raise ValueError( + f"Expected scalar function evaluation, " + f"got object with shape {getattr(value, 'shape', None)}" + ) + return value.item() if isinstance(x, DistributedArray): # Compute local function evaluation f = self.proxop(x.local_array) @@ -87,13 +101,10 @@ def __call__(self, x: DistributedArray) -> DistributedArray: op=reduce_op, engine=x.engine) - # Ensure that a bool/float/int is returned - if not isinstance(recv_buf, bool): - recv_buf = to_numpy(recv_buf) - return recv_buf + return _as_scalar(recv_buf) else: - # For broadcasted arrays, simply return the local f - return f + # For broadcasted arrays, simply return the local evaluation + return _as_scalar(f) else: # StackedDistributedArray reduce_op = _call_reduce_op[str(type(self.proxop).__name__)][1] fs = [self(x[iarr]) for iarr in range(x.narrays)] diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py index bb84b817..3ac30a73 100644 --- a/pylops_mpi/proximal/optimization/primal.py +++ b/pylops_mpi/proximal/optimization/primal.py @@ -329,7 +329,7 @@ def ADMML2( if show: if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0: - pf, pg = to_numpy(0.5 * (Op @ x - b).norm() ** 2), proxg(Ax) + pf, pg = to_numpy(0.5 * (Op @ x - b).norm().item() ** 2), proxg(Ax) if rank == 0: msg = "%6g %12.5e %10.3e %10.3e %10.3e" % ( iiter + 1, From df2a936176923c7974099a0ada4fe3e2bcb05553 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 20:39:00 +0530 Subject: [PATCH 38/47] Minor change in rtol --- tests/test_proxsolver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_proxsolver.py b/tests/test_proxsolver.py index 55f6c64f..a15dfe5e 100644 --- a/tests/test_proxsolver.py +++ b/tests/test_proxsolver.py @@ -287,4 +287,4 @@ def test_admml2_scatter(par): # Pretty high tolerance because a different # linear solver is used internally in the # serial vs distributed versions of ADMML2 - assert_allclose(xinv_array, xinv1, rtol=1e-4) + assert_allclose(xinv_array, xinv1, rtol=1e-3) From c0eaff04263453c70b6b85b8df3013d82fc9c38b Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 21:20:13 +0530 Subject: [PATCH 39/47] Update github action testing --- .github/workflows/build.yml | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 59fc9adf..6e6563b8 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -56,8 +56,14 @@ jobs: run: pip install .[all] - name: Testing using pytest-mpi run: | - if [ "${{ matrix.mpi }}" = "openmpi" ]; then - mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi - else - mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi - fi + case "${{ matrix.mpi }}" in + openmpi) + mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi + ;; + mpich) + mpiexec -n ${{ matrix.rank }} --bind-to none pytest tests/ --with-mpi + ;; + intelmpi) + mpiexec -n ${{ matrix.rank }} -genv I_MPI_PIN off pytest tests/ --with-mpi + ;; + esac From 62de62a7eef43d4b30fe0b67330e4f208d4577a9 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 22:57:36 +0530 Subject: [PATCH 40/47] Remove rank = 8 from GA --- .github/workflows/build.yml | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 6e6563b8..ff8ef0fc 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,7 +16,7 @@ jobs: os: [ubuntu-latest, macos-latest] python-version: ['3.11', '3.12', '3.13', '3.14'] mpi: ['mpich', 'openmpi', 'intelmpi'] - rank: ['2', '4', '8'] + rank: ['2', '4'] exclude: - os: macos-latest mpi: 'intelmpi' @@ -56,14 +56,8 @@ jobs: run: pip install .[all] - name: Testing using pytest-mpi run: | - case "${{ matrix.mpi }}" in - openmpi) - mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi - ;; - mpich) - mpiexec -n ${{ matrix.rank }} --bind-to none pytest tests/ --with-mpi - ;; - intelmpi) - mpiexec -n ${{ matrix.rank }} -genv I_MPI_PIN off pytest tests/ --with-mpi - ;; - esac + if [ "${{ matrix.mpi }}" = "openmpi" ]; then + mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi + else + mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi + fi From be0fd5c25625d6af2bae2287e9251f79ca77cb16 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 23:02:22 +0530 Subject: [PATCH 41/47] Remove rank = 8 from GA, update command-mpiexec --- .github/workflows/build.yml | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index ff8ef0fc..5a8539cc 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -56,8 +56,14 @@ jobs: run: pip install .[all] - name: Testing using pytest-mpi run: | - if [ "${{ matrix.mpi }}" = "openmpi" ]; then - mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi - else - mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi - fi + case "${{ matrix.mpi }}" in + openmpi) + mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi + ;; + mpich) + mpiexec -n ${{ matrix.rank }} --bind-to none pytest tests/ --with-mpi + ;; + intelmpi) + mpiexec -n ${{ matrix.rank }} -genv I_MPI_PIN off pytest tests/ --with-mpi + ;; + esac From 1de7bba405c5ff2f03abc9adea75eed226485e27 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Sun, 26 Jul 2026 23:14:14 +0530 Subject: [PATCH 42/47] Change fail-fast: false --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5a8539cc..f4baef58 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -12,6 +12,7 @@ on: jobs: build: strategy: + fail-fast: false matrix: os: [ubuntu-latest, macos-latest] python-version: ['3.11', '3.12', '3.13', '3.14'] From 475e0244ab0a4aeda5e41d476c85ea0d6a8d24ca Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Mon, 27 Jul 2026 00:06:59 +0530 Subject: [PATCH 43/47] Add rank=8 --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f4baef58..d232b945 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -17,7 +17,7 @@ jobs: os: [ubuntu-latest, macos-latest] python-version: ['3.11', '3.12', '3.13', '3.14'] mpi: ['mpich', 'openmpi', 'intelmpi'] - rank: ['2', '4'] + rank: ['2', '4', '8'] exclude: - os: macos-latest mpi: 'intelmpi' From ab65c4af65ae0823f1509a0ff670248b4ba13f83 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Mon, 27 Jul 2026 11:37:33 +0530 Subject: [PATCH 44/47] Remove rank=8 --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d232b945..f4baef58 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -17,7 +17,7 @@ jobs: os: [ubuntu-latest, macos-latest] python-version: ['3.11', '3.12', '3.13', '3.14'] mpi: ['mpich', 'openmpi', 'intelmpi'] - rank: ['2', '4', '8'] + rank: ['2', '4'] exclude: - os: macos-latest mpi: 'intelmpi' From 817712d1ead2a35c755300550cfd4158ed740b18 Mon Sep 17 00:00:00 2001 From: rohanbabbar04 Date: Mon, 27 Jul 2026 12:14:47 +0530 Subject: [PATCH 45/47] Update mpiexec commands to before --- .github/workflows/build.yml | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f4baef58..57c77015 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -57,14 +57,8 @@ jobs: run: pip install .[all] - name: Testing using pytest-mpi run: | - case "${{ matrix.mpi }}" in - openmpi) - mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi - ;; - mpich) - mpiexec -n ${{ matrix.rank }} --bind-to none pytest tests/ --with-mpi - ;; - intelmpi) - mpiexec -n ${{ matrix.rank }} -genv I_MPI_PIN off pytest tests/ --with-mpi - ;; - esac + if [ "${{ matrix.mpi }}" = "openmpi" ]; then + mpiexec --mca btl ^openib -n ${{ matrix.rank }} pytest tests/ --with-mpi + else + mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi + fi From 0fe26e484bce96ea8fde075104925d391a3fc232 Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Wed, 29 Jul 2026 21:12:28 +0000 Subject: [PATCH 46/47] minor: fix servedoc target --- Makefile | 10 +++++----- pylops_mpi/proximal/ProxOperator.py | 2 +- tests/test_proxsolver.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/Makefile b/Makefile index 10d121bd..fa5d47e1 100644 --- a/Makefile +++ b/Makefile @@ -70,22 +70,22 @@ doc: doc_cupy: cp tutorials_cupy/* tutorials/ cd docs && rm -rf source/api/generated && rm -rf source/gallery &&\ - rm -rf source/tutorials && rm -rf source/tutorials && rm -rf build &&\ - cd .. && sphinx-build -b html docs/source docs/build + rm -rf source/tutorials && rm -rf build &&\ + cd .. && sphinx-build -b html docs/source docs/build &&\ rm tutorials/*_cupy.py doc_nccl: cp tutorials_cupy/* tutorials_nccl/* tutorials/ cd docs && rm -rf source/api/generated && rm -rf source/gallery &&\ - rm -rf source/tutorials && rm -rf source/tutorials && rm -rf build &&\ - cd .. && sphinx-build -b html docs/source docs/build + rm -rf source/tutorials && rm -rf build &&\ + cd .. && sphinx-build -b html docs/source docs/build &&\ rm tutorials/*_cupy.py tutorials/*_nccl.py docupdate: cd docs && NCCL_PYLOPS_MPI=0 make html && cd .. servedoc: - $(PYTHON) -m http.server --directory docs/build/ + $(PYTHON) -m http.server --directory docs/build/html/ # Run examples using mpi run_examples: diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index 67026d6c..a71807ea 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -71,7 +71,6 @@ def __call__(self, x: DistributedArray) -> DistributedArray: Function evaluation """ - def _as_scalar(value): """Convert NumPy/CuPy/Python scalar-like objects to a Python scalar.""" # Ensure that a bool/int/float is returned @@ -85,6 +84,7 @@ def _as_scalar(value): f"got object with shape {getattr(value, 'shape', None)}" ) return value.item() + if isinstance(x, DistributedArray): # Compute local function evaluation f = self.proxop(x.local_array) diff --git a/tests/test_proxsolver.py b/tests/test_proxsolver.py index a15dfe5e..3fe0d9e3 100644 --- a/tests/test_proxsolver.py +++ b/tests/test_proxsolver.py @@ -286,5 +286,5 @@ def test_admml2_scatter(par): # Pretty high tolerance because a different # linear solver is used internally in the - # serial vs distributed versions of ADMML2 + # serial vs distributed versions of ADMML2 assert_allclose(xinv_array, xinv1, rtol=1e-3) From 098dbf7e6f5f1d52d2dab18f16e69d7ef26d579a Mon Sep 17 00:00:00 2001 From: mrava87SW Date: Wed, 29 Jul 2026 21:17:51 +0000 Subject: [PATCH 47/47] minor: fix flake8 --- pylops_mpi/proximal/ProxOperator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pylops_mpi/proximal/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py index a71807ea..676286cc 100644 --- a/pylops_mpi/proximal/ProxOperator.py +++ b/pylops_mpi/proximal/ProxOperator.py @@ -84,7 +84,7 @@ def _as_scalar(value): f"got object with shape {getattr(value, 'shape', None)}" ) return value.item() - + if isinstance(x, DistributedArray): # Compute local function evaluation f = self.proxop(x.local_array)