feature: added bilinear update to ProximalGradient

Author: mrava87

pyproximal/optimization/primal.py

@@ -6,6 +6,7 @@
 from pylops.optimization.leastsquares import regularized_inversion
 from pylops.utils.backend import to_numpy
 from pyproximal.proximal import L2
+from pyproximal.utils.bilinear import BilinearOperator
 
 
 def _backtracking(x, tau, proxf, proxg, epsg, beta=0.5, niterback=10):
@@ -239,7 +240,11 @@ def ProximalGradient(proxf, proxg, x0, tau=None, beta=0.5,
         else:
             x, tau = _backtracking(y, tau, proxf, proxg, epsg,
                                    beta=beta, niterback=niterback)
-        
+
+        # update internal parameters for bilinear operator
+        if isinstance(proxf, BilinearOperator):
+            proxf.updatexy(x)
+
         # update y
         if acceleration == 'vandenberghe':
             omega = iiter / (iiter + 3)

pyproximal/utils/bilinear.py

@@ -16,6 +16,11 @@ class BilinearOperator():
     - ``lx``: Lipschitz constant of :math:`\nabla_x H`
     - ``ly``: Lipschitz constant of :math:`\nabla_y H`
 
+    Two additional methods (``updatex`` and ``updatey``) are provided to
+    update the :math:`\mathbf{x}` and :math:`\mathbf{x}` internal
+    variables. It is user responsability to choose when to invoke such
+    method (i.e., when to update the internal variables).
+
     Notes
     -----
     A bilinear operator is defined as a differentiable nonlinear function
@@ -45,15 +50,17 @@ def ly(self, y):
         pass
 
     def updatex(self, x):
-        """Update x variable (used when evaluating the gradient over y
+        """Update x variable (to be used to update the internal variable x)
         """
         self.x = x
 
     def updatey(self, y):
-        """Update y variable (used when evaluating the gradient over y
+        """Update y variable (to be used to update the internal variable y)
         """
         self.y = y
 
+    def updatexy(self, xy):
+        pass
 
 class LowRankFactorizedMatrix(BilinearOperator):
     r"""Low-Rank Factorized Matrix operator.
@@ -83,16 +90,21 @@ class LowRankFactorizedMatrix(BilinearOperator):
 
     .. math::
 
-        \nabla_x H = \mathbf{Op}^H(\mathbf{Op}(\mathbf{X}\mathbf{Y})
+        \nabla_x H(\mathbf{x};\ mathbf{y}) =
+        \mathbf{Op}^H(\mathbf{Op}(\mathbf{X}\mathbf{Y})
         - \mathbf{d})\mathbf{Y}^H
 
     and gradient with respect to y equal to:
 
     .. math::
 
-        \nabla_y H = \mathbf{X}^H \mathbf{Op}^H(\mathbf{Op}
+        \nabla_y H(\mathbf{y}; \mathbf{x}) =
+        \mathbf{X}^H \mathbf{Op}^H(\mathbf{Op}
         (\mathbf{X}\mathbf{Y}) - \mathbf{d})
 
+    Note that in both cases, the currently stored x/y is used for
+    the second variable within parenthesis (after ;)
+
     """
     def __init__(self, X, Y, d, Op=None):
         self.n, self.k = X.shape
@@ -160,7 +172,7 @@ def gradx(self, x):
             r = (self.Op.H @ r).reshape(self.n, self.m)
         else:
             r = r.reshape(self.n, self.m)
-        g = -r @ self.y.reshape(self.k, self.m).T
+        g = -r @ np.conj(self.y.reshape(self.k, self.m).T)
         return g.ravel()
 
     def grady(self, y):
@@ -169,13 +181,15 @@ def grady(self, y):
             r = (self.Op.H @ r.ravel()).reshape(self.n, self.m)
         else:
             r = r.reshape(self.n, self.m)
-        g = -self.x.reshape(self.n, self.k).T @ r
+        g = -np.conj(self.x.reshape(self.n, self.k).T) @ r
         return g.ravel()
 
     def grad(self, x):
-        self.updatex(x[:self.n * self.k])
-        self.updatey(x[self.n * self.k:])
         gx = self.gradx(x[:self.n * self.k])
         gy = self.grady(x[self.n * self.k:])
         g = np.hstack([gx, gy])
         return g
+
+    def updatexy(self, x):
+        self.updatex(x[:self.n * self.k])
+        self.updatey(x[self.n * self.k:])

pytests/test_bilinear.py

@@ -0,0 +1,50 @@
+import pytest
+
+import numpy as np
+from numpy.testing import assert_array_equal
+from pylops import Diagonal
+
+from pyproximal.utils.bilinear import LowRankFactorizedMatrix
+
+
+par1 = {'n': 21, 'm': 11, 'k': 5, 'imag': 0,
+        'dtype': 'float32'}  # real
+par2 = {'n': 21, 'm': 11, 'k': 5, 'imag': 1j,
+        'dtype': 'complex64'}  # complex
+
+np.random.seed(10)
+
+
+@pytest.mark.parametrize("par", [(par1), (par2)])
+def test_lrfactorized(par):
+    """Check equivalence of matvec operations in LowRankFactorizedMatrix
+    """
+    U = (np.random.normal(0., 1., (par['n'], par['k'])) + \
+        par['imag'] * np.random.normal(0., 1., (par['n'], par['k']))).astype(par['dtype'])
+    V = (np.random.normal(0., 1., (par['m'], par['k'])) + \
+         par['imag'] * np.random.normal(0., 1., (par['m'], par['k']))).astype(par['dtype'])
+
+    X = U @ V.T
+    LOp = LowRankFactorizedMatrix(U, V.T, X)
+
+    assert_array_equal(X.ravel(), LOp._matvecx(U.ravel()))
+    assert_array_equal(X.ravel(), LOp._matvecy(V.T.ravel()))
+
+
+@pytest.mark.parametrize("par", [(par1), (par2)])
+def test_lrfactorizedoperator(par):
+    """Check equivalence of matvec operations in LowRankFactorizedMatrix with operator Op
+    """
+    U = (np.random.normal(0., 1., (par['n'], par['k'])) + \
+        par['imag'] * np.random.normal(0., 1., (par['n'], par['k']))).astype(par['dtype'])
+    V = (np.random.normal(0., 1., (par['m'], par['k'])) + \
+         par['imag'] * np.random.normal(0., 1., (par['m'], par['k']))).astype(par['dtype'])
+    Op = Diagonal(np.arange(par['n'] * par['m']) + 1.)
+
+    X = U @ V.T
+    y = Op @ X.ravel()
+
+    LOp = LowRankFactorizedMatrix(U, V.T, y, Op)
+
+    assert_array_equal(y, LOp._matvecx(U.ravel()))
+    assert_array_equal(y, LOp._matvecy(V.T.ravel()))