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

cuqi/sampler/_laplace_approximation.py

@@ -1,11 +1,9 @@
 import scipy as sp
 import numpy as np
 import cuqi
-from cuqi.distribution import Normal
 from cuqi.solver import CGLS
 from cuqi.sampler import Sampler
 
-
 class UGLA(Sampler):
     """ Unadjusted (Gaussian) Laplace Approximation sampler
     
@@ -16,103 +14,70 @@ class UGLA(Sampler):
 
     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).
+    For more details see: Uribe, Felipe, et al. A hybrid Gibbs sampler for edge-preserving 
+    tomographic reconstruction with uncertain view angles. SIAM/ASA Journal on UQ,
+    https://doi.org/10.1137/21M1412268 (2022).
 
     Parameters
     ----------
     target : `cuqi.distribution.Posterior`
         The target posterior distribution to sample.
 
-    x0 : ndarray
-        Initial parameters. *Optional*
+    initial_point : ndarray, *Optional*
+        Initial parameters.
+        If not provided, it defaults to zeros.
 
     maxit : int
         Maximum number of inner iterations for solver when generating one sample.
+        If not provided, it defaults to 50.
 
     tol : float
-        Tolerance for inner solver. Will stop before maxit if the inner solvers convergence check reaches tol.
+        Tolerance for inner solver.
+        The inner solvers will stop before maxit if convergence check reaches tol.
+        If not provided, it defaults to 1e-4.
 
     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.
-
+        Smoothing parameter for the Gaussian approximation of the Laplace distribution.
+        A small value in the range of 1e-7 to 1e-3 is recommended, though values out of this 
+        range might give better results in some cases. Generally, a larger beta value makes 
+        sampling easier but results in a worse approximation. See details in Section 3.3 of the paper.
+        If not provided, it defaults to 1e-5.
+
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
     """
+    def __init__(self, target=None, initial_point=None, maxit=50, tol=1e-4, beta=1e-5, **kwargs):
 
-    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 Linear model
-        if not isinstance(self.target.likelihood.model, cuqi.model.LinearModel):
-            raise TypeError("Model needs to be linear")
+        super().__init__(target=target, initial_point=initial_point, **kwargs)
 
-        # 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
+        # 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.
+    
+    def _initialize(self):
+        self._precompute()
 
-        Returns
-        -------
-        samples : ndarray
-            Samples from the approximate posterior.
+    @property
+    def prior(self):
+        return self.target.prior
 
-        target_eval : ndarray
-            Log-likelihood of each sample.
+    @property
+    def likelihood(self):
+        return self.target.likelihood
 
-        acc : ndarray
-            Acceptance rate of each sample.
+    @property
+    def model(self):
+        return self.target.model     
+    
+    @property
+    def _data(self):
+        return self.target.data - self.target.model._shift
 
-        """
+    def _precompute(self):
 
-        # Extract diff_op from target prior
-        D = self.target.prior._diff_op
+        D = self.prior._diff_op
         n = D.shape[0]
 
         # Gaussian approximation of LMRF prior as function of x_k
@@ -120,65 +85,70 @@ class UGLA(Sampler):
             dd =  1/np.sqrt((D @ x_k)**2 + self.beta*np.ones(n))
             W = sp.sparse.diags(dd)
             return W.sqrt() @ D
+        self.Lk_fun = Lk_fun
 
-        # 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._m = len(self._data)
-        self._L1 = self.target.likelihood.distribution.sqrtprec
+        self._L1 = self.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)
+        if len(self.prior.location) == 1:
+            self._priorloc = np.repeat(self.prior.location, self.dim)
         else:
-            self._priorloc = self.target.prior.location
+            self._priorloc = self.prior.location
 
         # Initial Laplace approx
-        self._L2 = Lk_fun(self.x0)
+        self._L2 = Lk_fun(self.initial_point)
         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(x)
-                out2 = np.sqrt(1/self.target.prior.scale)*(self._L2 @ x)
+                out1 = self._L1 @ self.model._forward_func_no_shift(x) # Use forward function which excludes shift
+                out2 = np.sqrt(1/self.prior.scale)*(self._L2 @ x)
                 out  = np.hstack([out1, out2])
             elif flag == 2:
                 idx = int(self._m)
-                out1 = self._model.adjoint(self._L1.T@x[:idx])
-                out2 = np.sqrt(1/self.target.prior.scale)*(self._L2.T @ x[idx:])
+                out1 = self.model._adjoint_func_no_shift(self._L1.T@x[:idx])
+                out2 = np.sqrt(1/self.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()
+            return out
+        self.M = M
 
-            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)
+    def step(self):
+        # Update Laplace approximation
+        self._L2 = self.Lk_fun(self.current_point)
+        self._L2mu = self._L2@self._priorloc
+        self._b_tild = np.hstack([self._L1@self._data, self._L2mu]) 
+    
+        # Sample from approximate posterior
+        e = np.random.randn(len(self._b_tild))
+        y = self._b_tild + e # Perturb data
+        sim = CGLS(self.M, y, self.current_point, self.maxit, self.tol)                     
+        self.current_point, _ = sim.solve()
+        acc = 1
+        return acc
+
+    def tune(self, skip_len, update_count):
+        pass
+    
+    def validate_target(self):
+        # 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.likelihood.model, cuqi.model.AffineModel):
+            raise TypeError("Model needs to be affine or linear")
 
-        # remove burn-in
-        samples = samples[:, Nb:]
+        # Check Gaussian likelihood
+        if not hasattr(self.likelihood.distribution, "sqrtprec"):
+            raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
+
+        # Check that prior is LMRF
+        if not isinstance(self.prior, cuqi.distribution.LMRF):
+            raise ValueError('Unadjusted Gaussian Laplace approximation (UGLA) requires LMRF prior')
         
-        return samples, None, None
+    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)

cuqi/sampler/_mh.py

@@ -4,187 +4,77 @@ from cuqi.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.
+    """ Metropolis-Hastings (MH) sampler.
 
     Parameters
     ----------
+    target : cuqi.density.Density
+        Target density or distribution.
 
-    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*.
+    proposal : cuqi.distribution.Distribution or callable
+        Proposal distribution. If None, a random walk MH is used (i.e., Gaussian proposal with identity covariance).
 
     scale : float
-        Scale parameter used to define correlation between previous and proposed sample in random-walk.  *Optional*.
-
-    x0 : ndarray
-        Initial parameters. *Optional*
+        Scaling parameter for the proposal distribution.
 
-    dim : int
-        Dimension of parameter space. Required if target and proposal are callable functions. *Optional*.
+    kwargs : dict
+        Additional keyword arguments to be passed to the base class :class:`ProposalBasedSampler`.
 
-    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
+    _STATE_KEYS = ProposalBasedSampler._STATE_KEYS.union({'scale', '_scale_temp'})
 
-        # Logpdf function
-        logpdf_func = lambda x: -1/(std**2)*np.sum((x-mu)**2)
+    def __init__(self, target=None, proposal=None, scale=1, **kwargs):
+        super().__init__(target, proposal=proposal, scale=scale, **kwargs)
 
-        # Define distribution from logpdf as UserDefinedDistribution (sample and gradients also supported)
-        target = cuqi.distribution.UserDefinedDistribution(dim=dim, logpdf_func=logpdf_func)
+    def _initialize(self):
+        # Due to a bug? in old MH, we must keep track of this extra variable to match behavior.
+        self._scale_temp = self.scale
 
-        # Set up sampler
-        sampler = cuqi.sampler.MH(target, scale=1)
+    def validate_target(self):
+        # Fail only when there is no log density, which is currently assumed to be the case in case NaN is returned.
+        if np.isnan(self.target.logd(self._get_default_initial_point(self.dim))):
+            raise ValueError("Target does not have valid logd")
 
-        # Sample
-        samples = sampler.sample(2000)
+    def validate_proposal(self):
+        if not isinstance(self.proposal, cuqi.distribution.Distribution):
+            raise ValueError("Proposal must be a cuqi.distribution.Distribution object")
+        if not self.proposal.is_symmetric:
+            raise ValueError("Proposal must be symmetric")
 
-    """
-    #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):
+    def step(self):
         # propose state
         xi = self.proposal.sample(1)   # sample from the proposal
-        x_star = x_t + self.scale*xi.flatten()   # MH proposal
+        x_star = self.current_point + 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
+        ratio = target_eval_star - self.current_target_logd # 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 = 0
+        if (u_theta <= alpha) and \
+           (not np.isnan(target_eval_star)) and \
+           (not np.isinf(target_eval_star)):
+            self.current_point = x_star
+            self.current_target_logd = 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
+        return acc
+
+    def tune(self, skip_len, update_count):
+        hat_acc = np.mean(self._acc[-skip_len:])
+
+        # d. compute new scaling parameter
+        zeta = 1/np.sqrt(update_count+1)   # ensures that the variation of lambda(i) vanishes
+
+        # We use self._scale_temp here instead of self.scale in update. This might be a bug,
+        # but is equivalent to old MH
+        self._scale_temp = np.exp(np.log(self._scale_temp) + zeta*(hat_acc-0.234))
+
+        # update parameters
+        self.scale = min(self._scale_temp, 1)