diff --git a/cuqi/sampler/__init__.py b/cuqi/sampler/__init__.py index c1c89fd580..512ba15ddd 100644 --- a/cuqi/sampler/__init__.py +++ b/cuqi/sampler/__init__.py @@ -3,9 +3,10 @@ from ._conjugate_approx import ConjugateApprox from ._cwmh import CWMH from ._gibbs import Gibbs -from ._hmc import NUTS +from ._hmc import HMC from ._langevin_algorithm import ULA, MALA from ._laplace_approximation import UGLA from ._mh import MH +from ._nuts import NUTS from ._pcn import pCN from ._rto import LinearRTO, RegularizedLinearRTO diff --git a/cuqi/sampler/_hmc.py b/cuqi/sampler/_hmc.py index 6f070e1852..0b1228e94d 100644 --- a/cuqi/sampler/_hmc.py +++ b/cuqi/sampler/_hmc.py @@ -2,25 +2,28 @@ from cuqi.sampler import Sampler -# another implementation is in https://github.com/mfouesneau/NUTS -class NUTS(Sampler): - """No-U-Turn Sampler (Hoffman and Gelman, 2014). +class HMC(Sampler): + """Hamiltonian Monte Carlo (Duane et al., 1987). 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. + For more details see: + Neal, R. M. (2011) - MCMC Using Hamiltonian Dynamics. Handbook of Markov chain Monte Carlo. Ed. by S. Brooks et al. Chapman & Hall/CRC, 2011. Chap. 5, pp. 113-162. + Duane, S., Kennedy, A. D., Pendleton, B. J., and Roweth, D. 1987. Hybrid Monte Carlo. Physics Letters B, 195:216-222. 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_traject_length : int + This is the trajectory length or steps in the numerical integrator (leapfrog). + If True, the step size is adapted. + If set to a scalar, the step size will be given by user and not adapted. adapt_step_size : Bool or float Whether to adapt the step size. @@ -30,7 +33,7 @@ class NUTS(Sampler): 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). + Suggested value is 0.65 (default). callback : callable, *Optional* If set this function will be called after every sample. @@ -78,50 +81,47 @@ class NUTS(Sampler): sampler.iteration_list """ - def __init__(self, target, x0=None, max_depth=15, adapt_step_size=True, opt_acc_rate=0.6, **kwargs): + def __init__(self, target, x0=None, adapt_traject_length=10, adapt_step_size=True, opt_acc_rate=0.8, **kwargs): super().__init__(target, x0=x0, **kwargs) - self.max_depth = max_depth + self.adapt_traject_length = adapt_traject_length 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 + # run diagnostic + # Create lists to store HMC run diagnostics self._create_run_diagnostic_attributes() def _create_run_diagnostic_attributes(self): - """A method to create attributes to store NUTS run diagnostic.""" + """A method to create attributes to store 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 + """A method to reset attributes to store run diagnostic.""" + # 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 + # List of step size used in each 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 = [] + self.traject_length_list = [] - def _update_run_diagnostic_attributes(self, k, n_tree, eps, eps_bar): - """A method to update attributes to store NUTS run diagnostic.""" + def _update_run_diagnostic_attributes(self, k, eps, eps_bar, L): + """A method to update attributes to store 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) + # Store the trajectory length during adaptation in iteration k + self.traject_length_list.append(L) - def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient + def _hmc_target(self, x): # returns logposterior tuple evaluation-gradient return self.target.logd(x), self.target.gradient(x) def _sample_adapt(self, N, Nb): @@ -135,201 +135,165 @@ def _sample(self, N, Nb): 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 + Ns = Nb+N # total number of chains theta = np.empty((self.dim, Ns)) joint_eval = np.empty(Ns) step_sizes = np.empty(Ns) + traject_lengths = np.empty(Ns, dtype=int) + acc = np.zeros(Ns, dtype=int) # Initial state theta[:, 0] = self.x0 - joint_eval[0], grad = self._nuts_target(self.x0) + joint_eval[0], grad = self._hmc_target(self.x0) + + # set the trajectory length as it has to be accessed many times + if (self.adapt_traject_length == True): + epsilon = 1e-2 # arbitrary initialization + # the best simulation length 'lambd' was 0.14-18; reported in some studies, + # but this is a tricky parameter and finding a good 'lambd' for HMC requires + # some number of preliminary runs (that is why NUTS was created) + lambd = 10 + L = np.max(1, np.round(lambd/epsilon)) + else: + L = self.adapt_traject_length # if scalar then user specifies the step size + traject_lengths[0] = L - # Step size variables + # parameters dual averaging epsilon, epsilon_bar = None, None - - # parameters dual averaging if (self.adapt_step_size == True): - epsilon = self._FindGoodEpsilon(theta[:, 0], joint_eval[0], grad) + epsilon = self._FindGoodEpsilon(theta[:, 0], joint_eval[0], grad, L) mu = np.log(10*epsilon) + log_epsilon_bar, H_bar = 0, 0 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) + epsilon = self._FindGoodEpsilon(theta[:, 0], joint_eval[0], grad, L) else: epsilon = self.adapt_step_size # if scalar then user specifies the step size + step_sizes[0] = epsilon - # run NUTS + # run HMC 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 + q_k = theta[:, k-1] # initial position (parameters) + r_k = self._neg_kinetic_func(1, 'sample') # initial momentum vector + + # LEAPFROG: alternate full steps for position and momentum + q_star, r_star, joint_star, grad_star = self._Leapfrog(q_k, r_k, grad, epsilon, L) + + # evaluate neg-potential and neg-kinetic energies at start and end of trajectory + joint_k = joint_eval[k-1] + nHam_k = self._neg_Hamiltonian(joint_k, r_k) # Hamiltonian + nHam_star = self._neg_Hamiltonian(joint_star, r_star) # Hamiltonian + + # accept/reject + log_alpha = -1000 if np.isneginf(nHam_star) else (0 if np.isinf(nHam_star) else min(0, nHam_star - nHam_k)) + log_u = np.log(np.random.rand()) + if (log_u <= log_alpha) and not np.isnan(nHam_star): + theta[:, k] = q_star + joint_eval[k] = joint_star + acc[k] = 1 + grad = grad_star + else: + theta[:, k] = q_k + joint_eval[k] = joint_k # update run diagnostic attributes - self._update_run_diagnostic_attributes( - k, self._num_tree_node, epsilon, epsilon_bar) - + self._update_run_diagnostic_attributes(k, epsilon, epsilon_bar, L) + # 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 + eta1, eta2 = 1/(k+t_0), k**(-kappa) + # + H_bar = (1-eta1)*H_bar + eta1*(delta - np.exp(log_alpha)) + log_epsilon = mu - (np.sqrt(k)/gamma)*H_bar + log_epsilon_bar = eta2*log_epsilon + (1-eta2)*log_epsilon_bar + # + epsilon, epsilon_bar = np.exp(log_epsilon), np.exp(log_epsilon_bar) + elif (k > Nb) and (self.adapt_step_size == True): + # after warm-up we jitter the epsilons to avoid pathological behavior, see Neal's reference + epsilon = np.random.uniform(0.9*epsilon_bar, 1.1*epsilon_bar) step_sizes[k] = epsilon - + + # adapt path length + if (k <= Nb) and (self.adapt_traject_length == True): + L = np.max(1, np.round(lambd/epsilon)) + traject_lengths[k] = L + # msg - self._print_progress(k+1, Ns) #k+1 is the sample number, k is index assuming x0 is the first sample + self._print_progress(k, 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:] + acc_rate = np.mean(acc) + print('\tAcceptance rate:', acc_rate, '\n') + return theta, joint_eval, step_sizes #========================================================================= - # auxiliary standard Gaussian PDF: kinetic energy function + # kinetic energy function (negative log PDF): assumed to be standard Gaussian PDF # d_log_2pi = d*np.log(2*np.pi) - def _Kfun(self, r, flag): + # here we implement the negative kinetic fun + def _neg_kinetic_func(self, r, flag): if flag == 'eval': # evaluate - return 0.5*(r.T @ r) #+ d_log_2pi + 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) + # Hamiltonian function (negative log) + def _neg_Hamiltonian(self, nU, r): + # here nU is the log-posterior (so it is the negative potential) + # and we work with negative kinetic fun, but the Hamiltonian is: + # H(q,r) = U(q) + K(r) + # U(q): potential energy: negative log-posterior + # K(r): kinetic energy: negative log-assumed-density + nK = self._neg_kinetic_func(r, 'eval') + + return nU + nK + + #========================================================================= + def _FindGoodEpsilon(self, theta, joint, grad, L, epsilon=1): + r = self._neg_kinetic_func(1, 'sample') # resample a momentum + nHam = self._neg_Hamiltonian(joint, r) # initial Hamiltonian + _, r_prime, joint_prime, grad_prime = self._Leapfrog(theta, r, grad, epsilon, L) # 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(): + while np.isnan(joint_prime) or np.isnan(grad_prime).any() or 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) + _, r_prime, joint_prime, grad_prime = self._Leapfrog(theta, r, grad, epsilon*k, L) 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 + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) + log_ratio = nHam_prime - nHam 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 + _, r_prime, joint_prime, _ = self._Leapfrog(theta, r, grad, epsilon, L) + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) + log_ratio = nHam_prime - nHam + return epsilon #========================================================================= - def _Leapfrog(self, theta_old, r_old, grad_old, epsilon): + def _Leapfrog(self, q0, r0, grad, epsilon, L): # symplectic integrator: trajectories preserve phase space volumen + q, r = np.copy(q0), np.copy(r0) + for n in range(L): + q, r, joint, grad = self._Leapfrog_single(q, r, grad, epsilon) + + return q, -r, joint, grad # negate momentum to make proposal symmetric + + #========================================================================= + def _Leapfrog_single(self, theta_old, r_old, grad_old, epsilon): + # single-step update 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 + joint_new, grad_new = self._hmc_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 + return theta_new, r_new, joint_new, grad_new diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py new file mode 100644 index 0000000000..e994720f42 --- /dev/null +++ b/cuqi/sampler/_nuts.py @@ -0,0 +1,350 @@ +import numpy as np +from cuqi.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: + 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.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) + log_epsilon_bar, H_bar = 0, 0 + gamma, t_0, kappa = 0.05, 10, 0.75 # kappa in (0.5, 1] + 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._neg_kinetic_func(1, 'sample') # resample momentum vector + nHam = self._neg_Hamiltonian(joint_k, r_k) # Hamiltonian + + # slice variable + log_u = np.log(np.random.uniform(0, np.exp(nHam))) + + # initialization + j, n, s = 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, nHam, 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, nHam, 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) and not np.isnan(joint_prime) and not np.isinf(joint_prime): + 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, eta2 = 1/(k+t_0), k**(-kappa) + # + H_bar = (1-eta1)*H_bar + eta1*(delta - (alpha/n_alpha)) + log_epsilon = mu - (np.sqrt(k)/gamma)*H_bar + log_epsilon_bar = eta2*log_epsilon + (1-eta2)*log_epsilon_bar + # + epsilon, epsilon_bar = np.exp(log_epsilon), np.exp(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 + + #========================================================================= + # kinetic energy function (negative log PDF): assumed to be standard Gaussian PDF + # d_log_2pi = d*np.log(2*np.pi) + # here we implement the negative kinetic fun + def _neg_kinetic_func(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) + + #========================================================================= + # Hamiltonian function (negative log) + def _neg_Hamiltonian(self, nU, r): + # here nU is the log-posterior (so it is the negative potential) + # and we work with negative kinetic fun, but the Hamiltonian is: + # H(q,r) = U(q) + K(r) + # U(q): potential energy: negative log-posterior + # K(r): kinetic energy: negative log-assumed-density + nK = self._neg_kinetic_func(r, 'eval') + return nU + nK + + #========================================================================= + def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1): + r = self._neg_kinetic_func(1, 'sample') # resample a momentum + nHam = self._neg_Hamiltonian(joint, r) # initial Hamiltonian + _, r_prime, joint_prime, grad_prime = self._Leapfrog_single(theta, r, grad, epsilon) + + # trick to make sure the step is not huge, leading to infinite values of the likelihood + k = 1 + while np.isnan(joint_prime) or np.isnan(grad_prime).any() or np.isinf(joint_prime) or np.isinf(grad_prime).any(): + k *= 0.5 + _, r_prime, joint_prime, grad_prime = self._Leapfrog_single(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 + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) + log_ratio = nHam_prime - nHam + 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_single(theta, r, grad, epsilon) + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) + log_ratio = nHam_prime - nHam + return epsilon + + #========================================================================= + def _Leapfrog_single(self, theta_old, r_old, grad_old, epsilon): + # symplectic integrator: trajectories preserve phase space volumen + # single-step update + 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, nHam, 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_single(theta, r, grad, v*epsilon) + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) # Hamiltonian eval + n_prime = int(log_u <= nHam_prime) # if particle is in the slice + s_prime = int(log_u < Delta_max + nHam_prime) # check U-turn + + # Compute the acceptance probability + # alpha_prime = min(1, np.exp(nHam_prime - nHam)) + # written in a stable way to avoid overflow when computing + # exp(nHam_prime - nHam) for large values of diff_Ham + diff_Ham = nHam_prime - nHam + alpha_prime = 1 if diff_Ham > 0 else np.exp(diff_Ham) + n_alpha_prime = 1 + + return theta_prime, r_prime, grad_prime, theta_prime, r_prime, grad_prime, \ + theta_prime, joint_prime, grad_prime, n_prime, s_prime, alpha_prime, n_alpha_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, nHam, 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, nHam, 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, nHam, 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 diff --git a/demos/demo38_HMC_and_NUTS_check.py b/demos/demo38_HMC_and_NUTS_check.py new file mode 100644 index 0000000000..b96798b6b5 --- /dev/null +++ b/demos/demo38_HMC_and_NUTS_check.py @@ -0,0 +1,96 @@ +# ================================================================= +# Created by: +# Felipe Uribe @ DTU +# ================================================================= +# Version 2022 +# ================================================================= +import sys +sys.path.append("..") +import numpy as np +import scipy as sp +import scipy.stats as sps +# +import cuqi +import matplotlib.pyplot as plt + +# ================================================================= +# sample a 100D Gaussian +d = 100 +mu_q = np.zeros(d) +sigma_q = np.linspace(0.01, 1, d) +Lambd_q = sp.sparse.spdiags(1/(sigma_q**2), 0, d, d) + +# =================================================================== +# unormalized target +def logpi_target(x): + return 0.5*( - (x.T @ Lambd_q @ x) ) +def grad_logpi_target(x): + return -(Lambd_q @ x) + +target = cuqi.distribution.UserDefinedDistribution(d, logpi_target, grad_logpi_target) + +# ================================================================= +# sample +np.random.seed(1) +x0 = mu_q #np.random.rand(d) +Ns = int(2e3) +Nb = int(5e2) + +# HMC +traject_length = 150 +MCMC = cuqi.sampler.HMC(target, x0, adapt_traject_length=traject_length) +solution_hmc = MCMC.sample(Ns, Nb) +x_chain_hmc = solution_hmc.samples.T +steps_hmc = solution_hmc.acc_rate +mu_hmc = np.mean(x_chain_hmc, axis=0) +sigma_hmc = np.std(x_chain_hmc, axis=0, ddof=1) + +# NUTS +MCMC = cuqi.sampler.NUTS(target, x0) +solution_nuts = MCMC.sample(Ns, Nb) +x_chain_nuts = solution_nuts.samples.T +steps_nuts = solution_nuts.acc_rate +mu_nuts = np.mean(x_chain_nuts, axis=0) +sigma_nuts = np.std(x_chain_nuts, axis=0, ddof=1) + +# ================================================================= +# plots +# ================================================================ +plt.figure() +plt.plot(steps_nuts, 'bo', markersize=1.5, label='NUTS') +plt.plot(steps_hmc, 'r.', markersize=1.5, label='HMC') +plt.title('step sizes') +plt.xlabel('iterations') +plt.ylabel('epsilon') +plt.legend() +plt.tight_layout() + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5)) +ax1.plot(x_chain_nuts[:, -1], 'b.', label='NUTS') +ax1.set_title('NUTS chain') +ax1.set_xlabel('iterations') +ax1.set_ylabel('theta_100') +ax1.legend() +ax2.plot(x_chain_hmc[:, -1], 'r.', label='HMC') +ax2.set_title('HMC chain') +ax2.set_xlabel('iterations') +ax2.set_ylabel('theta_100') +ax2.legend() +plt.tight_layout() + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) +ax1.plot(sigma_q, mu_nuts, 'bo', label='NUTS') +ax1.plot(sigma_q, mu_hmc, 'r.', label='HMC') +ax1.legend() +ax1.set_ylabel('sample mean at each coordinate') +ax1.set_xlim(0, 1) +ax1.set_ylim(-0.6, 0.6) +ax2.plot(sigma_q, sigma_nuts, 'bo', label='NUTS') +ax2.plot(sigma_q, sigma_hmc, 'r.', label='HMC') +ax2.set_ylabel('sample std at each coordinate') +ax2.set_xlim(0, 1) +ax2.set_ylim(0, 1.1) +fig.suptitle('NUTS') +fig.suptitle('stats comparison') +plt.tight_layout() +plt.show()