UserCode/RootMacros/overlayHists.py

#!/usr/bin/env python
"""Overlay the histograms from several root files with identical structure"""
__version__ = "1.0"

## Created by Jeff Klukas (klukas@wisc.edu), November 2009
## Updated February 2010

######## Import python libraries #############################################

import sys
import optparse
import shutil
import os
import re

## If we actually plan to do something other than show the help menu, 
## import the PyROOT package
if '-h' not in sys.argv and len(sys.argv) > 1:
    import ROOT
    # ROOT parses options when the first ROOT command is called, so we must
    # add '-b' before that to get batch mode, but we must immediately remove
    # it to avoid interference with option parsing for this script.
    sys.argv.append('-b')
    ROOT.gErrorIgnoreLevel = ROOT.kWarning
    sys.argv.remove('-b')


######## Feel free to change these style options to suit you  ################

def add_style_options(options):
    """Define a set of global variables storing style information, etc."""
    GetColor = ROOT.TColor.GetColor
    options.colors = [
        ## a default set of contrasting colors the author happens to like
        GetColor( 82, 124, 219), # blue
        GetColor(212,  58, 143), # red
        GetColor(231, 139,  77), # orange
        GetColor(145,  83, 207), # purple
        GetColor(114, 173, 117), # green
        GetColor( 67,  77,  83), # dark grey
        ]
    options.marker_styles = [
        ## some of the more clear markers in root
         3, # asterisk
         4, # circle 
         5, # x
        25, # square
        26, # triangle
        27, # diamond
        28, # cross
        30, # five-pointed star
        ]
    return options


######## Define classes and generators #######################################

class RootFile:
    """A wrapper for TFiles, allowing quick access to the name and Get."""
    def __init__(self, file_name):
        self.name = file_name[0:-5]
        self.file = ROOT.TFile(file_name, "read")
        if self.file.IsZombie():
            print "Error opening %s, exiting..." % file_name
            sys.exit(1)
    def Get(self, object_name):
        return self.file.Get(object_name)

def counter_generator():
    """Incremement the counter used to number plots."""
    k = 0
    while True:
        k += 1
        yield k
next_counter = counter_generator().next


######## These functions are the meat of this program #########################

#### A recursive function to drill down through directories
def process_directory(path, files):
    """Loop through all histograms in the directory and plot them."""
    dir_to_make = "%s/%s" % (options.plot_dir, path)
    if not os.path.exists(dir_to_make):
        os.mkdir(dir_to_make)
    keys = files[0].file.GetDirectory(path).GetListOfKeys()
    key = keys[0]
    while key:
        obj = key.ReadObj()
        key = keys.After(key)
        new_path = "%s/%s" % (path, obj.GetName())
        if obj.IsA().InheritsFrom("TDirectory"):
            process_directory(new_path, files)
        #### If obj is a desired histogram, process it
        if (options.regex.search(new_path) and
            obj.IsA().InheritsFrom("TH1") and
            not obj.IsA().InheritsFrom("TH2") and
            not obj.IsA().InheritsFrom("TH3")):
            process_hist(path, new_path, files, obj)


#### This is where all the plotting actually happens
def process_hist(path, new_path, files, obj):
    """Overlay all the instances of this plot and apply the options."""
    counter = next_counter() # used for page numbers
    name = obj.GetName()
    hist = files[0].file.GetDirectory(path).Get(name)
    title = hist.GetTitle()
    x_title = hist.GetXaxis().GetTitle()
    y_title = hist.GetYaxis().GetTitle()
    if options.normalize or (options.sticky and "Norm" in name):
        y_title = "Fraction of Events in Bin"
    if options.normalize_to_file:
        file_name = files[int(options.normalize_to_file) - 1].name
        y_title = "Events Normalized to %s" % file_name
    hists = []
    #### Apply options to hist from each file
    for i, file in enumerate(files):
        hist = file.file.GetDirectory(path).Get(name)
        if not hist: continue
        hist.SetTitle(file.name)
        color = options.colors[i % len(options.colors)]
        hist.SetLineColor(color)
#         if options.fill:
#             r, g, b = plot_colors_rgb[i % len(colors)]
#             #fill_color = ROOT.TColor.GetColor(r * 1.2, g * 1.2, b * 1.2)
#             fill_color = color
#             hist.SetFillColor(fill_color)
#             hist.SetFillStyle(1001)
#             print "Hist ", hist.GetFillColor()
        if options.markers:
            hist.SetMarkerColor(color)
            hist.SetMarkerStyle(marker_styles[i])
        else:
            hist.SetMarkerSize(0)
        if options.overflow or (options.sticky and "Overflow" in name):
            nbins = hist.GetNbinsX()
            overflow = hist.GetBinContent(nbins + 1)
            hist.AddBinContent(nbins, overflow)
        if options.underflow or (options.sticky and "Underflow" in name):
            underflow = hist.GetBinContent(0)
            hist.AddBinContent(1, underflow)
        if options.normalize or (options.sticky and "Norm" in name):
            integral = hist.Integral()
            if integral: hist.Scale(1. / integral)
        hists.append(hist)
    if options.normalize_to_file:
        integral = hists[int(options.normalize_to_file) - 1].Integral()
        if integral:
            for hist in hists:
                hist.Scale(hist.Integral() / integral)
    #### Combine hists in a THStack and draw
    pads = [canvas]
    stack = ROOT.THStack("st%.3i" % int(counter), title)
    legend_height = 0.04 * len(files) + 0.02
    legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
    for hist in hists:
        stack.Add(hist)
        legend.AddEntry(hist)
    stack.Draw(options.opt)
    stack.GetXaxis().SetTitle(x_title)
    stack.GetYaxis().SetTitle(y_title)
    if options.ratio or (options.sticky and "Ratio" in name):
        pads, stack, stack_ratio = add_ratio_plot(hists, stack, counter)
        pads[1].cd()
        stack_ratio.Draw(options.opt)
        pads[0].cd()
    pads[0].SetLogx(options.logx or (options.sticky and "Logx" in name))
    pads[0].SetLogy(options.logy or (options.sticky and "Logy" in name))
    stack.Draw(options.opt)
    if options.numbering:
        display_page_number(counter)
    if options.efficiency or (options.sticky and "Eff" in name):
        stack.Draw(options.opt + "e")
        stack.SetMaximum(1.)
        stack.SetMinimum(0.)
    if options.overflow or (options.sticky and "Overflow" in name):
        display_overflow(stack, hist)
    if options.underflow or (options.sticky and "Underflow" in name):
        display_underflow(stack, hist)
    legend.Draw()
    save_plot(stack, options.plot_dir, path, name, counter)


######## Define some supporting functions #####################################

def save_plot(stack, plot_dir, path, name, counter):
    """Save the canvas to the output format defined by --ext."""
    output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
    canvas.SaveAs(output_file_name)
    if options.ext == "pdf":
        numbered_pdf_name = "%.3i.pdf" % counter
        shutil.copy(output_file_name, numbered_pdf_name)
    report_progress(counter, 1)

def report_progress(counter, divisor):
    """Print the current number of finished plots."""
    if counter % divisor == 0:
        print "\r%i plots written to %s" % (counter, options.output),
        sys.stdout.flush()

def merge_pdf():
    """Merge together all the produced plots into one pdf file."""
    print "Writing merged pdf..."
    os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
              "-dAutoRotatePages=/All "
              "-sOutputFile=%s.pdf " % options.output +
              "[0-9][0-9][0-9].pdf")
    os.system("rm [0-9]*.pdf")

def display_page_number(page_number):
    """Add a page number to the top corner of the canvas."""
    page_text = ROOT.TText()
    page_text.SetTextSize(0.03)
    page_text.SetTextAlign(33)
    page_text.DrawTextNDC(0.97, 0.985, "%i" % page_number)

def display_overflow(stack, hist):
    """Add the overflow to the last bin and print 'Overflow' on the bin."""
    nbins = hist.GetNbinsX()
    x = 0.5 * (hist.GetBinLowEdge(nbins) +
               hist.GetBinLowEdge(nbins + 1))
    y = stack.GetMinimum("nostack")
    display_bin_text(x, y, nbins, "Overflow")

def display_underflow(stack, hist):
    """Add the underflow to the first bin and print 'Underflow' on the bin."""
    nbins = hist.GetNbinsX()
    x = 0.5 * (hist.GetBinLowEdge(1) +
               hist.GetBinLowEdge(2))
    y = stack.GetMinimum("nostack")
    display_bin_text(x, y, nbins, "Underflow")

def display_bin_text(x, y, nbins, text):
    """Overlay TEXT on this bin."""
    bin_text = ROOT.TText()
    bin_text.SetTextSize(min(1. / nbins, 0.04))
    bin_text.SetTextAlign(12)
    bin_text.SetTextAngle(90)
    bin_text.SetTextColor(13)
    bin_text.SetTextFont(42)
    bin_text.DrawText(x, y, text)

def add_ratio_plot(hists, stack, counter):
    """Divide canvas into two parts, and plot the ratio on the bottom."""
    ## Both pads are set to the full canvas size to maintain font sizes
    ## Fill style 4000 used to ensure pad transparency because of this
    div = 0.3 # portion of canvas to use for ratio plot
    margins = [ROOT.gStyle.GetPadTopMargin(), ROOT.gStyle.GetPadBottomMargin()]
    useable_height = 1 - (margins[0] + margins[1])
    canvas.Clear()
    pad = ROOT.TPad("mainPad", "mainPad", 0., 0., 1., 1.)
    pad.SetFillStyle(4000)
    pad.Draw()
    pad.SetBottomMargin(margins[1] + div * useable_height)
    pad_ratio = ROOT.TPad("ratioPad", "ratioPad", 0., 0., 1., 1.);
    pad_ratio.SetFillStyle(4000)
    pad_ratio.Draw()
    pad_ratio.SetTopMargin(margins[0] + (1 - div) * useable_height)
    pad.cd()
    stack.Draw()
    stack_ratio = ROOT.THStack("stRatio%.3i" % int(counter),
                               ";%s;Ratio" % stack.GetXaxis().GetTitle())
    for hist in hists[1:]:
        ratio_hist = hist.Clone()
        ratio_hist.Divide(hists[0])
        stack_ratio.Add(ratio_hist)
    stack_ratio.Draw()
    stack_ratio.GetYaxis().SetNdivisions(507) # Avoids crowded labels
    stack.GetXaxis().SetBinLabel(1, "") # Don't show numbers below top plot
    stack.GetXaxis().SetTitle("")
    if stack.GetYaxis().GetTitle() == "":
        stack.GetYaxis().SetTitle("Content")
    # Avoid overlap of y-axis numbers by supressing zero
    if stack.GetMinimum() / stack.GetMaximum() < 0.25:
        stack.SetMinimum(stack.GetMaximum() / 10000)
    return [pad, pad_ratio], stack, stack_ratio


######## Define the main program #############################################

def main():
    usage="""usage: %prog [options] file1.root file2.root file3.root ...

    function: overlays corresponding histograms from several files, dumping the
          images into an identical directory structure in the local directory
          and also merging all images into a single file (if output is pdf);
          most style options can be controlled from your rootlogon.C macro"""
    
    parser = optparse.OptionParser(usage=usage)
    parser.add_option('-e', '--ext', default="pdf", 
                      help="choose an output extension; default is pdf")
    parser.add_option('-o', '--opt', default="nostack p H",
                      help="pass OPT to the Draw command; default is "
                      "'nostack p H', add 'e' for error bars")
    parser.add_option('-m', '--markers', action="store_true", default=False,
                      help="add markers to histograms")
    parser.add_option('-s', '--sticky', action="store_true", default=False,
                      help="enable name-based special plotting options "
                      "(see below)")
#     parser.add_option('-f', '--fill', action="store_true", default=False,
#                       help="Fill histograms with a color")
    parser.add_option('--output', default="overlaidHists", metavar="NAME",
                      help="name of output directory; default is 'overlaidHists'")
    parser.add_option('--numbering', action="store_true", default=False,
                      help="add a page number in the upper right of each plot")
    parser.add_option('--match', default="", metavar="REGEX",
                      help="only make plots for paths containing the specified "
                      "regular expression (use '.*' for wildcard)")
    parser.add_option('--normalize-to-file', default="", metavar="FILENUM",
                      help="normalize to the FILENUMth file")
    group1 = optparse.OptionGroup(
        parser,
        "special plotting options",
        "Use the command line options given below to apply changes to all "
        "plots.  If you only wish to apply an option to a specific plot, "
        "you can use '-s' "
        "to turn on sticky keywords (such as 'Norm').  Any plot that includes "
        "the given keyword in its ROOT name will have the option applied "
        "regardless of its presence or absence on the command line."
        )
    group1.add_option('-n', '--normalize', action="store_true", default=False,
                      help="'Norm': area normalize the histograms")
    group1.add_option('--efficiency', action="store_true", default=False,
                      help="'Eff' : force y axis scale to run from 0 to 1")
    group1.add_option('--logx', action="store_true", default=False,
                      help="'Logx': force log scale for x axis")
    group1.add_option('--logy', action="store_true", default=False,
                      help="'Logy': force log scale for y axis")
    group1.add_option('--overflow', action="store_true", default=False,
                      help="'Overflow' : display overflow content in "
                      "highest bin")
    group1.add_option('--underflow', action="store_true", default=False,
                      help="'Underflow': display underflow content in "
                      "lowest bin")
    group1.add_option('--ratio', action="store_true", default=False,
                      help="'Ratio': display a ratio plot below the normal "
                      "plot")
    parser.add_option_group(group1)
    global options
    options, arguments = parser.parse_args()
    options.plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
    options.regex = re.compile(options.match)
    files = [RootFile(filename) for filename in arguments]
    ## if no arguments provided, just display the help message
    if len(files) == 0:
        parser.print_help()
        sys.exit(0)
    ## add style options and create the canvas
    options = add_style_options(options)
    global canvas
    canvas = ROOT.TCanvas()
    ## here, we decend into the files to start plotting
    process_directory("", files)
    print ""
    if options.ext == "pdf":
        merge_pdf()


if __name__ == "__main__":
    main()

Revision:	1.13
Committed:	Wed Feb 3 19:09:56 2010 UTC (15 years, 2 months ago) by klukas
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	HEAD
Changes since 1.12:	+169 -130 lines
Log Message:	Massive clean-up and documentation; added normalize-to-file option; removed timing option due to reorganization
#	User	Rev	Content
1	klukas	1.1	#!/usr/bin/env python
2	klukas	1.13	"""Overlay the histograms from several root files with identical structure"""
3			__version__ = "1.0"
4	klukas	1.1
5			## Created by Jeff Klukas (klukas@wisc.edu), November 2009
6	klukas	1.13	## Updated February 2010
7	klukas	1.1
8	klukas	1.13	######## Import python libraries #############################################
9	klukas	1.1
10			import sys
11			import optparse
12	klukas	1.9	import shutil
13	klukas	1.1	import os
14			import re
15
16	klukas	1.13	## If we actually plan to do something other than show the help menu,
17			## import the PyROOT package
18	klukas	1.7	if '-h' not in sys.argv and len(sys.argv) > 1:
19	klukas	1.13	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	klukas	1.1	sys.argv.append('-b')
24	klukas	1.13	ROOT.gErrorIgnoreLevel = ROOT.kWarning
25	klukas	1.1	sys.argv.remove('-b')
26
27	klukas	1.10
28	klukas	1.1
29	klukas	1.13	######## 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	klukas	1.1
60			class RootFile:
61	klukas	1.10	"""A wrapper for TFiles, allowing quick access to the name and Get."""
62	klukas	1.1	def __init__(self, file_name):
63	klukas	1.9	self.name = file_name[0:-5]
64	klukas	1.1	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	klukas	1.10	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	klukas	1.1
80
81	klukas	1.13	######## These functions are the meat of this program #########################
82	klukas	1.1
83	klukas	1.13	#### A recursive function to drill down through directories
84	klukas	1.1	def process_directory(path, files):
85	klukas	1.10	"""Loop through all histograms in the directory and plot them."""
86	klukas	1.13	dir_to_make = "%s/%s" % (options.plot_dir, path)
87	klukas	1.1	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	klukas	1.10	#### If obj is a desired histogram, process it
98	klukas	1.13	if (options.regex.search(new_path) and
99	klukas	1.1	obj.IsA().InheritsFrom("TH1") and
100			not obj.IsA().InheritsFrom("TH2") and
101			not obj.IsA().InheritsFrom("TH3")):
102	klukas	1.10	process_hist(path, new_path, files, obj)
103
104	klukas	1.1
105	klukas	1.10	#### This is where all the plotting actually happens
106			def process_hist(path, new_path, files, obj):
107	klukas	1.13	"""Overlay all the instances of this plot and apply the options."""
108	klukas	1.12	counter = next_counter() # used for page numbers
109	klukas	1.10	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	klukas	1.13	if options.normalize or (options.sticky and "Norm" in name):
115	klukas	1.10	y_title = "Fraction of Events in Bin"
116	klukas	1.13	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	klukas	1.12	hists = []
120			#### Apply options to hist from each file
121	klukas	1.10	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	klukas	1.13	color = options.colors[i % len(options.colors)]
126	klukas	1.10	hist.SetLineColor(color)
127	klukas	1.13	# 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	klukas	1.10	if options.markers:
135			hist.SetMarkerColor(color)
136			hist.SetMarkerStyle(marker_styles[i])
137			else:
138			hist.SetMarkerSize(0)
139	klukas	1.13	if options.overflow or (options.sticky and "Overflow" in name):
140	klukas	1.10	nbins = hist.GetNbinsX()
141			overflow = hist.GetBinContent(nbins + 1)
142			hist.AddBinContent(nbins, overflow)
143	klukas	1.13	if options.underflow or (options.sticky and "Underflow" in name):
144	klukas	1.10	underflow = hist.GetBinContent(0)
145			hist.AddBinContent(1, underflow)
146	klukas	1.13	if options.normalize or (options.sticky and "Norm" in name):
147	klukas	1.10	integral = hist.Integral()
148	klukas	1.11	if integral: hist.Scale(1. / integral)
149	klukas	1.12	hists.append(hist)
150	klukas	1.13	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	klukas	1.12	#### 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	klukas	1.10	stack.Add(hist)
162			legend.AddEntry(hist)
163	klukas	1.13	stack.Draw(options.opt)
164	klukas	1.12	stack.GetXaxis().SetTitle(x_title)
165			stack.GetYaxis().SetTitle(y_title)
166	klukas	1.13	if options.ratio or (options.sticky and "Ratio" in name):
167	klukas	1.12	pads, stack, stack_ratio = add_ratio_plot(hists, stack, counter)
168			pads[1].cd()
169	klukas	1.13	stack_ratio.Draw(options.opt)
170	klukas	1.12	pads[0].cd()
171	klukas	1.13	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	klukas	1.10	if options.numbering:
175			display_page_number(counter)
176	klukas	1.13	if options.efficiency or (options.sticky and "Eff" in name):
177			stack.Draw(options.opt + "e")
178	klukas	1.10	stack.SetMaximum(1.)
179			stack.SetMinimum(0.)
180	klukas	1.13	if options.overflow or (options.sticky and "Overflow" in name):
181	klukas	1.10	display_overflow(stack, hist)
182	klukas	1.13	if options.underflow or (options.sticky and "Underflow" in name):
183	klukas	1.10	display_underflow(stack, hist)
184			legend.Draw()
185	klukas	1.13	save_plot(stack, options.plot_dir, path, name, counter)
186
187	klukas	1.1
188	klukas	1.10
189	klukas	1.13	######## Define some supporting functions #####################################
190
191	klukas	1.8	def save_plot(stack, plot_dir, path, name, counter):
192	klukas	1.13	"""Save the canvas to the output format defined by --ext."""
193	klukas	1.9	output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
194			canvas.SaveAs(output_file_name)
195	klukas	1.8	if options.ext == "pdf":
196	klukas	1.9	numbered_pdf_name = "%.3i.pdf" % counter
197			shutil.copy(output_file_name, numbered_pdf_name)
198	klukas	1.8	report_progress(counter, 1)
199
200			def report_progress(counter, divisor):
201	klukas	1.13	"""Print the current number of finished plots."""
202	klukas	1.8	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	klukas	1.13	"""Merge together all the produced plots into one pdf file."""
208	klukas	1.8	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	klukas	1.10	def display_page_number(page_number):
216	klukas	1.13	"""Add a page number to the top corner of the canvas."""
217	klukas	1.10	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	klukas	1.13	"""Add the overflow to the last bin and print 'Overflow' on the bin."""
224	klukas	1.10	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	klukas	1.13	"""Add the underflow to the first bin and print 'Underflow' on the bin."""
232	klukas	1.10	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	klukas	1.13	"""Overlay TEXT on this bin."""
240	klukas	1.10	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	klukas	1.1
248	klukas	1.12	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	klukas	1.13
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	klukas	1.1	if __name__ == "__main__":
366	klukas	1.13	main()
367	klukas	1.1