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

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

cuqi/legacy/sampler/_gibbs.py

@@ -0,0 +1,231 @@
+from cuqi.distribution import JointDistribution
+from cuqi.legacy.sampler import Sampler
+from cuqi.samples import Samples
+from typing import Dict, Union
+import numpy as np
+import sys
+import warnings
+
+class Gibbs:
+    """
+    Gibbs sampler for sampling a joint distribution.
+
+    Gibbs sampling samples the variables of the distribution sequentially,
+    one variable at a time. When a variable represents a random vector, the
+    whole vector is sampled simultaneously.
+    
+    The sampling of each variable is done by sampling from the conditional
+    distribution of that variable given the values of the other variables.
+    This is often a very efficient way of sampling from a joint distribution
+    if the conditional distributions are easy to sample from.
+
+    Parameters
+    ----------
+    target : cuqi.distribution.JointDistribution
+        Target distribution to sample from.
+    
+    sampling_strategy : dict
+        Dictionary of sampling strategies for each parameter.
+        Keys are parameter names.
+        Values are sampler objects.
+
+    Example
+    -------
+    .. code-block:: python
+
+        import cuqi
+        import numpy as np
+
+        # Model and data
+        A, y_obs, probinfo = cuqi.testproblem.Deconvolution1D(phantom='square').get_components()
+        n = A.domain_dim
+
+        # Define distributions
+        d = cuqi.distribution.Gamma(1, 1e-4)
+        l = cuqi.distribution.Gamma(1, 1e-4)
+        x = cuqi.distribution.GMRF(np.zeros(n), lambda d: d)
+        y = cuqi.distribution.Gaussian(A, lambda l: 1/l)
+
+        # Combine into a joint distribution and create posterior
+        joint = cuqi.distribution.JointDistribution(d, l, x, y)
+        posterior = joint(y=y_obs)
+
+        # Define sampling strategy
+        sampling_strategy = {
+            'x': cuqi.legacy.sampler.LinearRTO,
+            ('d', 'l'): cuqi.legacy.sampler.Conjugate,
+        }
+
+        # Define Gibbs sampler
+        sampler = cuqi.legacy.sampler.Gibbs(posterior, sampling_strategy)
+
+        # Run sampler
+        samples = sampler.sample(Ns=1000, Nb=200)
+
+        # Plot results
+        samples['x'].plot_ci(exact=probinfo.exactSolution)
+        samples['d'].plot_trace(figsize=(8,2))
+        samples['l'].plot_trace(figsize=(8,2))
+            
+    """
+
+    def __init__(self, target: JointDistribution, sampling_strategy: Dict[Union[str,tuple], Sampler]):
+
+        warnings.warn(f"\nYou are using the legacy sampler '{self.__class__.__name__}'.\n"
+                      f"This will be removed in a future release of CUQIpy.\n"
+                      f"Please consider using the new samplers in the 'cuqi.sampler' module.\n", UserWarning, stacklevel=2)
+
+        # Store target and allow conditioning to reduce to a single density
+        self.target = target() # Create a copy of target distribution (to avoid modifying the original)
+
+        # Parse samplers and split any keys that are tuple into separate keys
+        self.samplers = {}
+        for par_name in sampling_strategy.keys():
+            if isinstance(par_name, tuple):
+                for par_name_ in par_name:
+                    self.samplers[par_name_] = sampling_strategy[par_name]
+            else:
+                self.samplers[par_name] = sampling_strategy[par_name]
+
+        # Store parameter names
+        self.par_names = self.target.get_parameter_names()
+
+    # ------------ Public methods ------------
+    def sample(self, Ns, Nb=0):
+        """ Sample from target distribution """
+
+        # Initial points
+        current_samples = self._get_initial_points()
+
+        # Compute how many samples were already taken previously
+        at_Nb = self._Nb
+        at_Ns = self._Ns
+
+        # Allocate memory for samples
+        self._allocate_samples_warmup(Nb)
+        self._allocate_samples(Ns)
+
+        # Sample tuning phase
+        for i in range(at_Nb, at_Nb+Nb):
+            current_samples = self.step_tune(current_samples)
+            self._store_samples(self.samples_warmup, current_samples, i)
+            self._print_progress(i+1+at_Nb, at_Nb+Nb, 'Warmup')
+
+        # Sample phase
+        for i in range(at_Ns, at_Ns+Ns):
+            current_samples = self.step(current_samples)
+            self._store_samples(self.samples, current_samples, i)
+            self._print_progress(i+1, at_Ns+Ns, 'Sample')
+
+        # Convert to samples objects and return
+        return self._convert_to_Samples(self.samples)
+
+    def step(self, current_samples):
+        """ Sequentially go through all parameters and sample them conditionally on each other """
+
+        # Extract par names
+        par_names = self.par_names
+
+        # Sample from each conditional distribution
+        for par_name in par_names:
+
+            # Dict of all other parameters to condition on
+            other_params = {par_name_: current_samples[par_name_] for par_name_ in par_names if par_name_ != par_name}
+
+            # Set up sampler for current conditional distribution
+            sampler = self.samplers[par_name](self.target(**other_params))
+
+            # Take a MCMC step
+            current_samples[par_name] = sampler.step(current_samples[par_name])
+
+            # Ensure even 1-dimensional samples are 1D arrays
+            current_samples[par_name] = current_samples[par_name].reshape(-1)
+        
+        return current_samples
+
+    def step_tune(self, current_samples):
+        """ Perform a single MCMC step for each parameter and tune the sampler """
+        # Not implemented. No tuning happening here yet. Requires samplers to be able to be modified after initialization.
+        return self.step(current_samples)
+
+    # ------------ Private methods ------------
+    def _allocate_samples(self, Ns):
+        """ Allocate memory for samples """
+        # Allocate memory for samples
+        samples = {}
+        for par_name in self.par_names:
+            samples[par_name] = np.zeros((self.target.get_density(par_name).dim, Ns))
+        
+        # Store samples in self
+        if hasattr(self, 'samples'):
+            # Append to existing samples (This makes a copy)
+            for par_name in self.par_names:
+                samples[par_name] = np.hstack((self.samples[par_name], samples[par_name]))
+        self.samples = samples
+
+    def _allocate_samples_warmup(self, Nb):
+        """ Allocate memory for samples """
+        
+        # If we already have warmup samples and more are requested raise error
+        if hasattr(self, 'samples_warmup') and Nb != 0:
+            raise ValueError('Sampler already has run warmup phase. Cannot run warmup phase again.')
+
+        # Allocate memory for samples
+        samples = {}
+        for par_name in self.par_names:
+            samples[par_name] = np.zeros((self.target.get_density(par_name).dim, Nb))
+        self.samples_warmup = samples
+
+    def _get_initial_points(self):
+        """ Get initial points for each parameter """
+        initial_points = {}
+        for par_name in self.par_names:
+            if hasattr(self, 'samples'):
+                initial_points[par_name] = self.samples[par_name][:, -1]
+            elif hasattr(self, 'samples_warmup'):
+                initial_points[par_name] = self.samples_warmup[par_name][:, -1]
+            elif hasattr(self.target.get_density(par_name), 'init_point'):
+                initial_points[par_name] = self.target.get_density(par_name).init_point
+            else:
+                initial_points[par_name] = np.ones(self.target.get_density(par_name).dim)
+        return initial_points
+
+    def _store_samples(self, samples, current_samples, i):
+        """ Store current samples at index i of samples dict """
+        for par_name in self.par_names:
+            samples[par_name][:, i] = current_samples[par_name]
+
+    def _convert_to_Samples(self, samples):
+        """ Convert each parameter in samples dict to cuqi.samples.Samples object with correct geometry """
+        samples_object = {}
+        for par_name in self.par_names:
+            samples_object[par_name] = Samples(samples[par_name], self.target.get_density(par_name).geometry)
+        return samples_object
+
+    def _print_progress(self, s, Ns, phase):
+        """Prints sampling progress"""
+        if Ns < 2: # Don't print progress if only one sample
+            return
+        if (s % (max(Ns//100,1))) == 0:
+            msg = f'{phase} {s} / {Ns}'
+            sys.stdout.write('\r'+msg)
+        if s==Ns:
+            msg = f'{phase} {s} / {Ns}'
+            sys.stdout.write('\r'+msg+'\n')
+
+    # ------------ Private properties ------------
+    @property
+    def _Ns(self):
+        """ Number of samples already taken """
+        if hasattr(self, 'samples'):
+            return self.samples[self.par_names[0]].shape[-1]
+        else:
+            return 0
+    
+    @property
+    def _Nb(self):
+        """ Number of samples already taken in warmup phase """
+        if hasattr(self, 'samples_warmup'):
+            return self.samples_warmup[self.par_names[0]].shape[-1]
+        else:
+            return 0