analysis/NNScaleMap/massClosure.C

0001 #include <TCanvas.h>
0002 #include <TFile.h>
0003 #include <TH1.h>
0004 #include <TLatex.h>
0005 #include <TLegend.h>
0006 #include <TLine.h>
0007 #include <TObject.h>
0008 #include <TPad.h>
0009 #include <TROOT.h>
0010 #include <TStyle.h>
0011 #include <TSystem.h>
0012
0013 #include <algorithm>
0014 #include <cctype>
0015 #include <cmath>
0016 #include <iostream>
0017 #include <memory>
0018 #include <string>
0019 #include <vector>
0020
0021 TGaxis::SetMaxDigits(3);
0022
0023 namespace
0024 {
0025 struct SpeciesSpec
0026 {
0027     std::string name;
0028     std::string label;
0029     std::string rawKey;
0030     std::string correctedKey;
0031     double pdgMass = 0.0;
0032 };
0033
0034 std::string directoryName(const std::string &path)
0035 {
0036     const std::size_t slash = path.find_last_of('/');
0037     return slash == std::string::npos ? "." : path.substr(0, slash);
0038 }
0039
0040 std::string baseName(const std::string &path)
0041 {
0042     const std::size_t slash = path.find_last_of('/');
0043     return slash == std::string::npos ? path : path.substr(slash + 1);
0044 }
0045
0046 std::string stripExtension(const std::string &name)
0047 {
0048     const std::size_t dot = name.find_last_of('.');
0049     return dot == std::string::npos ? name : name.substr(0, dot);
0050 }
0051
0052 std::string safeFileStem(std::string name)
0053 {
0054     for (char &c : name)
0055     {
0056         if (!(std::isalnum(static_cast<unsigned char>(c)) || c == '_' || c == '-' || c == '.'))
0057             c = '_';
0058     }
0059     return name;
0060 }
0061
0062 TH1 *getHistogram(TFile *input, const std::string &key)
0063 {
0064     TH1 *hist = dynamic_cast<TH1 *>(input->Get(key.c_str()));
0065     if (!hist)
0066     {
0067         std::cerr << "Missing TH1 histogram: " << key << std::endl;
0068         return nullptr;
0069     }
0070
0071     TH1 *clone = static_cast<TH1 *>(hist->Clone(Form("%s_plot_clone", key.c_str())));
0072     if (!clone)
0073     {
0074         std::cerr << "Could not clone histogram: " << key << std::endl;
0075         return nullptr;
0076     }
0077     clone->SetDirectory(nullptr);
0078     return clone;
0079 }
0080
0081 double histogramMax(const TH1 *a, const TH1 *b)
0082 {
0083     double ymax = 0.0;
0084     for (const TH1 *hist : {a, b})
0085     {
0086         for (int bin = 1; bin <= hist->GetNbinsX(); ++bin)
0087             ymax = std::max(ymax, hist->GetBinContent(bin));
0088     }
0089     return ymax > 0.0 ? ymax : 1.0;
0090 }
0091
0092 void styleHistogram(TH1 *hist, const int color, const int lineStyle)
0093 {
0094     hist->SetStats(false);
0095     hist->SetLineColor(color);
0096     hist->SetMarkerColor(color);
0097     hist->SetLineStyle(lineStyle);
0098     hist->SetLineWidth(3);
0099     hist->GetXaxis()->SetTitle("Invariant mass [GeV]");
0100     hist->GetYaxis()->SetTitle("Candidates");
0101     // hist->GetXaxis()->SetTitleOffset(1.1);
0102     hist->GetYaxis()->SetTitleOffset(1.25);
0103 }
0104
0105 void drawPdgLine(const double pdgMass, const double ymax)
0106 {
0107     TLine *line = new TLine(pdgMass, 0.0, pdgMass, ymax);
0108     line->SetLineColor(kGray + 2);
0109     line->SetLineStyle(2);
0110     line->SetLineWidth(2);
0111     line->SetBit(kCanDelete);
0112     line->Draw();
0113 }
0114
0115 void drawComparisonPad(TH1 *raw, TH1 *corrected, const SpeciesSpec &species, const bool drawLegend)
0116 {
0117     styleHistogram(raw, kAzure + 2, 1);
0118     styleHistogram(corrected, kOrange + 7, 1);
0119
0120     const double ymax = 1.25 * histogramMax(raw, corrected);
0121     raw->SetMaximum(ymax);
0122     corrected->SetMaximum(ymax);
0123     raw->SetTitle(Form("%s mass closure;Invariant mass [GeV];Candidates", species.label.c_str()));
0124     raw->GetXaxis()->SetNdivisions(505);
0125     raw->Draw("hist");
0126     corrected->Draw("hist same");
0127     drawPdgLine(species.pdgMass, ymax);
0128
0129     // TLatex label;
0130     // label.SetNDC();
0131     // label.SetTextFont(42);
0132     // label.SetTextSize(0.040);
0133     // label.DrawLatex(0.16, 0.86, species.label.c_str());
0134     // label.SetTextSize(0.032);
0135     // label.DrawLatex(0.16, 0.805, Form("PDG mass = %.6f GeV", species.pdgMass));
0136
0137     if (drawLegend)
0138     {
0139         TLegend *legend = new TLegend(0.54, 0.68, 0.88, 0.88);
0140         legend->SetHeader(Form("Resonance species: %s", species.label.c_str()), "C");
0141         legend->SetBorderSize(0);
0142         legend->SetFillStyle(0);
0143         legend->SetTextSize(0.04);
0144         legend->SetBit(kCanDelete);
0145         legend->AddEntry(raw, "Before correction", "l");
0146         legend->AddEntry(corrected, "After correction", "l");
0147         legend->AddEntry(static_cast<TObject *>(nullptr), "Dashed: PDG mass", "");
0148         legend->Draw();
0149     }
0150 }
0151 } // namespace
0152
0153 void massClosure(std::string inputfile = "./calib_out_pythiaInjection_20260805/stage1_mass_histograms.root")
0154 {
0155     const std::string inputdir = directoryName(inputfile);
0156     const std::string inputname = baseName(inputfile);
0157     const std::string inputname_noext = stripExtension(inputname);
0158     const std::string outdir = inputdir + "/" + inputname_noext + "_closure";
0159     gSystem->mkdir(outdir.c_str(), true);
0160
0161     gStyle->SetOptStat(0);
0162     gStyle->SetTitleFont(42, "XYZ");
0163     gStyle->SetLabelFont(42, "XYZ");
0164     gStyle->SetLegendFont(42);
0165
0166     std::unique_ptr<TFile> input(TFile::Open(inputfile.c_str(), "READ"));
0167     if (!input || input->IsZombie())
0168     {
0169         std::cerr << "Could not open mass histogram file: " << inputfile << std::endl;
0170         return;
0171     }
0172
0173     const std::vector<SpeciesSpec> speciesList = {
0174         {"kshort", "K^{0}_{S}", "kshort_mass_raw", "kshort_mass_corrected", 0.497611},
0175         {"lambda", "#Lambda", "lambda_mass_raw", "lambda_mass_corrected", 1.115683},
0176         {"anti_lambda", "#bar{#Lambda}", "anti_lambda_mass_raw", "anti_lambda_mass_corrected", 1.115683},
0177     };
0178
0179     std::vector<TH1 *> rawHists;
0180     std::vector<TH1 *> correctedHists;
0181     std::vector<SpeciesSpec> foundSpecies;
0182     rawHists.reserve(speciesList.size());
0183     correctedHists.reserve(speciesList.size());
0184     foundSpecies.reserve(speciesList.size());
0185
0186     for (const SpeciesSpec &species : speciesList)
0187     {
0188         TH1 *raw = getHistogram(input.get(), species.rawKey);
0189         TH1 *corrected = getHistogram(input.get(), species.correctedKey);
0190         if (!raw || !corrected)
0191         {
0192             delete raw;
0193             delete corrected;
0194             continue;
0195         }
0196
0197         rawHists.push_back(raw);
0198         correctedHists.push_back(corrected);
0199         foundSpecies.push_back(species);
0200
0201         TCanvas canvas(Form("c_%s_mass_closure", species.name.c_str()), Form("%s mass closure", species.label.c_str()), 800, 750);
0202         canvas.SetTicks(1, 1);
0203         canvas.SetLeftMargin(0.13);
0204         canvas.SetRightMargin(0.05);
0205         // canvas.SetBottomMargin(0.12);
0206         canvas.SetTopMargin(0.08);
0207         drawComparisonPad(raw, corrected, species, true);
0208         canvas.RedrawAxis();
0209
0210         const std::string stem = outdir + "/" + safeFileStem(species.name + "_mass_comparison");
0211         canvas.SaveAs((stem + ".png").c_str());
0212         canvas.SaveAs((stem + ".pdf").c_str());
0213     }
0214
0215     if (foundSpecies.empty())
0216     {
0217         std::cerr << "No complete raw/corrected mass histogram pairs found in " << inputfile << std::endl;
0218         return;
0219     }
0220
0221     {
0222         TCanvas gridCanvas("c_mass_closure_all", "Mass closure comparisons", 1500, 520);
0223         gridCanvas.Divide(static_cast<int>(foundSpecies.size()), 1, 0.001, 0.001);
0224         for (std::size_t i = 0; i < foundSpecies.size(); ++i)
0225         {
0226             gridCanvas.cd(static_cast<int>(i) + 1);
0227             gPad->SetTicks(1, 1);
0228             gPad->SetLeftMargin(0.14);
0229             gPad->SetRightMargin(0.04);
0230             gPad->SetBottomMargin(0.13);
0231             gPad->SetTopMargin(0.10);
0232             gPad->RedrawAxis();
0233             drawComparisonPad(rawHists[i], correctedHists[i], foundSpecies[i], i == 0);
0234         }
0235         gridCanvas.SaveAs((outdir + "/all_mass_comparisons.png").c_str());
0236         gridCanvas.SaveAs((outdir + "/all_mass_comparisons.pdf").c_str());
0237     }
0238
0239     std::cout << "Saved mass-closure comparison plots to " << outdir << std::endl;
0240
0241     for (TH1 *hist : rawHists)
0242         delete hist;
0243     for (TH1 *hist : correctedHists)
0244         delete hist;
0245 }