analysis/NNScaleMap/makeScaleMapPlot.C

0001 #include <TAxis.h>
0002 #include <TCanvas.h>
0003 #include <TColor.h>
0004 #include <TFile.h>
0005 #include <TH2.h>
0006 #include <TIterator.h>
0007 #include <TKey.h>
0008 #include <TLatex.h>
0009 #include <TObject.h>
0010 #include <TPad.h>
0011 #include <TROOT.h>
0012 #include <TString.h>
0013 #include <TStyle.h>
0014 #include <TSystem.h>
0015
0016 #include <algorithm>
0017 #include <cctype>
0018 #include <cmath>
0019 #include <iostream>
0020 #include <memory>
0021 #include <string>
0022 #include <utility>
0023 #include <vector>
0024
0025 TGaxis::SetMaxDigits(3);
0026
0027 namespace
0028 {
0029 struct PlotInfo
0030 {
0031     TH2 *hist = nullptr;
0032     std::string name;
0033     std::string title;
0034     std::string ztitle;
0035     int kind = 100;
0036     double pt = 0.0;
0037     bool hasPt = false;
0038 };
0039
0040 struct ZRange
0041 {
0042     double center = 0.0;
0043     double halfWidth = 0.0;
0044     bool initialized = false;
0045 };
0046
0047 std::string directoryName(const std::string &path)
0048 {
0049     const std::size_t slash = path.find_last_of('/');
0050     return slash == std::string::npos ? "." : path.substr(0, slash);
0051 }
0052
0053 std::string baseName(const std::string &path)
0054 {
0055     const std::size_t slash = path.find_last_of('/');
0056     return slash == std::string::npos ? path : path.substr(slash + 1);
0057 }
0058
0059 std::string stripExtension(const std::string &name)
0060 {
0061     const std::size_t dot = name.find_last_of('.');
0062     return dot == std::string::npos ? name : name.substr(0, dot);
0063 }
0064
0065 std::string safeFileStem(std::string name)
0066 {
0067     for (char &c : name)
0068     {
0069         if (!(std::isalnum(static_cast<unsigned char>(c)) || c == '_' || c == '-' || c == '.'))
0070             c = '_';
0071     }
0072     return name;
0073 }
0074
0075 bool parsePtFromName(const std::string &name, double &pt)
0076 {
0077     const std::size_t pos = name.rfind("_pt");
0078     if (pos == std::string::npos)
0079         return false;
0080
0081     std::string token = name.substr(pos + 3);
0082     std::replace(token.begin(), token.end(), 'p', '.');
0083     try
0084     {
0085         pt = std::stod(token);
0086     }
0087     catch (...)
0088     {
0089         return false;
0090     }
0091     return true;
0092 }
0093
0094 PlotInfo makePlotInfo(TH2 *hist)
0095 {
0096     PlotInfo info;
0097     info.hist = hist;
0098     info.name = hist->GetName();
0099     info.hasPt = parsePtFromName(info.name, info.pt);
0100
0101     if (info.name.rfind("eps_", 0) == 0)
0102     {
0103         info.kind = 0;
0104         info.title = "Charge-even curvature-scale correction";
0105         info.ztitle = "Charge-even correction #epsilon";
0106     }
0107     else if (info.name.rfind("delta_", 0) == 0)
0108     {
0109         info.kind = 1;
0110         info.title = "Charge-odd curvature-scale correction";
0111         info.ztitle = "Charge-odd correction #delta";
0112     }
0113     else if (info.name.rfind("kappa_qplus_", 0) == 0)
0114     {
0115         info.kind = 2;
0116         info.title = "Final p_{T} scale correction, q = +1";
0117         info.ztitle = "Final scale correction #kappa(q=+1)";
0118     }
0119     else if (info.name.rfind("kappa_qminus_", 0) == 0)
0120     {
0121         info.kind = 3;
0122         info.title = "Final p_{T} scale correction, q = -1";
0123         info.ztitle = "Final scale correction #kappa(q=-1)";
0124     }
0125     else
0126     {
0127         info.title = info.name;
0128         info.ztitle = "scale correction";
0129     }
0130
0131     return info;
0132 }
0133
0134 std::string histogramTitle(const PlotInfo &info, const bool grid)
0135 {
0136     std::string title = info.title;
0137     if (grid)
0138     {
0139         title = info.ztitle;
0140         if (info.name.rfind("kappa_qplus_", 0) == 0)
0141             title = "q = +1";
0142         else if (info.name.rfind("kappa_qminus_", 0) == 0)
0143             title = "q = -1";
0144
0145         return Form("%s;#eta;#phi;%s", title.c_str(), info.ztitle.c_str());
0146     }
0147
0148     if (info.hasPt)
0149         return Form("%s, p_{T} = %.2f GeV;#eta;#phi;%s", title.c_str(), info.pt, info.ztitle.c_str());
0150     return Form("%s;#eta;#phi;%s", title.c_str(), info.ztitle.c_str());
0151 }
0152
0153 void styleHistogram(PlotInfo &info, const ZRange &range, const bool grid)
0154 {
0155     TH2 *hist = info.hist;
0156     hist->SetTitle(histogramTitle(info, grid).c_str());
0157     hist->SetMinimum(range.center - range.halfWidth);
0158     hist->SetMaximum(range.center + range.halfWidth);
0159     // hist->SetContour(1000);
0160
0161     hist->GetXaxis()->SetTitle("#eta");
0162     hist->GetYaxis()->SetTitle("#phi");
0163     hist->GetZaxis()->SetTitle(info.ztitle.c_str());
0164
0165     // hist->GetXaxis()->CenterTitle();
0166     // hist->GetYaxis()->CenterTitle();
0167     // hist->GetZaxis()->CenterTitle();
0168     hist->GetXaxis()->SetTitleOffset(grid ? 1.2 : 1.1);
0169     hist->GetYaxis()->SetTitleOffset(grid ? 1.3 : 1.1);
0170     hist->GetZaxis()->SetTitleOffset(grid ? 1.4 : 1.8);
0171
0172     // if (grid)
0173     // {
0174     //     hist->SetTitleSize(0.045, "XYZ");
0175     //     hist->SetLabelSize(0.035, "XYZ");
0176     // }
0177     // else
0178     // {
0179     //     hist->SetTitleSize(0.045, "XYZ");
0180     //     hist->SetLabelSize(0.04, "XYZ");
0181     // }
0182 }
0183
0184 std::vector<ZRange> getZRanges(const std::vector<PlotInfo> &plots)
0185 {
0186     const int nKinds = 4;
0187     std::vector<ZRange> ranges(nKinds);
0188     for (int kind = 0; kind < nKinds; ++kind)
0189         ranges[kind].center = kind < 2 ? 0.0 : 1.0;
0190
0191     for (const PlotInfo &plot : plots)
0192     {
0193         if (plot.kind < 0 || plot.kind >= nKinds)
0194             continue;
0195
0196         const TH2 *hist = plot.hist;
0197         ZRange &range = ranges[plot.kind];
0198         for (int ix = 1; ix <= hist->GetNbinsX(); ++ix)
0199         {
0200             for (int iy = 1; iy <= hist->GetNbinsY(); ++iy)
0201             {
0202                 const double value = hist->GetBinContent(ix, iy);
0203                 if (std::isfinite(value))
0204                 {
0205                     range.halfWidth = std::max(range.halfWidth, std::abs(value - range.center));
0206                     range.initialized = true;
0207                 }
0208             }
0209         }
0210     }
0211
0212     for (ZRange &range : ranges)
0213     {
0214         if (range.halfWidth <= 0.0)
0215             range.halfWidth = 1.0e-12;
0216     }
0217
0218     return ranges;
0219 }
0220
0221 ZRange getRangeForPlot(const PlotInfo &plot, const std::vector<ZRange> &ranges)
0222 {
0223     if (plot.kind >= 0 && plot.kind < static_cast<int>(ranges.size()) && ranges[plot.kind].initialized)
0224         return ranges[plot.kind];
0225
0226     ZRange range;
0227     range.center = 0.0;
0228     range.halfWidth = 1.0e-12;
0229     for (int ix = 1; ix <= plot.hist->GetNbinsX(); ++ix)
0230     {
0231         for (int iy = 1; iy <= plot.hist->GetNbinsY(); ++iy)
0232         {
0233             const double value = plot.hist->GetBinContent(ix, iy);
0234             if (std::isfinite(value))
0235             {
0236                 range.halfWidth = std::max(range.halfWidth, std::abs(value));
0237                 range.initialized = true;
0238             }
0239         }
0240     }
0241     return range;
0242 }
0243
0244 std::string mapKindName(const int kind)
0245 {
0246     if (kind == 0)
0247         return "epsilon";
0248     if (kind == 1)
0249         return "delta";
0250     if (kind == 2)
0251         return "kappa q=+1";
0252     if (kind == 3)
0253         return "kappa q=-1";
0254     return Form("kind %d", kind);
0255 }
0256
0257 void drawPadLabel(const PlotInfo &info, const bool grid)
0258 {
0259     TLatex label;
0260     label.SetNDC();
0261     label.SetTextFont(42);
0262     label.SetTextSize(grid ? 0.045 : 0.035);
0263
0264     std::string charge_label;
0265     if (info.name.rfind("kappa_qplus_", 0) == 0)
0266         charge_label = "q = +1";
0267     else if (info.name.rfind("kappa_qminus_", 0) == 0)
0268         charge_label = "q = -1";
0269
0270     const double x = gPad->GetLeftMargin();
0271     double y = 1 - gPad->GetTopMargin() + 0.025;
0272     if (info.hasPt)
0273     {
0274         if (!charge_label.empty())
0275         {
0276             label.DrawLatex(x, y, Form("p_{T} = %.2f GeV, 1/p_{T} = %.2f GeV^{-1}, %s", info.pt, 1.0 / info.pt, charge_label.c_str()));
0277             y -= grid ? 0.055 : 0.045;
0278         }
0279         else
0280         {
0281             label.DrawLatex(x, y, Form("p_{T} = %.2f GeV, 1/p_{T} = %.2f GeV^{-1}", info.pt, 1.0 / info.pt));
0282             y -= grid ? 0.055 : 0.045;
0283         }
0284     }
0285 }
0286
0287 std::pair<int, int> gridShape(const std::vector<PlotInfo> &plots)
0288 {
0289     std::vector<double> pts;
0290     bool allHavePt = !plots.empty();
0291     for (const PlotInfo &plot : plots)
0292     {
0293         allHavePt = allHavePt && plot.hasPt;
0294         if (!plot.hasPt)
0295             continue;
0296         const bool seen = std::any_of(pts.begin(), pts.end(), [&](double value) { return std::abs(value - plot.pt) < 1.0e-6; });
0297         if (!seen)
0298             pts.push_back(plot.pt);
0299     }
0300
0301     if (allHavePt && pts.size() > 1 && pts.size() <= 12 && (plots.size() + pts.size() - 1) / pts.size() <= 6)
0302     {
0303         const int ncols = static_cast<int>(pts.size());
0304         const int nrows = static_cast<int>((plots.size() + ncols - 1) / ncols);
0305         return {ncols, nrows};
0306     }
0307
0308     const int ncols = static_cast<int>(std::ceil(std::sqrt(static_cast<double>(plots.size()))));
0309     const int nrows = static_cast<int>((plots.size() + ncols - 1) / ncols);
0310     return {ncols, nrows};
0311 }
0312 } // namespace
0313
0314 void makeScaleMapPlot( //
0315     std::string scalemapfile = "./calib_out_upgrade20260731/kappa_maps.root")
0316 {
0317     const std::string scalemapdir = directoryName(scalemapfile);
0318     const std::string scalemapname = baseName(scalemapfile);
0319     const std::string scalemapname_noext = stripExtension(scalemapname);
0320     const std::string outdir = scalemapdir + "/" + scalemapname_noext + "_plots";
0321     system(("mkdir -p " + outdir).c_str());
0322
0323     gStyle->SetPalette(kLightTemperature);
0324
0325     std::unique_ptr<TFile> input(TFile::Open(scalemapfile.c_str(), "READ"));
0326     if (!input || input->IsZombie())
0327     {
0328         std::cerr << "Could not open scale map file: " << scalemapfile << std::endl;
0329         return;
0330     }
0331
0332     std::vector<PlotInfo> plots;
0333     TIter next(input->GetListOfKeys());
0334     while (TKey *key = static_cast<TKey *>(next()))
0335     {
0336         std::unique_ptr<TObject> object(key->ReadObj());
0337         TH2 *hist = dynamic_cast<TH2 *>(object.get());
0338         if (!hist)
0339             continue;
0340
0341         hist->SetDirectory(nullptr);
0342         TH2 *clone = static_cast<TH2 *>(hist->Clone(Form("%s_plot_clone", hist->GetName())));
0343         if (!clone)
0344             continue;
0345         clone->SetDirectory(nullptr);
0346         clone->SetName(hist->GetName());
0347         plots.push_back(makePlotInfo(clone));
0348     }
0349
0350     if (plots.empty())
0351     {
0352         std::cerr << "No TH2 histograms found in scale map file: " << scalemapfile << std::endl;
0353         return;
0354     }
0355
0356     std::sort(plots.begin(), plots.end(),
0357         [](const PlotInfo &a, const PlotInfo &b)
0358         {
0359             if (a.kind != b.kind)
0360                 return a.kind < b.kind;
0361             if (a.hasPt != b.hasPt)
0362                 return a.hasPt > b.hasPt;
0363             if (a.hasPt && std::abs(a.pt - b.pt) > 1.0e-6)
0364                 return a.pt < b.pt;
0365             return a.name < b.name;
0366         });
0367
0368     const std::vector<ZRange> ranges = getZRanges(plots);
0369     for (int kind = 0; kind < static_cast<int>(ranges.size()); ++kind)
0370     {
0371         if (!ranges[kind].initialized)
0372             continue;
0373         std::cout << "Using z-axis range [" << ranges[kind].center - ranges[kind].halfWidth << ", " << ranges[kind].center + ranges[kind].halfWidth << "] for " << mapKindName(kind) << "."
0374                   << std::endl;
0375     }
0376
0377     for (PlotInfo &plot : plots)
0378     {
0379         styleHistogram(plot, getRangeForPlot(plot, ranges), false);
0380
0381         TCanvas canvas(Form("c_%s", plot.name.c_str()), plot.title.c_str(), 800, 700);
0382         canvas.SetTicks(1, 1);
0383         canvas.SetLeftMargin(0.13);
0384         canvas.SetRightMargin(0.23);
0385         canvas.SetBottomMargin(0.13);
0386         canvas.SetTopMargin(0.08);
0387         // plot.hist->SetContour(1000);
0388         plot.hist->Draw("COLZ");
0389         // plot.hist->SetLineColor(kBlack);
0390         // plot.hist->SetLineWidth(1);
0391         // plot.hist->Draw("CONT3 same");
0392         drawPadLabel(plot, false);
0393
0394         const std::string stem = outdir + "/" + safeFileStem(plot.name);
0395         canvas.SaveAs((stem + ".png").c_str());
0396         canvas.SaveAs((stem + ".pdf").c_str());
0397     }
0398
0399     const std::pair<int, int> shape = gridShape(plots);
0400     const int ncols = shape.first;
0401     const int nrows = shape.second;
0402     TCanvas gridCanvas("c_all_scale_maps", "All scale maps", 360 * ncols, 330 * nrows);
0403     gridCanvas.Divide(ncols, nrows, 0.001, 0.001);
0404
0405     for (std::size_t i = 0; i < plots.size(); ++i)
0406     {
0407         gridCanvas.cd(static_cast<int>(i) + 1);
0408         gPad->SetTicks(1, 1);
0409         gPad->SetLeftMargin(0.12);
0410         gPad->SetRightMargin(0.18);
0411         // gPad->SetBottomMargin(0.12);
0412         gPad->SetTopMargin(0.12);
0413
0414         styleHistogram(plots[i], getRangeForPlot(plots[i], ranges), true);
0415         plots[i].hist->Draw("COLZ");
0416         drawPadLabel(plots[i], true);
0417     }
0418
0419     gridCanvas.SaveAs((outdir + "/all_scale_maps_grid.png").c_str());
0420     gridCanvas.SaveAs((outdir + "/all_scale_maps_grid.pdf").c_str());
0421
0422     std::cout << "Saved scale-map plots to " << outdir << std::endl;
0423
0424     for (PlotInfo &plot : plots)
0425         delete plot.hist;
0426 }