CUQIpy 1.3.0.post0.dev298__py3-none-any.whl → 1.4.0.post0.dev92__py3-none-any.whl

cuqi/experimental/mcmc/_conjugate_approx.py

@@ -1,76 +0,0 @@
-import numpy as np
-from cuqi.experimental.mcmc import Conjugate
-from cuqi.experimental.mcmc._conjugate import _ConjugatePair, _get_conjugate_parameter, _check_conjugate_parameter_is_scalar_reciprocal
-from cuqi.distribution import LMRF, Gamma
-import scipy as sp
-
-class ConjugateApprox(Conjugate):
-    """ 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) where Gamma is defined on the inverse of the scale parameter of the LMRF distribution.
-
-    Gamma distribution must be univariate.
-
-    LMRF likelihood must have zero mean.
-
-    For more details on conjugacy see :class:`Conjugate`.
-
-    """
-
-    def _set_conjugatepair(self):
-        """ Set the conjugate pair based on the likelihood and prior. This requires target to be set. """
-        if isinstance(self.target.likelihood.distribution, LMRF) and isinstance(self.target.prior, Gamma):
-            self._conjugatepair = _LMRFGammaPair(self.target)
-        else:
-            raise ValueError(f"Conjugacy is not defined for likelihood {type(self.target.likelihood.distribution)} and prior {type(self.target.prior)}, in CUQIpy")
-
-
-class _LMRFGammaPair(_ConjugatePair):
-    """ Implementation of the conjugate pair (LMRF, Gamma) """
-
-    def validate_target(self):
-        if not self.target.prior.dim == 1:
-            raise ValueError("Approximate conjugate sampler only works with univariate Gamma prior")
-        
-        if np.sum(self.target.likelihood.distribution.location) != 0:
-            raise ValueError("Approximate conjugate sampler only works with zero mean LMRF likelihood")
-        
-        key_value_pairs = _get_conjugate_parameter(self.target)
-        if len(key_value_pairs) != 1:
-            raise ValueError(f"Multiple references to conjugate parameter {self.target.prior.name} found in likelihood. Only one occurance is supported.")
-        for key, value in key_value_pairs:
-            if key == "scale":
-                if not _check_conjugate_parameter_is_scalar_reciprocal(value):
-                    raise ValueError("Approximate conjugate sampler only works with Gamma prior on the inverse of the scale parameter of the LMRF likelihood")
-            else:
-                raise ValueError(f"No approximate conjugacy defined for likelihood {type(self.target.likelihood.distribution)} and prior {type(self.target.prior)}, in CUQIpy")
-        
-    def conjugate_distribution(self):
-        # Extract variables
-        # Here we approximate the LMRF 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
-        return Gamma(shape=d+alpha, rate=np.linalg.norm(Lx)**2+beta)

cuqi/experimental/mcmc/_cwmh.py

@@ -1,190 +0,0 @@
-import numpy as np
-import cuqi
-from cuqi.experimental.mcmc import ProposalBasedSampler
-from cuqi.array import CUQIarray
-from numbers import Number
-
-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 or ndarray
-        Scale parameter used to define correlation between previous and proposed
-        sample in random-walk.  *Optional*. If float, the same scale is used for
-        all dimensions. If ndarray, a (possibly) different scale is used for
-        each dimension.
-
-    initial_point : ndarray
-        Initial parameters. *Optional*
-
-    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)`.
-
-    kwargs : dict
-        Additional keyword arguments to be passed to the base class 
-        :class:`ProposalBasedSampler`.
-
-    Example
-    -------
-    .. code-block:: python
-        import numpy as np
-        import cuqi
-        # 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.experimental.mcmc.CWMH(target, scale=1)
-
-        # Sample
-        samples = sampler.sample(2000).get_samples()
-
-    """
-
-    _STATE_KEYS = ProposalBasedSampler._STATE_KEYS.union(['_scale_temp'])
-
-    def __init__(self, target:cuqi.density.Density=None, proposal=None, scale=1,
-                 initial_point=None, **kwargs):
-        super().__init__(target, proposal=proposal, scale=scale,
-                         initial_point=initial_point, **kwargs)
-        
-    def _initialize(self):
-        if isinstance(self.scale, Number):
-            self.scale = np.ones(self.dim)*self.scale
-        self._acc = [np.ones((self.dim))] # Overwrite acc from ProposalBasedSampler with list of arrays
-
-        # Handling of temporary scale parameter due to possible bug in old CWMH
-        self._scale_temp = self.scale.copy()
-
-    @property
-    def scale(self):
-        """ Get the scale parameter. """
-        return self._scale
-
-    @scale.setter
-    def scale(self, value):
-        """ Set the scale parameter. """
-        if self._is_initialized and isinstance(value, Number):
-            value = np.ones(self.dim)*value
-        self._scale = value
-
-    def validate_target(self):
-        if not isinstance(self.target, cuqi.density.Density):
-            raise ValueError(
-                "Target should be an instance of "+\
-                f"{cuqi.density.Density.__class__.__name__}")
-        # Fail 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")
-        
-    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")
-        
-    @property
-    def proposal(self):
-        if self._proposal is None:
-            self._proposal = cuqi.distribution.Normal(
-                mean=lambda location: location,
-                std=lambda scale: scale,
-                geometry=self.dim,
-            )
-        return self._proposal
-    
-    @proposal.setter
-    def proposal(self, value):
-        self._proposal = value
-
-    def step(self):
-        # Initialize x_t which is used to store the current CWMH sample
-        x_t = self.current_point.copy()
-
-        # Initialize x_star which is used to store the proposed sample by
-        # updating the current sample component-by-component
-        x_star = self.current_point.copy()
-
-        # Propose a sample x_all_components from the proposal distribution
-        # for all the components
-        target_eval_t = self.current_target_logd
-        if isinstance(self.proposal,cuqi.distribution.Distribution):
-            x_all_components = self.proposal(
-                location= self.current_point, scale=self.scale).sample()
-        else:
-            x_all_components = self.proposal(self.current_point, self.scale)
-
-        # Initialize acceptance rate
-        acc = np.zeros(self.dim)
-
-        # Loop over all the components of the sample and accept/reject
-        # each component update.
-        for j in range(self.dim):
-            # propose state x_star by updating the j-th component
-            x_star[j] = x_all_components[j]
-
-            # evaluate target
-            target_eval_star = self.target.logd(x_star)
-
-            # compute Metropolis acceptance ratio
-            alpha = min(0, target_eval_star - target_eval_t)
-
-            # accept/reject
-            u_theta = np.log(np.random.rand())
-            if (u_theta <= alpha) and \
-               (not np.isnan(target_eval_star)) and \
-               (not np.isinf(target_eval_star)):
-                x_t[j] = x_all_components[j]
-                target_eval_t = target_eval_star
-                acc[j] = 1
-
-            x_star = x_t.copy()
-
-        self.current_target_logd = target_eval_t
-        self.current_point = x_t
-
-        return acc
-
-    def tune(self, skip_len, update_count):
-        # Store update_count in variable i for readability
-        i = update_count
-
-        # Optimal acceptance rate for CWMH
-        star_acc = 0.21/self.dim + 0.23
-
-        # Mean of acceptance rate over the last skip_len samples
-        hat_acc = np.mean(self._acc[i*skip_len:(i+1)*skip_len], axis=0)
-
-        # Compute new intermediate scaling parameter scale_temp
-        # Factor zeta ensures that the variation of the scale update vanishes
-        zeta = 1/np.sqrt(update_count+1)  
-        scale_temp = np.exp(
-            np.log(self._scale_temp) + zeta*(hat_acc-star_acc))
-
-        # Update the scale parameter
-        self.scale = np.minimum(scale_temp, np.ones(self.dim))
-        self._scale_temp = scale_temp

cuqi/experimental/mcmc/_gibbs.py

@@ -1,374 +0,0 @@
-from cuqi.distribution import JointDistribution
-from cuqi.experimental.mcmc import Sampler
-from cuqi.samples import Samples, JointSamples
-from cuqi.experimental.mcmc import NUTS
-from typing import Dict
-import numpy as np
-import warnings
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    def tqdm(iterable, **kwargs):
-        warnings.warn("Module mcmc: tqdm not found. Install tqdm to get sampling progress.")
-        return iterable
-
-# Not subclassed from Sampler as Gibbs handles multiple samplers and samples multiple parameters
-# Similar approach as for JointDistribution
-class HybridGibbs: 
-    """
-    Hybrid 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. 
-
-    Hybrid Gibbs sampler is a generalization of the Gibbs sampler where the
-    conditional distributions are sampled using different MCMC samplers.
-    
-    When the conditionals are sampled exactly, the samples from the Gibbs 
-    sampler converge to the joint distribution. See e.g.
-    Gelman et al. "Bayesian Data Analysis" (2014), Third Edition
-    for more details.
-
-    In each Gibbs step, the corresponding sampler has the initial_point 
-    and initial_scale (if applicable) set to the value of the previous step
-    and the sampler is reinitialized. This means that the sampling is not 
-    fully stateful at this point. This means samplers like NUTS will lose
-    their internal state between Gibbs steps.
-
-    The order in which the conditionals are sampled is the order of the
-    variables in the sampling strategy, unless a different sampling order
-    is specified by the parameter `scan_order`
-
-    Parameters
-    ----------
-    target : cuqi.distribution.JointDistribution
-        Target distribution to sample from.
-    
-    sampling_strategy : dict
-        Dictionary of sampling strategies for each variable.
-        Keys are variable names.
-        Values are sampler objects.
-
-    num_sampling_steps : dict, *optional*
-        Dictionary of number of sampling steps for each variable.
-        The sampling steps are defined as the number of times the sampler
-        will call its step method in each Gibbs step.
-        Default is 1 for all variables.
-
-    scan_order : list or str, *optional*
-        Order in which the conditional distributions are sampled.
-        If set to "random", use a random ordering at each step.
-        If not specified, it will be the order in the sampling_strategy.
-
-    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)`.
-
-    Example
-    -------
-    .. code-block:: python
-
-        import cuqi
-        import numpy as np
-
-        # Model and data
-        A, y_obs, probinfo = cuqi.testproblem.Deconvolution1D(phantom='sinc').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.experimental.mcmc.LinearRTO(maxit=15),
-            'd': cuqi.experimental.mcmc.Conjugate(),
-            'l': cuqi.experimental.mcmc.Conjugate(),
-        }
-
-        # Define Gibbs sampler
-        sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
-
-        # Run sampler
-        sampler.warmup(200)
-        sampler.sample(1000)
-
-        # Get samples removing burn-in
-        samples = sampler.get_samples().burnthin(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[str, Sampler], num_sampling_steps: Dict[str, int] = None, scan_order = None, callback=None):
-
-        # 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)
-
-        # Store sampler instances (again as a copy to avoid modifying the original)
-        self.samplers = sampling_strategy.copy()
-
-        # Store number of sampling steps for each parameter
-        self.num_sampling_steps = num_sampling_steps
-
-        # Store parameter names
-        self.par_names = self.target.get_parameter_names()
-
-        # Store the scan order
-        self._scan_order = scan_order
-
-        # Check that the parameters of the target align with the sampling_strategy and scan_order
-        if set(self.par_names) != set(self.scan_order):
-            raise ValueError("Parameter names in JointDistribution do not equal the names in the scan order.")
-
-        # Initialize sampler (after target is set)
-        self._initialize()
-
-        # Set the callback function
-        self.callback = callback
-
-    def _initialize(self):
-        """ Initialize sampler """
-
-        # Initial points
-        self.current_samples = self._get_initial_points()
-
-        # Initialize sampling steps
-        self._initialize_num_sampling_steps()
-
-        # Allocate samples
-        self._allocate_samples()
-
-        # Set targets
-        self._set_targets()
-
-        # Initialize the samplers
-        self._initialize_samplers()
-
-        # Validate all targets for samplers.
-        self.validate_targets()
-
-    @property
-    def scan_order(self):
-        if self._scan_order is None:
-            return list(self.samplers.keys())
-        if self._scan_order == "random":
-            arr = list(self.samplers.keys())
-            np.random.shuffle(arr) # Shuffle works in-place
-            return arr
-        return self._scan_order
-
-    # ------------ Public methods ------------
-    def validate_targets(self):
-        """ Validate each of the conditional targets used in the Gibbs steps """
-        if not isinstance(self.target, JointDistribution):
-            raise ValueError('Target distribution must be a JointDistribution.')
-        for sampler in self.samplers.values():
-            sampler.validate_target()
-
-    def sample(self, Ns) -> 'HybridGibbs':
-        """ Sample from the joint distribution using Gibbs sampling
-
-        Parameters
-        ----------
-        Ns : int
-            The number of samples to draw.
-
-        """
-        for idx in tqdm(range(Ns), "Sample: "):
-
-            self.step()
-
-            self._store_samples()
-
-            # Call callback function if specified
-            self._call_callback(idx, Ns)
-
-        return self
-
-    def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
-        """ Warmup (tune) the samplers in the Gibbs sampling scheme
-
-        Parameters
-        ----------
-        Nb : int
-            The number of samples to draw during warmup.
-
-        tune_freq : float, optional
-            Frequency of tuning the samplers. Tuning is performed every tune_freq*Nb steps.
-
-        """
-
-        tune_interval = max(int(tune_freq * Nb), 1)
-
-        for idx in tqdm(range(Nb), "Warmup: "):
-
-            self.step()
-
-            # Tune the sampler at tuning intervals (matching behavior of Sampler class)
-            if (idx + 1) % tune_interval == 0:
-                self.tune(tune_interval, idx // tune_interval) 
-                
-            self._store_samples()
-
-            # Call callback function if specified
-            self._call_callback(idx, Nb)
-
-        return self
-
-    def get_samples(self) -> Dict[str, Samples]:
-        samples_object = JointSamples()
-        for par_name in self.par_names:
-            samples_array = np.array(self.samples[par_name]).T
-            samples_object[par_name] = Samples(samples_array, self.target.get_density(par_name).geometry)
-        return samples_object
-    
-    def step(self):
-        """ Sequentially go through all parameters and sample them conditionally on each other """
-
-        # Sample from each conditional distribution
-        for par_name in self.scan_order:
-
-            # Set target for current parameter
-            self._set_target(par_name)
-
-            # Get sampler
-            sampler = self.samplers[par_name]
-
-            # Instead of simply changing the target of the sampler, we reinitialize it.
-            # This is to ensure that all internal variables are set to match the new target.
-            # To return the sampler to the old state and history, we first extract the state and history
-            # before reinitializing the sampler and then set the state and history back to the sampler
-
-            # Extract state and history from sampler
-            if isinstance(sampler, NUTS): # Special case for NUTS as it is not playing nice with get_state and get_history
-                sampler.initial_point = sampler.current_point
-            else:
-                sampler_state = sampler.get_state()
-                sampler_history = sampler.get_history()
-
-            # Reinitialize sampler
-            sampler.reinitialize()
-
-            # Set state and history back to sampler
-            if not isinstance(sampler, NUTS): # Again, special case for NUTS.
-                sampler.set_state(sampler_state)
-                sampler.set_history(sampler_history)
-
-            # Allow for multiple sampling steps in each Gibbs step
-            for _ in range(self.num_sampling_steps[par_name]):
-                # Sampling step
-                acc = sampler.step()
-
-                # Store acceptance rate in sampler (matching behavior of Sampler class Sample method)
-                sampler._acc.append(acc)
-
-            # Extract samples (Ensure even 1-dimensional samples are 1D arrays)
-            if isinstance(sampler.current_point, np.ndarray):
-                self.current_samples[par_name] = sampler.current_point.reshape(-1)
-            else:
-                self.current_samples[par_name] = sampler.current_point
-
-    def tune(self, skip_len, update_count):
-        """ Run a single tuning step on each of the samplers in the Gibbs sampling scheme
-
-        Parameters
-        ----------
-        skip_len : int
-            Defines the number of steps in between tuning (i.e. the tuning interval).
-
-        update_count : int
-            The number of times tuning has been performed. Can be used for internal bookkeeping.
-
-        """
-        for par_name in self.par_names:
-            self.samplers[par_name].tune(skip_len=skip_len, update_count=update_count)
-
-    # ------------ Private methods ------------
-    def _call_callback(self, sample_index, num_of_samples):
-        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
-        if self.callback is not None:
-            self.callback(self, sample_index, num_of_samples)
-
-    def _initialize_samplers(self):
-        """ Initialize samplers """
-        for sampler in self.samplers.values():
-            if isinstance(sampler, NUTS):
-                print(f'Warning: NUTS sampler is not fully stateful in HybridGibbs. Sampler will be reinitialized in each Gibbs step.')
-            sampler.initialize()
-
-    def _initialize_num_sampling_steps(self):
-        """ Initialize the number of sampling steps for each sampler. Defaults to 1 if not set by user """
-
-        if self.num_sampling_steps is None:
-            self.num_sampling_steps = {par_name: 1 for par_name in self.par_names}
-
-        for par_name in self.par_names:
-            if par_name not in self.num_sampling_steps:
-                self.num_sampling_steps[par_name] = 1
-
-
-    def _set_targets(self):
-        """ Set targets for all samplers using the current samples """
-        par_names = self.par_names
-        for par_name in par_names:
-            self._set_target(par_name)
-
-    def _set_target(self, par_name):
-        """ Set target conditional distribution for a single parameter using the current samples """
-        # Get all other conditional parameters other than the current parameter and update the target
-        # This defines - from a joint p(x,y,z) - the conditional distribution p(x|y,z) or p(y|x,z) or p(z|x,y)
-        conditional_params = {par_name_: self.current_samples[par_name_] for par_name_ in self.par_names if par_name_ != par_name}
-        self.samplers[par_name].target = self.target(**conditional_params)
-
-    def _allocate_samples(self):
-        """ Allocate memory for samples """
-        samples = {}
-        for par_name in self.par_names:
-            samples[par_name] = []
-        self.samples = samples
-
-    def _get_initial_points(self):
-        """ Get initial points for each parameter """
-        initial_points = {}
-        for par_name in self.par_names:
-            sampler = self.samplers[par_name]
-            if sampler.initial_point is None:
-                sampler.initial_point = sampler._get_default_initial_point(self.target.get_density(par_name).dim)
-            initial_points[par_name] = sampler.initial_point
-            
-        return initial_points
-
-    def _store_samples(self):
-        """ Store current samples at index i of samples dict """
-        for par_name in self.par_names:
-            self.samples[par_name].append(self.current_samples[par_name])
-
-    def __repr__(self):
-        """ Return a string representation of the sampler. """
-        msg = f"Sampler: {self.__class__.__name__} \n"
-        if self.target is None:
-            msg += f" Target: None \n"
-        else:
-            msg += f" Target: \n \t {self.target} \n\n"
-            
-        for key, value in zip(self.samplers.keys(), self.samplers.values()):
-            msg += f" Variable '{key}' with {value} \n"
-
-        return msg