| 1 |
#!/usr/bin/env python
|
| 2 |
|
| 3 |
## Created by Jeff Klukas (klukas@wisc.edu), November 2009
|
| 4 |
|
| 5 |
## For more information, use the -h option:
|
| 6 |
## ./overlayHists.py -h
|
| 7 |
|
| 8 |
## Define usage string for help option
|
| 9 |
usage="""usage: %prog [options] file1.root file2.root file3.root ...
|
| 10 |
|
| 11 |
function: overlays corresponding histograms from several files, dumping the
|
| 12 |
images into an identical directory structure in the local directory
|
| 13 |
and also merging all images into a single file (if output is pdf)
|
| 14 |
|
| 15 |
naming: histograms whose names contain certain key terms will be handled
|
| 16 |
specially. Use this to your advantage!
|
| 17 |
'Eff' : y-axis will be scaled from 0 to 1
|
| 18 |
'Norm': plot will be area normalized
|
| 19 |
'Logx': x-axis will be on log scale
|
| 20 |
'Logy': y-axis will be on log scale"""
|
| 21 |
|
| 22 |
## Define colors and styles
|
| 23 |
rgbvals = [[82, 124, 219],
|
| 24 |
[212,58,143],
|
| 25 |
[231, 139, 77],
|
| 26 |
[145, 83, 207],
|
| 27 |
[114, 173, 117],
|
| 28 |
[67, 77, 83]]
|
| 29 |
marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]
|
| 30 |
|
| 31 |
## Import python libraries
|
| 32 |
import sys
|
| 33 |
import optparse
|
| 34 |
import shutil
|
| 35 |
import os
|
| 36 |
import re
|
| 37 |
|
| 38 |
## Import ROOT in batch mode
|
| 39 |
if '-h' not in sys.argv and len(sys.argv) > 1:
|
| 40 |
sys.argv.append('-b')
|
| 41 |
import ROOT
|
| 42 |
if os.path.exists('rootlogon.C'):
|
| 43 |
ROOT.gROOT.Macro('rootlogon.C')
|
| 44 |
else:
|
| 45 |
os.system('echo -e "{\n}\n" >> rootlogon.C')
|
| 46 |
ROOT.gROOT.Macro('rootlogon.C')
|
| 47 |
os.remove('rootlogon.C')
|
| 48 |
sys.argv.remove('-b')
|
| 49 |
ROOT.gErrorIgnoreLevel = ROOT.kWarning
|
| 50 |
colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
|
| 51 |
canvas = ROOT.TCanvas()
|
| 52 |
|
| 53 |
## Parse options
|
| 54 |
parser = optparse.OptionParser(usage=usage)
|
| 55 |
parser.add_option('-n', '--normalize', action="store_true", default=False,
|
| 56 |
help="area normalize all histograms")
|
| 57 |
parser.add_option('-m', '--markers', action="store_true", default=False,
|
| 58 |
help="add markers to histograms")
|
| 59 |
parser.add_option('--logx', action="store_true", default=False,
|
| 60 |
help="Force log scale for X axis on all histograms")
|
| 61 |
parser.add_option('--logy', action="store_true", default=False,
|
| 62 |
help="Force log scale for Y axis on all histograms")
|
| 63 |
parser.add_option('-e', '--ext', default="pdf",
|
| 64 |
help="choose an output extension; default is pdf")
|
| 65 |
parser.add_option('-o', '--opt', default="",
|
| 66 |
help="argument to be passed to THStack::Draw")
|
| 67 |
parser.add_option('--output', default="overlaidHists", metavar="NAME",
|
| 68 |
help="name of output directory; default is 'overlaidHists'")
|
| 69 |
parser.add_option('--match', default="", metavar="REGEX",
|
| 70 |
help="only make plots for paths containing the specified "
|
| 71 |
"regular expression (use '.*' for wildcard)")
|
| 72 |
parser.add_option('--timing', action="store_true", default=False,
|
| 73 |
help="output timing information")
|
| 74 |
options, arguments = parser.parse_args()
|
| 75 |
plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
|
| 76 |
regex = re.compile(options.match)
|
| 77 |
|
| 78 |
|
| 79 |
|
| 80 |
class RootFile:
|
| 81 |
def __init__(self, file_name):
|
| 82 |
self.name = file_name[0:-5]
|
| 83 |
self.file = ROOT.TFile(file_name, "read")
|
| 84 |
if self.file.IsZombie():
|
| 85 |
print "Error opening %s, exiting..." % file_name
|
| 86 |
sys.exit(1)
|
| 87 |
def Get(self, object_name):
|
| 88 |
return self.file.Get(object_name)
|
| 89 |
|
| 90 |
|
| 91 |
|
| 92 |
def main():
|
| 93 |
files = [RootFile(filename) for filename in arguments]
|
| 94 |
if len(files) == 0:
|
| 95 |
parser.print_help()
|
| 96 |
sys.exit(0)
|
| 97 |
process_directory("", files)
|
| 98 |
print ""
|
| 99 |
if options.ext == "pdf":
|
| 100 |
merge_pdf()
|
| 101 |
|
| 102 |
|
| 103 |
|
| 104 |
def process_directory(path, files):
|
| 105 |
dir_to_make = "%s/%s" % (plot_dir, path)
|
| 106 |
if not os.path.exists(dir_to_make):
|
| 107 |
os.mkdir(dir_to_make)
|
| 108 |
keys = files[0].file.GetDirectory(path).GetListOfKeys()
|
| 109 |
key = keys[0]
|
| 110 |
while key:
|
| 111 |
obj = key.ReadObj()
|
| 112 |
key = keys.After(key)
|
| 113 |
new_path = "%s/%s" % (path, obj.GetName())
|
| 114 |
if obj.IsA().InheritsFrom("TDirectory"):
|
| 115 |
process_directory(new_path, files)
|
| 116 |
if (regex.search(new_path) and
|
| 117 |
obj.IsA().InheritsFrom("TH1") and
|
| 118 |
not obj.IsA().InheritsFrom("TH2") and
|
| 119 |
not obj.IsA().InheritsFrom("TH3")):
|
| 120 |
counter = next_counter()
|
| 121 |
name = obj.GetName()
|
| 122 |
hist = files[0].file.GetDirectory(path).Get(name)
|
| 123 |
title = hist.GetTitle()
|
| 124 |
x_title = hist.GetXaxis().GetTitle()
|
| 125 |
y_title = hist.GetYaxis().GetTitle()
|
| 126 |
if "Norm" in name or options.normalize:
|
| 127 |
y_title = "Fraction of Events in Bin"
|
| 128 |
stack = ROOT.THStack("st%.3i" % int(counter), title)
|
| 129 |
legend_height = 0.04 * len(files) + 0.02
|
| 130 |
legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
|
| 131 |
canvas.SetLogx("Logx" in name or options.logx)
|
| 132 |
canvas.SetLogy("Logy" in name or options.logy)
|
| 133 |
for i, file in enumerate(files):
|
| 134 |
hist = file.file.GetDirectory(path).Get(name)
|
| 135 |
if not hist: continue
|
| 136 |
hist.SetTitle(file.name)
|
| 137 |
color = colors[i % len(colors)]
|
| 138 |
hist.SetLineColor(color)
|
| 139 |
if options.markers:
|
| 140 |
hist.SetMarkerColor(color)
|
| 141 |
hist.SetMarkerStyle(marker_styles[i])
|
| 142 |
else:
|
| 143 |
hist.SetMarkerSize(0)
|
| 144 |
if "Overflow" in name:
|
| 145 |
nbins = hist.GetNbinsX()
|
| 146 |
overflow = hist.GetBinContent(nbins + 1)
|
| 147 |
hist.AddBinContent(nbins, overflow)
|
| 148 |
if "Underflow" in name:
|
| 149 |
underflow = hist.GetBinContent(0)
|
| 150 |
hist.AddBinContent(1, underflow)
|
| 151 |
if "Norm" in name or options.normalize:
|
| 152 |
integral = hist.Integral()
|
| 153 |
hist.Scale(1. / integral)
|
| 154 |
stack.Add(hist)
|
| 155 |
legend.AddEntry(hist)
|
| 156 |
stack.Draw("nostack p H" + options.opt)
|
| 157 |
stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
|
| 158 |
if "Eff" in name:
|
| 159 |
stack.Draw("nostack e p")
|
| 160 |
stack.SetMaximum(1.)
|
| 161 |
stack.SetMinimum(0.)
|
| 162 |
if "Overflow" in name:
|
| 163 |
nbins = hist.GetNbinsX()
|
| 164 |
x = 0.5 * (hist.GetBinLowEdge(nbins) +
|
| 165 |
hist.GetBinLowEdge(nbins + 1))
|
| 166 |
y = stack.GetMinimum("nostack")
|
| 167 |
overflow_text = ROOT.TText(x, y, "Overflow")
|
| 168 |
overflow_text.SetTextSize(min(1. / nbins, 0.04))
|
| 169 |
overflow_text.SetTextAlign(12)
|
| 170 |
overflow_text.SetTextAngle(90)
|
| 171 |
overflow_text.SetTextColor(13)
|
| 172 |
overflow_text.SetTextFont(42)
|
| 173 |
overflow_text.Draw()
|
| 174 |
if "Underflow" in name:
|
| 175 |
nbins = hist.GetNbinsX()
|
| 176 |
x = hist.GetBinLowEdge(1)
|
| 177 |
y = hist.GetMinimum()
|
| 178 |
underflow_text = ROOT.TText(x, y, "Underflow")
|
| 179 |
underflow_text.SetTextSize(1. / nbins)
|
| 180 |
underflow_text.SetTextAlign(13)
|
| 181 |
underflow_text.SetTextAngle(90)
|
| 182 |
underflow_text.SetTextColor(13)
|
| 183 |
underflow_text.Draw()
|
| 184 |
legend.Draw()
|
| 185 |
save_plot(stack, plot_dir, path, name, counter)
|
| 186 |
|
| 187 |
|
| 188 |
|
| 189 |
def save_plot(stack, plot_dir, path, name, counter):
|
| 190 |
output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
|
| 191 |
canvas.SaveAs(output_file_name)
|
| 192 |
if options.ext == "pdf":
|
| 193 |
numbered_pdf_name = "%.3i.pdf" % counter
|
| 194 |
shutil.copy(output_file_name, numbered_pdf_name)
|
| 195 |
report_progress(counter, 1)
|
| 196 |
|
| 197 |
|
| 198 |
|
| 199 |
def report_progress(counter, divisor):
|
| 200 |
if counter % divisor == 0:
|
| 201 |
print "\r%i plots written to %s" % (counter, options.output),
|
| 202 |
sys.stdout.flush()
|
| 203 |
|
| 204 |
|
| 205 |
|
| 206 |
def merge_pdf():
|
| 207 |
print "Writing merged pdf..."
|
| 208 |
os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
|
| 209 |
"-dAutoRotatePages=/All "
|
| 210 |
"-sOutputFile=%s.pdf " % options.output +
|
| 211 |
"[0-9][0-9][0-9].pdf")
|
| 212 |
os.system("rm [0-9]*.pdf")
|
| 213 |
|
| 214 |
|
| 215 |
|
| 216 |
def counter_generator():
|
| 217 |
k = 0
|
| 218 |
while True:
|
| 219 |
k += 1
|
| 220 |
yield k
|
| 221 |
next_counter = counter_generator().next
|
| 222 |
|
| 223 |
|
| 224 |
|
| 225 |
if __name__ == "__main__":
|
| 226 |
if options.timing:
|
| 227 |
## import cProfile
|
| 228 |
## cProfile.run('main()', 'fooprof')
|
| 229 |
import profile
|
| 230 |
profile.run('main()', 'fooprof')
|
| 231 |
import pstats
|
| 232 |
p = pstats.Stats('fooprof')
|
| 233 |
p.sort_stats('cumulative').print_stats(15)
|
| 234 |
else:
|
| 235 |
sys.exit(main())
|
| 236 |
|
