TEST(McmcNutsBaseNuts, build_tree_test) {
rng_t base_rng(0);
int model_size = 1;
double init_momentum = 1.5;
stan::mcmc::ps_point z_init(model_size);
z_init.q(0) = 0;
z_init.p(0) = init_momentum;
stan::mcmc::ps_point z_propose(model_size);
Eigen::VectorXd p_sharp_left = Eigen::VectorXd::Zero(model_size);
Eigen::VectorXd p_sharp_right = Eigen::VectorXd::Zero(model_size);
Eigen::VectorXd rho = z_init.p;
double log_sum_weight = -std::numeric_limits<double>::infinity();
double H0 = -0.1;
int n_leapfrog = 0;
double sum_metro_prob = 0;
stan::mcmc::mock_model model(model_size);
stan::mcmc::mock_nuts sampler(model, base_rng);
sampler.set_nominal_stepsize(1);
sampler.set_stepsize_jitter(0);
sampler.sample_stepsize();
sampler.z() = z_init;
std::stringstream debug, info, warn, error, fatal;
stan::callbacks::stream_logger logger(debug, info, warn, error, fatal);
bool valid_subtree = sampler.build_tree(3, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, 1, n_leapfrog, log_sum_weight,
sum_metro_prob, logger);
EXPECT_TRUE(valid_subtree);
EXPECT_EQ(init_momentum * (n_leapfrog + 1), rho(0));
EXPECT_EQ(1, p_sharp_left(0));
EXPECT_EQ(1, p_sharp_right(0));
EXPECT_EQ(8 * init_momentum, sampler.z().q(0));
EXPECT_EQ(init_momentum, sampler.z().p(0));
EXPECT_EQ(8, n_leapfrog);
EXPECT_FLOAT_EQ(H0 + std::log(n_leapfrog), log_sum_weight);
EXPECT_FLOAT_EQ(std::exp(H0) * n_leapfrog, sum_metro_prob);
EXPECT_EQ("", debug.str());
EXPECT_EQ("", info.str());
EXPECT_EQ("", warn.str());
EXPECT_EQ("", error.str());
EXPECT_EQ("", fatal.str());
}
TEST(McmcUnitENuts, tree_boundary_test) {
rng_t base_rng(4839294);
stan::mcmc::unit_e_point z_init(3);
z_init.q(0) = 1;
z_init.q(1) = -1;
z_init.q(2) = 1;
z_init.p(0) = -1;
z_init.p(1) = 1;
z_init.p(2) = -1;
std::stringstream output;
stan::interface_callbacks::writer::stream_writer writer(output);
std::stringstream error_stream;
stan::interface_callbacks::writer::stream_writer error_writer(error_stream);
std::fstream empty_stream("", std::fstream::in);
stan::io::dump data_var_context(empty_stream);
typedef gauss3D_model_namespace::gauss3D_model model_t;
model_t model(data_var_context);
// Compute expected tree boundaries
typedef stan::mcmc::unit_e_metric<model_t, rng_t> metric_t;
metric_t metric(model);
stan::mcmc::expl_leapfrog<metric_t> unit_e_integrator;
double epsilon = 0.1;
stan::mcmc::unit_e_point z_test = z_init;
metric.init(z_test, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_forward_1 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_forward_2 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_forward_3 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_forward_4 = metric.dtau_dp(z_test);
z_test = z_init;
metric.init(z_test, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_backward_1 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_backward_2 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_backward_3 = metric.dtau_dp(z_test);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
unit_e_integrator.evolve(z_test, metric, -epsilon, writer, error_writer);
Eigen::VectorXd p_sharp_backward_4 = metric.dtau_dp(z_test);
// Check expected tree boundaries to those dynamically geneated by NUTS
stan::mcmc::unit_e_nuts<model_t, rng_t> sampler(model, base_rng);
sampler.set_nominal_stepsize(epsilon);
sampler.set_stepsize_jitter(0);
sampler.sample_stepsize();
stan::mcmc::ps_point z_propose = z_init;
Eigen::VectorXd p_sharp_left = Eigen::VectorXd::Zero(z_init.p.size());
Eigen::VectorXd p_sharp_right = Eigen::VectorXd::Zero(z_init.p.size());
Eigen::VectorXd rho = z_init.p;
double log_sum_weight = -std::numeric_limits<double>::infinity();
double H0 = -0.1;
int n_leapfrog = 0;
double sum_metro_prob = 0;
// Depth 0 forward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(0, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, 1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_right(n));
// Depth 1 forward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(1, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, 1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_2(n), p_sharp_right(n));
// Depth 2 forward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(2, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, 1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_3(n), p_sharp_right(n));
// Depth 3 forward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(3, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, 1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_forward_4(n), p_sharp_right(n));
// Depth 0 backward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(0, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, -1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_right(n));
// Depth 1 backward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(1, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, -1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_2(n), p_sharp_right(n));
// Depth 2 backward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(2, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, -1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_3(n), p_sharp_right(n));
// Depth 3 backward
sampler.z() = z_init;
sampler.init_hamiltonian(writer, error_writer);
sampler.build_tree(3, z_propose,
p_sharp_left, p_sharp_right, rho,
H0, -1, n_leapfrog, log_sum_weight,
sum_metro_prob,
writer, error_writer);
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_1(n), p_sharp_left(n));
for (int n = 0; n < rho.size(); ++n)
EXPECT_FLOAT_EQ(p_sharp_backward_4(n), p_sharp_right(n));
}
