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

cuqi/implicitprior/_restorator.py

@@ -0,0 +1,269 @@
+from abc import ABC, abstractmethod
+from cuqi.distribution import Distribution
+import numpy as np
+    
+class RestorationPrior(Distribution):
+    """    
+    This class defines an implicit distribution associated with a restoration operator
+    (eg denoiser). They are several works relating restorations operators with
+    priors, see 
+        -Laumont et al. https://arxiv.org/abs/2103.04715
+        -Hu et al. https://openreview.net/pdf?id=x7d1qXEn1e
+    We cannot sample from this distribution, neither compute its logpdf except in
+    some cases. It allows us to apply algorithms such as MYULA and PnPULA.
+    
+    Parameters
+    ---------- 
+    restorator : callable f(x, restoration_strength)
+        Function f that accepts input x to be restored and returns a two-element 
+        tuple of the restored version of x and extra information about the 
+        restoration operation. The second element can be of any type, including
+        `None` in case there is no information.
+            
+    restorator_kwargs : dictionary
+        Dictionary containing information about the restorator.
+        It contains keyword argument parameters that will be passed to the
+        restorator f. An example could be algorithm parameters such as the number
+        of iterations or the stopping criteria.
+    
+    potential : callable function, optional
+        The potential corresponds to the negative logpdf when it is accessible.
+        This function is a mapping from the parameter domain to the real set. 
+        It can be provided if the user knows how to relate it to the restorator.
+        Ex: restorator is the proximal operator of the total variation (TV), then
+        potential is the TV function. 
+    """
+    def __init__(self, restorator, restorator_kwargs
+                =None, potential=None, **kwargs):
+        if restorator_kwargs is None:
+            restorator_kwargs = {}
+        self.restorator = restorator 
+        self.restorator_kwargs = restorator_kwargs
+        self.potential = potential
+        super().__init__(**kwargs)
+
+    def restore(self, x, restoration_strength):
+        """This function allows us to restore the input x with the user-supplied
+        restorator. Extra information about the restoration operation is stored
+        in the `RestorationPrior` info attribute.
+        
+        Parameters
+        ---------- 
+        x : ndarray
+            parameter we want to restore.
+        
+        restoration_strength: positive float
+            Strength of the restoration operation. In the case where the
+            restorator is a denoiser, this parameter might correspond to the
+            noise level.
+        """
+        restorator_return = self.restorator(x, restoration_strength=restoration_strength,
+                                         **self.restorator_kwargs)
+
+        if type(restorator_return) == tuple and len(restorator_return) == 2:
+            solution, self.info = restorator_return
+        else:
+            raise ValueError("Unsupported return type from the user-supplied restorator function. "+ 
+                             "Please ensure that the restorator function returns a two-element tuple with the "+
+                             "restored solution as the first element and additional information about the "+
+                             "restoration as the second element. The second element can be of any type, "+
+                             "including `None` in case there is no particular information.")
+
+        return solution
+    
+    def logpdf(self, x):
+        """The logpdf function. It returns nan because we don't know the 
+        logpdf of the implicit prior."""
+        if self.potential is None:
+            return np.nan
+        else:
+            return -self.potential(x)
+    
+    def _sample(self, N, rng=None):
+        raise NotImplementedError("The sample method is not implemented for the"
+                                  + "RestorationPrior class.")
+
+    @property
+    def _mutable_vars(self):
+        """ Returns the mutable variables of the distribution. """
+        # Currently mutable variables are not supported for user-defined
+        # distributions.
+        return []
+
+    def get_conditioning_variables(self):
+        """ Returns the conditioning variables of the distribution. """
+        # Currently conditioning variables are not supported for user-defined
+        # distributions.
+        return []
+
+
+class MoreauYoshidaPrior(Distribution):
+    """    
+    This class defines (implicit) smoothed priors for which we can apply 
+    gradient-based algorithms. The smoothing is performed using
+    the Moreau-Yoshida envelope of the target prior potential.
+    
+    In the following we give a detailed explanation of the
+    Moreau-Yoshida smoothing.
+    
+    We consider a density such that - \log\pi(x) = -g(x) with g convex, lsc,
+    proper but not differentiable. Consequently, we cannot apply any
+    algorithm requiring the gradient of g.
+    Idea:
+    We consider the Moreau envelope of g defined as
+    
+    g_{smoothing_strength} (x) = inf_z 0.5*\| x-z \|_2^2/smoothing_strength + g(z).
+    
+    g_{smoothing_strength} has some nice properties
+        - g_{smoothing_strength}(x)-->g(x) as smoothing_strength-->0 for all x
+        - \nabla g_{smoothing_strength} is 1/smoothing_strength-Lipschitz
+        - \nabla g_{smoothing_strength}(x) = (x - prox_g^{smoothing_strength}(x))/smoothing_strength for all x with 
+        
+        prox_g^{smoothing_strength}(x) = argmin_z 0.5*\| x-z \|_2^2/smoothing_strength + g(z) .
+
+    Consequently, we can apply any gradient-based algorithm with
+    g_{smoothing_strength} in lieu of g. These algorithms do not require the
+    full knowledge of g_{smoothing_strength} but only its gradient. The gradient
+    of g_{smoothing_strength} is fully determined by prox_g^{smoothing_strength}
+    and smoothing_strength.
+    It is important as, although there exists an explicit formula for
+    g_{smoothing_strength}, it is rarely used in practice, as it would require
+    us to solve an optimization problem each time we want to 
+    estimate g_{smoothing_strength}. Furthermore, there exist cases where we dont't
+    the regularization g with which the mapping prox_g^{smoothing_strength} is
+    associated.
+
+    Remark (Proximal operators are denoisers):
+        We consider the denoising inverse problem x = u + n, with
+        n \sim \mathcal{N}(0, smoothing_strength I).
+        A mapping solving a denoising inverse problem is called denoiser. It takes
+        the noisy observation x as an input and returns a less noisy version of x
+        which is an estimate of u.
+        We assume a prior density \pi(u) \propto exp(- g(u)).
+        Then the MAP estimate is given by 
+            x_MAP = \argmin_z 0.5 \| x - z \|_2^2/smoothing_strength + g(z) = prox_g^smoothing_strength(x)
+        Then proximal operators are denoisers. 
+    
+    Remark (Denoisers are not necessarily proximal operators): Data-driven
+    denoisers are not necessarily proximal operators
+    (see https://arxiv.org/pdf/2201.13256)
+    
+    Parameters
+    ---------- 
+    prior : RestorationPrior
+        Prior of the RestorationPrior type. In order to stay within the MYULA
+        framework the restorator of RestorationPrior must be a proximal operator.
+        
+    smoothing_strength : float
+        Smoothing strength of the Moreau-Yoshida envelope of the prior potential.
+    """
+
+    def __init__(self, prior:RestorationPrior, smoothing_strength=0.1, 
+                 **kwargs):
+        self.prior = prior
+        self.smoothing_strength = smoothing_strength
+
+        # if kwargs does not contain the geometry,
+        # we set it to the geometry of the prior, if it exists
+        if "geometry" in kwargs:
+            raise ValueError(
+                "The geometry parameter is not supported for the"
+                + "MoreauYoshidaPrior class. The geometry is"
+                + "automatically set to the geometry of the prior.")
+        try:
+            geometry = prior.geometry
+        except:
+            geometry = None
+
+        super().__init__(geometry=geometry, **kwargs)
+
+    @property
+    def geometry(self):
+        return self.prior.geometry
+
+    @geometry.setter
+    def geometry(self, value):
+        self.prior.geometry = value
+
+    @property
+    def smoothing_strength(self):
+        """ smoothing_strength of the distribution"""
+        return self._smoothing_strength
+
+    @smoothing_strength.setter
+    def smoothing_strength(self, value):
+        self._smoothing_strength = value
+        
+    @property
+    def prior(self):
+        """Getter for the MoreauYoshida prior."""
+        return self._prior
+
+    @prior.setter
+    def prior(self, value):
+        self._prior = value
+    
+    def gradient(self, x):
+        """This is the gradient of the regularizer ie gradient of the negative
+        logpdf of the implicit prior."""
+        return -(x - self.prior.restore(x, self.smoothing_strength))/self.smoothing_strength
+        
+    def logpdf(self, x):
+        """The logpdf function. It returns nan because we don't know the 
+        logpdf of the implicit prior."""
+        if self.prior.potential == None:
+            return np.nan
+        else:
+            return -(self.prior.potential(self.prior.restore(x, self.smoothing_strength))*self.smoothing_strength +
+                     0.5*((x-self.prior.restore(x, self.smoothing_strength))**2).sum())
+    
+    def _sample(self, N, rng=None):
+        raise NotImplementedError("The sample method is not implemented for the"
+                                  + f"{self.__class__.__name__} class.")
+
+    @property
+    def _mutable_vars(self):
+        """ Returns the mutable variables of the distribution. """
+        # Currently mutable variables are not supported for user-defined
+        # distributions.
+        return []
+
+    def get_conditioning_variables(self):
+        """ Returns the conditioning variables of the distribution. """
+        # Currently conditioning variables are not supported for user-defined
+        # distributions.
+        return []
+
+class TweediePrior(MoreauYoshidaPrior):
+    """
+    Alias for MoreauYoshidaPrior following Tweedie's formula framework. TweediePrior 
+    defines priors where gradients are computed based on Tweedie's identity that links 
+    MMSE (Minimum Mean Square Error) denoisers with the underlying smoothed prior, see:
+        - Laumont et al. https://arxiv.org/abs/2103.04715 or https://doi.org/10.1137/21M1406349
+
+    Tweedie's Formula
+    -------------------------
+    In the context of denoising, Tweedie's identity states that for a signal x
+    corrupted by Gaussian noise:
+
+        ∇_x log p_e(x) = (D_e(x) - x) / e
+
+    where D_e(x) is the MMSE denoiser output and e is the noise variance.
+    This enables us to perform gradient-based sampling with algorithms like ULA.
+
+    At implementation level, TweediePrior shares identical functionality with MoreauYoshidaPrior. 
+    Thus, it is implemented as an alias of MoreauYoshidaPrior, meaning all methods, 
+    properties, and behavior are identical. The separate name provides clarity when 
+    working specifically with Tweedie's formula-based approaches.
+
+    Parameters
+    ----------
+    prior : RestorationPrior
+        Prior of the RestorationPrior type containing a denoiser/restorator.
+        
+    smoothing_strength : float, default=0.1
+        Corresponds to the noise variance e in Tweedie's formula context.
+
+    See MoreauYoshidaPrior for the underlying implementation with complete documentation.
+    """
+    pass

cuqi/legacy/__init__.py

@@ -0,0 +1,2 @@
+""" Legacy module for functionalities that are no longer supported or developed. """
+from . import sampler

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

@@ -1,15 +1,11 @@
-""" Re-implementation of sampler module in a more object oriented way. """
-
 from ._sampler import Sampler, ProposalBasedSampler
+from ._conjugate import Conjugate
+from ._conjugate_approx import ConjugateApprox
+from ._cwmh import CWMH
+from ._gibbs import Gibbs
+from ._hmc import NUTS
 from ._langevin_algorithm import ULA, MALA
+from ._laplace_approximation import UGLA
 from ._mh import MH
-from ._pcn import PCN
+from ._pcn import pCN
 from ._rto import LinearRTO, RegularizedLinearRTO
-from ._cwmh import CWMH
-from ._laplace_approximation import UGLA
-from ._hmc import NUTS
-from ._gibbs import HybridGibbs
-from ._conjugate import Conjugate
-from ._conjugate_approx import ConjugateApprox
-from ._direct import Direct
-from ._utilities import find_valid_samplers

cuqi/legacy/sampler/_conjugate.py

@@ -0,0 +1,55 @@
+from cuqi.distribution import Posterior, Gaussian, Gamma, GMRF
+from cuqi.implicitprior import RegularizedGaussian, RegularizedGMRF
+import numpy as np
+
+class Conjugate: # TODO: Subclass from Sampler once updated
+    """ Conjugate sampler
+
+    Sampler for sampling a posterior distribution where the likelihood and prior are conjugate.
+
+    Currently supported conjugate pairs are:
+    - (Gaussian, Gamma)
+    - (GMRF, Gamma)
+    - (RegularizedGaussian, Gamma) with nonnegativity constraints only
+
+    For more information on conjugate pairs, see https://en.wikipedia.org/wiki/Conjugate_prior.
+
+    For implicit regularized Gaussians see:
+    
+    [1] Everink, Jasper M., Yiqiu Dong, and Martin S. Andersen. "Bayesian inference with projected densities." SIAM/ASA Journal on Uncertainty Quantification 11.3 (2023): 1025-1043.
+
+    """
+
+    def __init__(self, target: Posterior):
+        if not isinstance(target.likelihood.distribution, (Gaussian, GMRF, RegularizedGaussian, RegularizedGMRF)):
+            raise ValueError("Conjugate sampler only works with a Gaussian-type likelihood function")
+        if not isinstance(target.prior, Gamma):
+            raise ValueError("Conjugate sampler only works with Gamma prior")
+        if not target.prior.dim == 1:
+            raise ValueError("Conjugate sampler only works with univariate Gamma prior")
+            
+        if isinstance(target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)) and (target.likelihood.distribution.preset["constraint"] not in ["nonnegativity"] or target.likelihood.distribution.preset["regularization"] is not None) :
+               raise ValueError("Conjugate sampler only works implicit regularized Gaussian likelihood with nonnegativity constraints")
+        
+        self.target = target
+
+    def step(self, x=None):
+        # Extract variables
+        b = self.target.likelihood.data                                 #mu
+        m = self._calc_m_for_Gaussians(b)                               #n
+        Ax = self.target.likelihood.distribution.mean                   #x_i
+        L = self.target.likelihood.distribution(np.array([1])).sqrtprec #L
+        alpha = self.target.prior.shape                                 #alpha
+        beta = self.target.prior.rate                                   #beta
+
+        # Create Gamma distribution and sample
+        dist = Gamma(shape=m/2+alpha,rate=.5*np.linalg.norm(L@(Ax-b))**2+beta)
+
+        return dist.sample()
+
+    def _calc_m_for_Gaussians(self, b):
+        """ Helper method to calculate m parameter for Gaussian-Gamma conjugate pair. """
+        if isinstance(self.target.likelihood.distribution, (Gaussian, GMRF)):
+            return len(b)
+        elif isinstance(self.target.likelihood.distribution, (RegularizedGaussian, RegularizedGMRF)):
+            return np.count_nonzero(b) # See 

cuqi/legacy/sampler/_conjugate_approx.py

@@ -0,0 +1,52 @@
+from cuqi.distribution import Posterior, LMRF, Gamma
+import numpy as np
+import scipy as sp
+
+class ConjugateApprox: # TODO: Subclass from Sampler once updated
+    """ Approximate Conjugate sampler
+
+    Sampler for sampling a posterior distribution where the likelihood and prior can be approximated
+    by a conjugate pair.
+
+    Currently supported pairs are:
+    - (LMRF, Gamma): Approximated by (Gaussian, Gamma)
+
+    For more information on conjugate pairs, see https://en.wikipedia.org/wiki/Conjugate_prior.
+
+    """
+
+    
+    def __init__(self, target: Posterior):
+        if not isinstance(target.likelihood.distribution, LMRF):
+            raise ValueError("Conjugate sampler only works with Laplace diff likelihood function")
+        if not isinstance(target.prior, Gamma):
+            raise ValueError("Conjugate sampler only works with Gamma prior")
+        self.target = target
+
+    def step(self, x=None):
+        # Extract variables
+        # Here we approximate the Laplace diff with a Gaussian
+
+        # Extract diff_op from target likelihood
+        D = self.target.likelihood.distribution._diff_op
+        n = D.shape[0]
+
+        # Gaussian approximation of LMRF prior as function of x_k
+        # See Uribe et al. (2022) for details
+        # Current has a zero mean assumption on likelihood! TODO
+        beta=1e-5
+        def Lk_fun(x_k):
+            dd =  1/np.sqrt((D @ x_k)**2 + beta*np.ones(n))
+            W = sp.sparse.diags(dd)
+            return W.sqrt() @ D
+
+        x = self.target.likelihood.data                 #x
+        d = len(x)                                      #d
+        Lx = Lk_fun(x)@x                                #Lx
+        alpha = self.target.prior.shape                 #alpha
+        beta = self.target.prior.rate                   #beta
+
+        # Create Gamma distribution and sample
+        dist = Gamma(shape=d+alpha, rate=np.linalg.norm(Lx)**2+beta)
+
+        return dist.sample()

cuqi/legacy/sampler/_cwmh.py

@@ -0,0 +1,196 @@
+import numpy as np
+import cuqi
+from cuqi.legacy.sampler import ProposalBasedSampler
+
+
+class CWMH(ProposalBasedSampler):
+    """Component-wise Metropolis Hastings sampler.
+
+    Allows sampling of a target distribution by a component-wise 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 as inputs to UserDefinedDistribution)
+        target = cuqi.distribution.UserDefinedDistribution(dim=dim, logpdf_func=logpdf_func)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.CWMH(target, scale=1)
+
+        # Sample
+        samples = sampler.sample(2000)
+
+    """
+    def __init__(self, target,  proposal=None, scale=1, x0=None, dim = None, **kwargs):
+        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, cuqi.distribution.Distribution conditioned only on 'location' and 'scale', lambda function, or cuqi.distribution.Normal conditioned only on 'mean' and 'std'"
+
+        if value is None:
+            self._proposal = cuqi.distribution.Normal(mean = lambda location:location,std = lambda scale:scale, geometry=self.dim)
+
+        elif isinstance(value, cuqi.distribution.Distribution) and sorted(value.get_conditioning_variables())==['location','scale']:
+            self._proposal = value
+
+        elif isinstance(value, cuqi.distribution.Normal) and sorted(value.get_conditioning_variables())==['mean','std']:
+            self._proposal = value(mean = lambda location:location, std = lambda scale:scale)
+
+        elif not isinstance(value, cuqi.distribution.Distribution) and callable(value):
+            self._proposal = value
+
+        else:
+            raise ValueError(fail_msg)
+
+
+    def _sample(self, N, Nb):
+        Ns = N+Nb   # number of simulations
+
+        # allocation
+        samples = np.empty((self.dim, Ns))
+        target_eval = np.empty(Ns)
+        acc = np.zeros((self.dim, Ns), dtype=int)
+
+        # initial state    
+        samples[:, 0] = self.x0
+        target_eval[0] = self.target.logd(self.x0)
+        acc[:, 0] = np.ones(self.dim)
+
+        # 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:]
+        acccomp = acc[:, Nb:].mean(axis=1)   
+        print('\nAverage acceptance rate all components:', acccomp.mean(), '\n')
+        
+        return samples, target_eval, acccomp
+
+    def _sample_adapt(self, N, Nb):
+        # this follows the vanishing adaptation Algorithm 4 in:
+        # Andrieu and Thoms (2008) - A tutorial on adaptive MCMC
+        Ns = N+Nb   # number of simulations
+
+        # allocation
+        samples = np.empty((self.dim, Ns))
+        target_eval = np.empty(Ns)
+        acc = np.zeros((self.dim, Ns), dtype=int)
+
+        # initial state
+        samples[:, 0] = self.x0
+        target_eval[0] = self.target.logd(self.x0)
+        acc[:, 0] = np.ones(self.dim)
+
+        # initial adaptation params 
+        Na = int(0.1*N)                                        # iterations to adapt
+        hat_acc = np.empty((self.dim, int(np.floor(Ns/Na))))     # average acceptance rate of the chains
+        lambd = np.empty((self.dim, int(np.floor(Ns/Na)+1)))     # scaling parameter \in (0,1)
+        lambd[:, 0] = self.scale
+        star_acc = 0.21/self.dim + 0.23    # 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 of each component using acc of past samples
+            if ((s+1) % Na == 0):
+                # evaluate average acceptance rate
+                hat_acc[:, i] = np.mean(acc[:, idx:idx+Na], axis=1)
+
+                # compute new scaling parameter
+                zeta = 1/np.sqrt(i+1)   # ensures that the variation of lambda(i) vanishes
+                lambd[:, i+1] = np.exp(np.log(lambd[:, i]) + zeta*(hat_acc[:, i]-star_acc))  
+
+                # update parameters
+                self.scale = np.minimum(lambd[:, i+1], np.ones(self.dim))
+
+                # 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
+            self._call_callback(samples[:, s+1], s+1)
+            
+        # remove burn-in
+        samples = samples[:, Nb:]
+        target_eval = target_eval[Nb:]
+        acccomp = acc[:, Nb:].mean(axis=1)
+        print('\nAverage acceptance rate all components:', acccomp.mean(), '\n')
+        
+        return samples, target_eval, acccomp
+
+    def single_update(self, x_t, target_eval_t):
+        if isinstance(self.proposal,cuqi.distribution.Distribution):
+            x_i_star = self.proposal(location= x_t, scale = self.scale).sample()
+        else:
+            x_i_star = self.proposal(x_t, self.scale) 
+        x_star = x_t.copy()
+        acc = np.zeros(self.dim)
+
+        for j in range(self.dim):
+            # propose state
+            x_star[j] = x_i_star[j]
+
+            # 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_t[j] = x_i_star[j]
+                target_eval_t = target_eval_star
+                acc[j] = 1
+            else:
+                pass
+                # x_t[j]       = x_t[j]
+                # target_eval_t = target_eval_t
+            x_star = x_t.copy()
+        #
+        return x_t, target_eval_t, acc