| 1 |
#!/usr/bin/env python
|
| 2 |
"""Overlay the histograms from several root files with identical structure"""
|
| 3 |
__version__ = "1.0"
|
| 4 |
|
| 5 |
## Created by Jeff Klukas (klukas@wisc.edu), November 2009
|
| 6 |
## Updated February 2010
|
| 7 |
|
| 8 |
######## Import python libraries #############################################
|
| 9 |
|
| 10 |
import sys
|
| 11 |
import optparse
|
| 12 |
import shutil
|
| 13 |
import os
|
| 14 |
import re
|
| 15 |
|
| 16 |
## If we actually plan to do something other than show the help menu,
|
| 17 |
## import the PyROOT package
|
| 18 |
if '-h' not in sys.argv and len(sys.argv) > 1:
|
| 19 |
import ROOT
|
| 20 |
# ROOT parses options when the first ROOT command is called, so we must
|
| 21 |
# add '-b' before that to get batch mode, but we must immediately remove
|
| 22 |
# it to avoid interference with option parsing for this script.
|
| 23 |
sys.argv.append('-b')
|
| 24 |
ROOT.gErrorIgnoreLevel = ROOT.kWarning
|
| 25 |
sys.argv.remove('-b')
|
| 26 |
|
| 27 |
|
| 28 |
|
| 29 |
######## Feel free to change these style options to suit you ################
|
| 30 |
|
| 31 |
def add_style_options(options):
|
| 32 |
"""Define a set of global variables storing style information, etc."""
|
| 33 |
GetColor = ROOT.TColor.GetColor
|
| 34 |
options.colors = [
|
| 35 |
## a default set of contrasting colors the author happens to like
|
| 36 |
GetColor( 82, 124, 219), # blue
|
| 37 |
GetColor(212, 58, 143), # red
|
| 38 |
GetColor(231, 139, 77), # orange
|
| 39 |
GetColor(145, 83, 207), # purple
|
| 40 |
GetColor(114, 173, 117), # green
|
| 41 |
GetColor( 67, 77, 83), # dark grey
|
| 42 |
]
|
| 43 |
options.marker_styles = [
|
| 44 |
## some of the more clear markers in root
|
| 45 |
3, # asterisk
|
| 46 |
4, # circle
|
| 47 |
5, # x
|
| 48 |
25, # square
|
| 49 |
26, # triangle
|
| 50 |
27, # diamond
|
| 51 |
28, # cross
|
| 52 |
30, # five-pointed star
|
| 53 |
]
|
| 54 |
return options
|
| 55 |
|
| 56 |
|
| 57 |
|
| 58 |
######## Define classes and generators #######################################
|
| 59 |
|
| 60 |
class RootFile:
|
| 61 |
"""A wrapper for TFiles, allowing quick access to the name and Get."""
|
| 62 |
def __init__(self, file_name):
|
| 63 |
self.name = file_name[0:-5]
|
| 64 |
self.file = ROOT.TFile(file_name, "read")
|
| 65 |
if self.file.IsZombie():
|
| 66 |
print "Error opening %s, exiting..." % file_name
|
| 67 |
sys.exit(1)
|
| 68 |
def Get(self, object_name):
|
| 69 |
return self.file.Get(object_name)
|
| 70 |
|
| 71 |
def counter_generator():
|
| 72 |
"""Incremement the counter used to number plots."""
|
| 73 |
k = 0
|
| 74 |
while True:
|
| 75 |
k += 1
|
| 76 |
yield k
|
| 77 |
next_counter = counter_generator().next
|
| 78 |
|
| 79 |
|
| 80 |
|
| 81 |
######## These functions are the meat of this program #########################
|
| 82 |
|
| 83 |
#### A recursive function to drill down through directories
|
| 84 |
def process_directory(path, files):
|
| 85 |
"""Loop through all histograms in the directory and plot them."""
|
| 86 |
dir_to_make = "%s/%s" % (options.plot_dir, path)
|
| 87 |
if not os.path.exists(dir_to_make):
|
| 88 |
os.mkdir(dir_to_make)
|
| 89 |
keys = files[0].file.GetDirectory(path).GetListOfKeys()
|
| 90 |
key = keys[0]
|
| 91 |
while key:
|
| 92 |
obj = key.ReadObj()
|
| 93 |
key = keys.After(key)
|
| 94 |
new_path = "%s/%s" % (path, obj.GetName())
|
| 95 |
if obj.IsA().InheritsFrom("TDirectory"):
|
| 96 |
process_directory(new_path, files)
|
| 97 |
#### If obj is a desired histogram, process it
|
| 98 |
if (options.regex.search(new_path) and
|
| 99 |
obj.IsA().InheritsFrom("TH1") and
|
| 100 |
not obj.IsA().InheritsFrom("TH2") and
|
| 101 |
not obj.IsA().InheritsFrom("TH3")):
|
| 102 |
process_hist(path, new_path, files, obj)
|
| 103 |
|
| 104 |
|
| 105 |
#### This is where all the plotting actually happens
|
| 106 |
def process_hist(path, new_path, files, obj):
|
| 107 |
"""Overlay all the instances of this plot and apply the options."""
|
| 108 |
counter = next_counter() # used for page numbers
|
| 109 |
name = obj.GetName()
|
| 110 |
hist = files[0].file.GetDirectory(path).Get(name)
|
| 111 |
title = hist.GetTitle()
|
| 112 |
x_title = hist.GetXaxis().GetTitle()
|
| 113 |
y_title = hist.GetYaxis().GetTitle()
|
| 114 |
if options.normalize or (options.sticky and "Norm" in name):
|
| 115 |
y_title = "Fraction of Events in Bin"
|
| 116 |
if options.normalize_to_file:
|
| 117 |
file_name = files[int(options.normalize_to_file) - 1].name
|
| 118 |
y_title = "Events Normalized to %s" % file_name
|
| 119 |
hists = []
|
| 120 |
#### Apply options to hist from each file
|
| 121 |
for i, file in enumerate(files):
|
| 122 |
hist = file.file.GetDirectory(path).Get(name)
|
| 123 |
if not hist: continue
|
| 124 |
hist.SetTitle(file.name)
|
| 125 |
color = options.colors[i % len(options.colors)]
|
| 126 |
hist.SetLineColor(color)
|
| 127 |
# if options.fill:
|
| 128 |
# r, g, b = plot_colors_rgb[i % len(colors)]
|
| 129 |
# #fill_color = ROOT.TColor.GetColor(r * 1.2, g * 1.2, b * 1.2)
|
| 130 |
# fill_color = color
|
| 131 |
# hist.SetFillColor(fill_color)
|
| 132 |
# hist.SetFillStyle(1001)
|
| 133 |
# print "Hist ", hist.GetFillColor()
|
| 134 |
if options.markers:
|
| 135 |
hist.SetMarkerColor(color)
|
| 136 |
hist.SetMarkerStyle(marker_styles[i])
|
| 137 |
else:
|
| 138 |
hist.SetMarkerSize(0)
|
| 139 |
if options.overflow or (options.sticky and "Overflow" in name):
|
| 140 |
nbins = hist.GetNbinsX()
|
| 141 |
overflow = hist.GetBinContent(nbins + 1)
|
| 142 |
hist.AddBinContent(nbins, overflow)
|
| 143 |
if options.underflow or (options.sticky and "Underflow" in name):
|
| 144 |
underflow = hist.GetBinContent(0)
|
| 145 |
hist.AddBinContent(1, underflow)
|
| 146 |
if options.normalize or (options.sticky and "Norm" in name):
|
| 147 |
integral = hist.Integral()
|
| 148 |
if integral: hist.Scale(1. / integral)
|
| 149 |
hists.append(hist)
|
| 150 |
if options.normalize_to_file:
|
| 151 |
integral = hists[int(options.normalize_to_file) - 1].Integral()
|
| 152 |
if integral:
|
| 153 |
for hist in hists:
|
| 154 |
hist.Scale(hist.Integral() / integral)
|
| 155 |
#### Combine hists in a THStack and draw
|
| 156 |
pads = [canvas]
|
| 157 |
stack = ROOT.THStack("st%.3i" % int(counter), title)
|
| 158 |
legend_height = 0.04 * len(files) + 0.02
|
| 159 |
legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
|
| 160 |
for hist in hists:
|
| 161 |
stack.Add(hist)
|
| 162 |
legend.AddEntry(hist)
|
| 163 |
stack.Draw(options.opt)
|
| 164 |
stack.GetXaxis().SetTitle(x_title)
|
| 165 |
stack.GetYaxis().SetTitle(y_title)
|
| 166 |
if options.ratio or (options.sticky and "Ratio" in name):
|
| 167 |
pads, stack, stack_ratio = add_ratio_plot(hists, stack, counter)
|
| 168 |
pads[1].cd()
|
| 169 |
stack_ratio.Draw(options.opt)
|
| 170 |
pads[0].cd()
|
| 171 |
pads[0].SetLogx(options.logx or (options.sticky and "Logx" in name))
|
| 172 |
pads[0].SetLogy(options.logy or (options.sticky and "Logy" in name))
|
| 173 |
stack.Draw(options.opt)
|
| 174 |
if options.numbering:
|
| 175 |
display_page_number(counter)
|
| 176 |
if options.efficiency or (options.sticky and "Eff" in name):
|
| 177 |
stack.Draw(options.opt + "e")
|
| 178 |
stack.SetMaximum(1.)
|
| 179 |
stack.SetMinimum(0.)
|
| 180 |
if options.overflow or (options.sticky and "Overflow" in name):
|
| 181 |
display_overflow(stack, hist)
|
| 182 |
if options.underflow or (options.sticky and "Underflow" in name):
|
| 183 |
display_underflow(stack, hist)
|
| 184 |
legend.Draw()
|
| 185 |
save_plot(stack, options.plot_dir, path, name, counter)
|
| 186 |
|
| 187 |
|
| 188 |
|
| 189 |
######## Define some supporting functions #####################################
|
| 190 |
|
| 191 |
def save_plot(stack, plot_dir, path, name, counter):
|
| 192 |
"""Save the canvas to the output format defined by --ext."""
|
| 193 |
output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
|
| 194 |
canvas.SaveAs(output_file_name)
|
| 195 |
if options.ext == "pdf":
|
| 196 |
numbered_pdf_name = "%.3i.pdf" % counter
|
| 197 |
shutil.copy(output_file_name, numbered_pdf_name)
|
| 198 |
report_progress(counter, 1)
|
| 199 |
|
| 200 |
def report_progress(counter, divisor):
|
| 201 |
"""Print the current number of finished plots."""
|
| 202 |
if counter % divisor == 0:
|
| 203 |
print "\r%i plots written to %s" % (counter, options.output),
|
| 204 |
sys.stdout.flush()
|
| 205 |
|
| 206 |
def merge_pdf():
|
| 207 |
"""Merge together all the produced plots into one pdf file."""
|
| 208 |
print "Writing merged pdf..."
|
| 209 |
os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
|
| 210 |
"-dAutoRotatePages=/All "
|
| 211 |
"-sOutputFile=%s.pdf " % options.output +
|
| 212 |
"[0-9][0-9][0-9].pdf")
|
| 213 |
os.system("rm [0-9]*.pdf")
|
| 214 |
|
| 215 |
def display_page_number(page_number):
|
| 216 |
"""Add a page number to the top corner of the canvas."""
|
| 217 |
page_text = ROOT.TText()
|
| 218 |
page_text.SetTextSize(0.03)
|
| 219 |
page_text.SetTextAlign(33)
|
| 220 |
page_text.DrawTextNDC(0.97, 0.985, "%i" % page_number)
|
| 221 |
|
| 222 |
def display_overflow(stack, hist):
|
| 223 |
"""Add the overflow to the last bin and print 'Overflow' on the bin."""
|
| 224 |
nbins = hist.GetNbinsX()
|
| 225 |
x = 0.5 * (hist.GetBinLowEdge(nbins) +
|
| 226 |
hist.GetBinLowEdge(nbins + 1))
|
| 227 |
y = stack.GetMinimum("nostack")
|
| 228 |
display_bin_text(x, y, nbins, "Overflow")
|
| 229 |
|
| 230 |
def display_underflow(stack, hist):
|
| 231 |
"""Add the underflow to the first bin and print 'Underflow' on the bin."""
|
| 232 |
nbins = hist.GetNbinsX()
|
| 233 |
x = 0.5 * (hist.GetBinLowEdge(1) +
|
| 234 |
hist.GetBinLowEdge(2))
|
| 235 |
y = stack.GetMinimum("nostack")
|
| 236 |
display_bin_text(x, y, nbins, "Underflow")
|
| 237 |
|
| 238 |
def display_bin_text(x, y, nbins, text):
|
| 239 |
"""Overlay TEXT on this bin."""
|
| 240 |
bin_text = ROOT.TText()
|
| 241 |
bin_text.SetTextSize(min(1. / nbins, 0.04))
|
| 242 |
bin_text.SetTextAlign(12)
|
| 243 |
bin_text.SetTextAngle(90)
|
| 244 |
bin_text.SetTextColor(13)
|
| 245 |
bin_text.SetTextFont(42)
|
| 246 |
bin_text.DrawText(x, y, text)
|
| 247 |
|
| 248 |
def add_ratio_plot(hists, stack, counter):
|
| 249 |
"""Divide canvas into two parts, and plot the ratio on the bottom."""
|
| 250 |
## Both pads are set to the full canvas size to maintain font sizes
|
| 251 |
## Fill style 4000 used to ensure pad transparency because of this
|
| 252 |
div = 0.3 # portion of canvas to use for ratio plot
|
| 253 |
margins = [ROOT.gStyle.GetPadTopMargin(), ROOT.gStyle.GetPadBottomMargin()]
|
| 254 |
useable_height = 1 - (margins[0] + margins[1])
|
| 255 |
canvas.Clear()
|
| 256 |
pad = ROOT.TPad("mainPad", "mainPad", 0., 0., 1., 1.)
|
| 257 |
pad.SetFillStyle(4000)
|
| 258 |
pad.Draw()
|
| 259 |
pad.SetBottomMargin(margins[1] + div * useable_height)
|
| 260 |
pad_ratio = ROOT.TPad("ratioPad", "ratioPad", 0., 0., 1., 1.);
|
| 261 |
pad_ratio.SetFillStyle(4000)
|
| 262 |
pad_ratio.Draw()
|
| 263 |
pad_ratio.SetTopMargin(margins[0] + (1 - div) * useable_height)
|
| 264 |
pad.cd()
|
| 265 |
stack.Draw()
|
| 266 |
stack_ratio = ROOT.THStack("stRatio%.3i" % int(counter),
|
| 267 |
";%s;Ratio" % stack.GetXaxis().GetTitle())
|
| 268 |
for hist in hists[1:]:
|
| 269 |
ratio_hist = hist.Clone()
|
| 270 |
ratio_hist.Divide(hists[0])
|
| 271 |
stack_ratio.Add(ratio_hist)
|
| 272 |
stack_ratio.Draw()
|
| 273 |
stack_ratio.GetYaxis().SetNdivisions(507) # Avoids crowded labels
|
| 274 |
stack.GetXaxis().SetBinLabel(1, "") # Don't show numbers below top plot
|
| 275 |
stack.GetXaxis().SetTitle("")
|
| 276 |
if stack.GetYaxis().GetTitle() == "":
|
| 277 |
stack.GetYaxis().SetTitle("Content")
|
| 278 |
# Avoid overlap of y-axis numbers by supressing zero
|
| 279 |
if stack.GetMinimum() / stack.GetMaximum() < 0.25:
|
| 280 |
stack.SetMinimum(stack.GetMaximum() / 10000)
|
| 281 |
return [pad, pad_ratio], stack, stack_ratio
|
| 282 |
|
| 283 |
|
| 284 |
|
| 285 |
######## Define the main program #############################################
|
| 286 |
|
| 287 |
def main():
|
| 288 |
usage="""usage: %prog [options] file1.root file2.root file3.root ...
|
| 289 |
|
| 290 |
function: overlays corresponding histograms from several files, dumping the
|
| 291 |
images into an identical directory structure in the local directory
|
| 292 |
and also merging all images into a single file (if output is pdf);
|
| 293 |
most style options can be controlled from your rootlogon.C macro"""
|
| 294 |
|
| 295 |
parser = optparse.OptionParser(usage=usage)
|
| 296 |
parser.add_option('-e', '--ext', default="pdf",
|
| 297 |
help="choose an output extension; default is pdf")
|
| 298 |
parser.add_option('-o', '--opt', default="nostack p H",
|
| 299 |
help="pass OPT to the Draw command; default is "
|
| 300 |
"'nostack p H', add 'e' for error bars")
|
| 301 |
parser.add_option('-m', '--markers', action="store_true", default=False,
|
| 302 |
help="add markers to histograms")
|
| 303 |
parser.add_option('-s', '--sticky', action="store_true", default=False,
|
| 304 |
help="enable name-based special plotting options "
|
| 305 |
"(see below)")
|
| 306 |
# parser.add_option('-f', '--fill', action="store_true", default=False,
|
| 307 |
# help="Fill histograms with a color")
|
| 308 |
parser.add_option('--output', default="overlaidHists", metavar="NAME",
|
| 309 |
help="name of output directory; default is 'overlaidHists'")
|
| 310 |
parser.add_option('--numbering', action="store_true", default=False,
|
| 311 |
help="add a page number in the upper right of each plot")
|
| 312 |
parser.add_option('--match', default="", metavar="REGEX",
|
| 313 |
help="only make plots for paths containing the specified "
|
| 314 |
"regular expression (use '.*' for wildcard)")
|
| 315 |
parser.add_option('--normalize-to-file', default="", metavar="FILENUM",
|
| 316 |
help="normalize to the FILENUMth file")
|
| 317 |
group1 = optparse.OptionGroup(
|
| 318 |
parser,
|
| 319 |
"special plotting options",
|
| 320 |
"Use the command line options given below to apply changes to all "
|
| 321 |
"plots. If you only wish to apply an option to a specific plot, "
|
| 322 |
"you can use '-s' "
|
| 323 |
"to turn on sticky keywords (such as 'Norm'). Any plot that includes "
|
| 324 |
"the given keyword in its ROOT name will have the option applied "
|
| 325 |
"regardless of its presence or absence on the command line."
|
| 326 |
)
|
| 327 |
group1.add_option('-n', '--normalize', action="store_true", default=False,
|
| 328 |
help="'Norm': area normalize the histograms")
|
| 329 |
group1.add_option('--efficiency', action="store_true", default=False,
|
| 330 |
help="'Eff' : force y axis scale to run from 0 to 1")
|
| 331 |
group1.add_option('--logx', action="store_true", default=False,
|
| 332 |
help="'Logx': force log scale for x axis")
|
| 333 |
group1.add_option('--logy', action="store_true", default=False,
|
| 334 |
help="'Logy': force log scale for y axis")
|
| 335 |
group1.add_option('--overflow', action="store_true", default=False,
|
| 336 |
help="'Overflow' : display overflow content in "
|
| 337 |
"highest bin")
|
| 338 |
group1.add_option('--underflow', action="store_true", default=False,
|
| 339 |
help="'Underflow': display underflow content in "
|
| 340 |
"lowest bin")
|
| 341 |
group1.add_option('--ratio', action="store_true", default=False,
|
| 342 |
help="'Ratio': display a ratio plot below the normal "
|
| 343 |
"plot")
|
| 344 |
parser.add_option_group(group1)
|
| 345 |
global options
|
| 346 |
options, arguments = parser.parse_args()
|
| 347 |
options.plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
|
| 348 |
options.regex = re.compile(options.match)
|
| 349 |
files = [RootFile(filename) for filename in arguments]
|
| 350 |
## if no arguments provided, just display the help message
|
| 351 |
if len(files) == 0:
|
| 352 |
parser.print_help()
|
| 353 |
sys.exit(0)
|
| 354 |
## add style options and create the canvas
|
| 355 |
options = add_style_options(options)
|
| 356 |
global canvas
|
| 357 |
canvas = ROOT.TCanvas()
|
| 358 |
## here, we decend into the files to start plotting
|
| 359 |
process_directory("", files)
|
| 360 |
print ""
|
| 361 |
if options.ext == "pdf":
|
| 362 |
merge_pdf()
|
| 363 |
|
| 364 |
|
| 365 |
if __name__ == "__main__":
|
| 366 |
main()
|
| 367 |
|
