From 5732f94931602696a6b4bb891e8a7604ff61ba25 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 1 Jul 2024 18:31:54 +0300 Subject: [PATCH 1/7] renamed hmc to nuts, and created the actual hmc sampler; example and checks todo --- cuqi/sampler/_hmc.py | 240 ++++++++++++------------------ cuqi/sampler/_nuts.py | 336 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 428 insertions(+), 148 deletions(-) create mode 100644 cuqi/sampler/_nuts.py diff --git a/cuqi/sampler/_hmc.py b/cuqi/sampler/_hmc.py index 6f070e1852..09a04d32df 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,44 @@ 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, traject_length=15, adapt_step_size=True, opt_acc_rate=0.65, **kwargs): super().__init__(target, x0=x0, **kwargs) - self.max_depth = max_depth + self.traject_length = 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 = [] - 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): + """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) - 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,112 +132,110 @@ 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) # Step size variables epsilon, epsilon_bar = None, None # parameters dual averaging + delta = self.opt_acc_rate # https://mc-stan.org/docs/2_18/reference-manual/hmc-algorithm-parameters.html 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 + # set the trajectory length as it has to be accessed many times + if (self.adapt_traject_length == True): + # 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 + + # run HMC + for k in range(Ns-1): + q_k = theta[:, k] # initial position (parameters) + p_k = self._Kfun(1, 'sample') # initial momentum vector + + # LEAPFROG: alternate full steps for position and momentum + q_star, p_star = self._Leapfrog(q_k, p_k, grad, epsilon, L) + + # evaluate neg-potential and neg-kinetic energies at start and end of trajectory + U_k = joint_eval[k] + U_star, grad_star = self._hmc_target(q_star) + K_star, K_k = self._Kfun(p_star, 'eval'), self._Kfun(p_k, 'eval') + + # accept/reject + log_alpha = min( 0, (U_star+K_star)-(U_k+K_k) ) + log_u = np.log(np.random.rand()) + if (log_u <= log_alpha): + theta[:, k+1] = q_star + joint_eval[k+1] = U_star + acc[k+1] = 1 + grad = grad_star + else: + theta[:, k+1] = q_k + joint_eval[k+1] = U_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) + # 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)) + H_bar = (1-eta1)*H_bar + eta1*(delta - np.exp(log_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 + # 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._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 + acc_rate = np.mean(acc) + + return theta, joint_eval, step_sizes, traject_lengths, acc_rate #========================================================================= - # auxiliary standard Gaussian PDF: kinetic energy function + # auxiliary standard Gaussian PDF: neg-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 + return -0.5*(r.T @ r) #- d_log_2pi if flag == 'sample': # sample return np.random.standard_normal(size=self.dim) @@ -269,67 +264,16 @@ def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1): return epsilon #========================================================================= - def _Leapfrog(self, theta_old, r_old, grad_old, epsilon): + def _Leapfrog(self, q0, p0, grad_old, epsilon, L): # 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 - + # faster: do not store trajectory + q, p = np.copy(q0), np.copy(p0) + p += 0.5*epsilon*grad_old # initial half step for momentum + for n in range(L): + q += epsilon*p # full step for the position + grad_q = self._hmc_target(q)[1] + if (n != L-1): + p += epsilon*grad_q # full step for the momentum, skip last one + p += (epsilon/2)*grad_q # final half step for momentum + + return q, p diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py new file mode 100644 index 0000000000..3a007ba79e --- /dev/null +++ b/cuqi/sampler/_nuts.py @@ -0,0 +1,336 @@ +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) + 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 + From 4b361d2a00275ee4cd4448b17e63904886f1bbff Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 1 Jul 2024 18:46:31 +0300 Subject: [PATCH 2/7] fixing the sampler init --- cuqi/sampler/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 From 4416828c9e4633ed40191fb6671a4ee5dc0eef20 Mon Sep 17 00:00:00 2001 From: Felipe Date: Tue, 2 Jul 2024 11:53:25 +0300 Subject: [PATCH 3/7] re-check _nuts implementation and keep working on the hmc --- cuqi/sampler/_hmc.py | 46 ++++++++++++++------- cuqi/sampler/_nuts.py | 94 ++++++++++++++++++++++++------------------- 2 files changed, 84 insertions(+), 56 deletions(-) diff --git a/cuqi/sampler/_hmc.py b/cuqi/sampler/_hmc.py index 09a04d32df..c4ec05e8e1 100644 --- a/cuqi/sampler/_hmc.py +++ b/cuqi/sampler/_hmc.py @@ -173,18 +173,20 @@ def _sample(self, N, Nb): # run HMC for k in range(Ns-1): q_k = theta[:, k] # initial position (parameters) - p_k = self._Kfun(1, 'sample') # initial momentum vector + r_k = self._Kfun(1, 'sample') # initial momentum vector # LEAPFROG: alternate full steps for position and momentum - q_star, p_star = self._Leapfrog(q_k, p_k, grad, epsilon, L) + q_star, r_star = self._Leapfrog(q_k, r_k, grad, epsilon, L) # evaluate neg-potential and neg-kinetic energies at start and end of trajectory U_k = joint_eval[k] + nHam_k = self._neg_Hamiltonian(U_k, r_k) # Hamiltonian + # U_star, grad_star = self._hmc_target(q_star) - K_star, K_k = self._Kfun(p_star, 'eval'), self._Kfun(p_k, 'eval') + nHam_star = self._neg_Hamiltonian(U_star, r_star) # Hamiltonian # accept/reject - log_alpha = min( 0, (U_star+K_star)-(U_k+K_k) ) + log_alpha = min( 0, nHam_star-nHam_k ) log_u = np.log(np.random.rand()) if (log_u <= log_alpha): theta[:, k+1] = q_star @@ -231,36 +233,50 @@ def _sample(self, N, Nb): return theta, joint_eval, step_sizes, traject_lengths, acc_rate #========================================================================= - # auxiliary standard Gaussian PDF: neg-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 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._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) + 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.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_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 - 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_single(theta, r, grad, epsilon) + nHam_prime = self._neg_Hamiltonian(joint_prime, r_prime) + log_ratio = nHam_prime - nHam + return epsilon #========================================================================= diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py index 3a007ba79e..40364fb3dd 100644 --- a/cuqi/sampler/_nuts.py +++ b/cuqi/sampler/_nuts.py @@ -167,14 +167,14 @@ def _sample(self, N, Nb): 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 + r_k = self._neg_kinetic_func(1, 'sample') # resample momentum vector + nHam = self._neg_Hamiltonian(joint_k, r_k) # Hamiltonian # slice variable - log_u = Ham - np.random.exponential(1, size=1) # u = np.log(np.random.uniform(0, np.exp(H))) + log_u = np.log(np.random.uniform(0, np.exp(nHam))) # initialization - j, s, n = 0, 1, 1 + 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) @@ -189,11 +189,11 @@ def _sample(self, N, Nb): 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) + self._BuildTree(theta_minus, r_minus, grad_minus, 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) + self._BuildTree(theta_plus, r_plus, grad_plus, log_u, v, j, epsilon) # Metropolis step alpha2 = min(1, (n_prime/n)) #min(0, np.log(n_p) - np.log(n)) @@ -211,7 +211,7 @@ def _sample(self, N, Nb): # 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) @@ -222,14 +222,14 @@ def _sample(self, N, Nb): 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:] @@ -237,41 +237,54 @@ def _sample(self, N, Nb): 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) + #========================================================================= + # 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._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) + 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.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_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 - 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_single(theta, r, grad, epsilon) + 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_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 @@ -280,43 +293,42 @@ def _Leapfrog(self, theta_old, r_old, grad_old, epsilon): #========================================================================= # @functools.lru_cache(maxsize=128) - def _BuildTree(self, theta, r, grad, Ham, log_u, v, j, epsilon, Delta_max=1000): + 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(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 + 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(diff_Ham)) + # alpha_prime = min(1, np.exp(nHam_prime - nHam)) # written in a stable way to avoid overflow when computing - # exp(diff_Ham) for large values of diff_Ham + # 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 - # - theta_minus, theta_plus = theta_prime, theta_prime - r_minus, r_plus = r_prime, r_prime - grad_minus, grad_plus = grad_prime, grad_prime + + 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, Ham, log_u, v, j-1, epsilon) + self._BuildTree(theta, r, grad, 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) + 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, Ham, log_u, v, j-1, epsilon) + 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)) @@ -328,9 +340,9 @@ def _BuildTree(self, theta, r, grad, Ham, log_u, v, j, epsilon, Delta_max=1000): # 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, \ + 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 - From 7ef58af38c4a6a87165cbdf8da68bf25006f4485 Mon Sep 17 00:00:00 2001 From: Felipe Date: Tue, 2 Jul 2024 12:01:38 +0300 Subject: [PATCH 4/7] fixing typo in nuts --- cuqi/sampler/_nuts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py index 40364fb3dd..820ee11302 100644 --- a/cuqi/sampler/_nuts.py +++ b/cuqi/sampler/_nuts.py @@ -189,11 +189,11 @@ def _sample(self, N, Nb): 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, log_u, v, j, epsilon) + 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, log_u, v, j, epsilon) + 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)) From 4503e9ebbe847c21821602a144825ffb0a220edc Mon Sep 17 00:00:00 2001 From: Felipe Date: Tue, 2 Jul 2024 12:12:18 +0300 Subject: [PATCH 5/7] fixing nuts typos --- cuqi/sampler/_nuts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py index 820ee11302..0086602606 100644 --- a/cuqi/sampler/_nuts.py +++ b/cuqi/sampler/_nuts.py @@ -319,7 +319,7 @@ def _BuildTree(self, theta, r, grad, nHam, log_u, v, j, epsilon, Delta_max=1000) # 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, log_u, v, j-1, epsilon) + 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, _, _, _, \ From 7a07573ed8f69244a689a053fd4ce4dab28008b3 Mon Sep 17 00:00:00 2001 From: Felipe Date: Tue, 2 Jul 2024 14:32:48 +0300 Subject: [PATCH 6/7] HMC and NUTS checks done --- cuqi/sampler/_hmc.py | 139 ++++++++++++++++------------- cuqi/sampler/_nuts.py | 17 ++-- demos/demo38_HMC_and_NUTS_check.py | 104 +++++++++++++++++++++ 3 files changed, 193 insertions(+), 67 deletions(-) create mode 100644 demos/demo38_HMC_and_NUTS_check.py diff --git a/cuqi/sampler/_hmc.py b/cuqi/sampler/_hmc.py index c4ec05e8e1..ba31371c7c 100644 --- a/cuqi/sampler/_hmc.py +++ b/cuqi/sampler/_hmc.py @@ -81,9 +81,9 @@ class HMC(Sampler): sampler.iteration_list """ - def __init__(self, target, x0=None, traject_length=15, adapt_step_size=True, opt_acc_rate=0.65, **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.traject_length = traject_length + 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 @@ -108,8 +108,9 @@ def _reset_run_diagnostic_attributes(self): # 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, eps, eps_bar): + 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) @@ -117,6 +118,8 @@ def _update_run_diagnostic_attributes(self, k, eps, eps_bar): 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 _hmc_target(self, x): # returns logposterior tuple evaluation-gradient return self.target.logd(x), self.target.gradient(x) @@ -143,24 +146,9 @@ def _sample(self, N, Nb): theta[:, 0] = self.x0 joint_eval[0], grad = self._hmc_target(self.x0) - # Step size variables - epsilon, epsilon_bar = None, None - - # parameters dual averaging - delta = self.opt_acc_rate # https://mc-stan.org/docs/2_18/reference-manual/hmc-algorithm-parameters.html - 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 - 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 - # 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) @@ -170,44 +158,58 @@ def _sample(self, N, Nb): L = self.adapt_traject_length # if scalar then user specifies the step size traject_lengths[0] = L + # parameters dual averaging + epsilon, epsilon_bar = None, None + if (self.adapt_step_size == True): + 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] + delta = self.opt_acc_rate # https://mc-stan.org/docs/2_18/reference-manual/hmc-algorithm-parameters.html + elif (self.adapt_step_size == False): + 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 HMC - for k in range(Ns-1): - q_k = theta[:, k] # initial position (parameters) - r_k = self._Kfun(1, 'sample') # initial momentum vector + for k in range(1, Ns): + 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 = self._Leapfrog(q_k, r_k, grad, epsilon, L) + q_star, r_star, joint_star, grad_star = self._leapfrog_all(q_k, r_k, grad, epsilon, L) # evaluate neg-potential and neg-kinetic energies at start and end of trajectory - U_k = joint_eval[k] - nHam_k = self._neg_Hamiltonian(U_k, r_k) # Hamiltonian - # - U_star, grad_star = self._hmc_target(q_star) - nHam_star = self._neg_Hamiltonian(U_star, r_star) # Hamiltonian + 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 = min( 0, nHam_star-nHam_k ) + 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): - theta[:, k+1] = q_star - joint_eval[k+1] = U_star - acc[k+1] = 1 + 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+1] = q_k - joint_eval[k+1] = U_k + theta[:, k] = q_k + joint_eval[k] = joint_k # update run diagnostic attributes - self._update_run_diagnostic_attributes(k, 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) + eta1, eta2 = 1/(k+t_0), k**(-kappa) + # H_bar = (1-eta1)*H_bar + eta1*(delta - np.exp(log_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): + 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 @@ -218,19 +220,17 @@ def _sample(self, N, Nb): 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, traject_lengths, acc_rate + return theta, joint_eval, step_sizes #========================================================================= # kinetic energy function (negative log PDF): assumed to be standard Gaussian PDF @@ -255,16 +255,16 @@ def _neg_Hamiltonian(self, nU, r): return nU + nK #========================================================================= - def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1): + 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_single(theta, r, grad, epsilon) + _, 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_single(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 @@ -273,23 +273,42 @@ def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1): 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) + _, 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, q0, p0, grad_old, epsilon, L): + def _leapfrog_all(self, q0, r0, grad, epsilon, L): # symplectic integrator: trajectories preserve phase space volumen - # faster: do not store trajectory - q, p = np.copy(q0), np.copy(p0) - p += 0.5*epsilon*grad_old # initial half step for momentum + q, r = np.copy(q0), np.copy(r0) + r += (epsilon/2)*grad # initial half step for momentum for n in range(L): - q += epsilon*p # full step for the position - grad_q = self._hmc_target(q)[1] + q += epsilon*r # full step for the position if (n != L-1): - p += epsilon*grad_q # full step for the momentum, skip last one - p += (epsilon/2)*grad_q # final half step for momentum + joint, grad = self._hmc_target(q) # new gradient + r += epsilon*grad # full step for the momentum, skip last one + joint, grad = self._hmc_target(q) # new gradient + r += (epsilon/2)*grad # final half step for momentum + + return q, -r, joint, grad # negate momentum to make proposal symmetric - return q, p + #========================================================================= + 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 + + #========================================================================= + 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._hmc_target(theta_new) # new gradient + r_new += 0.5*epsilon*grad_new # half-step + + return theta_new, r_new, joint_new, grad_new diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py index 0086602606..3d08015e03 100644 --- a/cuqi/sampler/_nuts.py +++ b/cuqi/sampler/_nuts.py @@ -152,8 +152,8 @@ def _sample(self, N, Nb): 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] - 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): @@ -197,7 +197,8 @@ def _sample(self, N, Nb): # 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): + if (s_prime == 1) and (np.random.rand() <= alpha2) and not np.isnan(joint_prime) \ + and not np.isinf(joint_prime) and not np.isneginf(joint_prime): theta[:, k] = theta_prime joint_eval[k] = joint_prime grad = np.copy(grad_prime) @@ -214,11 +215,13 @@ def _sample(self, N, Nb): # adapt epsilon during burn-in using dual averaging if (k <= Nb) and (self.adapt_step_size == True): - eta1 = 1/(k + t_0) + eta1, eta2 = 1/(k+t_0), k**(-kappa) + # 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)) + 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 @@ -265,7 +268,7 @@ def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1): # 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_single(theta, r, grad, epsilon*k) epsilon = 0.5*k*epsilon diff --git a/demos/demo38_HMC_and_NUTS_check.py b/demos/demo38_HMC_and_NUTS_check.py new file mode 100644 index 0000000000..50f8551804 --- /dev/null +++ b/demos/demo38_HMC_and_NUTS_check.py @@ -0,0 +1,104 @@ +# ================================================================= +# Created by: +# Felipe Uribe @ DTU +# ================================================================= +# Version 2022 +# ================================================================= +import sys +sys.path.append("/home/felipe/github/cuqi/CUQIpy/") +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) + +# 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) + +# 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) + +# ================================================================= +# 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, 'b.') +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, 'b.') +ax2.set_ylabel('sample std at each coordinate') +ax2.set_xlim(0, 1) +ax2.set_ylim(0, 1.1) +fig.suptitle('NUTS') +plt.tight_layout() +# +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) +ax1.plot(sigma_q, mu_hmc, 'r.') +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_hmc, 'r.') +ax2.set_ylabel('sample std at each coordinate') +ax2.set_xlim(0, 1) +ax2.set_ylim(0, 1.1) +fig.suptitle('HMC') +plt.tight_layout() +plt.show() From 6987cf51a043069144b76a0bcff9a50937a6d137 Mon Sep 17 00:00:00 2001 From: Felipe Date: Tue, 2 Jul 2024 15:01:13 +0300 Subject: [PATCH 7/7] some minor final touches --- cuqi/sampler/_hmc.py | 19 ++------------- cuqi/sampler/_nuts.py | 3 +-- demos/demo38_HMC_and_NUTS_check.py | 38 ++++++++++++------------------ 3 files changed, 18 insertions(+), 42 deletions(-) diff --git a/cuqi/sampler/_hmc.py b/cuqi/sampler/_hmc.py index ba31371c7c..0b1228e94d 100644 --- a/cuqi/sampler/_hmc.py +++ b/cuqi/sampler/_hmc.py @@ -178,7 +178,7 @@ def _sample(self, N, Nb): 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_all(q_k, r_k, grad, epsilon, L) + 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] @@ -279,21 +279,6 @@ def _FindGoodEpsilon(self, theta, joint, grad, L, epsilon=1): return epsilon - #========================================================================= - def _leapfrog_all(self, q0, r0, grad, epsilon, L): - # symplectic integrator: trajectories preserve phase space volumen - q, r = np.copy(q0), np.copy(r0) - r += (epsilon/2)*grad # initial half step for momentum - for n in range(L): - q += epsilon*r # full step for the position - if (n != L-1): - joint, grad = self._hmc_target(q) # new gradient - r += epsilon*grad # full step for the momentum, skip last one - joint, grad = self._hmc_target(q) # new gradient - r += (epsilon/2)*grad # final half step for momentum - - return q, -r, joint, grad # negate momentum to make proposal symmetric - #========================================================================= def _Leapfrog(self, q0, r0, grad, epsilon, L): # symplectic integrator: trajectories preserve phase space volumen @@ -301,7 +286,7 @@ def _Leapfrog(self, q0, r0, grad, epsilon, L): for n in range(L): q, r, joint, grad = self._Leapfrog_single(q, r, grad, epsilon) - return q, r, joint, grad + return q, -r, joint, grad # negate momentum to make proposal symmetric #========================================================================= def _Leapfrog_single(self, theta_old, r_old, grad_old, epsilon): diff --git a/cuqi/sampler/_nuts.py b/cuqi/sampler/_nuts.py index 3d08015e03..e994720f42 100644 --- a/cuqi/sampler/_nuts.py +++ b/cuqi/sampler/_nuts.py @@ -197,8 +197,7 @@ def _sample(self, N, Nb): # 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) and not np.isneginf(joint_prime): + 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) diff --git a/demos/demo38_HMC_and_NUTS_check.py b/demos/demo38_HMC_and_NUTS_check.py index 50f8551804..b96798b6b5 100644 --- a/demos/demo38_HMC_and_NUTS_check.py +++ b/demos/demo38_HMC_and_NUTS_check.py @@ -5,7 +5,7 @@ # Version 2022 # ================================================================= import sys -sys.path.append("/home/felipe/github/cuqi/CUQIpy/") +sys.path.append("..") import numpy as np import scipy as sp import scipy.stats as sps @@ -36,14 +36,6 @@ def grad_logpi_target(x): Ns = int(2e3) Nb = int(5e2) -# 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) - # HMC traject_length = 150 MCMC = cuqi.sampler.HMC(target, x0, adapt_traject_length=traject_length) @@ -53,6 +45,14 @@ def grad_logpi_target(x): 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 # ================================================================ @@ -79,26 +79,18 @@ def grad_logpi_target(x): plt.tight_layout() fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) -ax1.plot(sigma_q, mu_nuts, 'b.') +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, 'b.') +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') -plt.tight_layout() -# -fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5)) -ax1.plot(sigma_q, mu_hmc, 'r.') -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_hmc, 'r.') -ax2.set_ylabel('sample std at each coordinate') -ax2.set_xlim(0, 1) -ax2.set_ylim(0, 1.1) -fig.suptitle('HMC') +fig.suptitle('stats comparison') plt.tight_layout() plt.show()