Merge pull request #107 from mrava87/dev

feature: started porting for PyLops v2

Author: mrava87

environment-dev.yml

@@ -7,7 +7,7 @@ dependencies:
   - python>=3.8.12
   - numpy>=1.15.0
   - scipy>=1.8.0
-  - pylops<2.0.0
+  - pylops>=2.0.0
   - scikit-image
   - matplotlib
   - ipython

environment.yml

@@ -5,4 +5,4 @@ dependencies:
   - python>=3.8.12
   - numpy>=1.15.0
   - scipy>=1.8.0
-  - pylops<2.0.0
+  - pylops>=2.0.0

pyproximal/optimization/primal.py

@@ -4,6 +4,7 @@
 
 from math import sqrt
 from pylops.optimization.leastsquares import RegularizedInversion
+from pylops.utils.backend import to_numpy
 from pyproximal.proximal import L2
 
 
@@ -258,7 +259,7 @@ def ProximalGradient(proxf, proxg, x0, tau=None, beta=0.5,
             if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0:
                 pf, pg = proxf(x), proxg(x)
                 msg = '%6g  %12.5e  %10.3e  %10.3e  %10.3e' % \
-                      (iiter + 1, x[0] if x.ndim == 1 else x[0, 0],
+                      (iiter + 1, np.real(to_numpy(x[0])) if x.ndim == 1 else np.real(to_numpy(x[0, 0])),
                        pf, pg[0] if epsg_print == 'Multi' else pg,
                        pf + np.sum(epsg * pg))
                 print(msg)
@@ -1029,7 +1030,7 @@ def TwIST(proxg, A, b, x0, alpha=None, beta=None, eigs=None, niter=10,
             if iiter < 10 or niter - iiter < 10 or iiter % (niter // 10) == 0:
                 pf, pg = proxf(x), proxg(x)
                 msg = '%6g  %12.5e  %10.3e  %10.3e  %10.3e' % \
-                      (iiter + 1, x[0], pf, pg, pf + pg)
+                      (iiter + 1, np.real(to_numpy(x[0])), pf, pg, pf + pg)
                 print(msg)
     if show:
         print('\nTotal time (s) = %.2f' % (time.time() - tstart))

pyproximal/optimization/primaldual.py

@@ -1,9 +1,11 @@
 import time
 import numpy as np
 
+from pylops.utils.backend import get_array_module
 
-def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
-               gfirst=True, callback=None, callbacky=False, show=False):
+
+def PrimalDual(proxf, proxg, A, x0, tau, mu, y0=None, z=None, theta=1., niter=10,
+               gfirst=True, callback=None, callbacky=False, returny=False, show=False):
     r"""Primal-dual algorithm
 
     Solves the following (possibly) nonlinear minimization problem using
@@ -45,6 +47,8 @@ def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
     mu : :obj:`float` or :obj:`np.ndarray`
         Stepsize of subgradient of :math:`g^*`. This can be constant 
         or function of iterations (in the latter cases provided as np.ndarray)
+    z0 : :obj:`numpy.ndarray`
+        Initial auxiliary vector
     z : :obj:`numpy.ndarray`, optional
         Additional vector
     theta : :obj:`float`
@@ -62,6 +66,8 @@ def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
         where ``x`` is the current model vector
     callbacky : :obj:`bool`, optional
         Modify callback signature to (``callback(x, y)``) when ``callbacky=True``
+    returny : :obj:`bool`, optional
+        Return also ``y``
     show : :obj:`bool`, optional
         Display iterations log
 
@@ -102,13 +108,15 @@ def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
         Imaging and Vision, 40, 8pp. 120-145. 2011.
 
     """
+    ncp = get_array_module(x0)
+
     # check if tau and mu are scalars or arrays
     fixedtau = fixedmu = False
     if isinstance(tau, (int, float)):
-        tau = tau * np.ones(niter)
+        tau = tau * ncp.ones(niter)
         fixedtau = True
     if isinstance(mu, (int, float)):
-        mu = mu * np.ones(niter)
+        mu = mu * ncp.ones(niter)
         fixedmu = True
 
     if show:
@@ -128,8 +136,7 @@ def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
 
     x = x0.copy()
     xhat = x.copy()
-    y = np.zeros(A.shape[0], dtype=x.dtype)
-
+    y = y0.copy() if y0 is not None else ncp.zeros(A.shape[0], dtype=x.dtype)
     for iiter in range(niter):
         xold = x.copy()
         if gfirst:
@@ -165,7 +172,10 @@ def PrimalDual(proxf, proxg, A, x0, tau, mu, z=None, theta=1., niter=10,
     if show:
         print('\nTotal time (s) = %.2f' % (time.time() - tstart))
         print('---------------------------------------------------------\n')
-    return x
+    if not returny:
+        return x
+    else:
+        return x, y
 
 
 def AdaptivePrimalDual(proxf, proxg, A, x0, tau, mu,

pyproximal/projection/AffineSet.py

@@ -1,6 +1,7 @@
-import numpy as np
-from scipy.sparse.linalg import lsqr
+from scipy.sparse.linalg import lsqr as sp_lsqr
 
+from pylops.optimization.basic import lsqr
+from pylops.utils.backend import get_array_module, get_module_name
 
 class AffineSetProj():
     r"""Affine set projection.
@@ -39,6 +40,9 @@ def __init__(self, Op, b, niter):
         self.niter = niter
 
     def __call__(self, x):
-        inv = lsqr(self.Op * self.Op.H, self.Op * x - self.b, iter_lim=self.niter)[0]
-        y = x - self.Op.H * inv
+        if get_module_name(get_array_module(x)) == 'numpy':
+            inv = sp_lsqr(self.Op * self.Op.H, self.Op * x - self.b, iter_lim=self.niter)[0]
+        else:
+            inv = lsqr(self.Op * self.Op.H, self.Op * x - self.b, niter=self.niter)[0]
+        y = x - self.Op.H * inv.ravel() # currently ravel is added to ensure that the output is always a vector
         return y

pyproximal/proximal/Nuclear.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pylops.optimization.sparsity import _softthreshold
+from pylops.optimization.cls_sparsity import _softthreshold
 from pyproximal.ProxOperator import _check_tau
 from pyproximal.projection import NuclearBallProj
 from pyproximal import ProxOperator

pyproximal/proximal/Orthogonal.py

@@ -63,12 +63,11 @@ def __init__(self, f, Q, partial=False, b=None, alpha=1.):
         self.Q = Q
         self.partial = partial
         self.alpha = alpha
-        self.b = b if b is not None else np.zeros(Q.shape[0], dtype=Q.dtype)
+        self.b = b if b is not None else 0
 
     def __call__(self, x):
         y = self.Q.matvec(x)
-        if self.b is not None:
-            y += self.b
+        y += self.b
         f = self.f(y)
         return f
 
@@ -81,10 +80,8 @@ def prox(self, x, tau):
                  self.Q.rmatvec(self.f.prox(y + self.b, self.alpha * tau) -
                                 self.b))
         else:
-            if self.b is not None:
-                y = y + self.b
+            y = y + self.b
             z = self.f.prox(y, tau)
-            if self.b is not None:
-                z = z - self.b
+            z = z - self.b
             z = self.Q.rmatvec(z)
         return z

pyproximal/proximal/QuadraticEnvelope.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pylops.optimization.sparsity import _hardthreshold
+from pylops.optimization.cls_sparsity import _hardthreshold
 from pyproximal.ProxOperator import _check_tau
 from pyproximal import ProxOperator
 

pyproximal/proximal/TV.py

@@ -48,13 +48,12 @@ def __init__(self, dims, sigma=1., niter=10, rtol=1e-4, **kwargs):
     def __call__(self, x):
         x = x.reshape(self.dims)
         if self.ndim == 1:
-            derivOp = FirstDerivative(self.dims[0], dims=None, dir=0, edge=False,
+            derivOp = FirstDerivative(dims=self.dims[0], axis=0, edge=False,
                                       dtype=x.dtype, kind="forward")
             dx = derivOp @ x
             y = np.sum(np.abs(dx), axis=0)
         elif self.ndim >= 2:
             y = 0
-            grads = []
             gradOp = Gradient(self.dims, edge=False, dtype=x.dtype, kind="forward")
             grads = gradOp.matvec(x.ravel())
             grads = grads.reshape((self.ndim, ) + self.dims)

requirements-dev.txt

@@ -1,6 +1,6 @@
 numpy>=1.15.0
 scipy>=1.8.0
-pylops<=1.18.3
+pylops>=2.0.0
 numba
 scikit-image
 matplotlib