Merge pull request #89 from mrava87/patch-factorize

feature: added factorize option to L2

Author: mrava87

pyproximal/proximal/L2.py

@@ -1,5 +1,5 @@
 import numpy as np
-
+from scipy.linalg import cho_factor, cho_solve
 from scipy.sparse.linalg import lsqr
 from pylops import MatrixMult, Identity
 from pyproximal.ProxOperator import _check_tau
@@ -39,9 +39,12 @@ class L2(ProxOperator):
         Warm start (``True``) or not (``False``). Uses estimate from previous
         call of ``prox`` method.
     densesolver : :obj:`str`, optional
-        Use ``numpy`` or ``scipy`` solver when dealing with explicit operators.
-        Choose ``densesolver=None`` when using PyLops versions earlier than
-        v1.18.1 or v2.0.0
+        Use ``numpy``, ``scipy``, or ``factorize`` when dealing with explicit
+        operators. The former two rely on dense solvers from either library,
+        whilst the last computes a factorization of the matrix to invert and
+        avoids to do so unless the :math:`\tau` or :math:`\sigma` paramets
+        have changed. Choose ``densesolver=None`` when using PyLops versions
+        earlier than v1.18.1 or v2.0.0
 
     Notes
     -----
@@ -96,6 +99,11 @@ def __init__(self, Op=None, b=None, q=None, sigma=1., alpha=1.,
         self.densesolver = densesolver
         self.count = 0
 
+        # when using factorize, store the first tau*sigma=0 so that the
+        # first time it will be recomputed (as tau cannot be 0)
+        if self.densesolver == 'factorize':
+            self.tausigma = 0
+
         # create data term
         if self.Op is not None and self.b is not None:
             self.OpTb = self.sigma * self.Op.H @ self.b
@@ -137,14 +145,23 @@ def prox(self, x, tau):
             if self.q is not None:
                 y -= tau * self.alpha * self.q
             if self.Op.explicit:
-                Op1 = MatrixMult(np.eye(self.Op.shape[1]) +
-                                 tau * self.sigma * self.ATA)
-                if self.densesolver is None:
-                    # to allow backward compatibility with PyLops versions earlier
-                    # than v1.18.1 and v2.0.0
-                    x = Op1.div(y)
+                if self.densesolver != 'factorize':
+                    Op1 = MatrixMult(np.eye(self.Op.shape[1]) +
+                                     tau * self.sigma * self.ATA)
+                    if self.densesolver is None:
+                        # to allow backward compatibility with PyLops versions earlier
+                        # than v1.18.1 and v2.0.0
+                        x = Op1.div(y)
+                    else:
+                        x = Op1.div(y, densesolver=self.densesolver)
                 else:
-                    x = Op1.div(y, densesolver=self.densesolver)
+                    if self.tausigma != tau * self.sigma:
+                        # recompute factorization
+                        self.tausigma = tau * self.sigma
+                        ATA = np.eye(self.Op.shape[1]) + \
+                              self.tausigma * self.ATA
+                        self.cl = cho_factor(ATA)
+                    x = cho_solve(self.cl, y)
             else:
                 Op1 = Identity(self.Op.shape[1], dtype=self.Op.dtype) + \
                       tau * self.sigma * self.Op.H * self.Op

pytests/test_norms.py

@@ -89,20 +89,21 @@ def test_L2_op(par):
 def test_L2_dense(par):
     """L2 norm of Op*x with dense Op and proximal/dual proximal
     """
-    b = np.zeros(par['nx'], dtype=par['dtype'])
-    d = np.random.normal(0., 1., par['nx']).astype(par['dtype'])
-    l2 = L2(Op=MatrixMult(np.diag(d), dtype=par['dtype']),
-            b=b, sigma=par['sigma'], densesolver='numpy')
-
-    # norm
-    x = np.random.normal(0., 1., par['nx']).astype(par['dtype'])
-    assert l2(x) == (par['sigma'] / 2.) * np.linalg.norm(d * x) ** 2
-
-    # prox: since Op is a Diagonal operator the denominator becomes
-    # 1 + sigma*tau*d[i] for every i
-    tau = 2.
-    den = 1. + par['sigma'] * tau * d ** 2
-    assert_array_almost_equal(l2.prox(x, tau), x / den, decimal=4)
+    for densesolver in ('numpy', 'scipy', 'factorize'):
+        b = np.zeros(par['nx'], dtype=par['dtype'])
+        d = np.random.normal(0., 1., par['nx']).astype(par['dtype'])
+        l2 = L2(Op=MatrixMult(np.diag(d), dtype=par['dtype']),
+                b=b, sigma=par['sigma'], densesolver=densesolver)
+
+        # norm
+        x = np.random.normal(0., 1., par['nx']).astype(par['dtype'])
+        assert l2(x) == (par['sigma'] / 2.) * np.linalg.norm(d * x) ** 2
+
+        # prox: since Op is a Diagonal operator the denominator becomes
+        # 1 + sigma*tau*d[i] for every i
+        tau = 2.
+        den = 1. + par['sigma'] * tau * d ** 2
+        assert_array_almost_equal(l2.prox(x, tau), x / den, decimal=4)
 
 
 @pytest.mark.parametrize("par", [(par1), (par2)])