// Returns number of valid points in the completed subtree
int build_tree(int depth, Eigen::VectorXd& rho,
ps_point* z_init, ps_point& z_propose, nuts_util& util)
{
// Base case
if (depth == 0)
{
this->_integrator.evolve(this->_z, this->_hamiltonian,
util.sign * this->_epsilon);
rho += this->_z.p;
if (z_init) *z_init = this->_z;
z_propose = static_cast<ps_point>(this->_z);
double h = this->_hamiltonian.H(this->_z);
if (h != h) h = std::numeric_limits<double>::infinity();
util.criterion = util.log_u + (h - util.H0) < this->_max_delta;
util.sum_prob += stan::math::min(1, std::exp(util.H0 - h));
util.n_tree += 1;
return (util.log_u + (h - util.H0) < 0);
}
// General recursion
else
{
Eigen::VectorXd subtree_rho = Eigen::VectorXd::Zero(rho.size());
ps_point z_init(this->_z);
int n1 = build_tree(depth - 1, subtree_rho, &z_init, z_propose, util);
rho += subtree_rho;
if (!util.criterion) return 0;
subtree_rho.setZero();
ps_point z_propose_right(z_init);
int n2 = build_tree(depth - 1, subtree_rho, 0, z_propose_right, util);
rho += subtree_rho;
double accept_prob = static_cast<double>(n2) /
static_cast<double>(n1 + n2);
if ( util.criterion && (this->_rand_uniform() < accept_prob) )
z_propose = z_propose_right;
util.criterion &= _compute_criterion(z_init, this->_z, rho);
return n1 + n2;
}
}
// Returns number of valid points in the completed subtree
int build_tree(int depth, Eigen::VectorXd& rho,
ps_point* z_init_parent, ps_point& z_propose,
nuts_util& util) {
// Base case
if (depth == 0) {
this->integrator_.evolve(this->z_, this->hamiltonian_,
util.sign * this->epsilon_);
rho += this->z_.p;
if (z_init_parent) *z_init_parent = this->z_;
z_propose = this->z_;
double h = this->hamiltonian_.H(this->z_);
if (boost::math::isnan(h))
h = std::numeric_limits<double>::infinity();
util.criterion = util.log_u + (h - util.H0) < this->max_delta_;
if (!util.criterion) ++(this->n_divergent_);
util.sum_prob += stan::math::min(1, std::exp(util.H0 - h));
util.n_tree += 1;
return (util.log_u + (h - util.H0) < 0);
} else {
// General recursion
Eigen::VectorXd left_subtree_rho(rho.size());
left_subtree_rho.setZero();
ps_point z_init(this->z_);
int n1 = build_tree(depth - 1, left_subtree_rho, &z_init,
z_propose, util);
if (z_init_parent) *z_init_parent = z_init;
if (!util.criterion) return 0;
Eigen::VectorXd right_subtree_rho(rho.size());
right_subtree_rho.setZero();
ps_point z_propose_right(z_init);
int n2 = build_tree(depth - 1, right_subtree_rho, 0,
z_propose_right, util);
double accept_prob = static_cast<double>(n2) /
static_cast<double>(n1 + n2);
if ( util.criterion && (this->rand_uniform_() < accept_prob) )
z_propose = z_propose_right;
Eigen::VectorXd& subtree_rho = left_subtree_rho;
subtree_rho += right_subtree_rho;
rho += subtree_rho;
util.criterion &= compute_criterion(z_init, this->z_, subtree_rho);
return n1 + n2;
}
}