Merge pull request #602 from CUQI-DTU/add_solvers_to_rrto

Use scipy.optimize.lsq_linear in RegularizedLinearRTO for nonnegativity

Author: chaozg

cuqi/experimental/mcmc/_rto.py

@@ -3,7 +3,7 @@
 from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
-from cuqi.solver import CGLS, FISTA, ADMM
+from cuqi.solver import CGLS, FISTA, ADMM, ScipyLinearLSQ
 from cuqi.experimental.mcmc import Sampler
 
 
@@ -168,6 +168,7 @@ class RegularizedLinearRTO(LinearRTO):
            Used when prior.proximal is callable.
     ADMM:  [2] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
            Used when prior.proximal is a list of penalty terms.
+    ScipyLinearLSQ: Wrapper for Scipy's lsq_linear for the Trust Region Reflective algorithm. Optionally used when the constraint is either "nonnegativity" or "box".
 
     Parameters
     ------------
@@ -178,7 +179,7 @@ class RegularizedLinearRTO(LinearRTO):
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        Maximum number of iterations of the inner FISTA/ADMM solver. *Optional*.
+        Maximum number of iterations of the FISTA/ADMM/ScipyLinearLSQ solver. *Optional*.
 
     inner_max_it : int
         Maximum number of iterations of the CGLS solver used within the ADMM solver. *Optional*.
@@ -188,14 +189,20 @@ class RegularizedLinearRTO(LinearRTO):
         If stepsize is a float, then this stepsize is used.
 
     penalty_parameter : int
-        Penalty parameter of the inner ADMM solver. *Optional*.
+        Penalty parameter of the ADMM solver. *Optional*.
         See [2] or `cuqi.solver.ADMM`
 
     abstol : float
-        Absolute tolerance of the inner FISTA solver. *Optional*.
+        Absolute tolerance of the FISTA/ScipyLinearLSQ solver. *Optional*.
+    
+    inner_abstol : float
+        Tolerance parameter for ScipyLinearLSQ's inner solve of the unbounded least-squares problem. *Optional*.
     
     adaptive : bool
-        If True, FISTA is used as inner solver, otherwise ISTA is used. *Optional*.
+        If True, FISTA is used as solver, otherwise ISTA is used. *Optional*.
+    
+    solver : string
+        If set to "ScipyLinearLSQ", solver is set to cuqi.solver.ScipyLinearLSQ, otherwise FISTA/ISTA or ADMM is used. Note "ScipyLinearLSQ" can only be used with `RegularizedGaussian` of `box` or `nonnegativity` constraint. *Optional*.
 
     callback : callable, *Optional*
         If set this function will be called after every sample.
@@ -204,23 +211,41 @@ class RegularizedLinearRTO(LinearRTO):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, **kwargs):
+    def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, solver=None, inner_abstol=None, **kwargs):
 
         super().__init__(target=target, initial_point=initial_point, **kwargs)
 
         # Other parameters
         self.stepsize = stepsize
-        self.abstol = abstol   
+        self.abstol = abstol
+        self.inner_abstol = inner_abstol
         self.adaptive = adaptive
         self.maxit = maxit
         self.inner_max_it = inner_max_it
         self.penalty_parameter = penalty_parameter
+        self.solver = solver
 
     def _initialize(self):
         super()._initialize()
-        if self._inner_solver == "FISTA":
+        if self.solver is None:
+            self.solver = "FISTA" if callable(self.proximal) else "ADMM"
+        if self.solver == "FISTA":
             self._stepsize = self._choose_stepsize()
 
+    @property
+    def solver(self):
+        return self._solver
+
+    @solver.setter
+    def solver(self, value):
+        if value == "ScipyLinearLSQ":
+            if (self.target.prior._preset == "nonnegativity" or self.target.prior._preset == "box"):
+                self._solver = value
+            else:
+                raise ValueError("ScipyLinearLSQ only supports RegularizedGaussian with box or nonnegativity constraint.")
+        else:
+            self._solver = value
+
     @property
     def proximal(self):
         return self.target.prior.proximal
@@ -229,7 +254,6 @@ def validate_target(self):
         super().validate_target()
         if not isinstance(self.target.prior, (cuqi.implicitprior.RegularizedGaussian, cuqi.implicitprior.RegularizedGMRF)):
             raise TypeError("Prior needs to be RegularizedGaussian or RegularizedGMRF")
-        self._inner_solver = "FISTA" if callable(self.proximal) else "ADMM"
 
     def _choose_stepsize(self):
         if isinstance(self.stepsize, str):
@@ -254,15 +278,25 @@ def prior(self):
     def step(self):
         y = self.b_tild + np.random.randn(len(self.b_tild))
 
-        if self._inner_solver == "FISTA":
+        if self.solver == "FISTA":
             sim = FISTA(self.M, y, self.proximal,
                         self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
-        elif self._inner_solver == "ADMM":
+        elif self.solver == "ADMM":
             sim = ADMM(self.M, y, self.proximal,
-                        self.current_point, self.penalty_parameter, maxit = self.maxit, inner_max_it = self.inner_max_it, adaptive = self.adaptive)  
+                        self.current_point, self.penalty_parameter, maxit = self.maxit, inner_max_it = self.inner_max_it, adaptive = self.adaptive)
+        elif self.solver == "ScipyLinearLSQ":
+                A_op = sp.sparse.linalg.LinearOperator((sum([llh.dim for llh in self.likelihoods])+self.target.prior.dim, self.target.prior.dim),
+                                        matvec=lambda x: self.M(x, 1),
+                                        rmatvec=lambda x: self.M(x, 2)
+                                        )
+                sim = ScipyLinearLSQ(A_op, y, self.target.prior._box_bounds, 
+                                     max_iter = self.maxit,
+                                     lsmr_maxiter = self.inner_max_it, 
+                                     tol = self.abstol,
+                                     lsmr_tol = self.inner_abstol)
         else:
             raise ValueError("Choice of solver not supported.")
 
         self.current_point, _ = sim.solve()
         acc = 1
-        return acc
+        return acc

cuqi/implicitprior/_regularizedGaussian.py

@@ -113,10 +113,12 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         elif (isinstance(constraint, str) and constraint.lower() == "nonnegativity"):
             self._proximal = lambda z, gamma: ProjectNonnegative(z)
             self._preset = "nonnegativity"
+            self._box_bounds = (np.ones(self.dim)*0, np.ones(self.dim)*np.inf)
         elif (isinstance(constraint, str) and constraint.lower() == "box"):
-            lower = optional_regularization_parameters["lower_bound"]
-            upper = optional_regularization_parameters["upper_bound"]
-            self._proximal = lambda z, _: ProjectBox(z, lower, upper)
+            self._box_lower = optional_regularization_parameters["lower_bound"]
+            self._box_upper = optional_regularization_parameters["upper_bound"]
+            self._box_bounds = (np.ones(self.dim)*self._box_lower, np.ones(self.dim)*self._box_upper)
+            self._proximal = lambda z, _: ProjectBox(z, self._box_lower, self._box_upper)
             self._preset = "box" # Not supported in Gibbs
         elif (isinstance(regularization, str) and regularization.lower() in ["l1"]):
             self._strength = optional_regularization_parameters["strength"]

cuqi/solver/__init__.py

@@ -2,7 +2,8 @@
     ScipyLBFGSB,
     ScipyMinimizer,
     ScipyMaximizer,
-    ScipyLeastSquares,
+    ScipyLSQ,
+    ScipyLinearLSQ,
     CGLS,
     LM,
     PDHG,

cuqi/solver/_solver.py

@@ -164,7 +164,7 @@ def ngradfunc(*args,**kwargs):
 
 
 
-class ScipyLeastSquares(object):
+class ScipyLSQ(object):
     """Wrapper for :meth:`scipy.optimize.least_squares`.
 
     Solve nonlinear least-squares problems with bounds:
@@ -227,6 +227,44 @@ def solve(self):
             sol = solution['x']
         return sol, info
 
+class ScipyLinearLSQ(object):
+    """Wrapper for :meth:`scipy.optimize.lsq_linear`.
+
+    Solve linear least-squares problems with bounds:
+
+    .. math::
+    
+        \min \|A x - b\|_2^2
+
+    subject to :math:`lb <= x <= ub`.
+    
+    Parameters
+    ----------
+    A : ndarray, LinearOperator
+        Design matrix (system matrix).
+    b : ndarray
+        The right-hand side of the linear system.
+    bounds : 2-tuple of array_like or scipy.optimize Bounds
+        Bounds for variables. 
+    kwargs : Other keyword arguments passed to Scipy's `lsq_linear`. See documentation of `scipy.optimize.lsq_linear` for details.
+    """
+    def __init__(self, A, b, bounds=(-np.inf, np.inf), **kwargs):
+        self.A = A
+        self.b = b
+        self.bounds = bounds
+        self.kwargs = kwargs
+    
+    def solve(self):
+        """Runs optimization algorithm and returns solution and optimization information.
+
+        Returns
+        ----------
+        solution : Tuple
+            Solution found (array_like) and optimization information (dictionary).
+        """
+        res = opt.lsq_linear(self.A, self.b, bounds=self.bounds, **self.kwargs)
+        x = res.pop('x')
+        return x, res
 
 
 class CGLS(object):

tests/test_solver.py

@@ -1,7 +1,7 @@
 import numpy as np
 import scipy as sp
 
-from cuqi.solver import ScipyLBFGSB, ScipyMinimizer, ScipyLeastSquares, CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative
+from cuqi.solver import ScipyLBFGSB, ScipyMinimizer, ScipyLSQ, ScipyLinearLSQ, CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative, ProjectBox
 from scipy.optimize import lsq_linear
 
 
@@ -54,28 +54,66 @@ def test_ScipyLBFGSB_with_gradient():
     sol_ref = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
     assert np.allclose(sol, sol_ref)
 
-def test_ScipyLeastSquares_without_Jac():
+def test_ScipyLSQ_without_Jac():
     def fun_rosenbrock(x):
         return np.array([10 * (x[1] - x[0]**2), (1 - x[0])])
     x0 = np.array([2, 2])
-    solver = ScipyLeastSquares(fun_rosenbrock, x0)
+    solver = ScipyLSQ(fun_rosenbrock, x0)
     sol, _ = solver.solve()
     sol_ref = np.array([1, 1])
     assert np.allclose(sol, sol_ref)
 
-def test_ScipyLeastSquares_with_Jac():
+def test_ScipyLSQ_with_Jac():
     def fun_rosenbrock(x):
         return np.array([10 * (x[1] - x[0]**2), (1 - x[0])])
     def jac_rosenbrock(x):
         return np.array([
             [-20 * x[0], 10],
             [-1, 0]])
     x0 = np.array([2, 2])
-    solver = ScipyLeastSquares(fun_rosenbrock, x0, jacfun=jac_rosenbrock)
+    solver = ScipyLSQ(fun_rosenbrock, x0, jacfun=jac_rosenbrock)
     sol, _ = solver.solve()
     sol_ref = np.array([1, 1])
     assert np.allclose(sol, sol_ref)
 
+def test_ScipyLinearLSQ_with_matrix():
+    rng = np.random.default_rng(seed = 1219)
+    m, n = 10, 5
+    A = rng.standard_normal((m, n))
+    b = rng.standard_normal(m)
+    res = lsq_linear(A, b, tol=1e-8)
+    ref_sol = res.x
+    sol, _ = ScipyLinearLSQ(A, b).solve()
+    assert np.allclose(sol, ref_sol, rtol=1e-10)
+
+def test_ScipyLinearLSQ_with_LinearOperator():
+    rng = np.random.default_rng(seed = 1219)
+    m, n = 10, 5
+    A = rng.standard_normal((m, n))
+    b = rng.standard_normal(m)
+    A_op = sp.sparse.linalg.LinearOperator((m, n),
+                                           matvec=lambda x: A @ x,
+                                           rmatvec=lambda x: A.T @ x
+                                           )
+    res = lsq_linear(A, b, tol=1e-8)
+    ref_sol = res.x
+    sol, _ = ScipyLinearLSQ(A_op, b).solve()
+    assert np.allclose(sol, ref_sol, rtol=1e-10)
+
+def test_ScipyLinearLSQ_against_FISTA():
+    A = np.array([[73,71,52],[87,74,46],[72,2,7],[80,89,71]])
+    b = np.array([49,67,68,20])
+    # solve with ScipyLinearLSQ
+    lb = np.zeros(3)
+    ub = lb + np.inf
+    sol_lsq, _ = ScipyLinearLSQ(A, b, (lb,ub)).solve()
+    # solve with FISTA
+    rng = np.random.default_rng(seed = 1219)
+    x0 = rng.standard_normal(3)
+    sol_fista, _ = FISTA(A, b, lambda x, _: ProjectNonnegative(x), x0, stepsize=1e-7, maxit=100000, abstol=1e-16, adaptive=True).solve()
+
+    assert np.allclose(sol_lsq, sol_fista, rtol=1e-8)
+
 def test_LM():
     # compare to MATLAB's original code solution
     t = np.arange(1, 10, 2)