Skip to content

MCMC uniform prior in linear space and NUTS #507

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
markrosslonergan merged 7 commits into from
Aug 14, 2026
Merged

MCMC uniform prior in linear space and NUTS #507

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
a22ca26
Temp lin space prior target
jacoblarkin Jun 8, 2026
f6c032e
Prototype HMC implementation
jacoblarkin Jun 29, 2026
bf15cf0
NUTS, not working yet
jacoblarkin Jul 22, 2026
a3cafee
Some more tests done, actually running completely now
jacoblarkin Jul 29, 2026
d76ee87
Add HMC flag to MCMC subcommand
jacoblarkin Aug 11, 2026
83cbcba
Get things in a state ready to merge. Still work to be done long term.
jacoblarkin Aug 12, 2026
742a9a7
Attempt to fix CI build issue
jacoblarkin Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 116 additions & 42 deletions bin/PROfit.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,9 @@ std::wostream *OSTREAM = &wcout;
std::wofstream LOG_FILE_STREAM;
bool LOGGING_TO_FILE = false;

void mcmc_worker(std::vector<Metropolis<simple_target, adaptive_proposal>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps);
void mcmc_worker(std::vector<std::unique_ptr<Metropolis<unilin_prior_target, adaptive_proposal>>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps);
//void mcmc_worker(std::vector<std::unique_ptr<Metropolis<simple_target, adaptive_proposal>>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps);
void nuts_worker(std::vector<std::vector<Eigen::VectorXf>> &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;
Expand Down Expand Up @@ -374,6 +376,7 @@ int main(int argc, char* argv[])
std::vector<std::string> afc_merge_inputs;
std::vector<float> afc_cleanup_quantiles = {0.025f, 0.975f};
int afc_cleanup_halo = 1;
bool hmc = false;


//Global Arguments for all PROfit enables subcommands.
Expand Down Expand Up @@ -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.");
Expand Down Expand Up @@ -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<LOG_INFO>(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);
Expand Down Expand Up @@ -3019,20 +3024,33 @@ int main(int argc, char* argv[])
std::vector<std::vector<float>> samples = latin_hypercube_sampling(mcmc_chains, nparams, latin_distribution, myseed.global_rng);
recenter_latin_samples(samples, global_ub, global_lb);
std::vector<Eigen::VectorXf> 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<std::vector<Metropolis<simple_target, adaptive_proposal>>> mets;
std::vector<std::vector<std::unique_ptr<Metropolis<unilin_prior_target, adaptive_proposal>>>> mets;
std::vector<std::vector<std::vector<Eigen::VectorXf>>> chains;
//std::vector<std::vector<std::unique_ptr<Metropolis<simple_target, adaptive_proposal>>>> mets;
mets.reserve(mcmc_threads);
std::vector<std::thread> 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();
Expand All @@ -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<std::string> 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<std::string> 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<std::string> 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;
Expand Down Expand Up @@ -3462,14 +3519,31 @@ int main(int argc, char* argv[])
}


void mcmc_worker(std::vector<Metropolis<simple_target, adaptive_proposal>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) {
void nuts_worker(std::vector<std::vector<Eigen::VectorXf>> &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<uint32_t> dseed(0, std::numeric_limits<uint32_t>::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<std::unique_ptr<Metropolis<unilin_prior_target, adaptive_proposal>>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) {
//void mcmc_worker(std::vector<std::unique_ptr<Metropolis<simple_target, adaptive_proposal>>> &mets, Eigen::VectorXf initial, PROmetric *metric, uint32_t seed, size_t nchains, size_t burnin, size_t steps) {
std::uniform_int_distribution<uint32_t> dseed(0, std::numeric_limits<uint32_t>::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<Metropolis<unilin_prior_target, adaptive_proposal>>(target, proposal, initial, dseed(rng));
met->run(burnin, steps);
mets.emplace_back(std::move(met));
}
}

Expand Down
Loading
Loading