CUQIpy 1.3.0__py3-none-any.whl → 1.4.0.post0.dev61__py3-none-any.whl

cuqi/{experimental/mcmc → legacy/sampler}/_rto.py

@@ -3,8 +3,8 @@ 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, ADMM, ScipyLinearLSQ
-from cuqi.experimental.mcmc import Sampler
+from cuqi.solver import CGLS, FISTA
+from cuqi.legacy.sampler import Sampler
 
 
 class LinearRTO(Sampler):
@@ -27,7 +27,7 @@ class LinearRTO(Sampler):
         P_mean: is the prior mean.
         P_sqrtprec: is the squareroot of the precision matrix of the Gaussian mean.
 
-    initial_point : `np.ndarray` 
+    x0 : `np.ndarray` 
         Initial point for the sampler. *Optional*.
 
     maxit : int
@@ -43,51 +43,60 @@ class LinearRTO(Sampler):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target=None, initial_point=None, maxit=10, tol=1e-6, **kwargs):
-
-        super().__init__(target=target, initial_point=initial_point, **kwargs)
+    def __init__(self, target, x0=None, maxit=10, tol=1e-6, shift=0, **kwargs):
+        
+        # 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.AffineModel):
+                raise TypeError("Model needs to be cuqi.model.AffineModel 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)
 
         # Other parameters
         self.maxit = maxit
-        self.tol = tol
-
-    def _initialize(self):
-        self._precompute()
-
-    @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 _precompute(self):
+        self.tol = tol        
+        self.shift = 0
+                
         L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
         L2 = self.prior.sqrtprec
         L2mu = self.prior.sqrtprecTimesMean
 
         # pre-computations
-        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
+        self.n = len(self.x0)
+        self.b_tild = np.hstack([L@(likelihood.data - model._shift) for (L, likelihood, model) in zip(L1, self.likelihoods, self.models)]+ [L2mu])
+
         callability = [callable(likelihood.model) for likelihood in self.likelihoods]
         notcallability = [not c for c in callability]
         if all(notcallability):
@@ -105,26 +114,64 @@ class LinearRTO(Sampler):
                     out1 = np.zeros(self.n)
                     for likelihood in self.likelihoods:
                         idx_end += len(likelihood.data)
-                        out1 += likelihood.model._adjoint_func_no_shift(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]) # Use adjoint function which excludes shift
                         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.")
+            raise TypeError("All likelihoods need to be callable or none need to be callable.") 
 
-    def step(self):
-        y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = CGLS(self.M, y, self.current_point, self.maxit, self.tol)            
-        self.current_point, _ = sim.solve()
-        acc = 1
-        return acc
+    @property
+    def prior(self):
+        return self.target.prior
 
-    def tune(self, skip_len, update_count):
-        pass
+    @property
+    def likelihood(self):
+        return self.target.likelihood
     
-    def validate_target(self):
+    @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 _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 _sample_adapt(self, N, Nb):
+        return self._sample(N,Nb)
+    
+    def _check_posterior(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'.")       
@@ -132,15 +179,15 @@ class LinearRTO(Sampler):
         # Check Linear model and Gaussian likelihood(s)
         if isinstance(self.target, cuqi.distribution.Posterior):
             if not isinstance(self.model, cuqi.model.AffineModel):
-                raise TypeError("Model needs to be linear or more generally affine")
+                raise TypeError("Model needs to be linear or 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.AffineModel):
-                    raise TypeError("Model needs to be linear or more generally affine")
+                if not isinstance(likelihood.model, cuqi.model.LinearModel):
+                    raise TypeError("Model needs to be linear")
 
                 if not hasattr(likelihood.distribution, "sqrtprec"):
                     raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
@@ -151,58 +198,31 @@ 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".
 
     Parameters
     ------------
     target : `cuqi.distribution.Posterior`
-        See `cuqi.sampler.LinearRTO`
+        See `cuqi.legacy.sampler.LinearRTO`
 
-    initial_point : `np.ndarray` 
+    x0 : `np.ndarray` 
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        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*.
+        Maximum number of iterations of the inner FISTA 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 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 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*.
+        Absolute tolerance of the inner FISTA solver. *Optional*.
 
     callback : callable, *Optional*
         If set this function will be called after every sample.
@@ -211,51 +231,27 @@ 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, solver=None, inner_abstol=None, **kwargs):
+    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")
         
-        super().__init__(target=target, initial_point=initial_point, **kwargs)
+        super().__init__(target, x0=x0, maxit=100, **kwargs)
 
         # Other parameters
         self.stepsize = stepsize
-        self.abstol = abstol
-        self.inner_abstol = inner_abstol
+        self.abstol = 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.solver is None:
-            self.solver = "FISTA" if callable(self.proximal) else "ADMM"
-        if self.solver == "FISTA":
-            self._stepsize = self._choose_stepsize()
+        self.proximal = target.prior.proximal
 
     @property
-    def solver(self):
-        return self._solver
-
-    @solver.setter
-    def solver(self, value):
-        if value == "ScipyLinearLSQ":
-            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.")
-        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 prior(self):
+        return self.target.prior.gaussian
 
-    def _choose_stepsize(self):
+    def _sample(self, N, Nb):   
+        Ns = N+Nb   # number of simulations        
+        samples = np.empty((self.n, Ns))
+                   
         if isinstance(self.stepsize, str):
             if self.stepsize in ["automatic"]:
                 if not callable(self.M):
@@ -269,34 +265,20 @@ class RegularizedLinearRTO(LinearRTO):
                 raise ValueError("Stepsize choice not supported")
         else:
             _stepsize = self.stepsize
-        return _stepsize
-
-    @property
-    def prior(self):
-        return self.target.prior.gaussian
-
-    def step(self):
-        y = self.b_tild + np.random.randn(len(self.b_tild))
-
-        if self.solver == "FISTA":
+            
+        # 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, self.proximal,
-                        self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
-        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)
-        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.")
+                        samples[:, s], 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 samples, None, None
+
 
-        self.current_point, _ = sim.solve()
-        acc = 1
-        return acc

cuqi/legacy/sampler/_sampler.py

@@ -0,0 +1,182 @@
+from abc import ABC, abstractmethod
+import sys
+import numpy as np
+import cuqi
+from cuqi.samples import Samples
+import warnings
+
+class Sampler(ABC):
+
+    def __init__(self, target, x0=None, dim=None, callback=None):
+
+        warnings.warn(f"\nYou are using the legacy sampler '{self.__class__.__name__}'.\n"
+                      f"This will be removed in a future release of CUQIpy.\n"
+                      f"Please consider using the new samplers in the 'cuqi.sampler' module.\n", UserWarning, stacklevel=2)
+
+        self._dim = dim
+        if hasattr(target,'dim'): 
+            if self._dim is None:
+                self._dim = target.dim 
+            elif self._dim != target.dim:
+                raise ValueError("'dim' need to be None or equal to 'target.dim'") 
+        elif x0 is not None:
+            self._dim = len(x0)
+
+        self.target = target
+
+        if x0 is None:
+            x0 = np.ones(self.dim)
+        self.x0 = x0
+
+        self.callback = callback
+
+    def step(self, x):
+        """
+        Perform a single MCMC step
+        """
+        # Currently a hack to get step method for any sampler
+        self.x0 = x
+        return self.sample(2).samples[:,-1]
+
+    def step_tune(self, x, *args, **kwargs):
+        """
+        Perform a single MCMC step and tune the sampler. This is used during burn-in.
+        """
+        # Currently a hack to get step method for any sampler
+        out = self.step(x)
+        self.tune(*args, *kwargs)
+        return out
+
+    def tune(self):
+        """
+        Tune the sampler parameters.
+        """
+        pass
+
+
+    @property
+    def geometry(self):
+        if hasattr(self, 'target') and hasattr(self.target, 'geometry'):
+            geom =  self.target.geometry
+        else:
+            geom = cuqi.geometry._DefaultGeometry1D(self.dim)
+        return geom
+
+    @property 
+    def target(self):
+        return self._target 
+
+    @target.setter 
+    def target(self, value):
+        if  not isinstance(value, cuqi.distribution.Distribution) and callable(value):
+            # obtain self.dim
+            if self.dim is not None:
+                dim = self.dim
+            else:
+                raise ValueError(f"If 'target' is a lambda function, the parameter 'dim' need to be specified when initializing {self.__class__}.")
+
+            # set target
+            self._target = cuqi.distribution.UserDefinedDistribution(logpdf_func=value, dim = dim)
+
+        elif isinstance(value, cuqi.distribution.Distribution):
+            self._target = value
+        else:
+            raise ValueError("'target' need to be either a lambda function or of type 'cuqi.distribution.Distribution'")
+
+
+    @property
+    def dim(self):
+        if hasattr(self,'target') and hasattr(self.target,'dim'):
+            self._dim = self.target.dim 
+        return self._dim
+    
+
+    def sample(self,N,Nb=0):
+        # Get samples from the samplers sample method
+        result = self._sample(N,Nb)
+        return self._create_Sample_object(result,N+Nb)
+
+    def sample_adapt(self,N,Nb=0):
+        # Get samples from the samplers sample method
+        result = self._sample_adapt(N,Nb)
+        return self._create_Sample_object(result,N+Nb)
+
+    def _create_Sample_object(self,result,N):
+        loglike_eval = None
+        acc_rate = None
+        if isinstance(result,tuple):
+            #Unpack samples+loglike+acc_rate
+            s = result[0]
+            if len(result)>1: loglike_eval = result[1]
+            if len(result)>2: acc_rate = result[2]
+            if len(result)>3: raise TypeError("Expected tuple of at most 3 elements from sampling method.")
+        else:
+            s = result
+                
+        #Store samples in cuqi samples object if more than 1 sample
+        if N==1:
+            if len(s) == 1 and isinstance(s,np.ndarray): #Extract single value from numpy array
+                s = s.ravel()[0]
+            else:
+                s = s.flatten()
+        else:
+            s = Samples(s, self.geometry)#, geometry = self.geometry)
+            s.loglike_eval = loglike_eval
+            s.acc_rate = acc_rate
+        return s
+
+    @abstractmethod
+    def _sample(self,N,Nb):
+        pass
+
+    @abstractmethod
+    def _sample_adapt(self,N,Nb):
+        pass
+
+    def _print_progress(self,s,Ns):
+        """Prints sampling progress"""
+        if Ns > 2:
+            if (s % (max(Ns//100,1))) == 0:
+                msg = f'Sample {s} / {Ns}'
+                sys.stdout.write('\r'+msg)
+            if s==Ns:
+                msg = f'Sample {s} / {Ns}'
+                sys.stdout.write('\r'+msg+'\n')
+
+    def _call_callback(self, sample, sample_index):
+        """ Calls the callback function. Assumes input is sample and sample index"""
+        if self.callback is not None:
+            self.callback(sample, sample_index)
+
+class ProposalBasedSampler(Sampler,ABC):
+    def __init__(self, target,  proposal=None, scale=1, x0=None, dim=None, **kwargs):
+        #TODO: after fixing None dim
+        #if dim is None and hasattr(proposal,'dim'):
+        #    dim = proposal.dim
+        super().__init__(target, x0=x0, dim=dim, **kwargs)
+
+        self.proposal =proposal
+        self.scale = scale
+
+
+    @property 
+    def proposal(self):
+        return self._proposal 
+
+    @proposal.setter 
+    def proposal(self, value):
+        self._proposal = value
+
+    @property
+    def geometry(self):
+        geom1, geom2 = None, None
+        if hasattr(self, 'proposal') and hasattr(self.proposal, 'geometry') and self.proposal.geometry.par_dim is not None:
+            geom1=  self.proposal.geometry
+        if hasattr(self, 'target') and hasattr(self.target, 'geometry') and self.target.geometry.par_dim is not None:
+            geom2 = self.target.geometry
+        if not isinstance(geom1,cuqi.geometry._DefaultGeometry) and geom1 is not None:
+            return geom1
+        elif not isinstance(geom2,cuqi.geometry._DefaultGeometry) and geom2 is not None: 
+            return geom2
+        else:
+            return cuqi.geometry._DefaultGeometry1D(self.dim)

cuqi/likelihood/_likelihood.py

@@ -212,4 +212,4 @@ class UserDefinedLikelihood(object):
         return get_non_default_args(self.logpdf_func)
 
     def __repr__(self) -> str:
-        return "CUQI {} function. Parameters {}.".format(self.__class__.__name__,self.get_parameter_names())
+        return "CUQI {} function. Parameters {}.".format(self.__class__.__name__,self.get_parameter_names())