diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 59fc9adf..57c77015 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -12,11 +12,12 @@ on:
jobs:
build:
strategy:
+ fail-fast: false
matrix:
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'
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/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
-----
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/examples/plot_prox.py b/examples/plot_prox.py
new file mode 100644
index 00000000..c42dd3a9
--- /dev/null
+++ b/examples/plot_prox.py
@@ -0,0 +1,176 @@
+r"""
+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
+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 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.
+
+n = 10
+arr = pylops_mpi.DistributedArray(global_shape=n * size,
+ partition=pylops_mpi.Partition.SCATTER)
+arr[:] = rank * np.arange(n)
+
+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()
+
+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))
+
+###############################################################################
+# 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 # 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()
+
+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))
+
+###############################################################################
+# 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`
+
+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)
+
+# Call
+f = l2d(arr)
+
+# Proximal
+prox = l2d.prox(arr, .1)
+proxdlocal = prox.asarray()
+
+# Gradient
+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))
+
+###############################################################################
+# Next we move onto the more general case
+# :math:`||\mathbf{Op} \mathbf{x} - \mathbf{b}||_2^2`
+
+solver="cgls"
+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 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()
+
+arrlocal = arr.asarray()
+if rank == 0:
+ flocal = l2(arrlocal)
+ proxlocal = l2.prox(arrlocal, .1)
+ gradlocal = l2.grad(arrlocal)
+ 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))
diff --git a/examples/plot_proxsolver.py b/examples/plot_proxsolver.py
new file mode 100644
index 00000000..4575f3f0
--- /dev/null
+++ b/examples/plot_proxsolver.py
@@ -0,0 +1,208 @@
+r"""
+Proximal solvers
+================
+
+This example demonstrates the use of the solvers in the
+``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
+
+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.
+# 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,))
+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
+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)
+
+ 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.tight_layout()
+
+
+###############################################################################
+# 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=False, iter_lim=5,
+ )[0]
+
+ 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()
+
+
+###############################################################################
+# 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=False, iter_lim=5,
+ )[0]
+
+ 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()
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/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/ProxOperator.py b/pylops_mpi/proximal/ProxOperator.py
new file mode 100644
index 00000000..676286cc
--- /dev/null
+++ b/pylops_mpi/proximal/ProxOperator.py
@@ -0,0 +1,140 @@
+from mpi4py import MPI
+from typing import Any
+
+from pyproximal import ProxOperator
+from pylops.utils.backend import get_module
+
+from pylops_mpi import DistributedArray, Partition
+
+
+_call_reduce_op = dict(
+ Box=(MPI.LAND, all),
+ L0=(MPI.SUM, sum),
+ L1=(MPI.SUM, 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 __repr__(self) -> str:
+ 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.
+
+ 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
+
+ """
+ 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)
+
+ 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 _as_scalar(recv_buf)
+ else:
+ # 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)]
+ f = reduce_op(fs)
+ return f
+
+ def prox(self, x: DistributedArray, tau: float, **kwargs: Any) -> DistributedArray:
+ """Proximal operator applied to a vector
+ """
+ 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:
+ """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..3a4e9c82
--- /dev/null
+++ b/pylops_mpi/proximal/__init__.py
@@ -0,0 +1,24 @@
+"""
+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.
+
+A list of proximal solvers present in pylops_mpi.proximal.optimization.primal.
+
+"""
+
+from .ProxOperator import *
+from .proximal import *
+from .optimization import *
+
+
+__all__ = [
+ "MPIProxOperator",
+]
diff --git a/pylops_mpi/proximal/optimization/__init__.py b/pylops_mpi/proximal/optimization/__init__.py
new file mode 100644
index 00000000..cf7c2bb7
--- /dev/null
+++ b/pylops_mpi/proximal/optimization/__init__.py
@@ -0,0 +1,21 @@
+"""
+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
+ ADMML2 ADMM with L2 misfit term
+
+"""
+
+from .primal import *
+
+
+__all__ = [
+ "ProximalGradient",
+ "ADMML2",
+]
diff --git a/pylops_mpi/proximal/optimization/primal.py b/pylops_mpi/proximal/optimization/primal.py
new file mode 100644
index 00000000..3ac30a73
--- /dev/null
+++ b/pylops_mpi/proximal/optimization/primal.py
@@ -0,0 +1,347 @@
+import sys
+import time
+from collections.abc import Callable
+from math import sqrt
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from pylops.utils.backend import to_numpy
+from pylops.utils.typing import NDArray
+
+from pyproximal.optimization.primal import _x0z0_init
+
+from pylops_mpi import DistributedArray, StackedDistributedArray
+from pylops_mpi.basicoperators import MPIStackedVStack
+from pylops_mpi.optimization.basic import cgls
+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,
+ eta: float = 1.0,
+ niter: int = 10,
+ niterback: int = 100,
+ acceleration: str | None = None,
+ tol: float | None = None,
+ callback: Callable[[DistributedArray], None] | None = None,
+ show: bool = False,
+) -> 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
+
+ # 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\tepsg = %s\n"
+ "niter = %d\ttol = %s\n"
+ ""
+ "niterback = %d\tacceleration = %s\n"
+ % (
+ proxf,
+ proxg,
+ str(tau),
+ epsg_print,
+ niter,
+ str(tol),
+ niterback,
+ acceleration,
+ )
+ )
+ head = " Itn x[0] f g J=f+eps*g tau"
+ print(head)
+ sys.stdout.flush()
+
+ # 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 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
+ )
+
+ # 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")
+ 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
+
+ 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
+
+ # 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 = 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,
+ 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
new file mode 100644
index 00000000..1bc6e231
--- /dev/null
+++ b/pylops_mpi/proximal/proximal/L2.py
@@ -0,0 +1,192 @@
+from math import sqrt
+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
+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
+
+
+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(to_numpy(f.item()))
+
+ 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:
+ 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,
+ 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:
+ 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..57d990fb
--- /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",
+]
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
diff --git a/tests/test_distributedarray.py b/tests/test_distributedarray.py
index a32b1354..7ea59226 100644
--- a/tests/test_distributedarray.py
+++ b/tests/test_distributedarray.py
@@ -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
diff --git a/tests/test_prox.py b/tests/test_prox.py
new file mode 100644
index 00000000..7d28a79f
--- /dev/null
+++ b/tests/test_prox.py
@@ -0,0 +1,201 @@
+"""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 Diagonal
+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
+} # 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
+ (does not support complex numbers)"""
+ 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()
+
+ 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)
+
+
+@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)
+
+ 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)
+
+
+@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)
+
+ 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), (par1j), (par1b)]
+)
+@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),
+ 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),
+ 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'] * (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 = Diagonal(
+ np.ones(par['n'] * (size if par["partition"] == pylops_mpi.Partition.SCATTER else 1), dtype=par['dtype']),
+ dtype=par['dtype'])
+ 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:
+ Op_local = Diagonal(np.ones(par['n'], dtype=par['dtype']), dtype=par['dtype'])
+ Opd = pylops_mpi.MPILinearOperator(Op_local)
+
+ 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, 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]):
+ 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)
diff --git a/tests/test_proxsolver.py b/tests/test_proxsolver.py
new file mode 100644
index 00000000..3fe0d9e3
--- /dev/null
+++ b/tests/test_proxsolver.py
@@ -0,0 +1,290 @@
+"""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, ADMML2
+
+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, dtype=par['dtype']), ])
+
+ 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, dtype=par['dtype'])
+ 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 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
+
+ # 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, 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
+
+ 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)
+
+
+@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-3)
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 cdccce8f..a794f512 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]
@@ -191,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.
@@ -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))
+###############################################################################
+# TV-Regularized inversion
+
+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,8 @@
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=(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")
@@ -277,6 +302,16 @@
axs[5][2].set_title('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
# tutorial `_
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.
diff --git a/tutorials_cupy/poststack_cupy.py b/tutorials_cupy/poststack_cupy.py
index 16f53a4f..5521173e 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),
@@ -125,25 +126,42 @@
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))
###############################################################################
-
# 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.
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.
@@ -156,9 +174,10 @@
# 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=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 +222,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')
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')