Feat: proximal module

.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'

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:

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
 -----
 

docs/source/installation.rst

@@ -14,6 +14,7 @@ The minimal set of dependencies for the PyLops-MPI project is:
 * `Matplotlib <https://matplotlib.org/>`_
 * `MPI4py <https://mpi4py.readthedocs.io/en/stable/>`_
 * `PyLops <https://pylops.readthedocs.io/en/stable/>`_
+* `PyProximal <https://pyproximal.readthedocs.io/en/stable/>`_
 
 Additionally, to use the CUDA-aware MPI engine, the following additional 
 dependencies are required:

environment-dev.yml

@@ -9,6 +9,7 @@ dependencies:
   - scipy>=1.8.0
   - mpi4py
   - pylops>=2.0.0
+  - pyproximal
   - matplotlib
   - ipython
   - pytest

environment.yml

@@ -7,5 +7,6 @@ dependencies:
   - numpy>=1.15.0
   - scipy>=1.8.0
   - pylops>=2.0.0
+  - pyproximal
   - matplotlib
   - mpi4py

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))