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

cuqi/legacy/sampler/_pcn.py

@@ -0,0 +1,244 @@
+import numpy as np
+import cuqi
+from cuqi.legacy.sampler import Sampler
+
+class pCN(Sampler):   
+    #Samples target*proposal
+    #TODO. Check proposal, needs to be Gaussian and zero mean.
+    """Preconditioned Crank-Nicolson sampler 
+    
+    Parameters
+    ----------
+    target : `cuqi.distribution.Posterior` or tuple of likelihood and prior objects
+        If target is of type cuqi.distribution.Posterior, it represents the posterior distribution.
+        If target is a tuple of (cuqi.likelihood.Likelihood, cuqi.distribution.Distribution) objects,
+        the first element is considered the likelihood and the second is considered the prior.
+
+    scale : int
+
+    x0 : `np.ndarray` 
+      Initial point for the sampler
+
+    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.
+
+    Example 
+    -------
+
+    This uses a custom logpdf and sample function.
+
+    .. code-block:: python
+
+        # Parameters
+        dim = 5 # Dimension of distribution
+        mu = np.arange(dim) # Mean of Gaussian
+        std = 1 # standard deviation of Gaussian
+
+        # Logpdf function of likelihood
+        logpdf_func = lambda x: -1/(std**2)*np.sum((x-mu)**2)
+
+        # sample function of prior N(0,I)
+        sample_func = lambda : 0 + 1*np.random.randn(dim,1)
+
+        # Define as UserDefinedDistributions
+        likelihood = cuqi.likelihood.UserDefinedLikelihood(dim=dim, logpdf_func=logpdf_func)
+        prior = cuqi.distribution.UserDefinedDistribution(dim=dim, sample_func=sample_func)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.pCN((likelihood,prior), scale = 0.1)
+
+        # Sample
+        samples = sampler.sample(5000)
+
+    Example
+    -------
+
+    This uses CUQIpy distributions.
+
+    .. code-block:: python
+
+        # Parameters
+        dim = 5 # Dimension of distribution
+        mu = np.arange(dim) # Mean of Gaussian
+        std = 1 # standard deviation of Gaussian
+
+        # Define as UserDefinedDistributions
+        model = cuqi.model.Model(lambda x: x, range_geometry=dim, domain_geometry=dim)
+        likelihood = cuqi.distribution.Gaussian(mean=model, cov=np.ones(dim)).to_likelihood(mu)
+        prior = cuqi.distribution.Gaussian(mean=np.zeros(dim), cov=1)
+
+        target = cuqi.distribution.Posterior(likelihood, prior)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.pCN(target, scale = 0.1)
+
+        # Sample
+        samples = sampler.sample(5000)
+        
+    """
+    def __init__(self, target, scale=None, x0=None, **kwargs):
+        super().__init__(target, x0=x0, dim=None, **kwargs) 
+        self.scale = scale
+    
+    @property
+    def prior(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return self.target.prior
+        elif isinstance(self.target,tuple) and len(self.target)==2:
+            return self.target[1]
+
+    @property
+    def likelihood(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return self.target.likelihood
+        elif isinstance(self.target,tuple) and len(self.target)==2:
+            return self.target[0]
+
+
+    @Sampler.target.setter 
+    def target(self, value):
+        if isinstance(value, cuqi.distribution.Posterior):
+            self._target = value
+            self._loglikelihood = lambda x : self.likelihood.logd(x)
+        elif isinstance(value,tuple) and len(value)==2 and \
+             (isinstance(value[0], cuqi.likelihood.Likelihood) or isinstance(value[0], cuqi.likelihood.UserDefinedLikelihood))  and \
+             isinstance(value[1], cuqi.distribution.Distribution):
+            self._target = value
+            self._loglikelihood = lambda x : self.likelihood.logd(x)
+        else:
+            raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")
+        
+        #TODO:
+        #if not isinstance(self.prior,(cuqi.distribution.Gaussian, cuqi.distribution.Normal)):
+        #    raise ValueError("The prior distribution of the target need to be Gaussian")
+
+    @property
+    def dim(self):
+        if hasattr(self,'target') and hasattr(self.target,'dim'):
+            self._dim = self.target.dim
+        elif hasattr(self,'target') and isinstance(self.target,tuple) and len(self.target)==2:
+            self._dim = self.target[0].dim
+        return self._dim
+
+    def _sample(self, N, Nb):
+        if self.scale is None:
+            raise ValueError("Scale must be set to sample without adaptation. Consider using sample_adapt instead.")
+
+        Ns = N+Nb   # number of simulations
+
+        # allocation
+        samples = np.empty((self.dim, Ns))
+        loglike_eval = np.empty(Ns)
+        acc = np.zeros(Ns, dtype=int)
+
+        # initial state    
+        samples[:, 0] = self.x0
+        loglike_eval[0] = self._loglikelihood(self.x0)
+        acc[0] = 1
+
+        # run MCMC
+        for s in range(Ns-1):
+            # run component by component
+            samples[:, s+1], loglike_eval[s+1], acc[s+1] = self.single_update(samples[:, s], loglike_eval[s])
+
+            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:]
+        loglike_eval = loglike_eval[Nb:]
+        accave = acc[Nb:].mean()   
+        print('\nAverage acceptance rate:', accave, '\n')
+        #
+        return samples, loglike_eval, accave
+
+    def _sample_adapt(self, N, Nb):
+        # Set intial scale if not set
+        if self.scale is None:
+            self.scale = 0.1
+
+        Ns = N+Nb   # number of simulations
+
+        # allocation
+        samples = np.empty((self.dim, Ns))
+        loglike_eval = np.empty(Ns)
+        acc = np.zeros(Ns)
+
+        # initial state    
+        samples[:, 0] = self.x0
+        loglike_eval[0] = self._loglikelihood(self.x0) 
+        acc[0] = 1
+
+        # initial adaptation params 
+        Na = int(0.1*N)                              # iterations to adapt
+        hat_acc = np.empty(int(np.floor(Ns/Na)))     # average acceptance rate of the chains
+        lambd = self.scale
+        star_acc = 0.44    # target acceptance rate RW
+        i, idx = 0, 0
+
+        # run MCMC
+        for s in range(Ns-1):
+            # run component by component
+            samples[:, s+1], loglike_eval[s+1], acc[s+1] = self.single_update(samples[:, s], loglike_eval[s])
+            
+            # adapt prop spread using acc of past samples
+            if ((s+1) % Na == 0):
+                # evaluate average acceptance rate
+                hat_acc[i] = np.mean(acc[idx:idx+Na])
+
+                # d. compute new scaling parameter
+                zeta = 1/np.sqrt(i+1)   # ensures that the variation of lambda(i) vanishes
+                lambd = np.exp(np.log(lambd) + zeta*(hat_acc[i]-star_acc))
+
+                # update parameters
+                self.scale = min(lambd, 1)
+
+                # update counters
+                i += 1
+                idx += Na
+
+            # display iterations
+            if ((s+1) % (max(Ns//100,1))) == 0 or (s+1) == Ns-1:
+                print("\r",'Sample', s+1, '/', Ns, end="")
+
+            self._call_callback(samples[:, s+1], s+1)
+
+        print("\r",'Sample', s+2, '/', Ns)
+
+        # remove burn-in
+        samples = samples[:, Nb:]
+        loglike_eval = loglike_eval[Nb:]
+        accave = acc[Nb:].mean()   
+        print('\nAverage acceptance rate:', accave, 'MCMC scale:', self.scale, '\n')
+        
+        return samples, loglike_eval, accave
+
+    def single_update(self, x_t, loglike_eval_t):
+        # propose state
+        xi = self.prior.sample(1).flatten()   # sample from the prior
+        x_star = np.sqrt(1-self.scale**2)*x_t + self.scale*xi   # pCN proposal
+
+        # evaluate target
+        loglike_eval_star =  self._loglikelihood(x_star) 
+
+        # ratio and acceptance probability
+        ratio = loglike_eval_star - loglike_eval_t  # proposal is symmetric
+        alpha = min(0, ratio)
+
+        # accept/reject
+        u_theta = np.log(np.random.rand())
+        if (u_theta <= alpha):
+            x_next = x_star
+            loglike_eval_next = loglike_eval_star
+            acc = 1
+        else:
+            x_next = x_t
+            loglike_eval_next = loglike_eval_t
+            acc = 0
+        
+        return x_next, loglike_eval_next, acc
+    
+    

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

@@ -4,14 +4,14 @@ from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
 from cuqi.solver import CGLS, FISTA
-from cuqi.experimental.mcmc import Sampler
+from cuqi.legacy.sampler import Sampler
 
 
 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.
 
-    initial_point : `np.ndarray` 
+    x0 : `np.ndarray` 
         Initial point for the sampler. *Optional*.
 
     maxit : int
@@ -43,48 +43,59 @@ 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 data(self):
-        return self.target.data
-    
-    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 for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
+        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]
@@ -94,7 +105,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:
@@ -103,34 +114,72 @@ 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]) # 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'.")       
 
         # 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 affine")
 
             if not hasattr(self.likelihood.distribution, "sqrtprec"):
                 raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
@@ -149,10 +198,7 @@ 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):
     """
@@ -163,9 +209,9 @@ class RegularizedLinearRTO(LinearRTO):
     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
@@ -177,9 +223,6 @@ class RegularizedLinearRTO(LinearRTO):
 
     abstol : float
         Absolute tolerance of the inner FISTA solver. *Optional*.
-    
-    adaptive : bool
-        If True, FISTA is used as inner solver, otherwise ISTA is used. *Optional*.
 
     callback : callable, *Optional*
         If set this function will be called after every sample.
@@ -188,32 +231,27 @@ class RegularizedLinearRTO(LinearRTO):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target=None, initial_point=None, maxit=100, stepsize="automatic", abstol=1e-10, adaptive=True, **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.adaptive = adaptive
-        self.maxit = maxit
-
-    def _initialize(self):
-        super()._initialize()
-        self._stepsize = self._choose_stepsize()
+        self.proximal = target.prior.proximal
 
     @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")
-        if not callable(self.proximal):
-            raise TypeError("Proximal needs to be callable")
-
-    def _choose_stepsize(self):
+    def prior(self):
+        return self.target.prior.gaussian
+
+    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):
@@ -227,16 +265,20 @@ class RegularizedLinearRTO(LinearRTO):
                 raise ValueError("Stepsize choice not supported")
         else:
             _stepsize = self.stepsize
-        return _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, self.proximal,
+                        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
 
-    @property
-    def prior(self):
-        return self.target.prior.gaussian
 
-    def step(self):
-        y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = FISTA(self.M, y, self.current_point, self.proximal,
-                    maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
-        self.current_point, _ = sim.solve()
-        acc = 1
-        return acc