CUQIpy 1.1.1.post0.dev36__py3-none-any.whl → 1.4.1.post0.dev124__py3-none-any.whl

cuqi/sampler/_hmc.py

@@ -1,42 +1,46 @@
 import numpy as np
+import numpy as np
 from cuqi.sampler import Sampler
+from cuqi.array import CUQIarray
+from numbers import Number
 
-
-# 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.
+    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: 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`.
+        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*
+    initial_point : ndarray
+        Initial parameters. *Optional*. If not provided, the initial point is 
+        an array of ones.
 
     max_depth : int
-        Maximum depth of the tree.
+        Maximum depth of the tree >=0 and the default is 15.
 
-    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.
+    step_size : None or float
+        If step_size is provided (as positive float), it will be used as initial
+        step size. If None, the step size will be estimated by the 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 values are 0.6 (default) or 0.8 (as in stan). In principle,
+        opt_acc_rate should be in (0, 1), however, choosing a value that is very
+        close to 1 or 0 might lead to poor performance of the sampler.
 
-    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.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     Example
     -------
@@ -53,7 +57,11 @@ class NUTS(Sampler):
         sampler = cuqi.sampler.NUTS(target)
 
         # Sample
-        samples = sampler.sample(10000, 5000)
+        sampler.warmup(5000)
+        sampler.sample(10000)
+
+        # Get samples
+        samples = sampler.get_samples()
 
         # Plot samples
         samples.plot_pair()
@@ -70,170 +78,234 @@ class NUTS(Sampler):
         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.
+        # only used after adaptation).
         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)
+
+    _STATE_KEYS = Sampler._STATE_KEYS.union({'_epsilon', '_epsilon_bar',
+                                                '_H_bar',
+                                                'current_target_logd',
+                                                'current_target_grad',
+                                                'max_depth'})
+
+    _HISTORY_KEYS = Sampler._HISTORY_KEYS.union({'num_tree_node_list',
+                                                    'epsilon_list',
+                                                    'epsilon_bar_list'})
+
+    def __init__(self, target=None, initial_point=None, max_depth=None,
+                 step_size=None, opt_acc_rate=0.6, **kwargs):
+        super().__init__(target, initial_point=initial_point, **kwargs)
+
+        # Assign parameters as attributes
         self.max_depth = max_depth
-        self.adapt_step_size = adapt_step_size
+        self.step_size = 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
+
+    def _initialize(self):
+
+        self._current_alpha_ratio = np.nan # Current alpha ratio will be set to some
+                                           # value (other than np.nan) before 
+                                           # being used
+
+        self.current_target_logd, self.current_target_grad = self._nuts_target(self.current_point)
+
+        # Parameters dual averaging
+        # Initialize epsilon and epsilon_bar
+        # epsilon is the step size used in the current iteration
+        # after warm up and one sampling step, epsilon is updated
+        # to epsilon_bar for the remaining sampling steps.
+        if self.step_size is None:
+            self._epsilon = self._FindGoodEpsilon()
+            self.step_size = self._epsilon
+        else:
+            self._epsilon = self.step_size
+
+        self._epsilon_bar = "unset"
+
+        # Parameter mu, does not change during the run
+        self._mu = np.log(10*self._epsilon)
+
+        self._H_bar = 0
+
+        # NUTS run diagnostics
         # 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()
+    #=========================================================================
+    #============================== Properties ===============================
+    #=========================================================================
+    @property
+    def max_depth(self):
+        return self._max_depth
+
+    @max_depth.setter
+    def max_depth(self, value):
+        if value is None:
+            value = 15 # default value
+        if not isinstance(value, int):
+            raise TypeError('max_depth must be an integer.')
+        if value < 0:
+            raise ValueError('max_depth must be >= 0.')
+        self._max_depth = value
+
+    @property
+    def step_size(self):
+        return self._step_size
+    
+    @step_size.setter
+    def step_size(self, value):
+        if value is None:
+            pass # NUTS will adapt the step size
+
+        # step_size must be a positive float, raise error otherwise
+        elif isinstance(value, bool)\
+            or not isinstance(value, Number)\
+            or value <= 0:
+            raise TypeError('step_size must be a positive float or None.')
+        self._step_size = value
+
+    @property
+    def opt_acc_rate(self):
+        return self._opt_acc_rate
+    
+    @opt_acc_rate.setter
+    def opt_acc_rate(self, value):
+        if not isinstance(value, Number) or value <= 0 or value >= 1:
+            raise ValueError('opt_acc_rate must be a float in (0, 1).')
+        self._opt_acc_rate = value
 
-    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 = []
+    #=========================================================================
+    #================== Implement methods required by Sampler =============
+    #=========================================================================
+    def validate_target(self):
+        # Check if the target has logd and gradient methods
+        try:
+            current_target_logd, current_target_grad =\
+            self._nuts_target(np.ones(self.dim))
+        except:
+            raise ValueError('Target must have logd and gradient methods.')
+
+    def reinitialize(self):
+        # Call the parent reset method
+        super().reinitialize()
+        # Reset NUTS run diagnostic attributes
+        self._reset_run_diagnostic_attributes()
 
-    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 step(self):
+        if isinstance(self._epsilon_bar, str) and self._epsilon_bar == "unset":
+            self._epsilon_bar = self._epsilon
 
-    def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient
-        return self.target.logd(x), self.target.gradient(x)
+        # Convert current_point, logd, and grad to numpy arrays
+        # if they are CUQIarray objects
+        if isinstance(self.current_point, CUQIarray):
+            self.current_point = self.current_point.to_numpy() 
+        if isinstance(self.current_target_logd, CUQIarray):
+            self.current_target_logd = self.current_target_logd.to_numpy()
+        if isinstance(self.current_target_grad, CUQIarray):
+            self.current_target_grad = self.current_target_grad.to_numpy()
 
-    def _sample_adapt(self, N, Nb):
-        return self._sample(N, Nb)
+        # reset number of tree nodes for each iteration
+        self._num_tree_node = 0
 
-    def _sample(self, N, Nb):
-        # Reset run diagnostic attributes
-        self._reset_run_diagnostic_attributes()
+        # copy current point, logd, and grad in local variables
+        point_k = self.current_point # initial position (parameters)
+        logd_k = self.current_target_logd
+        grad_k = self.current_target_grad # initial gradient
         
-        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
+        # compute r_k and Hamiltonian
+        r_k = self._Kfun(1, 'sample') # resample momentum vector
+        Ham = logd_k - self._Kfun(r_k, 'eval') # Hamiltonian
+
+        # slice variable
+        log_u = Ham - np.random.exponential(1, size=1)
+
+        # initialization
+        j, s, n = 0, 1, 1
+        point_minus, point_plus = point_k.copy(), point_k.copy()
+        grad_minus, grad_plus = grad_k.copy(), grad_k.copy()
+        r_minus, r_plus = r_k.copy(), r_k.copy()
 
         # 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)
+        acc = 0
+        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):
+                point_minus, r_minus, grad_minus, _, _, _, \
+                    point_prime, logd_prime, grad_prime,\
+                        n_prime, s_prime, alpha, n_alpha = \
+                            self._BuildTree(point_minus, r_minus, grad_minus,
+                                            Ham, log_u, v, j, self._epsilon)
+            else:
+                _, _, _, point_plus, r_plus, grad_plus, \
+                    point_prime, logd_prime, grad_prime,\
+                        n_prime, s_prime, alpha, n_alpha = \
+                            self._BuildTree(point_plus, r_plus, grad_plus,
+                                            Ham, log_u, v, j, self._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(logd_prime)) and \
+                (not np.isinf(logd_prime)):
+                self.current_point = point_prime.copy()
+                # copy if array, else assign if scalar
+                self.current_target_logd = (
+                        logd_prime.copy()
+                        if isinstance(logd_prime, np.ndarray)
+                        else logd_prime
+                    )
+                self.current_target_grad = grad_prime.copy()
+                acc = 1
+
+
+            # update number of particles, tree level, and stopping criterion
+            n += n_prime
+            dpoints = point_plus - point_minus
+            s = s_prime *\
+                int((dpoints @ r_minus.T) >= 0) * int((dpoints @ r_plus.T) >= 0)
+            j += 1
+            self._current_alpha_ratio = alpha/n_alpha
+
+        # update run diagnostic attributes
+        self._update_run_diagnostic_attributes(
+            self._num_tree_node, self._epsilon, self._epsilon_bar)
+        
+        self._epsilon = self._epsilon_bar 
+        if np.isnan(self.current_target_logd):
+            raise NameError('NaN potential func')
 
-                # 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
+        return acc
+
+    def tune(self, skip_len, update_count):
+        """ adapt epsilon during burn-in using dual averaging"""
+        if isinstance(self._epsilon_bar, str) and self._epsilon_bar == "unset":
+            self._epsilon_bar = 1
+
+        k = update_count+1
+
+        # Fixed parameters that do not change during the run
+        gamma, t_0, kappa = 0.05, 10, 0.75 # kappa in (0.5, 1]
+
+        eta1 = 1/(k + t_0)
+        self._H_bar = (1-eta1)*self._H_bar +\
+            eta1*(self.opt_acc_rate - (self._current_alpha_ratio))
+        self._epsilon = np.exp(self._mu - (np.sqrt(k)/gamma)*self._H_bar)
+        eta = k**(-kappa)
+        self._epsilon_bar =\
+            np.exp(eta*np.log(self._epsilon) +(1-eta)*np.log(self._epsilon_bar))
+
+    #=========================================================================
+    def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient
+        return self.target.logd(x), self.target.gradient(x)
 
     #=========================================================================
     # auxiliary standard Gaussian PDF: kinetic energy function
@@ -245,48 +317,61 @@ class NUTS(Sampler):
             return np.random.standard_normal(size=self.dim)
 
     #=========================================================================
-    def _FindGoodEpsilon(self, theta, joint, grad, epsilon=1):
+    def _FindGoodEpsilon(self, epsilon=1):
+        point_k = self.current_point
+        self.current_target_logd, self.current_target_grad = self._nuts_target(
+            point_k)
+        logd = self.current_target_logd
+        grad = self.current_target_grad
+
         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)
+        Ham = logd - self._Kfun(r, 'eval')     # initial Hamiltonian
+        _, r_prime, logd_prime, grad_prime = self._Leapfrog(
+            point_k, r, grad, epsilon)
 
-        # trick to make sure the step is not huge, leading to infinite values of the likelihood
+        # 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.isinf(logd_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, logd_prime, grad_prime = self._Leapfrog(
+                point_k, 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')
+        # doubles/halves the value of epsilon until the accprob of the Langevin
+        # proposal crosses 0.5
+        Ham_prime = logd_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')
+            _, r_prime, logd_prime, _ = self._Leapfrog(
+                point_k, r, grad, epsilon)
+            Ham_prime = logd_prime - self._Kfun(r_prime, 'eval')
             log_ratio = Ham_prime - Ham
         return epsilon
 
     #=========================================================================
-    def _Leapfrog(self, theta_old, r_old, grad_old, epsilon):
+    def _Leapfrog(self, point_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
+        point_new = point_old + epsilon*r_new     # full-step
+        logd_new, grad_new = self._nuts_target(point_new)     # new gradient
         r_new += 0.5*epsilon*grad_new     # half-step
-        return theta_new, r_new, joint_new, grad_new
+        return point_new, r_new, logd_new, grad_new
 
     #=========================================================================
-    # @functools.lru_cache(maxsize=128)
-    def _BuildTree(self, theta, r, grad, Ham, log_u, v, j, epsilon, Delta_max=1000):
+    def _BuildTree(
+            self, point_k, 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
+            point_prime, r_prime, logd_prime, grad_prime = self._Leapfrog(
+                point_k, r, grad, v*epsilon)
+            Ham_prime = logd_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
             #
@@ -299,37 +384,79 @@ class NUTS(Sampler):
             alpha_prime = 1 if diff_Ham > 0 else np.exp(diff_Ham)
             n_alpha_prime = 1
             #
-            theta_minus, theta_plus = theta_prime, theta_prime
+            point_minus, point_plus = point_prime, point_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
+            point_minus, r_minus, grad_minus, point_plus, r_plus, grad_plus, \
+                point_prime, logd_prime, grad_prime,\
+                    n_prime, s_prime, alpha_prime, n_alpha_prime = \
+                        self._BuildTree(point_k, 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)
+                    point_minus, r_minus, grad_minus, _, _, _, \
+                        point_2prime, logd_2prime, grad_2prime,\
+                            n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                                self._BuildTree(point_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)
+                    _, _, _, point_plus, r_plus, grad_plus, \
+                        point_2prime, logd_2prime, grad_2prime,\
+                            n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                                self._BuildTree(point_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)
+                    point_prime = point_2prime.copy()
+                    # copy if array, else assign if scalar
+                    logd_prime = (
+                        logd_2prime.copy()
+                        if isinstance(logd_2prime, np.ndarray)
+                        else logd_2prime
+                    )
+                    grad_prime = grad_2prime.copy()
 
                 # 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)
+                dpoints = point_plus - point_minus
+                s_prime = s_2prime *\
+                    int((dpoints@r_minus.T)>=0) * int((dpoints@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 point_minus, r_minus, grad_minus, point_plus, r_plus, grad_plus,\
+            point_prime, logd_prime, grad_prime,\
+                n_prime, s_prime, alpha_prime, n_alpha_prime
+
+    #=========================================================================
+    #======================== Diagnostic methods =============================
+    #=========================================================================
+
+    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."""
+        # 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, n_tree, eps, eps_bar):
+        """A method to update attributes to store NUTS run diagnostic."""
+        # 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)