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

cuqi/legacy/sampler/_hmc.py

@@ -0,0 +1,335 @@
+import numpy as np
+from cuqi.legacy.sampler import Sampler
+
+
+# another implementation is in https://github.com/mfouesneau/NUTS
+class NUTS(Sampler):
+    """No-U-Turn Sampler (Hoffman and Gelman, 2014).
+
+    Samples a distribution given its logpdf and gradient using a Hamiltonian Monte Carlo (HMC) algorithm with automatic parameter tuning.
+
+    For more details see: See Hoffman, M. D., & Gelman, A. (2014). The no-U-turn sampler: Adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15, 1593-1623.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample. Must have logpdf and gradient method. Custom logpdfs and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
+    
+    x0 : ndarray
+        Initial parameters. *Optional*
+
+    max_depth : int
+        Maximum depth of the tree.
+
+    adapt_step_size : Bool or float
+        Whether to adapt the step size.
+        If True, the step size is adapted automatically.
+        If False, the step size is fixed to the initially estimated value.
+        If set to a scalar, the step size will be given by user and not adapted.
+
+    opt_acc_rate : float
+        The optimal acceptance rate to reach if using adaptive step size.
+        Suggested values are 0.6 (default) or 0.8 (as in stan).
+
+    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
+
+        # Import cuqi
+        import cuqi
+
+        # Define a target distribution
+        tp = cuqi.testproblem.WangCubic()
+        target = tp.posterior
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.NUTS(target)
+
+        # Sample
+        samples = sampler.sample(10000, 5000)
+
+        # Plot samples
+        samples.plot_pair()
+
+    After running the NUTS sampler, run diagnostics can be accessed via the 
+    following attributes:
+
+    .. code-block:: python
+
+        # Number of tree nodes created each NUTS iteration
+        sampler.num_tree_node_list
+
+        # Step size used in each NUTS iteration
+        sampler.epsilon_list
+
+        # Suggested step size during adaptation (the value of this step size is
+        # only used after adaptation). The suggested step size is None if 
+        # adaptation is not requested.
+        sampler.epsilon_bar_list
+
+        # Additionally, iterations' number can be accessed via
+        sampler.iteration_list
+
+    """
+    def __init__(self, target, x0=None, max_depth=15, adapt_step_size=True, opt_acc_rate=0.6, **kwargs):
+        super().__init__(target, x0=x0, **kwargs)
+        self.max_depth = max_depth
+        self.adapt_step_size = adapt_step_size
+        self.opt_acc_rate = opt_acc_rate
+        # if this flag is True, the samples and the burn-in will be returned
+        # otherwise, the burn-in will be truncated
+        self._return_burnin = False
+
+        # NUTS run diagnostic
+        # number of tree nodes created each NUTS iteration
+        self._num_tree_node = 0
+        # Create lists to store NUTS run diagnostics
+        self._create_run_diagnostic_attributes()
+
+    def _create_run_diagnostic_attributes(self):
+        """A method to create attributes to store NUTS run diagnostic."""
+        self._reset_run_diagnostic_attributes()
+
+    def _reset_run_diagnostic_attributes(self):
+        """A method to reset attributes to store NUTS run diagnostic."""
+        # NUTS iterations
+        self.iteration_list = []
+        # List to store number of tree nodes created each NUTS iteration
+        self.num_tree_node_list = []
+        # List of step size used in each NUTS iteration 
+        self.epsilon_list = []
+        # List of burn-in step size suggestion during adaptation 
+        # only used when adaptation is done
+        # remains fixed after adaptation (after burn-in)
+        self.epsilon_bar_list = []
+
+    def _update_run_diagnostic_attributes(self, k, n_tree, eps, eps_bar):
+        """A method to update attributes to store NUTS run diagnostic."""
+        # Store the current iteration number k
+        self.iteration_list.append(k)
+        # Store the number of tree nodes created in iteration k
+        self.num_tree_node_list.append(n_tree)
+        # Store the step size used in iteration k
+        self.epsilon_list.append(eps)
+        # Store the step size suggestion during adaptation in iteration k
+        self.epsilon_bar_list.append(eps_bar)
+
+    def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient
+        return self.target.logd(x), self.target.gradient(x)
+
+    def _sample_adapt(self, N, Nb):
+        return self._sample(N, Nb)
+
+    def _sample(self, N, Nb):
+        # Reset run diagnostic attributes
+        self._reset_run_diagnostic_attributes()
+        
+        if self.adapt_step_size is True and Nb == 0:
+            raise ValueError("Adaptive step size is True but number of burn-in steps is 0. Please set Nb > 0.")
+
+        # Allocation
+        Ns = Nb+N # total number of chains
+        theta = np.empty((self.dim, Ns))
+        joint_eval = np.empty(Ns)
+        step_sizes = np.empty(Ns)
+
+        # Initial state
+        theta[:, 0] = self.x0
+        joint_eval[0], grad = self._nuts_target(self.x0)
+
+        # Step size variables
+        epsilon, epsilon_bar = None, None
+
+        # parameters dual averaging
+        if (self.adapt_step_size == True):
+            epsilon = self._FindGoodEpsilon(theta[:, 0], joint_eval[0], grad)
+            mu = np.log(10*epsilon)
+            gamma, t_0, kappa = 0.05, 10, 0.75 # kappa in (0.5, 1]
+            epsilon_bar, H_bar = 1, 0
+            delta = self.opt_acc_rate # https://mc-stan.org/docs/2_18/reference-manual/hmc-algorithm-parameters.html
+            step_sizes[0] = epsilon
+        elif (self.adapt_step_size == False):
+            epsilon = self._FindGoodEpsilon(theta[:, 0], joint_eval[0], grad)
+        else:
+            epsilon = self.adapt_step_size # if scalar then user specifies the step size
+
+        # run NUTS
+        for k in range(1, Ns):
+            # reset number of tree nodes for each iteration
+            self._num_tree_node = 0
+
+            theta_k, joint_k = theta[:, k-1], joint_eval[k-1] # initial position (parameters)
+            r_k = self._Kfun(1, 'sample') # resample momentum vector
+            Ham = joint_k - self._Kfun(r_k, 'eval') # Hamiltonian
+
+            # slice variable
+            log_u = Ham - np.random.exponential(1, size=1) # u = np.log(np.random.uniform(0, np.exp(H)))
+
+            # initialization
+            j, s, n = 0, 1, 1
+            theta[:, k], joint_eval[k] = theta_k, joint_k
+            theta_minus, theta_plus = np.copy(theta_k), np.copy(theta_k)
+            grad_minus, grad_plus = np.copy(grad), np.copy(grad)
+            r_minus, r_plus = np.copy(r_k), np.copy(r_k)
+
+            # run NUTS
+            while (s == 1) and (j <= self.max_depth):
+                # sample a direction
+                v = int(2*(np.random.rand() < 0.5)-1)
+
+                # build tree: doubling procedure
+                if (v == -1):
+                    theta_minus, r_minus, grad_minus, _, _, _, \
+                    theta_prime, joint_prime, grad_prime, n_prime, s_prime, alpha, n_alpha = \
+                        self._BuildTree(theta_minus, r_minus, grad_minus, Ham, log_u, v, j, epsilon)
+                else:
+                    _, _, _, theta_plus, r_plus, grad_plus, \
+                    theta_prime, joint_prime, grad_prime, n_prime, s_prime, alpha, n_alpha = \
+                        self._BuildTree(theta_plus, r_plus, grad_plus, Ham, log_u, v, j, epsilon)
+
+                # Metropolis step
+                alpha2 = min(1, (n_prime/n)) #min(0, np.log(n_p) - np.log(n))
+                if (s_prime == 1) and (np.random.rand() <= alpha2):
+                    theta[:, k] = theta_prime
+                    joint_eval[k] = joint_prime
+                    grad = np.copy(grad_prime)
+
+                # update number of particles, tree level, and stopping criterion
+                n += n_prime
+                dtheta = theta_plus - theta_minus
+                s = s_prime * int((dtheta @ r_minus.T) >= 0) * int((dtheta @ r_plus.T) >= 0)
+                j += 1
+
+            # update run diagnostic attributes
+            self._update_run_diagnostic_attributes(
+                k, self._num_tree_node, epsilon, epsilon_bar)
+            
+            # adapt epsilon during burn-in using dual averaging
+            if (k <= Nb) and (self.adapt_step_size == True):
+                eta1 = 1/(k + t_0)
+                H_bar = (1-eta1)*H_bar + eta1*(delta - (alpha/n_alpha))
+                epsilon = np.exp(mu - (np.sqrt(k)/gamma)*H_bar)
+                eta = k**(-kappa)
+                epsilon_bar = np.exp(eta*np.log(epsilon) + (1-eta)*np.log(epsilon_bar))
+            elif (k == Nb+1) and (self.adapt_step_size == True):
+                epsilon = epsilon_bar   # fix epsilon after burn-in
+            step_sizes[k] = epsilon
+            
+            # msg
+            self._print_progress(k+1, Ns) #k+1 is the sample number, k is index assuming x0 is the first sample
+            self._call_callback(theta[:, k], k)
+            
+            if np.isnan(joint_eval[k]):
+                raise NameError('NaN potential func')
+            
+        # apply burn-in
+        if not self._return_burnin: 
+            theta = theta[:, Nb:]
+            joint_eval = joint_eval[Nb:]
+        return theta, joint_eval, step_sizes
+
+    #=========================================================================
+    # auxiliary standard Gaussian PDF: kinetic energy function
+    # d_log_2pi = d*np.log(2*np.pi)
+    def _Kfun(self, r, flag):
+        if flag == 'eval': # evaluate
+            return 0.5*(r.T @ r) #+ d_log_2pi 
+        if flag == 'sample': # sample
+            return np.random.standard_normal(size=self.dim)
+
+    #=========================================================================
+    def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1):
+        r = self._Kfun(1, 'sample')    # resample a momentum
+        Ham = joint - self._Kfun(r, 'eval')     # initial Hamiltonian
+        _, r_prime, joint_prime, grad_prime = self._Leapfrog(theta, r, grad, epsilon)
+
+        # trick to make sure the step is not huge, leading to infinite values of the likelihood
+        k = 1
+        while np.isinf(joint_prime) or np.isinf(grad_prime).any():
+            k *= 0.5
+            _, r_prime, joint_prime, grad_prime = self._Leapfrog(theta, r, grad, epsilon*k)
+        epsilon = 0.5*k*epsilon
+
+        # doubles/halves the value of epsilon until the accprob of the Langevin proposal crosses 0.5
+        Ham_prime = joint_prime - self._Kfun(r_prime, 'eval')
+        log_ratio = Ham_prime - Ham
+        a = 1 if log_ratio > np.log(0.5) else -1
+        while (a*log_ratio > -a*np.log(2)):
+            epsilon = (2**a)*epsilon
+            _, r_prime, joint_prime, _ = self._Leapfrog(theta, r, grad, epsilon)
+            Ham_prime = joint_prime - self._Kfun(r_prime, 'eval')
+            log_ratio = Ham_prime - Ham
+        return epsilon
+
+    #=========================================================================
+    def _Leapfrog(self, theta_old, r_old, grad_old, epsilon):
+        # symplectic integrator: trajectories preserve phase space volumen
+        r_new = r_old + 0.5*epsilon*grad_old     # half-step
+        theta_new = theta_old + epsilon*r_new     # full-step
+        joint_new, grad_new = self._nuts_target(theta_new)     # new gradient
+        r_new += 0.5*epsilon*grad_new     # half-step
+        return theta_new, r_new, joint_new, grad_new
+
+    #=========================================================================
+    # @functools.lru_cache(maxsize=128)
+    def _BuildTree(self, theta, r, grad, Ham, log_u, v, j, epsilon, Delta_max=1000):
+        # Increment the number of tree nodes counter
+        self._num_tree_node += 1
+
+        if (j == 0):     # base case
+            # single leapfrog step in the direction v
+            theta_prime, r_prime, joint_prime, grad_prime = self._Leapfrog(theta, r, grad, v*epsilon)
+            Ham_prime = joint_prime - self._Kfun(r_prime, 'eval')     # Hamiltonian eval
+            n_prime = int(log_u <= Ham_prime)     # if particle is in the slice
+            s_prime = int(log_u < Delta_max + Ham_prime)     # check U-turn
+            #
+            diff_Ham = Ham_prime - Ham
+
+            # Compute the acceptance probability
+            # alpha_prime = min(1, np.exp(diff_Ham))
+            # written in a stable way to avoid overflow when computing
+            # exp(diff_Ham) for large values of diff_Ham
+            alpha_prime = 1 if diff_Ham > 0 else np.exp(diff_Ham)
+            n_alpha_prime = 1
+            #
+            theta_minus, theta_plus = theta_prime, theta_prime
+            r_minus, r_plus = r_prime, r_prime
+            grad_minus, grad_plus = grad_prime, grad_prime
+        else: 
+            # recursion: build the left/right subtrees
+            theta_minus, r_minus, grad_minus, theta_plus, r_plus, grad_plus, \
+            theta_prime, joint_prime, grad_prime, n_prime, s_prime, alpha_prime, n_alpha_prime = \
+                self._BuildTree(theta, r, grad, Ham, log_u, v, j-1, epsilon)
+            if (s_prime == 1): # do only if the stopping criteria does not verify at the first subtree
+                if (v == -1):
+                    theta_minus, r_minus, grad_minus, _, _, _, \
+                    theta_2prime, joint_2prime, grad_2prime, n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                        self._BuildTree(theta_minus, r_minus, grad_minus, Ham, log_u, v, j-1, epsilon)
+                else:
+                    _, _, _, theta_plus, r_plus, grad_plus, \
+                    theta_2prime, joint_2prime, grad_2prime, n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                        self._BuildTree(theta_plus, r_plus, grad_plus, Ham, log_u, v, j-1, epsilon)
+
+                # Metropolis step
+                alpha2 = n_2prime / max(1, (n_prime + n_2prime))
+                if (np.random.rand() <= alpha2):
+                    theta_prime = np.copy(theta_2prime)
+                    joint_prime = np.copy(joint_2prime)
+                    grad_prime = np.copy(grad_2prime)
+
+                # update number of particles and stopping criterion
+                alpha_prime += alpha_2prime
+                n_alpha_prime += n_alpha_2prime
+                dtheta = theta_plus - theta_minus
+                s_prime = s_2prime * int((dtheta@r_minus.T)>=0) * int((dtheta@r_plus.T)>=0)
+                n_prime += n_2prime
+        return theta_minus, r_minus, grad_minus, theta_plus, r_plus, grad_plus, \
+                theta_prime, joint_prime, grad_prime, n_prime, s_prime, alpha_prime, n_alpha_prime
+

cuqi/legacy/sampler/_langevin_algorithm.py

@@ -0,0 +1,198 @@
+import numpy as np
+import cuqi
+from cuqi.legacy.sampler import 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
+    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt,  where L_t is 
+    the Langevin diffusion and W_t is the `dim`-dimensional standard Brownian motion.
+
+    For more details see: Roberts, G. O., & Tweedie, R. L. (1996). Exponential convergence
+    of Langevin distributions and their discrete approximations. Bernoulli, 341-363.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample. Must have logd and gradient method. Custom logpdfs 
+        and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
+    
+    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)`,
+        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)
+        gradient_func = lambda x: -2/(std**2)*(x - mu)
+
+        # Define distribution from logpdf and gradient as UserDefinedDistribution
+        target = cuqi.distribution.UserDefinedDistribution(dim=dim, logpdf_func=logpdf_func,
+            gradient_func=gradient_func)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.ULA(target, scale=1/dim**2)
+
+        # Sample
+        samples = sampler.sample(2000)
+
+    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)
+
+    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)
+
+        # msg
+        if np.isnan(logpi_eval_star):
+            raise NameError('NaN potential func. Consider using smaller scale parameter')
+
+        return x_star, logpi_eval_star, g_logpi_star, 1 # sample always accepted without Metropolis correction
+
+
+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
+    Langevin diffusion dL_t = dW_t + 1/2*Nabla target.logd(L_t)dt,  where L_t is 
+    the Langevin diffusion and W_t is the `dim`-dimensional standard Brownian motion. 
+    The sample is then accepted or rejected according to Metropolis–Hastings algorithm.
+
+    For more details see: Roberts, G. O., & Tweedie, R. L. (1996). Exponential convergence
+    of Langevin distributions and their discrete approximations. Bernoulli, 341-363.
+
+    Parameters
+    ----------
+
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample. Must have logpdf and gradient method. Custom logpdfs 
+        and gradients are supported by using a :class:`cuqi.distribution.UserDefinedDistribution`.
+    
+    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)`,
+        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)
+        gradient_func = lambda x: -2/(std**2)*(x-mu)
+
+        # Define distribution from logpdf as UserDefinedDistribution (sample and gradients also supported)
+        target = cuqi.distribution.UserDefinedDistribution(dim=dim, logpdf_func=logpdf_func,
+            gradient_func=gradient_func)
+
+        # Set up sampler
+        sampler = cuqi.legacy.sampler.MALA(target, scale=1/5**2)
+
+        # Sample
+        samples = sampler.sample(2000)
+
+    A Deblur example can be found in demos/demo28_MALA.py
+    """
+    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
+        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 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))
+