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)); }