ANASEN_analysis/Armory/AutoHist2D.h
2026-08-24 10:42:51 -04:00

103 lines
3.8 KiB
C++

#ifndef AUTOHIST2D_H
#define AUTOHIST2D_H
#include <TCanvas.h>
#include <TH2F.h>
#include <algorithm>
#include <iostream>
#include <limits>
#include <string>
#include <unordered_map>
#include <vector>
#include <filesystem>
// Buffers (name -> x,y) pairs during the event loop and, on WriteAll(),
// books a TH2F per name with a fixed number of bins per axis and a range
// derived from the buffered data, then Write()s it to the current
// TDirectory/TFile so it shows up alongside the tree in a TBrowser.
class AutoHist2D {
public:
static const int kDefaultBins = 200;
// Buffers one (xValue, yValue) pair under the histogram 'name'. Axis
// titles are optional and only need to be given once (first call wins).
static void Fill(const std::string& name, double xValue, double yValue,
const std::string& xTitle = "", const std::string& yTitle = "") {
Buffer& buf = Registry()[name];
buf.x.push_back(xValue);
buf.y.push_back(yValue);
if (buf.xTitle.empty() && !xTitle.empty()) buf.xTitle = xTitle;
if (buf.yTitle.empty() && !yTitle.empty()) buf.yTitle = yTitle;
}
// Books, fills and Write()s every registered histogram (nbins x nbins,
// range taken from min/max of the buffered data) to gDirectory, then
// clears the buffers.
static void WriteAll(int nbins = kDefaultBins) {
for (auto& kv : Registry()) {
const std::string& name = kv.first;
Buffer& buf = kv.second;
if (buf.x.empty()) continue;
double xlo = *std::min_element(buf.x.begin(), buf.x.end());
double xhi = *std::max_element(buf.x.begin(), buf.x.end());
double ylo = *std::min_element(buf.y.begin(), buf.y.end());
double yhi = *std::max_element(buf.y.begin(), buf.y.end());
PadRange(xlo, xhi);
PadRange(ylo, yhi);
std::string title = name + ";" + (buf.xTitle.empty() ? "x" : buf.xTitle)
+ ";" + (buf.yTitle.empty() ? "y" : buf.yTitle);
TH2F h(name.c_str(), title.c_str(), nbins, xlo, xhi, nbins, ylo, yhi);
for (size_t i = 0; i < buf.x.size(); ++i) h.Fill(buf.x[i], buf.y[i]);
h.Write();
SavePNG(h, name);
std::cout << "AutoHist2D: wrote \"" << name << "\" (" << buf.x.size()
<< " entries, " << nbins << "x" << nbins << " bins, range ["
<< xlo << "," << xhi << "] x [" << ylo << "," << yhi << "])" << std::endl;
}
Registry().clear();
}
private:
static constexpr const char* kPlotDir = "Plots";
struct Buffer {
std::vector<double> x, y;
std::string xTitle, yTitle;
};
// Draws 'h' on a throwaway canvas and saves it under kPlotDir/<name>.png,
// creating the directory first if it doesn't already exist; overwrites
// any existing file of that name.
static void SavePNG(TH2F& h, const std::string& name) {
std::filesystem::create_directories(kPlotDir);
TCanvas c(("c_" + name).c_str(), "", 800, 600);
h.Draw("COLZ");
c.SaveAs((std::string(kPlotDir) + "/" + name + ".png").c_str());
}
// Registered histograms, keyed by name. Function-local static avoids
// needing a separate translation unit for a header-only class.
static std::unordered_map<std::string, Buffer>& Registry() {
static std::unordered_map<std::string, Buffer> registry;
return registry;
}
// Pads a [lo, hi] range by 2% on each side so points at the extremes
// aren't binned into the edge/overflow bin; handles the lo == hi case.
static void PadRange(double& lo, double& hi) {
if (lo == hi) {
lo -= 1.0;
hi += 1.0;
return;
}
double pad = 0.02 * (hi - lo);
lo -= pad;
hi += pad;
}
};
#endif