irt/scripts/plot_irt.py

0001 #!/usr/bin/env python3
0002 """Plot the histograms produced by the ``irt_analysis`` executable.
0003
0004 Several input files (e.g. one per generated particle species) can be given, in
0005 which case the histograms are summed before plotting.
0006 """
0007
0008 import argparse
0009 import os
0010
0011 import ROOT
0012
0013 ROOT.gROOT.SetBatch(True)
0014 ROOT.TH1.AddDirectory(False)
0015 ROOT.gErrorIgnoreLevel = ROOT.kWarning
0016
0017
0018 def load_histograms(input_files):
0019     """Sum the histograms of all input files, keyed by histogram name."""
0020     histograms = {}
0021     for input_file in input_files:
0022         f = ROOT.TFile.Open(input_file)
0023         if not f or f.IsZombie():
0024             raise OSError(f"Can not open {input_file}")
0025         for key in f.GetListOfKeys():
0026             obj = key.ReadObj()
0027             if not obj.InheritsFrom("TH1"):
0028                 continue
0029             name = obj.GetName()
0030             if name in histograms:
0031                 # TH1::Add checks binning consistency
0032                 if not histograms[name].Add(obj):
0033                     raise ValueError(f"Inconsistent binning for histogram {name}")
0034             else:
0035                 histograms[name] = obj.Clone(name)
0036         f.Close()
0037     return histograms
0038
0039
0040 def save(canvas, path):
0041     canvas.SaveAs(path)
0042     canvas.Close()
0043
0044
0045 def plot_th1(hist, path, title, xlabel):
0046     canvas = ROOT.TCanvas("c", "", 800, 600)
0047     hist.SetTitle(f"{title};{xlabel};Entries")
0048     hist.SetFillColorAlpha(ROOT.kAzure + 1, 0.7)
0049     hist.SetLineColor(ROOT.kAzure + 2)
0050     hist.SetStats(False)
0051     hist.Draw("HIST")
0052     save(canvas, path)
0053
0054
0055 def plot_th2(hist, path, title, xlabel, ylabel):
0056     canvas = ROOT.TCanvas("c", "", 800, 600)
0057     canvas.SetRightMargin(0.15)
0058     if hist.GetMaximum() > 0:
0059         canvas.SetLogz()
0060     hist.SetTitle(f"{title};{xlabel};{ylabel};Entries")
0061     hist.SetStats(False)
0062     hist.Draw("COLZ")
0063     save(canvas, path)
0064
0065
0066 def plot_confusion(hist, path, title):
0067     # normalize per true species (y axis) to get identification probabilities
0068     fractions = hist.Clone("Confusion_fractions")
0069     for j in range(1, hist.GetNbinsY() + 1):
0070         total = sum(hist.GetBinContent(i, j) for i in range(1, hist.GetNbinsX() + 1))
0071         for i in range(1, hist.GetNbinsX() + 1):
0072             fractions.SetBinContent(
0073                 i, j, hist.GetBinContent(i, j) / total if total > 0 else 0.0
0074             )
0075             fractions.SetBinError(i, j, 0.0)
0076
0077     canvas = ROOT.TCanvas("c", "", 800, 700)
0078     canvas.SetRightMargin(0.15)
0079     ROOT.gStyle.SetPaintTextFormat(".2f")
0080     fractions.SetTitle(f"{title};Reconstructed PID;True PID;Fraction of true species")
0081     fractions.SetStats(False)
0082     fractions.SetMinimum(0.0)
0083     fractions.SetMaximum(1.0)
0084     fractions.GetXaxis().SetLabelSize(0.06)
0085     fractions.GetYaxis().SetLabelSize(0.06)
0086     fractions.SetMarkerSize(1.5)
0087     fractions.Draw("COLZ TEXT")
0088     save(canvas, path)
0089
0090
0091 def main():
0092     parser = argparse.ArgumentParser(description=__doc__)
0093     parser.add_argument("input_files", nargs="+", help="histogram files to plot")
0094     parser.add_argument("-o", "--output-dir", required=True, help="output directory")
0095     parser.add_argument("-t", "--title", default="", help="prefix for plot titles")
0096     args = parser.parse_args()
0097
0098     os.makedirs(args.output_dir, exist_ok=True)
0099     histograms = load_histograms(args.input_files)
0100
0101     def out(name):
0102         return os.path.join(args.output_dir, name)
0103
0104     def title(text):
0105         return f"{args.title}: {text}" if args.title else text
0106
0107     if "m_RICH_Mom" in histograms:
0108         plot_th1(
0109             histograms["m_RICH_Mom"],
0110             out("momentum.png"),
0111             title("Track momentum"),
0112             "p [GeV/c]",
0113         )
0114     if "m_RICH_ETA" in histograms:
0115         plot_th1(
0116             histograms["m_RICH_ETA"],
0117             out("eta.png"),
0118             title("Track pseudorapidity"),
0119             "#eta",
0120         )
0121
0122     for prefix, filename, xlabel, text in [
0123         ("m_Theta_Rad_", "theta_rad_", "#theta_{C} [mrad]", "Cherenkov angle, radiator"),
0124         ("m_Npe_Rad_", "npe_rad_", "N_{pe}", "Number of photoelectrons, radiator"),
0125         ("m_NHits_Rad_", "nhits_rad_", "N_{hits}", "Number of hits, radiator"),
0126     ]:
0127         for name in sorted(k for k in histograms if k.startswith(prefix)):
0128             index = name[len(prefix):]
0129             plot_th1(
0130                 histograms[name],
0131                 out(f"{filename}{index}.png"),
0132                 title(f"{text} {index}"),
0133                 xlabel,
0134             )
0135
0136     for name in sorted(k for k in histograms if k.startswith("m_ThetaMom_Rad_")):
0137         index = name[len("m_ThetaMom_Rad_"):]
0138         plot_th2(
0139             histograms[name],
0140             out(f"theta_vs_momentum_rad_{index}.png"),
0141             title(f"Cherenkov angle vs. momentum, radiator {index}"),
0142             "p [GeV/c]",
0143             "#theta_{C} [mrad]",
0144         )
0145
0146     if "Confusion" in histograms:
0147         plot_confusion(
0148             histograms["Confusion"],
0149             out("confusion_matrix.png"),
0150             title("PID confusion matrix"),
0151         )
0152
0153
0154 if __name__ == "__main__":
0155     main()