Merge pull request #612 from jeverink/Regularized-Gaussian-TVwNonnegativity

Regularized Gaussian: Combined regularized and constraint presets

Author: chaozg

cuqi/experimental/mcmc/_conjugate.py

@@ -147,7 +147,7 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedGaussian-Gamma conjugacy only works with univariate ModifiedHalfNormal prior")
 
-        if self.target.likelihood.distribution.preset not in ["nonnegativity"]:
+        if self.target.likelihood.distribution.preset["constraint"] not in ["nonnegativity"]:
             raise ValueError("RegularizedGaussian-Gamma conjugacy only works with implicit regularized Gaussian likelihood with nonnegativity constraints")
 
         key_value_pairs = _get_conjugate_parameter(self.target)
@@ -183,7 +183,7 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy only works with univariate Gamma prior")
 
-        if self.target.likelihood.distribution.preset not in ["l1", "tv"]:
+        if self.target.likelihood.distribution.preset["regularization"] not in ["l1", "tv"]:
             raise ValueError("RegularizedUnboundedUniform-Gamma conjugacy only works with implicit regularized Gaussian likelihood with l1 or tv regularization")
 
         key_value_pairs = _get_conjugate_parameter(self.target)
@@ -203,12 +203,7 @@ def conjugate_distribution(self):
 
         # Compute likelihood quantities
         x = self.target.likelihood.data
-        if self.target.likelihood.distribution.preset == "l1":
-            m = count_nonzero(x)
-        elif self.target.likelihood.distribution.preset == "tv" and isinstance(self.target.likelihood.distribution.geometry, Continuous1D):
-            m = count_constant_components_1D(x)
-        elif self.target.likelihood.distribution.preset == "tv" and isinstance(self.target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
-            m = count_constant_components_2D(self.target.likelihood.distribution.geometry.par2fun(x))
+        m = _compute_sparsity_level(self.target)
 
         reg_op = self.target.likelihood.distribution._regularization_oper
         reg_strength = self.target.likelihood.distribution(np.array([1])).strength
@@ -224,7 +219,7 @@ def validate_target(self):
         if self.target.prior.dim != 1:
             raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy only works with univariate ModifiedHalfNormal prior")
 
-        if self.target.likelihood.distribution.preset not in ["l1", "tv"]:
+        if self.target.likelihood.distribution.preset["regularization"] not in ["l1", "tv"]:
             raise ValueError("RegularizedGaussian-ModifiedHalfNormal conjugacy only works with implicit regularized Gaussian likelihood with l1 or tv regularization")
 
         key_value_pairs = _get_conjugate_parameter(self.target)
@@ -254,13 +249,8 @@ def conjugate_distribution(self):
         x = self.target.likelihood.data
         mu = self.target.likelihood.distribution.mean
         L = self.target.likelihood.distribution(np.array([1])).sqrtprec
-
-        if self.target.likelihood.distribution.preset == "l1":
-            m = count_nonzero(x)
-        elif self.target.likelihood.distribution.preset == "tv" and isinstance(self.target.likelihood.distribution.geometry, Continuous1D):
-            m = count_constant_components_1D(x)
-        elif self.target.likelihood.distribution.preset == "tv" and isinstance(self.target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
-            m = count_constant_components_2D(self.target.likelihood.distribution.geometry.par2fun(x))
+        
+        m = _compute_sparsity_level(self.target)
 
         reg_op = self.target.likelihood.distribution._regularization_oper
         reg_strength = self.target.likelihood.distribution(np.array([1])).strength
@@ -275,6 +265,26 @@ def conjugate_distribution(self):
         return ModifiedHalfNormal(conj_alpha, conj_beta, conj_gamma)
 
 
+def _compute_sparsity_level(target):
+    """Computes the sparsity level in accordance with Section 4 from [2],"""
+    x = target.likelihood.data
+    if target.likelihood.distribution.preset["constraint"] == "nonnegativity":
+        if target.likelihood.distribution.preset["regularization"] == "l1":
+            m = count_nonzero(x)
+        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            m = count_constant_components_1D(x, lower = 0.0)
+        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            m = count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x), lower = 0.0)
+    else: # No constraints, only regularization
+        if target.likelihood.distribution.preset["regularization"] == "l1":
+            m = count_nonzero(x)
+        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, Continuous1D):
+            m = count_constant_components_1D(x)
+        elif target.likelihood.distribution.preset["regularization"] == "tv" and isinstance(target.likelihood.distribution.geometry, (Continuous2D, Image2D)):
+            m = count_constant_components_2D(target.likelihood.distribution.geometry.par2fun(x))
+    return m
+
+
 def _get_conjugate_parameter(target):
     """Extract the conjugate parameter name (e.g. d), and returns the mutable variable that is defined by the conjugate parameter, e.g. cov and its value e.g. lambda d:1/d"""
     par_name = target.prior.name

cuqi/experimental/mcmc/_rto.py

@@ -239,7 +239,7 @@ def solver(self):
     @solver.setter
     def solver(self, value):
         if value == "ScipyLinearLSQ":
-            if (self.target.prior._preset == "nonnegativity" or self.target.prior._preset == "box"):
+            if (self.target.prior.preset["constraint"] == "nonnegativity" or self.target.prior.preset["constraint"] == "box"):
                 self._solver = value
             else:
                 raise ValueError("ScipyLinearLSQ only supports RegularizedGaussian with box or nonnegativity constraint.")

cuqi/implicitprior/_regularizedGMRF.py

@@ -68,6 +68,7 @@ def __init__(self, mean=None, prec=None, bc_type='zero', order=1, proximal = Non
             # Init from abstract distribution class
             super(Distribution, self).__init__(**kwargs)
 
+            self._force_list = False
             self._parse_regularization_input_arguments(proximal, projector, constraint, regularization, args)
 
 

cuqi/implicitprior/_regularizedGaussian.py

@@ -2,9 +2,10 @@
 from cuqi.distribution import Distribution, Gaussian
 from cuqi.solver import ProjectNonnegative, ProjectBox, ProximalL1
 from cuqi.geometry import Continuous1D, Continuous2D, Image2D
-from cuqi.operator import FirstOrderFiniteDifference
+from cuqi.operator import FirstOrderFiniteDifference, Operator
 
 import numpy as np
+import scipy.sparse as sparse
 from copy import copy
 
 
@@ -48,6 +49,8 @@ class RegularizedGaussian(Distribution):
             min_z 0.5||x-z||_2^2+scale*g(x).
         If list of tuples (callable proximal operator of f_i, linear operator L_i):
             Each callable proximal operator of f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+            Each linear operator needs to have the '__matmul__', 'T' and 'shape' attributes;
+            this includes numpy.ndarray, scipy.sparse.sparray, scipy.sparse.linalg.LinearOperator and cuqi.operator.Operator.
             The corresponding regularization takes the form
                 sum_i f_i(L_i x),
             where the sum ranges from 1 to an arbitrary n.
@@ -88,59 +91,137 @@ def __init__(self, mean=None, cov=None, prec=None, sqrtcov=None, sqrtprec=None,
         # Init from abstract distribution class
         super().__init__(**kwargs)
 
+        self._force_list = False
         self._parse_regularization_input_arguments(proximal, projector, constraint, regularization, optional_regularization_parameters)
 
     def _parse_regularization_input_arguments(self, proximal, projector, constraint, regularization, optional_regularization_parameters):
         """ Parse regularization input arguments with guarding statements and store internal states """
 
-        # Check that only one of proximal, projector, constraint or regularization is provided        
-        if (proximal is not None) + (projector is not None) + (constraint is not None) + (regularization is not None) != 1:
-            raise ValueError("Precisely one of proximal, projector, constraint or regularization needs to be provided.")
+        # Guards checking whether the regularization inputs are valid
+        if (proximal is not None) + (projector is not None) + max((constraint is not None), (regularization is not None)) == 0:
+            raise ValueError("At least some constraint or regularization has to be specified for RegularizedGaussian")
+            
+        if (proximal is not None) + (projector is not None) == 2:
+            raise ValueError("Only one of proximal or projector can be used.")
+
+        if (proximal is not None) + (projector is not None) + max((constraint is not None), (regularization is not None)) > 1:
+            raise ValueError("User-defined proximals and projectors cannot be combined with pre-defined constraints and regularization.")
+
+        # Branch between user-defined and preset
+        if (proximal is not None) + (projector is not None) >= 1:
+            self._parse_user_specified_input(proximal, projector)
+        else:
+            # Set constraint and regularization presets for use with Gibbs
+            self._preset = {"constraint": None,
+                            "regularization": None}
 
+            self._parse_preset_constraint_input(constraint, optional_regularization_parameters)
+            self._parse_preset_regularization_input(regularization, optional_regularization_parameters)
+
+            # Merge
+            self._merge_predefined_option()
+
+    def _parse_user_specified_input(self, proximal, projector):
+        # Guard for checking partial validy of proximals or projectors
+        if proximal is not None:
+            if callable(proximal):
+                if len(get_non_default_args(proximal)) != 2:
+                    raise ValueError("Proximal should take 2 arguments.")
+            elif isinstance(proximal, list):
+                for val in proximal:
+                    if len(val) != 2:
+                        raise ValueError("Each value in the proximal list needs to consistent of two elements: a proximal operator and a linear operator.")
+                    if callable(val[0]):
+                        if len(get_non_default_args(val[0])) != 2:
+                            raise ValueError("Proximal should take 2 arguments.")
+                    else:
+                        raise ValueError("Proximal operators need to be callable.")
+                    if not (hasattr(val[1], '__matmul__') and hasattr(val[1], 'T') and hasattr(val[1], 'shape')):
+                        raise ValueError("Linear operator not supported, must have '__matmul__', 'T' and 'shape' attributes.")
+            else:
+                raise ValueError("Proximal needs to be callable or a list. See documentation.")
+            
         if projector is not None:
-            if not callable(projector):
-                raise ValueError("Projector needs to be callable.")
-            if len(get_non_default_args(projector)) != 1:
-                raise ValueError("Projector should take 1 argument.")
+            if callable(projector):
+                if len(get_non_default_args(projector)) != 1:
+                    raise ValueError("Projector should take 1 argument.")
+            else:
+                raise ValueError("Projector needs to be callable")
 
-        # Preset information, for use in Gibbs
-        self._preset = None
-        
+        # Set user-defined proximals or projectors
         if proximal is not None:
-            # No need to generate the proximal and associated information
-            self.proximal = proximal
-        elif projector is not None:
+            self._preset = None
+            self._proximal = proximal
+            return
+        
+        if projector is not None:
+            self._preset = None
             self._proximal = lambda z, gamma: projector(z)
-        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"):
-            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"]
-            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
-            self._preset = "l1"
-        elif (isinstance(regularization, str) and regularization.lower() in ["tv"]):
+            return
+
+    def _parse_preset_constraint_input(self, constraint, optional_regularization_parameters):
+        # Create data for constraints
+        self._constraint_prox = None
+        self._constraint_oper = None
+        if constraint is not None:
+            if not isinstance(constraint, str):
+                raise ValueError("Constraint needs to be specified as a string.")
+            
+            c_lower = constraint.lower()
+            if c_lower == "nonnegativity":
+                self._constraint_prox = lambda z, gamma: ProjectNonnegative(z)
+                self._box_bounds = (np.ones(self.dim)*0, np.ones(self.dim)*np.inf)
+                self._preset["constraint"] = "nonnegativity"
+            elif c_lower == "box":
+                _box_lower = optional_regularization_parameters["lower_bound"]
+                _box_upper = optional_regularization_parameters["upper_bound"]
+                self._proximal = lambda z, _: ProjectBox(z, _box_lower, _box_upper)
+                self._box_bounds = (np.ones(self.dim)*_box_lower, np.ones(self.dim)*_box_upper)
+                self._preset["constraint"] = "box"
+            else:
+                raise ValueError("Constraint not supported.")
+
+    def _parse_preset_regularization_input(self, regularization, optional_regularization_parameters):
+        # Create data for regularization
+        self._regularization_prox = None
+        self._regularization_oper = None
+        if regularization is not None:
+            if not isinstance(regularization, str):
+                raise ValueError("Regularization needs to be specified as a string.")
+                
             self._strength = optional_regularization_parameters["strength"]
-            if isinstance(self.geometry, (Continuous1D, Continuous2D, Image2D)):
-                self._transformation = FirstOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+            r_lower = regularization.lower()
+            if r_lower == "l1":
+                self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+                self._preset["regularization"] = "l1"
+            elif r_lower == "tv":
+                # Store the transformation to reuse when modifying the strength
+                if not isinstance(self.geometry, (Continuous1D, Continuous2D, Image2D)):
+                    raise ValueError("Geometry not supported for total variation")
+                self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+                self._regularization_oper = FirstOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+                self._preset["regularization"] = "tv"
             else:
-                raise ValueError("Geometry not supported for total variation")
-            
-            self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
-            self._regularization_oper = self._transformation
-
-            self._proximal = [(self._regularization_prox, self._regularization_oper)]
-            self._preset = "tv"
-        else:
-            raise ValueError("Regularization not supported")
-
+                raise ValueError("Regularization not supported.")
 
+    def _merge_predefined_option(self):
+        # Check whether it is a single proximal and hence FISTA could be used in RegularizedLinearRTO 
+        if ((not self._force_list) and
+            ((self._constraint_prox is not None) + (self._regularization_prox is not None) == 1) and
+            ((self._constraint_oper is not None) + (self._regularization_oper is not None) == 0)):
+                if self._constraint_prox is not None:
+                    self._proximal = self._constraint_prox
+                else:
+                    self._proximal = self._regularization_prox 
+                return
+
+        # Merge regularization choices in list for use in ADMM by RegularizedLinearRTO
+        self._proximal = []
+        if self._constraint_prox is not None:
+            self._proximal += [(self._constraint_prox, self._constraint_oper if self._constraint_oper is not None else sparse.eye(self.geometry.par_dim))]
+        if self._regularization_prox is not None:
+            self._proximal += [(self._regularization_prox, self._regularization_oper if self._regularization_oper is not None else sparse.eye(self.geometry.par_dim))]
+
     @property
     def transformation(self):
         return self._transformation
@@ -151,15 +232,15 @@ def strength(self):
 
     @strength.setter
     def strength(self, value):
-        if self._preset not in self.regularization_options():
+        if self._preset is None or self._preset["regularization"] is None:
             raise TypeError("Strength is only used when the regularization is set to l1 or TV.")
 
         self._strength = value
-        if self._preset == "tv":
+        if self._preset["regularization"] in ["l1", "tv"]:        
             self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
-            self._proximal = [(self._regularization_prox, self._regularization_oper)]
-        elif self._preset == "l1":
-            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+
+        # Create new list of proximals based on updated regularization
+        self._merge_predefined_option()
 
     # This is a getter only attribute for the underlying Gaussian
     # It also ensures that the name of the underlying Gaussian
@@ -266,7 +347,7 @@ def sqrtcov(self, value):
 
     def get_mutable_variables(self):
         mutable_vars = self.gaussian.get_mutable_variables().copy()
-        if self.preset in self.regularization_options():
+        if self.preset is not None and self.preset['regularization'] in ["l1", "tv"]:
             mutable_vars += ["strength"]
         return mutable_vars
 

cuqi/implicitprior/_regularizedUnboundedUniform.py

@@ -63,4 +63,5 @@ def __init__(self, geometry, proximal = None, projector = None, constraint = Non
                 # Init from abstract distribution class
                 super(Distribution, self).__init__(**kwargs)
 
+                self._force_list = False
                 self._parse_regularization_input_arguments(proximal, projector, constraint, regularization, args)