diff --git a/bin/PROfit.cxx b/bin/PROfit.cxx index 27f6c58d..2d96e4ba 100644 --- a/bin/PROfit.cxx +++ b/bin/PROfit.cxx @@ -169,7 +169,9 @@ std::wostream *OSTREAM = &wcout; std::wofstream LOG_FILE_STREAM; bool LOGGING_TO_FILE = false; -void mcmc_worker(std::vector> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps); +void mcmc_worker(std::vector>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps); +//void mcmc_worker(std::vector>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps); +void nuts_worker(std::vector> &chains, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps, const Eigen::VectorXf &lb, const Eigen::VectorXf &ub); struct GlobalFitResult { PROfitter fitter; @@ -374,6 +376,7 @@ int main(int argc, char* argv[]) std::vector afc_merge_inputs; std::vector afc_cleanup_quantiles = {0.025f, 0.975f}; int afc_cleanup_halo = 1; + bool hmc = false; //Global Arguments for all PROfit enables subcommands. @@ -617,6 +620,7 @@ int main(int argc, char* argv[]) CLI::App *promcmc_command = app.add_subcommand("mcmc", "Get bayesian posteriors using MCMC"); promcmc_command->add_option("--vars", mcmc_vars, "Variables to find posteriors of."); promcmc_command->add_option("--nchains", mcmc_chains, "Number of chains to run with MCMC.")->default_val(1); + promcmc_command->add_flag("--hmc", hmc, "Run Hamiltonian MC instead of Metropolis"); //PROtest, test things CLI::App *protest_command = app.add_subcommand("protest", "Testing ground for rapid quick tests."); @@ -2619,6 +2623,7 @@ int main(int argc, char* argv[]) } // Debug PDF for covariance_to_spline systematics — only emitted if at least one is present. + log(L"%1% || cov2spline debug info %2%") % __func__ % !variable_systs[config.i_prime].cov2spline_debug_info.empty(); if(!variable_systs[config.i_prime].cov2spline_debug_info.empty()) { const std::string cov2spline_pdf = final_output_tag + "_covariance_to_spline_checks.pdf"; plotCov2SplineChecks(config, variable_cvs[config.i_prime], variable_systs[config.i_prime], cov2spline_pdf, config.i_prime); @@ -3019,20 +3024,33 @@ int main(int argc, char* argv[]) std::vector> samples = latin_hypercube_sampling(mcmc_chains, nparams, latin_distribution, myseed.global_rng); recenter_latin_samples(samples, global_ub, global_lb); std::vector samples_eigen; - for(size_t i = 0; i < samples.size(); ++i) + for(size_t i = 0; i < samples.size(); ++i) { samples_eigen.push_back(Eigen::VectorXf::Map(samples[i].data(), samples[i].size())); + //samples_eigen.back()(0) = std::pow(10, samples_eigen.back()(0)); + //samples_eigen.back()(1) = std::pow(10, samples_eigen.back()(1)); + } size_t mcmc_threads = mcmc_chains >= nthread ? nthread : mcmc_chains; - std::vector>> mets; + std::vector>>> mets; + std::vector>> chains; + //std::vector>>> mets; mets.reserve(mcmc_threads); std::vector threads; size_t chains_per_thread = mcmc_chains / mcmc_threads; size_t addone = mcmc_threads - mcmc_chains%mcmc_threads; for(size_t i = 0; i < mcmc_threads; ++i) { - mets.emplace_back(); - threads.emplace_back( - [&, i](){ - mcmc_worker(mets[i], samples_eigen[i], metric->Clone(), myseed.getThreadSeeds()->at(i), chains_per_thread + (i >= addone), fitConfig.MCMCburn, fitConfig.MCMCiter); - }); + if(hmc) { + chains.emplace_back(); + threads.emplace_back( + [&, i](){ + nuts_worker(chains[i], samples_eigen[i], metric->Clone(), myseed.getThreadSeeds()->at(i), chains_per_thread + (i >= addone), fitConfig.MCMCburn, fitConfig.MCMCiter, global_lb, global_ub); + }); + } else { + mets.emplace_back(); + threads.emplace_back( + [&, i](){ + mcmc_worker(mets[i], samples_eigen[i], metric->Clone(), myseed.getThreadSeeds()->at(i), chains_per_thread + (i >= addone), fitConfig.MCMCburn, fitConfig.MCMCiter); + }); + } } for(auto&& t : threads) { t.join(); @@ -3049,39 +3067,78 @@ int main(int argc, char* argv[]) //oned.push_back(TH1D("one2", ";#Deltam^{2}_{41} [eV^{2}];Posterior PDF", 200, -2, 2)); TFile fout((final_output_tag+"_PROMCMC_chains.root").c_str(), "RECREATE"); size_t chain_counter = 0; - for(const auto &tmets : mets) { - for(const auto &met : tmets) { - chain_counter++; - std::string name = "chain"+std::to_string(chain_counter); - TTree tree(name.c_str(), name.c_str()); - Eigen::VectorXf v = Eigen::VectorXf::Zero(nparams); - std::vector param_names; - for(size_t i = 0; i < metric->GetModel().nparams; ++i) { - tree.Branch(metric->GetModel().param_names[i].c_str(), &v(i)); - param_names.push_back(metric->GetModel().pretty_param_names[i]); - } - for(size_t i = metric->GetModel().nparams; i < nparams; ++i) { - const std::string &sname = metric->GetSysts().spline_names[i-metric->GetModel().nparams]; - std::string::size_type l = sname.find(':'); - // TODO: This only handles names with a single colon in them. I don't think we ever have more than that, it's really just meant for the 'flat' and 'norm' systs. - if(l != std::string::npos) { - std::string bname = sname; - bname[l] = '_'; - tree.Branch(bname.c_str(), &v(i)); - } else { - tree.Branch(sname.c_str(), &v(i)); + if(hmc) { + for(const auto &tchain : chains) { + for(const auto &chain : tchain) { + chain_counter++; + std::string name = "chain"+std::to_string(chain_counter); + TTree tree(name.c_str(), name.c_str()); + Eigen::VectorXf v = Eigen::VectorXf::Zero(nparams); + std::vector param_names; + for(size_t i = 0; i < metric->GetModel().nparams; ++i) { + tree.Branch(metric->GetModel().param_names[i].c_str(), &v(i)); + param_names.push_back(metric->GetModel().pretty_param_names[i]); } - param_names.push_back(config.m_mcgen_variation_plotname_map.at(metric->GetSysts().spline_names[i-N_phys_params])); + for(size_t i = metric->GetModel().nparams; i < nparams; ++i) { + const std::string &sname = metric->GetSysts().spline_names[i-metric->GetModel().nparams]; + std::string::size_type l = sname.find(':'); + // TODO: This only handles names with a single colon in them. I don't think we ever have more than that, it's really just meant for the 'flat' and 'norm' systs. + if(l != std::string::npos) { + std::string bname = sname; + bname[l] = '_'; + tree.Branch(bname.c_str(), &v(i)); + } else { + tree.Branch(sname.c_str(), &v(i)); + } + param_names.push_back(config.m_mcgen_variation_plotname_map.at(metric->GetSysts().spline_names[i-N_phys_params])); + } + for(const auto &p : chain) { + // twod[0].Fill(p(1), p(0)); + // oned[0].Fill(p(1)); + // oned[1].Fill(p(0)); + v = p; + tree.Fill(); + } + tree.Write(); + //met->plot_autocorrelation((final_output_tag+"_PROMCMC_autocorrelation_chain"+std::to_string(chain_counter)+".pdf").c_str(), param_names); } - for(const auto &p : met.chain) { - // twod[0].Fill(p(1), p(0)); - // oned[0].Fill(p(1)); - // oned[1].Fill(p(0)); - v = p; - tree.Fill(); + } + + } else { + for(const auto &tmets : mets) { + for(const auto &met : tmets) { + chain_counter++; + std::string name = "chain"+std::to_string(chain_counter); + TTree tree(name.c_str(), name.c_str()); + Eigen::VectorXf v = Eigen::VectorXf::Zero(nparams); + std::vector param_names; + for(size_t i = 0; i < metric->GetModel().nparams; ++i) { + tree.Branch(metric->GetModel().param_names[i].c_str(), &v(i)); + param_names.push_back(metric->GetModel().pretty_param_names[i]); + } + for(size_t i = metric->GetModel().nparams; i < nparams; ++i) { + const std::string &sname = metric->GetSysts().spline_names[i-metric->GetModel().nparams]; + std::string::size_type l = sname.find(':'); + // TODO: This only handles names with a single colon in them. I don't think we ever have more than that, it's really just meant for the 'flat' and 'norm' systs. + if(l != std::string::npos) { + std::string bname = sname; + bname[l] = '_'; + tree.Branch(bname.c_str(), &v(i)); + } else { + tree.Branch(sname.c_str(), &v(i)); + } + param_names.push_back(config.m_mcgen_variation_plotname_map.at(metric->GetSysts().spline_names[i-N_phys_params])); + } + for(const auto &p : met->chain) { + // twod[0].Fill(p(1), p(0)); + // oned[0].Fill(p(1)); + // oned[1].Fill(p(0)); + v = p; + tree.Fill(); + } + tree.Write(); + met->plot_autocorrelation((final_output_tag+"_PROMCMC_autocorrelation_chain"+std::to_string(chain_counter)+".pdf").c_str(), param_names, {}); } - tree.Write(); - met.plot_autocorrelation((final_output_tag+"_PROMCMC_autocorrelation_chain"+std::to_string(chain_counter)+".pdf").c_str(), param_names, {}); } } //TCanvas c; @@ -3462,14 +3519,31 @@ int main(int argc, char* argv[]) } -void mcmc_worker(std::vector> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) { +void nuts_worker(std::vector> &chains, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps, const Eigen::VectorXf &lb, const Eigen::VectorXf &ub) { + std::uniform_int_distribution dseed(0, std::numeric_limits::max()); + std::mt19937 rng(seed); + metric->setBounds(lb, ub); + for(size_t i = 0; i < nchains; ++i) { + NUTS nuts; + nuts.M = burnin+steps; + nuts.Madapt = burnin; + nuts(initial, *metric, dseed(rng)); + chains.emplace_back(std::move(nuts.chain)); + } +} + +void mcmc_worker(std::vector>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) { +//void mcmc_worker(std::vector>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) { std::uniform_int_distribution dseed(0, std::numeric_limits::max()); std::mt19937 rng(seed); + mets.reserve(nchains); for(size_t i = 0; i < nchains; ++i) { - simple_target target{*metric}; + //simple_target target{*metric}; + unilin_prior_target target{*metric}; adaptive_proposal proposal(*metric, dseed(rng)); - mets.emplace_back(target, proposal, initial, dseed(rng)); - mets.back().run(burnin, steps); + auto met = std::make_unique>(target, proposal, initial, dseed(rng)); + met->run(burnin, steps); + mets.emplace_back(std::move(met)); } } diff --git a/inc/PROMCMC.h b/inc/PROMCMC.h index 693485f0..6b9c68c1 100644 --- a/inc/PROMCMC.h +++ b/inc/PROMCMC.h @@ -204,6 +204,23 @@ namespace PROfit { } }; + // This target distribution uses a prior which is uniform in linear space for models + // where the parameters are in log space. + struct unilin_prior_target { + PROmetric &metric; + + // Returns log-target (-0.5*chi^2). Metropolis::step does exp(target(p) - target(current)) + // so the exp argument stays in safe float32 range even when chi^2 is large. + float operator()(Eigen::VectorXf &value) { + Eigen::VectorXf empty = value; + Eigen::VectorXf corr = value; + //corr.segment(0,2) = corr.segment(0,2).array().log10(); + //return -0.5f*metric(value, empty, false)+value.segment(0,2).array().sum() - std::log(100-0.01); + return -0.5f*metric(corr, empty, false)+corr.segment(0,2).array().sum(); + //return -0.5f*metric(corr, empty, false)-corr(0)-corr(1); + } + }; + struct prior_only_target { PROmetric &metric; @@ -349,7 +366,7 @@ namespace PROfit { std::vector active; std::mt19937 rng; static constexpr bool has_tune = true; - std::vector proposed; + //std::vector proposed; Eigen::VectorXf last_proposed; Eigen::VectorXf last_accepted; Eigen::VectorXf mean; @@ -428,6 +445,7 @@ namespace PROfit { if(std::find(fixed.begin(), fixed.end(), i) != std::end(fixed)) continue; if(i < nparams) { if(value(i) > metric.GetModel().ub(i) || value(i) < std::max(metric.GetModel().lb(i),-5.0f)) + //if(value(i) > std::pow(10, metric.GetModel().ub(i)) || value(i) < std::max(std::pow(10, metric.GetModel().lb(i)),1e-5)) return false; } else { size_t si = i - nparams; @@ -489,6 +507,226 @@ namespace PROfit { } }; + +class HMC { + public: + float epsilon, L; + std::vector chain; + + void leapfrog(Eigen::VectorXf &theta, Eigen::VectorXf &r, PROmetric &metric) { + Eigen::VectorXf g = theta; + metric(theta, g, true); + r += epsilon/2 * g; + theta += epsilon * r; + metric(theta, g, true); + r += epsilon/2 * g; + } + + void operator()(Eigen::VectorXf theta0, PROmetric &metric, size_t M, uint32_t seed) { + std::mt19937 rng(seed); + std::normal_distribution normal(0, 1); + std::uniform_real_distribution uniform(0, 1); + Eigen::VectorXf empty; + chain.push_back(theta0); + Eigen::VectorXf r0 = theta0, r = r0; // Just set the size, we'll set the values with rng in the loop + for(size_t m = 0; m < M; ++m) { + Eigen::VectorXf theta = chain.back(); + for(long i = 0; i < r.size(); ++i) r(i) = normal(rng); + r0 = r; + for(size_t i = 0; i < L; ++i) leapfrog(theta, r, metric); + float alpha = std::min(1.0f, expf(metric(theta, empty, false) - 0.5 * r.dot(r) - metric(chain.back(), empty, false) + 0.5 * r0.dot(r0))); + if(alpha < uniform(rng)) { + chain.push_back(theta); + } else { + chain.push_back(theta0); + } + } + } +}; + +class NUTS { + std::mt19937 rng; + public: + float M, Madapt; + // These are the values recommended in the NUTS papaer, but they are tunable + float delta = 0.6, DeltaMax = 1000; + std::vector chain; + + void normalize_theta(Eigen::VectorXf &theta) { + while(theta(0) > 2 || theta(0) < -2) { + if(theta(0) > 2) { + theta(0) -= 2; + theta(0) *= -1; + theta(0) += 2; + } + if(theta(0) < -2) { + theta(0) += 2; + theta(0) *= -1; + theta(0) -= 2; + } + } + while(theta(1) > 0 || theta(1) < -3) { + if(theta(1) > 0) { + theta(1) *= -1; + } + if(theta(1) < -3) { + theta(1) += 3; + theta(1) *= -1; + theta(1) -= 3; + } + } + } + + float eval(PROmetric &metric, const Eigen::VectorXf &theta, Eigen::VectorXf &grad, bool run_grad) { + float v = -0.5 * metric(theta, grad, run_grad) + theta(0) + theta(1); + if(run_grad) { + grad(0) += 1; + grad(1) += 1; + } + return v; + } + + void leapfrog(Eigen::VectorXf &theta, Eigen::VectorXf &r, float epsilon, PROmetric &metric) { + Eigen::VectorXf g = theta; + eval(metric, theta, g, true); + r -= epsilon/4 * g; + theta += epsilon * r; + normalize_theta(theta); + eval(metric, theta, g, true); + r -= epsilon/4 * g; + } + + float find_reasonable_epsilon(const Eigen::VectorXf &theta, PROmetric &metric) { + float epsilon = 1; + Eigen::VectorXf empty; + Eigen::VectorXf r = theta; + std::normal_distribution normal(0, 1); + for(long i = 0; i < r.size(); ++i) r(i) = normal(rng); + Eigen::VectorXf thetap = theta; + Eigen::VectorXf rp = r; + leapfrog(thetap, rp, epsilon, metric); + float a = 2*(expf(eval(metric, thetap, empty, false) - eval(metric, theta, empty, false) + 0.5 * (r.dot(r) - rp.dot(rp)) ) > 0.5) - 1; + while(powf(expf(eval(metric, thetap, empty, false) - eval(metric, theta, empty, false) + 0.5 * (r.dot(r) - rp.dot(rp))), a) > powf(2, -a)) { + epsilon *= powf(2, a); + thetap = theta; + rp = r; + leapfrog(thetap, rp, epsilon, metric); + } + return epsilon; + } + + struct BTR { + Eigen::VectorXf theta_minus, theta_plus, theta_prime, + r_minus, r_plus; + float n, s, alpha, nalpha; + }; + + BTR BuildTree(const Eigen::VectorXf &theta, const Eigen::VectorXf &r, float u, int v, + int j, float epsilon, const Eigen::VectorXf &theta0, const Eigen::VectorXf &r0, PROmetric &metric) { + BTR ret; + if(j == 0) { + ret.theta_minus = theta; + ret.r_minus = r; + leapfrog(ret.theta_minus, ret.r_minus, v*epsilon, metric); + ret.theta_plus = ret.theta_minus; + ret.theta_prime = ret.theta_minus; + ret.r_plus = ret.r_minus; + Eigen::VectorXf empty; + ret.n = (u <= std::exp(eval(metric, ret.theta_prime, empty, false) - 0.5 * ret.r_minus.dot(ret.r_minus))); + ret.s = (u < std::exp(DeltaMax + eval(metric, ret.theta_prime, empty, false) - 0.5 * ret.r_minus.dot(ret.r_minus))); + ret.alpha = std::min(1.0f, std::exp(eval(metric, ret.theta_prime, empty, false) - 0.5f * ret.r_minus.dot(ret.r_minus) + - eval(metric, theta0, empty, false) + 0.5f * r0.dot(r0))); + ret.nalpha = 1; + return ret; + } + ret = BuildTree(theta, r, u, v, j-1, epsilon, theta0, r0, metric); + std::uniform_real_distribution uniform(0, 1); + if(ret.s == 1) { + Eigen::VectorXf t,r, tp = ret.theta_prime; + float np = ret.n, alphap = ret.alpha, nalphap = ret.nalpha; + if(v == -1) { + t = ret.theta_plus; + r = ret.r_plus; + ret = BuildTree(ret.theta_minus, ret.r_minus, u, v, j-1, epsilon, theta0, r0, metric); + ret.theta_plus = t; + ret.r_plus = r; + } else { + t = ret.theta_minus; + r = ret.r_minus; + ret = BuildTree(ret.theta_plus, ret.r_plus, u, v, j-1, epsilon, theta0, r0, metric); + ret.theta_minus = t; + ret.r_minus = r; + } + if(uniform(rng) > std::min((ret.n/(ret.n+np)),1.0f)) + ret.theta_prime = tp; + ret.alpha += alphap; + ret.nalpha += nalphap; + ret.s = ret.s * ((ret.theta_plus - ret.theta_minus).dot(ret.r_minus) >= 0) + * ((ret.theta_plus - ret.theta_minus).dot(ret.r_plus) >= 0); + ret.n += np; + } + return ret; + } + + void operator()(const Eigen::VectorXf &theta0, PROmetric &metric, uint32_t seed) { + rng.seed(seed); + std::normal_distribution normal(0, 1); + std::uniform_real_distribution uniform(0, 1); + Eigen::VectorXf theta = theta0, empty; + chain.push_back(theta0); + float epsilon = find_reasonable_epsilon(theta, metric), + mu = std::log(10 * epsilon), + epsilonbar = 1, + Hbar = 0, + gamma = 0.05, + t0 = 10, + kappa = 0.75; + Eigen::VectorXf r0 = theta; + for(size_t m = 0; m < M; ++m) { + for(long i = 0; i < r0.size(); ++i) r0(i) = normal(rng); + float u = std::uniform_real_distribution(0, std::exp(eval(metric, theta, empty, false) - 0.5 * r0.dot(r0)))(rng); + BTR t{theta, theta, theta, r0, r0, 1, 1, 0, 0}; + int j = 0; + while(t.s == 1) { + int v = 2 * (uniform(rng) < 0.5) - 1; + BTR tp; + if(v == -1) { + tp = BuildTree(t.theta_minus, t.r_minus, u, v, j, epsilon, chain.back(), r0, metric); + t.theta_minus = tp.theta_minus; + t.r_minus = tp.r_minus; + t.alpha = tp.alpha; + t.nalpha = tp.nalpha; + } else { + tp = BuildTree(t.theta_plus, t.r_plus, u, v, j, epsilon, chain.back(), r0, metric); + t.theta_plus = tp.theta_plus; + t.r_plus = tp.r_plus; + t.alpha = tp.alpha; + t.nalpha = tp.nalpha; + } + if(tp.s == 1) { + if(uniform(rng) < std::min(1.0f, tp.n/t.n)) + theta = tp.theta_prime; + } + t.n += tp.n; + t.s = tp.s * ((t.theta_plus - t.theta_minus).dot(t.r_minus) >= 0) * ((t.theta_plus - t.theta_minus).dot(t.r_plus) >= 0); + j++; + //log(L"%1% || Iteration %2%: %3%") % __func__ % j % theta; + if(j >= 10) break; + } + if(m < Madapt) { + Hbar = (1 - 1/(m + 1 + t0)) * Hbar + 1/(m + 1 + t0) * (delta - t.alpha/t.nalpha); + epsilon = std::exp(mu - sqrt(m+1)/gamma * Hbar); + epsilonbar = std::exp(std::pow(m+1, -kappa) * std::log(epsilon) + (1-std::pow(m+1, -kappa))*std::log(epsilonbar)); + } else { + epsilon = epsilonbar; + } + log(L"%1% || Adding %2% to chain of length %3%.") % __func__ % theta % chain.size(); + chain.push_back(theta); + } + } +}; + + }; #endif