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

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

@@ -1,9 +1,8 @@
 import numpy as np
 import cuqi
-from cuqi.experimental.mcmc import Sampler
-from cuqi.array import CUQIarray
+from cuqi.legacy.sampler import Sampler
 
-class ULA(Sampler): # Refactor to Proposal-based sampler?
+class ULA(Sampler):
     """Unadjusted Langevin algorithm (ULA) (Roberts and Tweedie, 1996)
 
     Samples a distribution given its logpdf and gradient (up to a constant) based on
@@ -20,13 +19,17 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
         The target distribution to sample. Must have logd and gradient method. Custom logpdfs 
         and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
     
-    initial_point : ndarray
+    x0 : ndarray
         Initial parameters. *Optional*
 
     scale : int
         The Langevin diffusion discretization time step (In practice, a scale of 1/dim**2 is
         recommended but not guaranteed to be the optimal choice).
 
+    dim : int
+        Dimension of parameter space. Required if target logpdf and gradient are callable 
+        functions. *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)`,
@@ -52,80 +55,61 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
             gradient_func=gradient_func)
 
         # Set up sampler
-        sampler = cuqi.experimental.mcmc.ULA(target, scale=1/dim**2)
+        sampler = cuqi.legacy.sampler.ULA(target, scale=1/dim**2)
 
         # Sample
-        sampler.sample(2000)
+        samples = sampler.sample(2000)
 
-    A Deblur example can be found in demos/demo27_ULA.py
-    # TODO: update demo once sampler merged
+    A Deblur example can be found in demos/demo27_ula.py
     """
+    def __init__(self, target, scale, x0=None, dim=None, rng=None, **kwargs):
+        super().__init__(target, x0=x0, dim=dim, **kwargs)
+        self.scale = scale
+        self.rng = rng
+
+    def _sample_adapt(self, N, Nb):
+        return self._sample(N, Nb)
+
+    def _sample(self, N, Nb):    
+        # allocation
+        Ns = Nb+N
+        samples = np.empty((self.dim, Ns))
+        target_eval = np.empty(Ns)
+        g_target_eval = np.empty((self.dim, Ns))
+        acc = np.zeros(Ns)
+
+        # initial state
+        samples[:, 0] = self.x0
+        target_eval[0], g_target_eval[:,0] = self.target.logd(self.x0), self.target.gradient(self.x0)
+        acc[0] = 1
+
+        # ULA
+        for s in range(Ns-1):
+            samples[:, s+1], target_eval[s+1], g_target_eval[:,s+1], acc[s+1] = \
+                self.single_update(samples[:, s], target_eval[s], g_target_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)
+    
+        # apply burn-in 
+        samples = samples[:, Nb:]
+        target_eval = target_eval[Nb:]
+        acc = acc[Nb:]
+        return samples, target_eval, np.mean(acc)
 
-    _STATE_KEYS = Sampler._STATE_KEYS.union({'current_target_logd', 'scale', 'current_target_grad'})
-
-    def __init__(self, target=None, scale=1.0, **kwargs):
-
-        super().__init__(target, **kwargs)
-
-        self.initial_scale = scale
-
-    def _initialize(self):
-        self.scale = self.initial_scale
-        self.current_target_logd = self.target.logd(self.current_point)
-        self.current_target_grad = self.target.gradient(self.current_point)
-
-    def validate_target(self):
-        try:
-            self.target.gradient(np.ones(self.dim))
-            pass
-        except (NotImplementedError, AttributeError):
-            raise ValueError("The target needs to have a gradient method")
-
-    def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
-        """
-        Accepts the proposed state and updates the sampler's state accordingly, i.e.,
-        current_point, current_target_eval, and current_target_grad_eval.
-
-        Parameters
-        ----------
-        x_star : 
-            The proposed state
-
-        target_eval_star: 
-            The log likelihood evaluated at x_star
-
-        target_grad_star: 
-            The gradient of log likelihood evaluated at x_star
-
-        Returns
-        -------
-        scalar
-            1 (accepted)
-        """
-        self.current_point = x_star
-        self.current_target_logd = target_eval_star
-        self.current_target_grad = target_grad_star
-        acc = 1
-        return acc
-
-    def step(self):
-        # propose state
-        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample()
-        x_star = self.current_point + 0.5*self.scale*self.current_target_grad + xi
-
-        # evaluate target
-        target_eval_star, target_grad_star = self.target.logd(x_star), self.target.gradient(x_star)
-
-        # accept or reject proposal
-        acc = self._accept_or_reject(x_star, target_eval_star, target_grad_star)
+    def single_update(self, x_t, target_eval_t, g_target_eval_t):
+        # approximate Langevin diffusion
+        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample(rng=self.rng)
+        x_star = x_t + 0.5*self.scale*g_target_eval_t + xi
+        logpi_eval_star, g_logpi_star = self.target.logd(x_star), self.target.gradient(x_star)
 
-        return acc
+        # msg
+        if np.isnan(logpi_eval_star):
+            raise NameError('NaN potential func. Consider using smaller scale parameter')
 
-    def tune(self, skip_len, update_count):
-        pass
+        return x_star, logpi_eval_star, g_logpi_star, 1 # sample always accepted without Metropolis correction
 
 
-class MALA(ULA): # Refactor to Proposal-based sampler?
+class MALA(ULA):
     """  Metropolis-adjusted Langevin algorithm (MALA) (Roberts and Tweedie, 1996)
 
     Samples a distribution given its logd and gradient (up to a constant) based on
@@ -143,12 +127,16 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
         The target distribution to sample. Must have logpdf and gradient method. Custom logpdfs 
         and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
     
-    initial_point : ndarray
+    x0 : ndarray
         Initial parameters. *Optional*
 
     scale : int
         The Langevin diffusion discretization time step.
 
+    dim : int
+        Dimension of parameter space. Required if target logpdf and gradient are callable 
+        functions. *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)`,
@@ -174,54 +162,37 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
             gradient_func=gradient_func)
 
         # Set up sampler
-        sampler = cuqi.experimental.mcmc.MALA(target, scale=1/5**2)
+        sampler = cuqi.legacy.sampler.MALA(target, scale=1/5**2)
 
         # Sample
-        sampler.sample(2000)
+        samples = sampler.sample(2000)
 
-    A Deblur example can be found in demos/demo28_MALA.py
-    # TODO: update demo once sampler merged
+    A Deblur example can be found in demos/demo28_mala.py
     """
-
-    def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
-        """
-        Accepts the proposed state according to a Metropolis step and updates the sampler's state accordingly, i.e., current_point, current_target_eval, and current_target_grad_eval.
-
-        Parameters
-        ----------
-        x_star : 
-            The proposed state
-
-        target_eval_star: 
-            The log likelihood evaluated at x_star
-
-        target_grad_star: 
-            The gradient of log likelihood evaluated at x_star
-
-        Returns
-        -------
-        scaler
-            1 if accepted, 0 otherwise
-        """
-        log_target_ratio = target_eval_star - self.current_target_logd
-        log_prop_ratio = self._log_proposal(self.current_point, x_star, target_grad_star) \
-            - self._log_proposal(x_star, self.current_point,  self.current_target_grad)
+    def __init__(self, target, scale, x0=None, dim=None, rng=None, **kwargs):
+        super().__init__(target, scale, x0=x0, dim=dim, rng=rng, **kwargs)
+
+    def single_update(self, x_t, target_eval_t, g_target_eval_t):
+        # approximate Langevin diffusion
+        xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample(rng=self.rng)
+        x_star = x_t + (self.scale/2)*g_target_eval_t + xi
+        logpi_eval_star, g_logpi_star = self.target.logd(x_star), self.target.gradient(x_star)
+
+        # Metropolis step
+        log_target_ratio = logpi_eval_star - target_eval_t
+        log_prop_ratio = self.log_proposal(x_t, x_star, g_logpi_star) \
+            - self.log_proposal(x_star, x_t,  g_target_eval_t)
         log_alpha = min(0, log_target_ratio + log_prop_ratio)
 
-        # accept/reject with Metropolis
-        acc = 0
-        log_u = np.log(np.random.rand())
-        if (log_u <= log_alpha) and (np.isnan(target_eval_star) == False):
-            self.current_point = x_star
-            self.current_target_logd = target_eval_star
-            self.current_target_grad = target_grad_star
-            acc = 1
-        return acc
+        # accept/reject
+        log_u = np.log(cuqi.distribution.Uniform(low=0, high=1).sample(rng=self.rng))
+        if (log_u <= log_alpha) and (np.isnan(logpi_eval_star) == False):
+            return x_star, logpi_eval_star, g_logpi_star, 1
+        else:
+            return x_t.copy(), target_eval_t, g_target_eval_t.copy(), 0
 
-    def tune(self, skip_len, update_count):
-        pass
-
-    def _log_proposal(self, theta_star, theta_k, g_logpi_k):
+    def log_proposal(self, theta_star, theta_k, g_logpi_k):
         mu = theta_k + ((self.scale)/2)*g_logpi_k
         misfit = theta_star - mu
         return -0.5*((1/(self.scale))*(misfit.T @ misfit))
+

cuqi/legacy/sampler/_laplace_approximation.py

@@ -0,0 +1,184 @@
+import scipy as sp
+import numpy as np
+import cuqi
+from cuqi.distribution import Normal
+from cuqi.solver import CGLS
+from cuqi.legacy.sampler import Sampler
+
+
+class UGLA(Sampler):
+    """ Unadjusted (Gaussian) Laplace Approximation sampler
+    
+    Samples an approximate posterior where the prior is approximated
+    by a Gaussian distribution. The likelihood must be Gaussian.
+
+    Currently only works for LMRF priors.
+
+    The inner solver is Conjugate Gradient Least Squares (CGLS) solver.
+
+    For more details see: Uribe, Felipe, et al. "A hybrid Gibbs sampler for edge-preserving 
+    tomographic reconstruction with uncertain view angles." arXiv preprint arXiv:2104.06919 (2021).
+
+    Parameters
+    ----------
+    target : `cuqi.distribution.Posterior`
+        The target posterior distribution to sample.
+
+    x0 : ndarray
+        Initial parameters. *Optional*
+
+    maxit : int
+        Maximum number of inner iterations for solver when generating one sample.
+
+    tol : float
+        Tolerance for inner solver. Will stop before maxit if the inner solvers convergence check reaches tol.
+
+    beta : float
+        Smoothing parameter for the Gaussian approximation of the Laplace distribution. Larger beta is easier to sample but is a worse approximation.
+
+    rng : np.random.RandomState
+        Random number generator used for sampling. *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.
+
+    Returns
+    -------
+    cuqi.samples.Samples
+        Samples from the posterior distribution.
+
+    """
+
+    def __init__(self, target, x0=None, maxit=50, tol=1e-4, beta=1e-5, rng=None, **kwargs):
+        
+        super().__init__(target, x0=x0, **kwargs)
+
+        # Check target type
+        if not isinstance(self.target, cuqi.distribution.Posterior):
+            raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")       
+
+        # Check Affine model
+        if not isinstance(self.target.likelihood.model, cuqi.model.AffineModel):
+            raise TypeError("Model needs to be affine or linear")
+
+        # Check Gaussian likelihood
+        if not hasattr(self.target.likelihood.distribution, "sqrtprec"):
+            raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
+
+        # Check that prior is LMRF
+        if not isinstance(self.target.prior, cuqi.distribution.LMRF):
+            raise ValueError('Unadjusted Gaussian Laplace approximation (UGLA) requires LMRF prior')
+
+        # Modify initial guess since Sampler sets it to ones.       
+        if x0 is not None:
+            self.x0 = x0
+        else:
+            self.x0 = np.zeros(self.target.prior.dim)
+        
+        # Store internal parameters
+        self.maxit = maxit
+        self.tol = tol
+        self.beta = beta
+        self.rng = rng
+
+    def _sample_adapt(self, Ns, Nb):
+        return self._sample(Ns, Nb)
+
+    def _sample(self, Ns, Nb):
+        """ Sample from the approximate posterior.
+
+        Parameters
+        ----------
+        Ns : int
+            Number of samples to draw.
+
+        Nb : int
+            Number of burn-in samples to discard.
+
+        Returns
+        -------
+        samples : ndarray
+            Samples from the approximate posterior.
+
+        target_eval : ndarray
+            Log-likelihood of each sample.
+
+        acc : ndarray
+            Acceptance rate of each sample.
+
+        """
+
+        # Extract diff_op from target prior
+        D = self.target.prior._diff_op
+        n = D.shape[0]
+
+        # Gaussian approximation of LMRF prior as function of x_k
+        def Lk_fun(x_k):
+            dd =  1/np.sqrt((D @ x_k)**2 + self.beta*np.ones(n))
+            W = sp.sparse.diags(dd)
+            return W.sqrt() @ D
+
+        # Now prepare "LinearRTO" type sampler. TODO: Use LinearRTO for this instead
+        self._shift = 0
+
+        # Pre-computations
+        self._model = self.target.likelihood.model   
+        self._data = self.target.likelihood.data - self.target.model._shift
+        self._m = len(self._data)
+        self._L1 = self.target.likelihood.distribution.sqrtprec
+
+        # If prior location is scalar, repeat it to match dimensions
+        if len(self.target.prior.location) == 1:
+            self._priorloc = np.repeat(self.target.prior.location, self.dim)
+        else:
+            self._priorloc = self.target.prior.location
+
+        # Initial Laplace approx
+        self._L2 = Lk_fun(self.x0)
+        self._L2mu = self._L2@self._priorloc
+        self._b_tild = np.hstack([self._L1@self._data, self._L2mu]) 
+        
+        #self.n = len(self.x0)
+        
+        # Least squares form
+        def M(x, flag):
+            if flag == 1:
+                out1 = self._L1 @ self._model._forward_func_no_shift(x) # Use forward function which excludes shift
+                out2 = np.sqrt(1/self.target.prior.scale)*(self._L2 @ x)
+                out  = np.hstack([out1, out2])
+            elif flag == 2:
+                idx = int(self._m)
+                out1 = self._model._adjoint_func_no_shift(self._L1.T@x[:idx])
+                out2 = np.sqrt(1/self.target.prior.scale)*(self._L2.T @ x[idx:])
+                out  = out1 + out2                
+            return out 
+        
+        # Initialize samples
+        N = Ns+Nb   # number of simulations        
+        samples = np.empty((self.target.dim, N))
+                     
+        # initial state   
+        samples[:, 0] = self.x0
+        for s in range(N-1):
+
+            # Update Laplace approximation
+            self._L2 = Lk_fun(samples[:, s])
+            self._L2mu = self._L2@self._priorloc
+            self._b_tild = np.hstack([self._L1@self._data, self._L2mu]) 
+        
+            # Sample from approximate posterior
+            e = Normal(mean=np.zeros(len(self._b_tild)), std=1).sample(rng=self.rng)
+            y = self._b_tild + e # Perturb data
+            sim = CGLS(M, y, samples[:, s], self.maxit, self.tol, self._shift)            
+            samples[:, s+1], _ = sim.solve()
+
+            self._print_progress(s+2,N) #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

cuqi/legacy/sampler/_mh.py

@@ -0,0 +1,190 @@
+import numpy as np
+import cuqi
+from cuqi.legacy.sampler import ProposalBasedSampler
+
+
+class MH(ProposalBasedSampler):
+    """Metropolis Hastings sampler.
+
+    Allows sampling of a target distribution by random-walk sampling of a proposal distribution along with an accept/reject step.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution` or lambda function
+        The target distribution to sample. Custom logpdfs are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
+    
+    proposal : `cuqi.distribution.Distribution` or callable method
+        The proposal to sample from. If a callable method it should provide a single independent sample from proposal distribution. Defaults to a Gaussian proposal.  *Optional*.
+
+    scale : float
+        Scale parameter used to define correlation between previous and proposed sample in random-walk.  *Optional*.
+
+    x0 : ndarray
+        Initial parameters. *Optional*
+
+    dim : int
+        Dimension of parameter space. Required if target and proposal are callable functions. *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.
+
+    Example
+    -------
+    .. code-block:: python
+
+        # Parameters
+        dim = 5 # Dimension of distribution
+        mu = np.arange(dim) # Mean of Gaussian
+        std = 1 # standard deviation of Gaussian
+
+        # Logpdf function
+        logpdf_func = lambda x: -1/(std**2)*np.sum((x-mu)**2)
+
+        # Define distribution from logpdf as UserDefinedDistribution (sample and gradients also supported)
+        target = cuqi.distribution.UserDefinedDistribution(dim=dim, logpdf_func=logpdf_func)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.MH(target, scale=1)
+
+        # Sample
+        samples = sampler.sample(2000)
+
+    """
+    #target,  proposal=None, scale=1, x0=None, dim=None
+    #    super().__init__(target, proposal=proposal, scale=scale,  x0=x0, dim=dim)
+    def __init__(self, target, proposal=None, scale=None, x0=None, dim=None, **kwargs):
+        """ Metropolis-Hastings (MH) sampler. Default (if proposal is None) is random walk MH with proposal that is Gaussian with identity covariance"""
+        super().__init__(target, proposal=proposal, scale=scale,  x0=x0, dim=dim, **kwargs)
+
+
+    @ProposalBasedSampler.proposal.setter 
+    def proposal(self, value):
+        fail_msg = "Proposal should be either None, symmetric cuqi.distribution.Distribution or a lambda function."
+
+        if value is None:
+            self._proposal = cuqi.distribution.Gaussian(np.zeros(self.dim), 1)
+        elif not isinstance(value, cuqi.distribution.Distribution) and callable(value):
+            raise NotImplementedError(fail_msg)
+        elif isinstance(value, cuqi.distribution.Distribution) and value.is_symmetric:
+            self._proposal = value
+        else:
+            raise ValueError(fail_msg)
+        self._proposal.geometry = self.target.geometry
+
+    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))
+        target_eval = np.empty(Ns)
+        acc = np.zeros(Ns, dtype=int)
+
+        # initial state    
+        samples[:, 0] = self.x0
+        target_eval[0] = self.target.logd(self.x0)
+        acc[0] = 1
+
+        # run MCMC
+        for s in range(Ns-1):
+            # run component by component
+            samples[:, s+1], target_eval[s+1], acc[s+1] = self.single_update(samples[:, s], target_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:]
+        target_eval = target_eval[Nb:]
+        accave = acc[Nb:].mean()   
+        print('\nAverage acceptance rate:', accave, '\n')
+        #
+        return samples, target_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))
+        target_eval = np.empty(Ns)
+        acc = np.zeros(Ns)
+
+        # initial state    
+        samples[:, 0] = self.x0
+        target_eval[0] = self.target.logd(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.234    # target acceptance rate RW
+        i, idx = 0, 0
+
+        # run MCMC
+        for s in range(Ns-1):
+            # run component by component
+            samples[:, s+1], target_eval[s+1], acc[s+1] = self.single_update(samples[:, s], target_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
+            self._print_progress(s+2,Ns) #s+2 is the sample number, s+1 is index assuming x0 is the first sample
+
+
+        # remove burn-in
+        samples = samples[:, Nb:]
+        target_eval = target_eval[Nb:]
+        accave = acc[Nb:].mean()   
+        print('\nAverage acceptance rate:', accave, 'MCMC scale:', self.scale, '\n')
+        
+        return samples, target_eval, accave
+
+
+    def single_update(self, x_t, target_eval_t):
+        # propose state
+        xi = self.proposal.sample(1)   # sample from the proposal
+        x_star = x_t + self.scale*xi.flatten()   # MH proposal
+
+        # evaluate target
+        target_eval_star = self.target.logd(x_star)
+
+        # ratio and acceptance probability
+        ratio = target_eval_star - target_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
+            target_eval_next = target_eval_star
+            acc = 1
+        else:
+            x_next = x_t
+            target_eval_next = target_eval_t
+            acc = 0
+        
+        return x_next, target_eval_next, acc