UmbrellaSolver.cpp File Reference

#include <filesystem>
#include "Fitters/Processing/MCMCProcessor.h"
#include "Manager/Manager.h"
#include "Fitters/Algorithms/MulticanonicalMCMCHandler.h"
#include "TSystem.h"
#include "TChain.h"
#include "TSystemDirectory.h"

Include dependency graph for UmbrellaSolver.cpp:

Go to the source code of this file.

Classes
struct	WindowConfig
struct	UmbrellaConfig

Functions
UmbrellaConfig	parseYAMLConfig (const std::string &filename)
	YAML-based config parser using yaml-cpp library. More...
double	gaussianWindow (double x, double center, double width)
double	vonMisesWindow (double x, double center, double kappa)
double	generalisedGaussian2 (double x, double mean, double width)
double	GetMulticanonicalWeightGenGaussian (double deltacp, double mean, double width, double multicanonicalBeta)
double	summedWindowsWeighted (double x, const std::vector< WindowConfig > &windows, const std::vector< double > &z_values)
size_t	calculateWindowCacheSizeBytes (const std::vector< WindowConfig > &windows, const std::vector< std::vector< double >> &samples)
std::vector< std::vector< std::vector< double > > >	buildWindowCache (const std::vector< WindowConfig > &windows, const std::vector< std::vector< double >> &samples)
std::vector< std::vector< double > >	calcFmatrix (const std::vector< double > &z_current, const std::vector< WindowConfig > &windows, const std::vector< std::vector< double >> &samples, const std::vector< std::vector< std::vector< double >>> &window_cache)
std::vector< double >	zSolver (const std::vector< double > &z_current, const std::vector< WindowConfig > &windows, const std::vector< std::vector< double >> &samples, const std::vector< std::vector< std::vector< double >>> &window_cache, [[maybe_unused]] bool verbose=false, [[maybe_unused]] int *total_lines=nullptr)
bool	checkConvergence (const std::vector< double > &z_current, const std::vector< double > &z_prev, double tolerance)
	A few different convergence checks. More...
std::vector< double >	getZDiffs (const std::vector< double > &z_current, const std::vector< double > &z_prev)
bool	checkConvergenceStalled (const std::vector< double > &z_current, const std::vector< double > &z_prev, double tolerance)
void	LoadFiles (std::vector< std::vector< double >> &samples, UmbrellaConfig &config, std::vector< TFile * > &input_files, std::vector< TTree * > &input_trees)
void	DiagnoseUmbrella (const UmbrellaConfig &config, TFile *output_file, const std::vector< std::vector< double >> &z_evolution, TTree *combined_tree)
void	UmbrellaSolver (const std::string &config_file)
int	main (int argc, char *argv[])

Variables
_MaCh3_Safe_Include_Start_ _MaCh3_Safe_Include_End_ bool	debug_mode = false

Detailed Description

Author

David Riley

Definition in file UmbrellaSolver.cpp.

Function Documentation

buildWindowCache()

std::vector<std::vector<std::vector<double> > > buildWindowCache	(	const std::vector< WindowConfig > &	windows,
		const std::vector< std::vector< double >> &	samples
	)

Definition at line 175 of file UmbrellaSolver.cpp.

                                                                                                            {
  TStopwatch clock;
  clock.Start();
  int n_windows = static_cast<int>(windows.size());
  const size_t cache_size_bytes = calculateWindowCacheSizeBytes(windows, samples);
  const double cache_size_gb = static_cast<double>(cache_size_bytes) / (1024.0 * 1024.0 * 1024.0);
  if (cache_size_gb > 10.0) {
    MACH3LOG_WARN("Window cache size: {:.2f} GB", cache_size_gb);
  } else {
    MACH3LOG_INFO("Window cache size: {:.2f} GB", cache_size_gb);
  }
 
  std::vector<std::vector<std::vector<double>>> cache(n_windows);
  for (int i = 0; i < n_windows; i++) {
    cache[i].resize(n_windows);
    for (int j = 0; j < n_windows; j++) {
      cache[i][j].resize(samples[i].size());
    }
  }
  MACH3LOG_INFO("Window cache built successfully.");
 
  #ifdef MULTITHREAD
  #pragma omp parallel for collapse(2) schedule(dynamic)
  #endif
  for (int i = 0; i < n_windows; i++) {
    for (int j = 0; j < n_windows; j++) {
      for (size_t s = 0; s < samples[i].size(); s++) {
        if (windows[j].umbrellaBiasFunction == M3::BiasFunction::kVonMises) {
          cache[i][j][s] = vonMisesWindow(samples[i][s], windows[j].center, windows[j].vonMises_kappa);
        } else if (windows[j].umbrellaBiasFunction == M3::BiasFunction::kGaussian) {
          cache[i][j][s] = gaussianWindow(samples[i][s], windows[j].center, windows[j].width);
        } else if (windows[j].umbrellaBiasFunction == M3::BiasFunction::kGeneralisedGaussian) {
          cache[i][j][s] = GetMulticanonicalWeightGenGaussian(samples[i][s], windows[j].center, windows[j].width, windows[j].multicanonical_beta);
        } else {
          MACH3LOG_ERROR("Unrecognised BiasFunction!!");
          throw MaCh3Exception(__FILE__, __LINE__);
        }
      }
    }
  }
 
  if (debug_mode) { // this should be TH1D of cache against sample value in the
                    // variable of interest for each window, to check the cache
                    // is being built correctly and the windows look correct
                    // across the range of samples
    // TODO break this into a standalone function
    // first print the size in memory of the cache to check it is reasonable and
    // not being built incorrectly with an extra dimension or something
    TFile *cache_file = TFile::Open("window_cache_debug_histograms.root", "RECREATE");
    for (int i = 0; i < n_windows; i++) {
      for (int j = 0; j < n_windows; j++) {
        std::string hist_name = "Sample" + std::to_string(i) + "_window" + std::to_string(j);
        TH1D *window_cache_hist = new TH1D(hist_name.c_str(), hist_name.c_str(), 100, -3.1415, 3.1415);
        for (size_t s = 0; s < cache[i][j].size(); s++) {
          window_cache_hist->AddBinContent(window_cache_hist->FindBin(samples[i][s]));
        }
        window_cache_hist->Write();
      }
    }
 
    // additionally save a TH2D of the values in the cache against the sample
    // values for a specific window to check the shape of the windows is correct
    // across the range of samples
    for (int i = 0; i < n_windows; i++) {
      for (int j = 0; j < n_windows; j++) {
        std::string hist_name = "Sample" + std::to_string(i) + "_window" + std::to_string(j) + "_cache_2D";
        TH2D *window_cache_2D_hist = new TH2D(hist_name.c_str(), hist_name.c_str(), 100, -3.1415, 3.1415, 100, -1000, 10);
        for (size_t s = 0; s < cache[i][j].size(); s++) {
          // take the log of the likelihood values
          window_cache_2D_hist->Fill(samples[i][s], log(cache[i][j][s]));
        }
        // show overflow numbers on the hist
        window_cache_2D_hist->SetStats(1);
        window_cache_2D_hist->Write();
      }
    }
 
    cache_file->Close();
  }
 
  clock.Stop();
  MACH3LOG_INFO("Finished building window cache, it took {:.2f}s to finish", clock.RealTime());
  return cache;
}

calcFmatrix()

std::vector<std::vector<double> > calcFmatrix	(	const std::vector< double > &	z_current,
		const std::vector< WindowConfig > &	windows,
		const std::vector< std::vector< double >> &	samples,
		const std::vector< std::vector< std::vector< double >>> &	window_cache
	)

Definition at line 262 of file UmbrellaSolver.cpp.

                                                                         {
  int n_windows = static_cast<int>(windows.size());
  std::vector<std::vector<double>> F(n_windows, std::vector<double>(n_windows, 0.0));
 
  std::vector<double> z_inv = z_current; // make a copy to avoid modifying the original z_current
  for (size_t i = 0; i < z_current.size(); i++) {
    if (z_current[i] > 0) {
      z_inv[i] = 1.0 / z_current[i];
    } else {
      z_inv[i] = 0.0; // Handle zero or negative z values gracefully
      if (debug_mode) {
        MACH3LOG_WARN("Warning: z_current[{}] is non-positive ({}). Setting its inverse to 0 in F matrix calculation.", i, z_current[i]);
      }
    }
  }
  #ifdef MULTITHREAD
  #pragma omp parallel for schedule(dynamic)
  #endif
  for (int i = 0; i < n_windows; i++) {
    std::vector<double> denominator_cache(samples[i].size(), 0.0);
    for (size_t s = 0; s < samples[i].size(); s++) {
      double denominator = 0.0;
      for (int k = 0; k < n_windows; k++) {
        denominator += window_cache[i][k][s] * z_inv[k];
      }
      denominator_cache[s] = 1 / denominator;
    }
 
    for (int j = 0; j < n_windows; j++) {
      double sum = 0.0;
      int count = 0;
 
      for (size_t s = 0; s < samples[i].size(); s++) {
        double sample = samples[i][s];
        double window_j = window_cache[i][j][s];
        double denominator = denominator_cache[s];
 
        if (denominator > 0) {
          double integrand = (window_j * z_inv[i]) * denominator;
          sum += integrand;
          count++;
        } else if (debug_mode) {
          MACH3LOG_WARN("Denominator is zero for sample {} in window {}, skipping...", sample, i);
        }
      }
 
      if (debug_mode) {
        MACH3LOG_INFO("F[{}][{}] sum: {}, count: {}", i, j, sum, count);
      }
 
      if (count > 0) {
        F[i][j] = sum / count;
      }
    }
  }
 
  return F;
}

calculateWindowCacheSizeBytes()

size_t calculateWindowCacheSizeBytes	(	const std::vector< WindowConfig > &	windows,
		const std::vector< std::vector< double >> &	samples
	)

Definition at line 163 of file UmbrellaSolver.cpp.

                                                                                                                          {
  const size_t n_windows = windows.size();
  size_t total_samples = 0;
  for (const auto &sample_set : samples) {
    total_samples += sample_set.size();
  }
  return n_windows * total_samples * sizeof(double);
}

checkConvergence()

bool checkConvergence	(	const std::vector< double > &	z_current,
		const std::vector< double > &	z_prev,
		double	tolerance
	)

A few different convergence checks.

Definition at line 375 of file UmbrellaSolver.cpp.

                                                                                                             {
  double sum_diffs = 0.0;
  for (size_t i = 0; i < z_current.size(); i++) {
    // if (std::abs(z_current[i] - z_prev[i]) > tolerance *
    // std::max(std::abs(z_current[i]), std::abs(z_prev[i]))) { // intention
    // here: is the difference between any of the z elements more than the
    // tolerance*the magnitude of the largest element in z? in this way the
    // tolerance sets the number of decimal points below the largest element we
    // are targeting.
    //     return false;
    // }
    sum_diffs += std::abs(z_current[i] - z_prev[i]);
    // check the average difference across all elements instead
    // of requiring every element to meet the condition, this
    // should be more robust to individual elements fluctuating
    // around their target value while the overall z
    // distribution is still converging
    if (sum_diffs / static_cast<double>(z_current.size()) > tolerance) {
      return false;
    }
  }
  return true;
}

checkConvergenceStalled()

bool checkConvergenceStalled	(	const std::vector< double > &	z_current,
		const std::vector< double > &	z_prev,
		double	tolerance
	)

Definition at line 409 of file UmbrellaSolver.cpp.

                                                                                                                    {
  (void)z_prev;
 
  static std::deque<std::vector<double>> z_history;
  static std::vector<double> previous_moving_average;
  static int stagnant_iterations = 0;
 
  constexpr int moving_average_window = 500;
  constexpr int stagnant_required = 500;
  const double bound = tolerance;
 
  if (z_current.empty()) {
    return false;
  }
 
  // Reset state safely if number of windows changes between solver runs.
  if (!z_history.empty() && z_history.front().size() != z_current.size()) {
    z_history.clear();
    previous_moving_average.clear();
    stagnant_iterations = 0;
  }
 
  z_history.push_back(z_current);
  if (static_cast<int>(z_history.size()) > moving_average_window) {
    z_history.pop_front();
  }
 
  // Wait until the moving-average window is fully populated.
  if (static_cast<int>(z_history.size()) < moving_average_window) {
    return false;
  }
 
  std::vector<double> moving_average(z_current.size(), 0.0);
  for (const auto &z_vec : z_history) {
    for (size_t i = 0; i < z_vec.size(); i++) {
      moving_average[i] += z_vec[i];
    }
  }
  for (size_t i = 0; i < moving_average.size(); i++) {
    moving_average[i] /= moving_average_window;
  }
 
  if (previous_moving_average.empty()) {
    previous_moving_average = moving_average;
    return false;
  }
 
  bool all_within_bound = true;
  for (size_t i = 0; i < moving_average.size(); i++) {
    if (std::abs(moving_average[i] - previous_moving_average[i]) > bound) {
      all_within_bound = false;
      break;
    }
  }
 
  if (all_within_bound) {
    stagnant_iterations++;
  } else {
    stagnant_iterations = 0;
  }
 
  previous_moving_average = moving_average;
 
  if (stagnant_iterations == stagnant_required) {
    MACH3LOG_WARN("Convergence appears stalled: moving-average change stayed within {} for {} iterations.",
                  bound, stagnant_required);
  }
 
  return stagnant_iterations >= stagnant_required;
}

DiagnoseUmbrella()

void DiagnoseUmbrella	(	const UmbrellaConfig &	config,
		TFile *	output_file,
		const std::vector< std::vector< double >> &	z_evolution,
		TTree *	combined_tree
	)

Definition at line 663 of file UmbrellaSolver.cpp.

                                            {
  TDirectory* UmbrellaDir = output_file->mkdir("Umbrella");
  UmbrellaDir->cd();
 
  TCanvas c1("c1", "Z Evolution", 800, 600);
  std::vector<std::unique_ptr<TGraph>> z_graphs(config.windows.size());
  TLegend legend(0.7, 0.7, 0.9, 0.9);
 
  double ymax = 0.0;
  double ymin = std::numeric_limits<double>::max();
 
  for (size_t i = 0; i < config.windows.size(); i++) {
    std::vector<double> iterations, z_vals;
    for (size_t j = 0; j < z_evolution.size(); j++) {
      iterations.push_back(static_cast<int>(j));
      z_vals.push_back(z_evolution[j][i]);
    }
 
    for (double val : z_vals) {
      if (val > ymax)
        ymax = val;
      if (val < ymin)
        ymin = val;
    }
 
    z_graphs[i] = std::make_unique<TGraph>((iterations.size()), &iterations[0], &z_vals[0]);
    z_graphs[i]->SetLineColor(static_cast<Color_t>(i + 1));
    z_graphs[i]->SetLineWidth(2);
    z_graphs[i]->SetName(Form("z_evolution_window_%lu", i));
    z_graphs[i]->SetTitle("Evolution of Z Values");
    z_graphs[i]->GetXaxis()->SetTitle("Iteration");
    z_graphs[i]->GetYaxis()->SetTitle("Z Value");
 
    if (i == 0) {
      z_graphs[i]->Draw("AL");
    } else {
      z_graphs[i]->Draw("L SAME");
    }
 
    legend.AddEntry(z_graphs[i].get(), Form("Window %lu", i), "l");
    z_graphs[i]->Write();
  }
 
  legend.Draw();
  z_graphs[0]->SetMaximum(ymax);
  z_graphs[0]->SetMinimum(ymin);
  c1.Update();
  c1.SetLogy();
  c1.Write();
 
  // Create summary histogram of umbrella_variable distribution
  TH1D *h_delta_cp = new TH1D("h_delta_cp_weighted", "Weighted Delta CP Distribution", 100, -TMath::Pi(), TMath::Pi());
  TH1D *h_delta_cp_unweighted = new TH1D("h_delta_cp_unweighted", "Unweighted Delta CP Distribution", 100, -TMath::Pi(), TMath::Pi());
 
  const std::string& variable = config.variable_of_interest;
  combined_tree->Draw((variable + ">>h_delta_cp_weighted").c_str(), "umbrella_weight", "goff");
  combined_tree->Draw((variable + ">>h_delta_cp_unweighted").c_str(), "", "goff");
 
  h_delta_cp->Write();
  h_delta_cp_unweighted->Write();
 
  UmbrellaDir->Close();
  delete UmbrellaDir;
}

gaussianWindow()

double gaussianWindow	(	double	x,
		double	center,
		double	width
	)

Definition at line 102 of file UmbrellaSolver.cpp.

                                                             {
  return exp(-0.5 * pow((x - center) / width, 2)) / (width * sqrt(2 * TMath::Pi()));
}

generalisedGaussian2()

double generalisedGaussian2	(	double	x,
		double	mean,
		double	width
	)

Definition at line 122 of file UmbrellaSolver.cpp.

                                                                 {
  constexpr int n = 2; // this controls the tightness of the gaussian fixed at 2
                       // for now due to normalisation
  const double normFactor = 1 / ((M3::UmbrellaGaussianNormFactor) * 2 * std::sqrt(2) * width); // the normalisation is a little ugly (uses gamma functions),
                   // im just going to hardcode them for now
  double likelihood = normFactor * std::exp(-std::pow((std::pow(x - mean, 2) / (2 * std::pow(width, 2))), n));
  return likelihood;
}

GetMulticanonicalWeightGenGaussian()

double GetMulticanonicalWeightGenGaussian	(	double	deltacp,
		double	mean,
		double	width,
		double	multicanonicalBeta
	)

Todo:

modify this to instead use atan2 implementation for wrapping

Definition at line 132 of file UmbrellaSolver.cpp.

                                                                                                                {
  // implementation of the generalised gaussian as a bias function
  // for now with a fixed n = 2 for simplicity
 
  double g0 = generalisedGaussian2(deltacp, mean, width);
  double g1 = generalisedGaussian2(deltacp, mean - 2 * TMath::Pi(),width); // these two repeats are required for wrapping the gaussian around -+pi
  double g2 = generalisedGaussian2(deltacp, mean + 2 * TMath::Pi(), width);
  return (g0 + g1 + g2) * (multicanonicalBeta);
}

getZDiffs()

std::vector<double> getZDiffs	(	const std::vector< double > &	z_current,
		const std::vector< double > &	z_prev
	)

Definition at line 399 of file UmbrellaSolver.cpp.

                                                                                                 {
  std::vector<double> diffs(z_current.size(), 0.0);
  for (size_t i = 0; i < z_current.size(); i++) {
    diffs[i] = std::abs(z_current[i] - z_prev[i]);
  }
  return diffs;
}

LoadFiles()

void LoadFiles	(	std::vector< std::vector< double >> &	samples,
		UmbrellaConfig &	config,
		std::vector< TFile * > &	input_files,
		std::vector< TTree * > &	input_trees
	)

Definition at line 480 of file UmbrellaSolver.cpp.

                                                                                                                                                {
  TStopwatch clock;
  clock.Start();
  if (!config.dynamic_files) {
    MACH3LOG_INFO("Using static input files from configuration.");
    for (size_t i = 0; i < config.windows.size(); i++) {
      MACH3LOG_INFO("Loading file: {}", config.windows[i].input_file);
      TFile *file = M3::Open(config.windows[i].input_file.c_str(), "READ", __FILE__, __LINE__);
 
      TTree *tree = static_cast<TTree*>(file->Get("posteriors"));
      if (!tree) {
        MACH3LOG_ERROR("Cannot find 'posteriors' tree in {}", config.windows[i].input_file);
        file->Close();
        continue;
      }
 
      input_files.push_back(file);
      input_trees.push_back(tree);
    }
  } else {
    MACH3LOG_INFO("Dynamic file loading enabled. Searching for files in directory: {}", config.dynamic_pattern);
    // Use ROOT's TSystem to find all files in the directory
    TSystemDirectory dir("", config.dynamic_pattern.c_str());
    TList *files = dir.GetListOfFiles();
    int file_count = 0;
    if (files) {
      TIter next(files);
      TSystemFile *file;
      while ((file = static_cast<TSystemFile*>(next()))) {
        std::string filename = file->GetName();
        if (filename.find(".root") == std::string::npos) {
          MACH3LOG_DEBUG("Skipping non-root file: {}", filename);
          continue;
        }
 
        std::string full_path = config.dynamic_pattern + "/" + filename;
        MACH3LOG_DEBUG("Found file: {}", full_path);
 
        TFile *root_file = M3::Open(full_path, "READ", __FILE__, __LINE__);
 
        TTree *tree = static_cast<TTree*>(root_file->Get("posteriors"));
        if (!tree) {
          MACH3LOG_ERROR("Cannot find 'posteriors' tree in {}", full_path);
          root_file->Close();
          throw MaCh3Exception(__FILE__, __LINE__, "Missing 'posteriors' tree in file: " + full_path);
        }
        file_count++;
 
        input_files.push_back(root_file);
        input_trees.push_back(tree);
        MACH3LOG_DEBUG("Loaded tree 'posteriors' from file: {}", full_path);
      }
    } else {
      MACH3LOG_ERROR("No files found matching pattern: {}", config.dynamic_pattern);
      throw MaCh3Exception(__FILE__, __LINE__, "No files found matching pattern: " + config.dynamic_pattern);
    }
 
    if (file_count != config.dynamic_n_windows) {
      MACH3LOG_ERROR("Number of files found ({}) does not match expected dynamic_n_windows ({}).",
                     file_count, config.dynamic_n_windows);
      throw MaCh3Exception(__FILE__, __LINE__, "File count mismatch for dynamic loading.");
    }
  }
 
  for (size_t i = 0; i < input_trees.size(); i++) {
    TTree *tree = input_trees[i];
    TFile *file = input_files[i];
 
    MACH3LOG_DEBUG("Processing file {}/{}: {}", i + 1, input_trees.size(), file->GetName());
    TMacro *macro = static_cast<TMacro*>(file->Get("MaCh3_Config"));
    if (macro) {
      MACH3LOG_DEBUG("Found MaCh3_Config macro in file.");
 
      // Convert TMacro to YAML: concatenate all lines and parse
      std::stringstream yaml_text;
      TList *lines = macro->GetListOfLines();
      for (int iline = 0; iline < lines->GetEntries(); ++iline) {
        TObjString *line = static_cast<TObjString*>(lines->At(iline));
        yaml_text << line->GetString().Data() << "\n";
      }
 
      try {
        YAML::Node macro_yaml = YAML::Load(yaml_text.str());
        YAML::Node umbrellaConfig =
        macro_yaml["General"]["MCMC"]["Multicanonical"];
        // Extract window parameters
        config.windows[i].center = Get<double>(umbrellaConfig["Umbrella"]["UmbrellaMean"], __FILE__, __LINE__);
        MACH3LOG_DEBUG("Window {} center updated to {}", i, config.windows[i].center);
 
        config.variable_of_interest = Get<std::string>(umbrellaConfig["VariableName"], __FILE__, __LINE__);
        config.windows[i].multicanonical_beta = Get<double>(umbrellaConfig["Beta"], __FILE__, __LINE__);
 
        // Check if using von Mises distribution
        auto biasString = Get<std::string>(umbrellaConfig["Umbrella"]["UmbrellaBiasFunction"], __FILE__, __LINE__);
        M3::BiasFunction biasMode = M3::ParseBiasFunction(biasString);
        config.windows[i].umbrellaBiasFunction = biasMode;
 
        if (config.windows[i].umbrellaBiasFunction == M3::BiasFunction::kVonMises) {
          // Extract von Mises sigma and compute kappa
          double vonMises_sigma = Get<double>(umbrellaConfig["Umbrella"]["UmbrellaWidth"], __FILE__, __LINE__);
          config.windows[i].vonMises_kappa = 1.0 / (vonMises_sigma * vonMises_sigma);
          config.windows[i].width = vonMises_sigma; // Store sigma in width for reference
          MACH3LOG_DEBUG("Window {} using von Mises: sigma = {}, kappa = {}",
                         i, vonMises_sigma, config.windows[i].vonMises_kappa);
        } else {
          // Extract Gaussian sigma
          config.windows[i].width = Get<double>(umbrellaConfig["Umbrella"]["UmbrellaWidth"], __FILE__, __LINE__);
          config.windows[i].vonMises_kappa = -1.0; // Not using von Mises
          MACH3LOG_DEBUG("Window {} using Gaussian: width = {}", i, config.windows[i].width);
        }
      } catch (const std::exception &e) {
        MACH3LOG_WARN("Could not parse macro as YAML: {}", e.what());
      }
    } else{
      MACH3LOG_ERROR("Found MaCh3 chain without embedded config, this can't be right");
      throw MaCh3Exception(__FILE__, __LINE__);
    }
 
    double var_value;
 
    tree->SetBranchStatus("*", false);
    const auto* branch = tree->GetBranch(config.variable_of_interest.c_str());
    if (branch == nullptr) {
      MACH3LOG_ERROR("Branch '{}' does not exist in the TTree.", config.variable_of_interest);
      throw MaCh3Exception(__FILE__, __LINE__);
    }
    tree->SetBranchStatus(config.variable_of_interest.c_str(), true);
    tree->SetBranchAddress(config.variable_of_interest.c_str(), &var_value);
    Long64_t nentries = tree->GetEntries();
    MACH3LOG_DEBUG("Window {}: {} entries", i, nentries);
    samples[i].resize(nentries);
    for (Long64_t entry = 0; entry < nentries; entry++) {
      tree->GetEntry(entry);
      samples[i][entry] = var_value;
    }
    // need to set all to true so later part of the code don't produce bonkers results
    tree->SetBranchStatus("*", true);
  }
 
  // Sort by window center and keep all per-window containers aligned.
  // The final weighting stage loops over input_trees by index, so those indices
  // must track the same sorted window order used by z_current.
  for (size_t i = 0; i < config.windows.size(); i++) {
    for (size_t j = i + 1; j < config.windows.size(); j++) {
      if (config.windows[i].center > config.windows[j].center) {
        std::swap(config.windows[i], config.windows[j]);
        std::swap(samples[i], samples[j]);
        std::swap(input_trees[i], input_trees[j]);
        std::swap(input_files[i], input_files[j]);
      }
    }
  }
  clock.Stop();
  MACH3LOG_INFO("Finished loading samples, it took {:.2f}s to finish", clock.RealTime());
 
  // verify the order and file associations after sorting
  MACH3LOG_INFO("Final window configurations after sorting:");
 
  M3::Utils::TablePrinter table({
    {"#",             5,  5,  M3::Utils::Alignment::Right},
    {"Filename",     10, 50, M3::Utils::Alignment::Left},
    {"Center",       10, 15, M3::Utils::Alignment::Right},
    {"Width",        10, 15, M3::Utils::Alignment::Right},
    {"BiasFunction", 15, 30, M3::Utils::Alignment::Left},
    {"Samples",      10, 20, M3::Utils::Alignment::Right},
  });
 
  for (size_t i = 0; i < config.windows.size(); ++i) {
    const auto& window = config.windows[i];
    TFile* file = input_files[i];
 
    table.AddRow(i, file->GetName(),
                 M3::Utils::Cell(window.center, "{:.2f}"),
                 M3::Utils::Cell(window.width, "{:.2f}"),
                 M3::BiasFunction_ToString(window.umbrellaBiasFunction),
                 samples[i].size());
  }
 
  for (const auto& line : table.Render()) {
    MACH3LOG_INFO("{}", line);
  }
}

main()

int main	(	int	argc,
		char *	argv[]
	)

Definition at line 1035 of file UmbrellaSolver.cpp.

                                 {
  SetMaCh3LoggerFormat();
  std::string config_file = "umbrella_config.yaml";
  if (argc > 1) {
    config_file = argv[1];
  }
 
  try {
    UmbrellaSolver(config_file);
  } catch (const std::exception &e) {
    MACH3LOG_ERROR("Error: {}", e.what());
    return 1;
  }
 
  return 0;
}

parseYAMLConfig()

UmbrellaConfig parseYAMLConfig

(

const std::string &

filename

)

YAML-based config parser using yaml-cpp library.

Definition at line 45 of file UmbrellaSolver.cpp.

                                                          {
  UmbrellaConfig config;
 
  try {
    YAML::Node yaml_diag_config = YAML::LoadFile(filename);
    YAML::Node yaml_config = yaml_diag_config["UmbrellaSolver"];
 
    // Parse other configuration
    config.output_file = Get<std::string>(yaml_config["output_file"], __FILE__ , __LINE__);
    config.max_iterations = Get<int>(yaml_config["max_iterations"], __FILE__ , __LINE__);
    config.tolerance = Get<double>(yaml_config["tolerance"], __FILE__ , __LINE__);
    config.print_frequency = GetFromManager<int>(yaml_config["print_frequency"], 0, __FILE__ , __LINE__);
    config.dynamic_files = GetFromManager<bool>(yaml_config["dynamic_files"], false, __FILE__ , __LINE__);
    config.dynamic_pattern = Get<std::string>(yaml_config["dynamic_pattern"], __FILE__ , __LINE__);
    config.dynamic_n_windows = Get<int>(yaml_config["dynamic_n_windows"], __FILE__ , __LINE__);
 
    if (!config.dynamic_files) {
      // TODO: these are mostly redundant now that the MaCh3_Config macro is
      // being parsed, work out if its safe to remove them at somepoint!!! Parse
      // windows
      if (yaml_config["windows"]) {
        const YAML::Node &windows = yaml_config["windows"];
        for (size_t i = 0; i < windows.size(); i++) {
          WindowConfig window;
          window.name = Get<std::string>(windows[i]["name"], __FILE__, __LINE__);
          window.center = Get<double>(windows[i]["center"], __FILE__, __LINE__);
          window.width = Get<double>(windows[i]["width"], __FILE__, __LINE__);
          config.windows.push_back(window);
        }
      }
 
      // Parse input files
      if (yaml_config["input_files"]) {
        const YAML::Node &input_files = yaml_config["input_files"];
        for (size_t i = 0; i < input_files.size() && i < config.windows.size(); i++) {
          config.windows[i].input_file = Get<std::string>(input_files[i], __FILE__, __LINE__);
        }
      }
    } else {
      // set up placeholders to be filled later
      for (int i = 0; i < config.dynamic_n_windows; i++) {
        WindowConfig window;
        window.name = "Window_" + std::to_string(i);
        window.center = 0.0; // Placeholder, will be updated from MaCh3_Config macro
        window.width = 1.0; // Placeholder, will be updated from MaCh3_Config macro
        config.windows.push_back(window);
      }
    }
  } catch (const YAML::Exception &e) {
    MACH3LOG_ERROR("Error parsing YAML file {}: {}", filename, e.what());
    throw MaCh3Exception(__FILE__, __LINE__);
  }
 
  return config;
}

summedWindowsWeighted()

double summedWindowsWeighted	(	double	x,
		const std::vector< WindowConfig > &	windows,
		const std::vector< double > &	z_values
	)

Definition at line 144 of file UmbrellaSolver.cpp.

                                                                                                                  {
  double sum = 0.0;
  for (size_t k = 0; k < windows.size(); k++) {
    double window_val;
    if (windows[k].umbrellaBiasFunction == M3::BiasFunction::kVonMises) {
      window_val = vonMisesWindow(x, windows[k].center, windows[k].vonMises_kappa);
    } else if (windows[k].umbrellaBiasFunction == M3::BiasFunction::kGaussian) {
      window_val = gaussianWindow(x, windows[k].center, windows[k].width);
    } else if (windows[k].umbrellaBiasFunction == M3::BiasFunction::kGeneralisedGaussian) {
      window_val = GetMulticanonicalWeightGenGaussian(x, windows[k].center, windows[k].width, windows[k].multicanonical_beta);
    } else {
      MACH3LOG_ERROR("Unrecognised BiasFunction!!");
      throw MaCh3Exception(__FILE__, __LINE__);
    }
    sum += window_val / z_values[k];
  }
  return sum;
}

UmbrellaSolver()

void UmbrellaSolver

(

const std::string &

config_file

)

Definition at line 730 of file UmbrellaSolver.cpp.

                                                  {
  MACH3LOG_INFO("=== Umbrella Sampling Z-Factor Solver ===");
 
  #ifdef MULTITHREAD
  MACH3LOG_INFO("Max threads available: {}", omp_get_max_threads());
  #else
  MACH3LOG_WARN("OpenMP not available");
  #endif
  MACH3LOG_INFO("Loading configuration from: {}", config_file);
 
  // Parse configuration
  UmbrellaConfig config = parseYAMLConfig(config_file);
 
  if (config.windows.empty() && !config.dynamic_files) {
    MACH3LOG_ERROR("No windows defined in configuration and dynamic file loading is disabled.");
    throw MaCh3Exception(__FILE__, __LINE__);
  }
  MACH3LOG_INFO("Variable of interest: {}", config.variable_of_interest);
  MACH3LOG_INFO("Output file: {}", config.output_file);
 
  // Load data from input files
  // Declare here to ensure it exists in the scope of the entire function
  std::vector<std::vector<double>> samples;
 
  if (!config.dynamic_files) {
    samples.resize(config.windows.size());
  } else {
    samples.resize(config.dynamic_n_windows);
  }
  std::vector<TFile*> input_files;
  std::vector<TTree*> input_trees;
 
  LoadFiles(samples, config, input_files, input_trees);
 
  // Initialize z values
  std::vector<double> z_current(config.windows.size(), 1.0);
  std::vector<double> z_prev(config.windows.size(), 1.0);
  std::vector<std::vector<double>> z_evolution;
 
  MACH3LOG_INFO("Starting iterative z-solver...");
  int max_threads = M3::GetNThreads();
  MACH3LOG_INFO("Max threads reported: {}", max_threads);
 
  MACH3LOG_INFO("Precomputing window cache...");
  // TODO: this is slow as hell and causes massive memory usage, what's a smarter
  // way to do this? break out to a file? RDataFrame?
  std::vector<std::vector<std::vector<double>>> window_cache = buildWindowCache(config.windows, samples);
 
  // TFile to hold the F matrix evolution for the first 15 iterations if needed
  bool save_matrix = true; // Set to true to enable saving F matrix evolution
  // Create F_file here to avoid reopening it multiple times in the loop
  TFile *F_file = nullptr;
  if (save_matrix) {
    size_t pos = config.output_file.find(".root");
    std::string base_name = (pos != std::string::npos) ? config.output_file.substr(0, pos) : config.output_file;
    F_file = TFile::Open((base_name + "_matrix_evolution.root").c_str(), "RECREATE"); // use name from config with .root subtracted with _matrix_evolution suffix added
    if (!F_file || F_file->IsZombie()) {
      MACH3LOG_ERROR("Cannot create file {}", base_name + "_matrix_evolution.root");
      save_matrix = false; // Disable saving if file cannot be created
    }
    // add an initial FMatrix with the initial z values for reference
    std::vector<std::vector<double>> initial_F = calcFmatrix(z_current, config.windows, samples, window_cache);
    int n_windows = static_cast<int>(config.windows.size());
    TH2D initial_F_TH2D("F_matrix_initial", "Initial F matrix;Window j;Window i", n_windows, 0, n_windows, n_windows, 0, n_windows);
    for (int i = 0; i < n_windows; i++) {
      for (int j = 0; j < n_windows; j++) {initial_F_TH2D.SetBinContent(j + 1, i + 1, initial_F[i][j]); // Note the order of i and j for correct axis labeling
      }
    }
    F_file->cd();
    MACH3LOG_INFO("Saving initial F matrix to file...");
    initial_F_TH2D.Write();
  }
 
  // Timing variables
  auto start_time = std::chrono::high_resolution_clock::now();
  auto last_print_time = start_time;
 
  bool converged_robustness_check = false; // Flag to indicate if convergence check has been passed at least
             // once, used to control when to start checking for stalled
             // convergence
 
  if (!converged_robustness_check) {
    MACH3LOG_INFO("Starting iterative solver with convergence checks...");
  }
  // Iterative solver
  int total_output_lines = 0; // Track total lines printed for clearing
  for (int iteration = 0; iteration < config.max_iterations; iteration++) {
    if (iteration % config.print_frequency == 0 || iteration == 1) {
      auto current_time = std::chrono::high_resolution_clock::now();
      // Calculate average relative change (precision metric)
      double avg_relative_change = 0.0;
      if (iteration > 0) {
        for (size_t i = 0; i < z_current.size(); i++) {
          double rel_change = std::abs(z_current[i] - z_prev[i]) / std::max(std::abs(z_current[i]), 1e-10);
          avg_relative_change += rel_change;
        }
        avg_relative_change /= static_cast<double>(z_current.size());
      }
 
      // Print z-values
      std::string z_values = "[";
      for (size_t i = 0; i < z_current.size(); i++) {
        z_values += fmt::format("{:10.5f}", z_current[i]);
        if (i < z_current.size() - 1) {
          z_values += ", ";
        }
      }
      z_values += "]";
      if (iteration > 0) {
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(current_time - last_print_time);
        double avg_time_per_iteration = static_cast<double>(duration.count()) /
                                        static_cast<double>(config.print_frequency);
        MACH3LOG_INFO("Iteration {:6}: z values: {} (avg: {:6.1f} ms/iter)", iteration, z_values, avg_time_per_iteration);
        MACH3LOG_INFO("Avg relative change: {:.3e} (target: {:.3e})", avg_relative_change, config.tolerance);
        total_output_lines = 2;
      } else {
        MACH3LOG_INFO("Iteration {:6}: z values: {}", iteration, z_values);
        total_output_lines = 1;
      }
      last_print_time = current_time;
    }
 
    z_prev = z_current;
    z_current = zSolver(z_current, config.windows, samples, window_cache, iteration % config.print_frequency == 0, &total_output_lines);
    z_evolution.push_back(z_current);
 
    // save the first 15 iterations of the F matrix to check convergence
    // behaviour and debug if needed should be a root file with Th2D for easy
    // plotting in root, with axes of iteration number and window index, and the
    // value being the F matrix element
    if (save_matrix && (iteration < 15 || iteration % config.print_frequency == 0)) {
      std::vector<std::vector<double>> F_matrix = calcFmatrix(z_current, config.windows, samples, window_cache);
      // convert F_matrix to Th2D for saving to root file
      int n_windows = static_cast<int>(config.windows.size());
      TH2D F_TH2D(Form("F_matrix_iter_%02d", iteration),Form("F matrix at iteration %02d;Window j;Window i", iteration), n_windows, 0, n_windows, n_windows, 0, n_windows);
      for (int i = 0; i < n_windows; i++) {
        for (int j = 0; j < n_windows; j++) {
          F_TH2D.SetBinContent(j + 1, i + 1, F_matrix[i][j]); // Note the order of i and j for correct axis labeling
        }
      }
      // for the purposes of picking back up a solve after it has been
      // interrupted also save the std::vector of z_current. You can put this into hacky start to pick up
      TTree *z_tree = new TTree(Form("z_saved_iter_%02d", iteration), Form("Z vector at iteration %02d", iteration));
      z_tree->Branch("z_saved", &z_current);
      z_tree->Fill();
      F_file->cd();
      MACH3LOG_INFO("Saving F matrix for iteration {} to file...", iteration);
      F_TH2D.Write();
      z_tree->Write();
    }
    // if (save_matrix && iteration == 15) {
    //     F_file->Close();
    // }
 
    // after the first convergence check has been passed, randomly perturb the z
    // values to see if they return to the same values, this is a robustness
    // check to see if the solution is stable or if it is just meeting the
    // convergence criteria by chance due to small changes in z values
    bool apply_robustness_check = true;
    if (iteration % 100 == 0 &&
        (checkConvergence(z_current, z_prev, config.tolerance) || checkConvergenceStalled(z_current, z_prev, config.tolerance))) {
      if (!converged_robustness_check && apply_robustness_check) {
        MACH3LOG_INFO("Convergence check passed at iteration {}. Starting robustness check with random perturbation...", iteration);
        converged_robustness_check = true;
 
        // Apply random perturbation to z_current
        std::vector<double> z_perturbed = z_current;
        for (size_t i = 0; i < z_perturbed.size(); i++) {
          // Random perturbation up to 10 times the tolerance
          double perturbation = M3::rand::Uniform(-0.5, 0.5) * z_perturbed[i];
          z_perturbed[i] += perturbation;
          if (z_perturbed[i] < 0) {
            z_perturbed[i] = abs(z_perturbed[i]); // Ensure no negative values
          }
          MACH3LOG_INFO("z[{:2}] | Original: {:.6e} | Perturbation: {:+.6e} | Updated: {:.6e}",
                          i, z_current[i], perturbation, z_perturbed[i]);
        }
        z_current = z_perturbed;
          MACH3LOG_INFO("Applied random perturbation to z values for robustness check.");
      } else {
        if (checkConvergence(z_current, z_prev, config.tolerance)) {
          MACH3LOG_INFO("Convergence achieved at iteration {}", iteration);
        } else {
          MACH3LOG_WARN("Convergence appears to be stalled at iteration {}", iteration);
        }
 
        if (iteration == config.max_iterations - 1) {
          MACH3LOG_WARN("Reached maximum iterations without convergence.");
        }
 
        auto end_time = std::chrono::high_resolution_clock::now();
        auto total_duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
        double avg_time_total = static_cast<double>(total_duration.count()) /
                                static_cast<double>(iteration + 1);
        if (save_matrix) {
          F_file->Close();
        }
        MACH3LOG_INFO("Terminating at iteration {}", iteration);
        MACH3LOG_INFO("Average time per iteration: {:.6f} ms", avg_time_total);
        break;
      }
    }
  }
 
  std::ostringstream oss;
  oss << "Final z values: [";
  for (size_t i = 0; i < z_current.size(); i++) {
    oss << std::fixed << std::setprecision(5) << z_current[i];
    if (i < z_current.size() - 1)
      oss << ", ";
  }
  oss << "]";
  MACH3LOG_INFO("{}", oss.str());
 
  std::filesystem::copy_file(input_files[0]->GetName(), config.output_file,
                             std::filesystem::copy_options::overwrite_existing);
  // Create output file
  TFile *output_file = M3::Open(config.output_file.c_str(), "UPDATE", __FILE__, __LINE__);
  output_file->cd();
 
  TTree *input_tree = dynamic_cast<TTree*>(output_file->Get("posteriors"));
  // Create combined tree with weights
  TTree *combined_tree = input_tree->CloneTree(0);
 
  // Variables for the combined tree
  double umbrella_weight;
  int window_id;
  double umbrella_variable;
 
  combined_tree->Branch("umbrella_weight", &umbrella_weight, "umbrella_weight/D");
  combined_tree->Branch("window_id", &window_id, "window_id/I");
 
  // Fill combined tree
  for (size_t i = 0; i < input_trees.size(); i++) {
    TTree *tree = input_trees[i];
 
    Long64_t nentries = tree->GetEntries();
 
    // KS: This is to avoid warnings about missing umbrella branches...
    int oldLevel = gErrorIgnoreLevel;
    gErrorIgnoreLevel = kError;
    combined_tree->CopyAddresses(tree);
    gErrorIgnoreLevel = oldLevel;
 
    tree->SetBranchAddress(config.variable_of_interest.c_str(), &umbrella_variable);
    // KS: SetBranchAddress above decouples the input branch address, so update the
    // copied output branch address to use the current umbrella_variable value.
    combined_tree->GetBranch(config.variable_of_interest.c_str())->SetAddress(&umbrella_variable);
    if (z_current[i] == 0) {
      MACH3LOG_WARN("Z value for window {} is zero, skipping weighting for this window to avoid division by zero.", i);
    }
    window_id = static_cast<int>(i);
 
    for (Long64_t entry = 0; entry < nentries; entry++) {
      tree->GetEntry(entry);
 
      if (z_current[i] == 0) {
        umbrella_weight = 0.0; // If z is zero, we cannot apply the umbrella weight, so we set it to 0 (completely downweigh this window's contribution)
      } else {
        // Calculate umbrella weight for this event
        // The umbrella weight corrects for the bias introduced by the window function Weight is 1 / sum of all window contributions (equation 4 from paper)
        double denominator = 1 / summedWindowsWeighted(umbrella_variable, config.windows, z_current);
 
        // umbrella_weight = z_current[i] / denominator; // with or without z_current[i] / denominator? why did I have this originally
        umbrella_weight = denominator; // This is the correct form based on the paper - the z_current[i] factor is already included in the summedWindowsWeighted function
      }
 
      if (combined_tree->Fill() < 0) {
        MACH3LOG_WARN("Failed writing output tree. Check disk quota/space and write permissions for: {}", config.output_file);
        throw MaCh3Exception(__FILE__, __LINE__);
      }
    }
  }
  // Write final results
  combined_tree->Write(input_tree->GetName(), TObject::kOverwrite);
 
  // Save diagnostics
  DiagnoseUmbrella(config, output_file, z_evolution, combined_tree);
 
  // Get entry count before closing the file
  Long64_t total_entries = combined_tree->GetEntries();
 
  output_file->cd();
 
  YAML::Node yaml_config = M3OpenConfig(config_file);
  YAML::Node umbrella_config;
  umbrella_config["UmbrellaSolver"] = yaml_config["UmbrellaSolver"];
 
  // Convert YAML -> TMacro
  TMacro UmbrellaHeader = YAMLtoTMacro(umbrella_config, "Umbrella_Config");
  UmbrellaHeader.Write();
 
  output_file->Close();
 
  // Close input files
  for (TFile *file : input_files) {
    file->Close();
  }
 
  MACH3LOG_INFO("Umbrella sampling outputs created (make sure to check for convergence issues)!");
  MACH3LOG_INFO("Output written to: {}", config.output_file);
  MACH3LOG_INFO("Combined tree contains {} entries with umbrella weights.", total_entries);
}

vonMisesWindow()

double vonMisesWindow	(	double	x,
		double	center,
		double	kappa
	)

Definition at line 107 of file UmbrellaSolver.cpp.

                                                             {
  // von Mises PDF: exp(kappa * cos(x - center)) / (2*pi*I0(kappa))
  // For numerical stability, compute in log space when possible
  if (kappa > 700) {
    // Use asymptotic approximation for large kappa to avoid overflow
    // log(I0(kappa)) ≈ kappa - 0.5*log(2*pi*kappa)
    double log_I0 = kappa - 0.5 * log(2 * TMath::Pi() * kappa);
    return exp(kappa * cos(x - center) - log_I0 - log(2 * TMath::Pi()));
  } else {
    double I0_kappa = TMath::BesselI0(kappa);
    return exp(kappa * cos(x - center)) / (2 * TMath::Pi() * I0_kappa);
  }
}

zSolver()

std::vector<double> zSolver	(	const std::vector< double > &	z_current,
		const std::vector< WindowConfig > &	windows,
		const std::vector< std::vector< double >> &	samples,
		const std::vector< std::vector< std::vector< double >>> &	window_cache,
		[[maybe_unused] ] bool	verbose = false,
		[[maybe_unused] ] int *	total_lines = nullptr
	)

Normalize z to prevent scale drift and maintain sum(z)=1

Todo:

: try this with magnitude = 1, the absolute scale doesn't matter but maybee this condition is preventing elements meeting their required values? also just try completely disabled double z_sum = 0.0; for (int i = 0; i < n_windows; i++) { z_sum += z_new[i]; } if (z_sum > 0) { for (int i = 0; i < n_windows; i++) { z_new[i] /= z_sum; }

Definition at line 325 of file UmbrellaSolver.cpp.

                                                     {
  int n_windows = static_cast<int>(windows.size());
 
  // F matrix and update z values
  std::vector<std::vector<double>> F = calcFmatrix(z_current, windows, samples, window_cache);
  // Compute z_new = z_current * F
  // could do faster matrix multiplication here?
  std::vector<double> z_new(n_windows, 0.0);
  for (int i = 0; i < n_windows; i++) {
    for (int j = 0; j < n_windows; j++) {
      z_new[i] += z_current[j] * F[j][i];
    }
  }
 
  //}
 
  // normalise the z values so that magnitude of the vector is 1
  double z_magnitude = 0.0;
  for (int i = 0; i < n_windows; i++) {
    z_magnitude += z_new[i] * z_new[i];
  }
  z_magnitude = sqrt(z_magnitude);
  if (z_magnitude > 0) {
    for (int i = 0; i < n_windows; i++) {
      z_new[i] /= z_magnitude;
    }
  }
 
  return z_new;
}

Variable Documentation

debug_mode

_MaCh3_Safe_Include_Start_ _MaCh3_Safe_Include_End_ bool debug_mode = false

Definition at line 17 of file UmbrellaSolver.cpp.