CUQIpy 1.1.1.post0.dev36__py3-none-any.whl → 1.4.1.post0.dev124__py3-none-any.whl

cuqi/sampler/_rto.py

@@ -3,7 +3,7 @@ from scipy.linalg.interpolative import estimate_spectral_norm
 from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
-from cuqi.solver import CGLS, FISTA
+from cuqi.solver import CGLS, FISTA, ADMM, ScipyLinearLSQ, ScipyMinimizer
 from cuqi.sampler import Sampler
 
 
@@ -11,7 +11,7 @@ class LinearRTO(Sampler):
     """
     Linear RTO (Randomize-Then-Optimize) sampler.
 
-    Samples posterior related to the inverse problem with Gaussian likelihood and prior, and where the forward model is Linear.
+    Samples posterior related to the inverse problem with Gaussian likelihood and prior, and where the forward model is linear or more generally affine.
 
     Parameters
     ------------
@@ -22,12 +22,12 @@ class LinearRTO(Sampler):
         
         Here:
         data: is a m-dimensional numpy array containing the measured data.
-        model: is a m by n dimensional matrix or LinearModel representing the forward model.
+        model: is a m by n dimensional matrix, AffineModel or LinearModel representing the forward model.
         L_sqrtprec: is the squareroot of the precision matrix of the Gaussian likelihood.
         P_mean: is the prior mean.
         P_sqrtprec: is the squareroot of the precision matrix of the Gaussian mean.
 
-    x0 : `np.ndarray` 
+    initial_point : `np.ndarray` 
         Initial point for the sampler. *Optional*.
 
     maxit : int
@@ -36,67 +36,87 @@ class LinearRTO(Sampler):
     tol : float
         Tolerance of the inner CGLS solver. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    inner_initial_point : string or np.ndarray or cuqi.array.CUQIArray
+        Initial point for the inner optimization problem. Can be "previous_sample" (default), "MAP", or a specific numpy or cuqi array. *Optional*.
         
-    """
-    def __init__(self, target, x0=None, maxit=10, tol=1e-6, shift=0, **kwargs):
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
-        # Accept tuple of inputs and construct posterior
-        if isinstance(target, tuple) and len(target) == 5:
-            # Structure (data, model, L_sqrtprec, P_mean, P_sqrtprec)
-            data = target[0]
-            model = target[1]
-            L_sqrtprec = target[2]
-            P_mean = target[3]
-            P_sqrtprec = target[4]
-
-            # If numpy matrix convert to CUQI model
-            if isinstance(model, np.ndarray) and len(model.shape) == 2:
-                model = cuqi.model.LinearModel(model)
-
-            # Check model input
-            if not isinstance(model, cuqi.model.LinearModel):
-                raise TypeError("Model needs to be cuqi.model.LinearModel or matrix")
-
-            # Likelihood
-            L = cuqi.distribution.Gaussian(model, sqrtprec=L_sqrtprec).to_likelihood(data)
-
-            # Prior TODO: allow multiple priors stacked
-            #if isinstance(P_mean, list) and isinstance(P_sqrtprec, list):
-            #    P = cuqi.distribution.JointGaussianSqrtPrec(P_mean, P_sqrtprec)
-            #else:
-            P = cuqi.distribution.Gaussian(P_mean, sqrtprec=P_sqrtprec)
-
-            # Construct posterior
-            target = cuqi.distribution.Posterior(L, P)
-
-        super().__init__(target, x0=x0, **kwargs)
-
-        self._check_posterior()
-
-        # Modify initial guess        
-        if x0 is not None:
-            self.x0 = x0
-        else:
-            self.x0 = np.zeros(self.prior.dim)
+    """
+    def __init__(self, target=None, initial_point=None, maxit=10, tol=1e-6, inner_initial_point="previous_sample", **kwargs):
+
+        super().__init__(target=target, initial_point=initial_point, **kwargs)
 
         # Other parameters
         self.maxit = maxit
-        self.tol = tol        
-        self.shift = 0
-                
+        self.tol = tol
+        self.inner_initial_point = inner_initial_point
+
+    def _initialize(self):
+        self._precompute()
+        self._compute_map()
+
+    @property
+    def inner_initial_point(self):
+        if isinstance(self._inner_initial_point, str):
+            if self._inner_initial_point == "previous_sample":
+                return self.current_point
+            elif self._inner_initial_point == "map":
+                return self._map
+        else:
+            return self._inner_initial_point
+
+    @inner_initial_point.setter
+    def inner_initial_point(self, value):
+        is_correct_string = (isinstance(value, str) and
+                             (value.lower() == "previous_sample" or
+                              value.lower() == "map"))
+        if is_correct_string:
+            self._inner_initial_point = value.lower()
+        elif isinstance(value, (np.ndarray, cuqi.array.CUQIarray)):
+            self._inner_initial_point = value
+        else:
+            raise ValueError("Invalid value for inner_initial_point. Choose either 'previous_sample', 'MAP', or provide a numpy array/cuqi array.")
+
+    @property
+    def prior(self):
+        return self.target.prior
+
+    @property
+    def likelihood(self):
+        return self.target.likelihood
+    
+    @property
+    def likelihoods(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return [self.target.likelihood]
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
+            return self.target.likelihoods
+
+    @property
+    def model(self):
+        return self.target.model 
+
+    @property
+    def models(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return [self.target.model]
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
+            return self.target.models    
+
+    def _compute_map(self):
+        sim = CGLS(self.M, self.b_tild, self.current_point, self.maxit, self.tol)            
+        self._map, _ = sim.solve()
+
+    def _precompute(self):
         L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
         L2 = self.prior.sqrtprec
         L2mu = self.prior.sqrtprecTimesMean
 
         # pre-computations
-        self.n = len(self.x0)
-        self.b_tild = np.hstack([L@likelihood.data for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
-
+        self.n = self.prior.dim
+        self.b_tild = np.hstack([L@(likelihood.data - model._shift) for (L, likelihood, model) in zip(L1, self.likelihoods, self.models)]+ [L2mu]) # With shift from AffineModel
         callability = [callable(likelihood.model) for likelihood in self.likelihoods]
         notcallability = [not c for c in callability]
         if all(notcallability):
@@ -105,7 +125,7 @@ class LinearRTO(Sampler):
             # in this case, model is a function doing forward and backward operations
             def M(x, flag):
                 if flag == 1:
-                    out1 = [L @ likelihood.model.forward(x) for (L, likelihood) in zip(L1, self.likelihoods)]
+                    out1 = [L @ likelihood.model._forward_func_no_shift(x) for (L, likelihood) in zip(L1, self.likelihoods)] # Use forward function which excludes shift
                     out2 = L2 @ x
                     out  = np.hstack(out1 + [out2])
                 elif flag == 2:
@@ -114,77 +134,42 @@ class LinearRTO(Sampler):
                     out1 = np.zeros(self.n)
                     for likelihood in self.likelihoods:
                         idx_end += len(likelihood.data)
-                        out1 += likelihood.model.adjoint(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end])
+                        out1 += likelihood.model._adjoint_func_no_shift(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end])
                         idx_start = idx_end
                     out2 = L2.T @ x[idx_end:]
                     out  = out1 + out2                
                 return out   
             self.M = M  
         else:
-            raise TypeError("All likelihoods need to be callable or none need to be callable.") 
-
-    @property
-    def prior(self):
-        return self.target.prior
-
-    @property
-    def likelihood(self):
-        return self.target.likelihood
-    
-    @property
-    def likelihoods(self):
-        if isinstance(self.target, cuqi.distribution.Posterior):
-            return [self.target.likelihood]
-        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
-            return self.target.likelihoods
+            raise TypeError("All likelihoods need to be callable or none need to be callable.")
 
-    @property
-    def model(self):
-        return self.target.model     
-    
-    @property
-    def data(self):
-        return self.target.data
-
-    def _sample(self, N, Nb):   
-        Ns = N+Nb   # number of simulations        
-        samples = np.empty((self.n, Ns))
-                     
-        # initial state   
-        samples[:, 0] = self.x0
-        for s in range(Ns-1):
-            y = self.b_tild + np.random.randn(len(self.b_tild))
-            sim = CGLS(self.M, y, samples[:, s], self.maxit, self.tol, self.shift)            
-            samples[:, s+1], _ = sim.solve()
-
-            self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample
-            self._call_callback(samples[:, s+1], s+1)
-
-        # remove burn-in
-        samples = samples[:, Nb:]
-        
-        return samples, None, None
+    def step(self):
+        y = self.b_tild + np.random.randn(len(self.b_tild))
+        sim = CGLS(self.M, y, self.inner_initial_point, self.maxit, self.tol)            
+        self.current_point, _ = sim.solve()
+        acc = 1
+        return acc
 
-    def _sample_adapt(self, N, Nb):
-        return self._sample(N,Nb)
+    def tune(self, skip_len, update_count):
+        pass
     
-    def _check_posterior(self):
+    def validate_target(self):
         # Check target type
         if not isinstance(self.target, (cuqi.distribution.Posterior, cuqi.distribution.MultipleLikelihoodPosterior)):
             raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior' or 'cuqi.distribution.MultipleLikelihoodPosterior'.")       
 
         # Check Linear model and Gaussian likelihood(s)
         if isinstance(self.target, cuqi.distribution.Posterior):
-            if not isinstance(self.model, cuqi.model.LinearModel):
-                raise TypeError("Model needs to be linear")
+            if not isinstance(self.model, cuqi.model.AffineModel):
+                raise TypeError("Model needs to be linear or more generally affine")
 
             if not hasattr(self.likelihood.distribution, "sqrtprec"):
                 raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
             
         elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior): # Elif used for further alternatives, e.g., stacked posterior
             for likelihood in self.likelihoods:
-                if not isinstance(likelihood.model, cuqi.model.LinearModel):
-                    raise TypeError("Model needs to be linear")
+                if not isinstance(likelihood.model, cuqi.model.AffineModel):
+                    raise TypeError("Model needs to be linear or more generally affine")
 
                 if not hasattr(likelihood.distribution, "sqrtprec"):
                     raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
@@ -195,60 +180,115 @@ class LinearRTO(Sampler):
 
         if not hasattr(self.prior, "sqrtprecTimesMean"):
             raise TypeError("Prior must contain a sqrtprecTimesMean attribute")
-
+    
+    def _get_default_initial_point(self, dim):
+        """ Get the default initial point for the sampler. Defaults to an array of zeros. """
+        return np.zeros(dim)
 
 class RegularizedLinearRTO(LinearRTO):
     """
     Regularized Linear RTO (Randomize-Then-Optimize) sampler.
 
     Samples posterior related to the inverse problem with Gaussian likelihood and implicit Gaussian prior, and where the forward model is Linear.
+    The sampler works by repeatedly solving regularized linear least squares problems for perturbed data.
+    The solver for these optimization problems is chosen based on how the regularized is provided in the implicit Gaussian prior.
+    Currently we use the following solvers:
+    FISTA: [1] Beck, Amir, and Marc Teboulle. "A fast iterative shrinkage-thresholding algorithm for linear inverse problems." SIAM journal on imaging sciences 2.1 (2009): 183-202.
+           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".
+    ScipyMinimizer: Wrapper for Scipy's minimize. Optionally used when the constraint is either "nonnegativity" or "box".
 
     Parameters
     ------------
     target : `cuqi.distribution.Posterior`
         See `cuqi.sampler.LinearRTO`
 
-    x0 : `np.ndarray` 
+    initial_point : `np.ndarray` 
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        Maximum number of iterations of the inner FISTA solver. *Optional*.
+        Maximum number of iterations of the FISTA/ADMM/ScipyLinearLSQ/ScipyMinimizer solver. *Optional*.
+
+    inner_max_it : int
+        Maximum number of iterations of the CGLS solver used within the ADMM solver. *Optional*.
         
     stepsize : string or float
         If stepsize is a string and equals either "automatic", then the stepsize is automatically estimated based on the spectral norm.
         If stepsize is a float, then this stepsize is used.
 
+    penalty_parameter : int
+        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/ScipyMinimizer 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 solver, otherwise ISTA is used. *Optional*.
+    
+    solver : string
+        Options are "FISTA" (default for a single constraint or regularization), "ADMM" (default and the only option for multiple constraints or regularizations), "ScipyLinearLSQ" and "ScipyMinimizer". Note "ScipyLinearLSQ" and "ScipyMinimizer" can only be used with `RegularizedGaussian` of a single `box` or `nonnegativity` constraint. *Optional*.
+
+    inner_initial_point : string or np.ndarray or cuqi.array.CUQIArray
+        Initial point for the inner optimization problem. Can be "previous_sample" (default), "MAP", or a specific numpy or cuqi array. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
     """
-    def __init__(self, target, x0=None, maxit=100, stepsize = "automatic", abstol=1e-10, adaptive = True, **kwargs):
-
-        if not callable(target.prior.proximal):
-            raise TypeError("Projector needs to be callable")
+    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, inner_initial_point="previous_sample", **kwargs):
         
-        super().__init__(target, x0=x0, maxit=100, **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.proximal = target.prior.proximal
+        self.maxit = maxit
+        self.inner_max_it = inner_max_it
+        self.penalty_parameter = penalty_parameter
+        self.solver = solver
+        self.inner_initial_point = inner_initial_point
+
+    def _initialize(self):
+        super()._initialize()
+        if self.solver is None:
+            self.solver = "FISTA" if callable(self.proximal) else "ADMM"
+        if self.solver == "FISTA":
+            self._stepsize = self._choose_stepsize()
+        self._compute_map_regularized()
 
     @property
-    def prior(self):
-        return self.target.prior.gaussian
+    def solver(self):
+        return self._solver
+
+    @solver.setter
+    def solver(self, value):
+        if value == "ScipyLinearLSQ" or value == "ScipyMinimizer":
+            if (self.target.prior.preset["constraint"] == "nonnegativity" or self.target.prior.preset["constraint"] == "box"):
+                self._solver = value
+            else:
+                raise ValueError("ScipyLinearLSQ and ScipyMinimizer only support RegularizedGaussian with box or nonnegativity constraint.")
+        else:
+            self._solver = value
+
+    @property
+    def proximal(self):
+        return self.target.prior.proximal
+    
+    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")
 
-    def _sample(self, N, Nb):   
-        Ns = N+Nb   # number of simulations        
-        samples = np.empty((self.n, Ns))
-                   
+    def _choose_stepsize(self):
         if isinstance(self.stepsize, str):
             if self.stepsize in ["automatic"]:
                 if not callable(self.M):
@@ -262,20 +302,49 @@ class RegularizedLinearRTO(LinearRTO):
                 raise ValueError("Stepsize choice not supported")
         else:
             _stepsize = self.stepsize
-            
-        # initial state   
-        samples[:, 0] = self.x0
-        for s in range(Ns-1):
-            y = self.b_tild + np.random.randn(len(self.b_tild))
-            sim = FISTA(self.M, y, samples[:, s], self.proximal,
-                        maxit = self.maxit, stepsize = _stepsize, abstol = self.abstol, adaptive = self.adaptive)         
-            samples[:, s+1], _ = sim.solve()
-            
-            self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample
-            self._call_callback(samples[:, s+1], s+1)
-        # remove burn-in
-        samples = samples[:, Nb:]
+        return _stepsize
+
+    @property
+    def prior(self):
+        return self.target.prior.gaussian
+
+    def _compute_map_regularized(self):
+        self._map = self._customized_step(self.b_tild, self.initial_point)
+
+    def _customized_step(self, y, x0):
+        if self.solver == "FISTA":
+            sim = FISTA(self.M, y, self.proximal,
+                        x0, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        elif self.solver == "ADMM":
+            sim = ADMM(self.M, y, self.proximal,
+                        x0, 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.distribution.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)
+        elif self.solver == "ScipyMinimizer":
+            # Adapt bounds format, as scipy.minimize requires a bounds format 
+            # different than that in scipy.lsq_linear.
+            bounds = [(self.target.prior._box_bounds[0][i], self.target.prior._box_bounds[1][i]) for i in range(self.target.prior.dim)]
+            # Note that the objective function is defined as 0.5*||Mx-y||^2, 
+            # and the corresponding gradient (gradfunc) is given by M^T(Mx-y).
+            sim = ScipyMinimizer(lambda x: 0.5*np.sum((self.M(x, 1)-y)**2), x0, gradfunc=lambda x: self.M(self.M(x, 1) - y, 2), bounds=bounds, tol=self.abstol, options={"maxiter": self.maxit})
+        else:
+            raise ValueError("Choice of solver not supported.")
         
-        return samples, None, None
+        sol, _ = sim.solve()
+        return sol
+
+    def step(self):
+        y = self.b_tild + np.random.randn(len(self.b_tild))
 
+        self.current_point = self._customized_step(y, self.inner_initial_point)
 
+        acc = 1
+        return acc