UserCode/RootMacros/overlayHists.py

#!/usr/bin/env python

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

## For more information, use the -h option:
##     ./overlayHists.py -h

## Define usage string for help option
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"""


#### Define colors and styles
rgbvals = [[82, 124, 219],
           [212,58,143],
           [231, 139, 77],
           [145, 83, 207],
           [114, 173, 117],
           [67, 77, 83]]
marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]


#### Import python libraries
import sys
import optparse
import shutil
import os
import re


#### Import ROOT in batch mode
if '-h' not in sys.argv and len(sys.argv) > 1:
    sys.argv.append('-b')
    import ROOT
    if os.path.exists('rootlogon.C'):
        ROOT.gROOT.Macro('rootlogon.C')
    else:
        os.system('echo -e "{\n}\n" >> rootlogon.C')
        ROOT.gROOT.Macro('rootlogon.C')
        os.remove('rootlogon.C')
    sys.argv.remove('-b')
    ROOT.gErrorIgnoreLevel = ROOT.kWarning
    colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
    canvas = ROOT.TCanvas()


#### Parse options
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="",
                  help="pass OPT to the Draw command (try 'e' for error bars)")
parser.add_option('-m', '--markers', action="store_true", default=False,
                  help="add markers to histograms")
parser.add_option('-s', '--special', action="store_true", default=False,
                  help="enable name-based special plotting options "
                  "(see below)")
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('--timing', action="store_true", default=False,
                  help="output timing information")
parser.add_option('--match', default="", metavar="REGEX",
                  help="only make plots for paths containing the specified "
                  "regular expression (use '.*' for wildcard)")
group1 = optparse.OptionGroup(
    parser,
    "special plotting options",
    "These options can be applied to all plots or to only specific plots "
    "you choose.  To affect all plots, simply include the option on the "
    "command line.  With the '-s' option, however, you may instead include "
    "a keyword (such as 'Norm') in the ROOT name of a histogram and it will "
    "always have that option turned on."
    )
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")
parser.add_option_group(group1)
options, arguments = parser.parse_args()
plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
regex = re.compile(options.match)


#### Define classes and utility functions
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


#### Define the main program functions
def main():
    """Initialize files and enter into process_directory loop"""
    files = [RootFile(filename) for filename in arguments]
    if len(files) == 0:
        parser.print_help()
        sys.exit(0)
    process_directory("", files)
    print ""
    if options.ext == "pdf":
        merge_pdf()


def process_directory(path, files):
    """Loop through all histograms in the directory and plot them."""
    dir_to_make = "%s/%s" % (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 (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 the plots and apply options."""
    counter = next_counter()
    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.special and "Norm" in name):
        y_title = "Fraction of Events in Bin"
    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)
    canvas.SetLogx(options.logx or (options.special and "Logx" in name))
    canvas.SetLogy(options.logy or (options.special and "Logy" in name))
    #### Fill the THStack with histograms 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 = colors[i % len(colors)]
        hist.SetLineColor(color)
        if options.markers:
            hist.SetMarkerColor(color)
            hist.SetMarkerStyle(marker_styles[i])
        else:
            hist.SetMarkerSize(0)
        if options.overflow or (options.special and "Overflow" in name):
            nbins = hist.GetNbinsX()
            overflow = hist.GetBinContent(nbins + 1)
            hist.AddBinContent(nbins, overflow)
        if options.underflow or (options.special and "Underflow" in name):
            underflow = hist.GetBinContent(0)
            hist.AddBinContent(1, underflow)
        if options.normalize or (options.special and "Norm" in name):
            integral = hist.Integral()
            if integral: hist.Scale(1. / integral)
        stack.Add(hist)
        legend.AddEntry(hist)
    #### Draw the stack and apply text overlays
    stack.Draw("nostack p H" + options.opt)
    stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
    if options.numbering:
        display_page_number(counter)
    if options.efficiency or (options.special and "Eff" in name):
        stack.Draw("nostack e p")
        stack.SetMaximum(1.)
        stack.SetMinimum(0.)
    if options.overflow or (options.special and "Overflow" in name):
        display_overflow(stack, hist)
    if options.underflow or (options.special and "Underflow" in name):
        display_underflow(stack, hist)
    legend.Draw()
    save_plot(stack, plot_dir, path, name, counter)


#### Define the extra functions used by the main routines
def save_plot(stack, plot_dir, path, name, counter):
    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):
    if counter % divisor == 0:
        print "\r%i plots written to %s" % (counter, options.output),
        sys.stdout.flush()

def merge_pdf():
    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):
    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):
    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):
    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):
    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)


#### Run main program
if __name__ == "__main__":
    if options.timing:
        import profile
        profile.run('main()', 'fooprof')
        import pstats
        p = pstats.Stats('fooprof')
        p.sort_stats('cumulative').print_stats(15)
    else:
        sys.exit(main())

Revision:	1.11
Committed:	Mon Jan 11 16:02:50 2010 UTC (15 years, 3 months ago) by klukas
Content type:	text/x-python
Branch:	MAIN
Changes since 1.10:	+5 -3 lines
Log Message:	Added safeguard in normalization for dividing by zero.
#	User	Rev	Content
1	klukas	1.1	#!/usr/bin/env python
2
3			## Created by Jeff Klukas (klukas@wisc.edu), November 2009
4	klukas	1.11	## Updated January 2010
5	klukas	1.1
6			## For more information, use the -h option:
7			## ./overlayHists.py -h
8
9			## Define usage string for help option
10			usage="""usage: %prog [options] file1.root file2.root file3.root ...
11
12	klukas	1.3	function: overlays corresponding histograms from several files, dumping the
13			images into an identical directory structure in the local directory
14	klukas	1.11	and also merging all images into a single file (if output is pdf);
15			most style options can be controlled from your rootlogon.C macro"""
16	klukas	1.1
17
18	klukas	1.10	#### Define colors and styles
19	klukas	1.4	rgbvals = [[82, 124, 219],
20	klukas	1.7	[212,58,143],
21			[231, 139, 77],
22	klukas	1.4	[145, 83, 207],
23			[114, 173, 117],
24			[67, 77, 83]]
25	klukas	1.7	marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]
26	klukas	1.1
27	klukas	1.10
28			#### Import python libraries
29	klukas	1.1	import sys
30			import optparse
31	klukas	1.9	import shutil
32	klukas	1.1	import os
33			import re
34
35	klukas	1.10
36			#### Import ROOT in batch mode
37	klukas	1.7	if '-h' not in sys.argv and len(sys.argv) > 1:
38	klukas	1.1	sys.argv.append('-b')
39			import ROOT
40	klukas	1.7	if os.path.exists('rootlogon.C'):
41			ROOT.gROOT.Macro('rootlogon.C')
42			else:
43			os.system('echo -e "{\n}\n" >> rootlogon.C')
44			ROOT.gROOT.Macro('rootlogon.C')
45			os.remove('rootlogon.C')
46	klukas	1.1	sys.argv.remove('-b')
47			ROOT.gErrorIgnoreLevel = ROOT.kWarning
48	klukas	1.4	colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
49	klukas	1.9	canvas = ROOT.TCanvas()
50	klukas	1.1
51	klukas	1.10
52			#### Parse options
53	klukas	1.1	parser = optparse.OptionParser(usage=usage)
54			parser.add_option('-e', '--ext', default="pdf",
55	klukas	1.3	help="choose an output extension; default is pdf")
56	klukas	1.9	parser.add_option('-o', '--opt', default="",
57	klukas	1.10	help="pass OPT to the Draw command (try 'e' for error bars)")
58			parser.add_option('-m', '--markers', action="store_true", default=False,
59			help="add markers to histograms")
60			parser.add_option('-s', '--special', action="store_true", default=False,
61			help="enable name-based special plotting options "
62			"(see below)")
63	klukas	1.9	parser.add_option('--output', default="overlaidHists", metavar="NAME",
64	klukas	1.3	help="name of output directory; default is 'overlaidHists'")
65	klukas	1.10	parser.add_option('--numbering', action="store_true", default=False,
66			help="add a page number in the upper right of each plot")
67			parser.add_option('--timing', action="store_true", default=False,
68			help="output timing information")
69	klukas	1.7	parser.add_option('--match', default="", metavar="REGEX",
70	klukas	1.3	help="only make plots for paths containing the specified "
71			"regular expression (use '.*' for wildcard)")
72	klukas	1.10	group1 = optparse.OptionGroup(
73			parser,
74			"special plotting options",
75			"These options can be applied to all plots or to only specific plots "
76			"you choose. To affect all plots, simply include the option on the "
77			"command line. With the '-s' option, however, you may instead include "
78			"a keyword (such as 'Norm') in the ROOT name of a histogram and it will "
79	klukas	1.11	"always have that option turned on."
80	klukas	1.10	)
81			group1.add_option('-n', '--normalize', action="store_true", default=False,
82			help="'Norm': area normalize the histograms")
83			group1.add_option('--efficiency', action="store_true", default=False,
84			help="'Eff' : force y axis scale to run from 0 to 1")
85			group1.add_option('--logx', action="store_true", default=False,
86			help="'Logx': force log scale for x axis")
87			group1.add_option('--logy', action="store_true", default=False,
88			help="'Logy': force log scale for y axis")
89			group1.add_option('--overflow', action="store_true", default=False,
90			help="'Overflow' : display overflow content in highest bin")
91			group1.add_option('--underflow', action="store_true", default=False,
92			help="'Underflow': display underflow content in lowest bin")
93			parser.add_option_group(group1)
94	klukas	1.1	options, arguments = parser.parse_args()
95			plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
96			regex = re.compile(options.match)
97
98
99	klukas	1.10	#### Define classes and utility functions
100	klukas	1.1	class RootFile:
101	klukas	1.10	"""A wrapper for TFiles, allowing quick access to the name and Get."""
102	klukas	1.1	def __init__(self, file_name):
103	klukas	1.9	self.name = file_name[0:-5]
104	klukas	1.1	self.file = ROOT.TFile(file_name, "read")
105			if self.file.IsZombie():
106			print "Error opening %s, exiting..." % file_name
107			sys.exit(1)
108			def Get(self, object_name):
109			return self.file.Get(object_name)
110
111
112	klukas	1.10	def counter_generator():
113			"""Incremement the counter used to number plots."""
114			k = 0
115			while True:
116			k += 1
117			yield k
118			next_counter = counter_generator().next
119
120	klukas	1.1
121	klukas	1.10	#### Define the main program functions
122	klukas	1.1	def main():
123	klukas	1.10	"""Initialize files and enter into process_directory loop"""
124	klukas	1.4	files = [RootFile(filename) for filename in arguments]
125	klukas	1.2	if len(files) == 0:
126			parser.print_help()
127			sys.exit(0)
128	klukas	1.1	process_directory("", files)
129	klukas	1.8	print ""
130	klukas	1.1	if options.ext == "pdf":
131	klukas	1.8	merge_pdf()
132	klukas	1.1
133
134			def process_directory(path, files):
135	klukas	1.10	"""Loop through all histograms in the directory and plot them."""
136	klukas	1.1	dir_to_make = "%s/%s" % (plot_dir, path)
137			if not os.path.exists(dir_to_make):
138			os.mkdir(dir_to_make)
139			keys = files[0].file.GetDirectory(path).GetListOfKeys()
140			key = keys[0]
141			while key:
142			obj = key.ReadObj()
143			key = keys.After(key)
144			new_path = "%s/%s" % (path, obj.GetName())
145			if obj.IsA().InheritsFrom("TDirectory"):
146			process_directory(new_path, files)
147	klukas	1.10	#### If obj is a desired histogram, process it
148	klukas	1.1	if (regex.search(new_path) and
149			obj.IsA().InheritsFrom("TH1") and
150			not obj.IsA().InheritsFrom("TH2") and
151			not obj.IsA().InheritsFrom("TH3")):
152	klukas	1.10	process_hist(path, new_path, files, obj)
153
154	klukas	1.1
155	klukas	1.10	#### This is where all the plotting actually happens
156			def process_hist(path, new_path, files, obj):
157			"""Overlay the plots and apply options."""
158			counter = next_counter()
159			name = obj.GetName()
160			hist = files[0].file.GetDirectory(path).Get(name)
161			title = hist.GetTitle()
162			x_title = hist.GetXaxis().GetTitle()
163			y_title = hist.GetYaxis().GetTitle()
164			if options.normalize or (options.special and "Norm" in name):
165			y_title = "Fraction of Events in Bin"
166			stack = ROOT.THStack("st%.3i" % int(counter), title)
167			legend_height = 0.04 * len(files) + 0.02
168			legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
169			canvas.SetLogx(options.logx or (options.special and "Logx" in name))
170			canvas.SetLogy(options.logy or (options.special and "Logy" in name))
171			#### Fill the THStack with histograms from each file
172			for i, file in enumerate(files):
173			hist = file.file.GetDirectory(path).Get(name)
174			if not hist: continue
175			hist.SetTitle(file.name)
176			color = colors[i % len(colors)]
177			hist.SetLineColor(color)
178			if options.markers:
179			hist.SetMarkerColor(color)
180			hist.SetMarkerStyle(marker_styles[i])
181			else:
182			hist.SetMarkerSize(0)
183			if options.overflow or (options.special and "Overflow" in name):
184			nbins = hist.GetNbinsX()
185			overflow = hist.GetBinContent(nbins + 1)
186			hist.AddBinContent(nbins, overflow)
187			if options.underflow or (options.special and "Underflow" in name):
188			underflow = hist.GetBinContent(0)
189			hist.AddBinContent(1, underflow)
190			if options.normalize or (options.special and "Norm" in name):
191			integral = hist.Integral()
192	klukas	1.11	if integral: hist.Scale(1. / integral)
193	klukas	1.10	stack.Add(hist)
194			legend.AddEntry(hist)
195			#### Draw the stack and apply text overlays
196			stack.Draw("nostack p H" + options.opt)
197			stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
198			if options.numbering:
199			display_page_number(counter)
200			if options.efficiency or (options.special and "Eff" in name):
201			stack.Draw("nostack e p")
202			stack.SetMaximum(1.)
203			stack.SetMinimum(0.)
204			if options.overflow or (options.special and "Overflow" in name):
205			display_overflow(stack, hist)
206			if options.underflow or (options.special and "Underflow" in name):
207			display_underflow(stack, hist)
208			legend.Draw()
209			save_plot(stack, plot_dir, path, name, counter)
210	klukas	1.1
211	klukas	1.10
212			#### Define the extra functions used by the main routines
213	klukas	1.8	def save_plot(stack, plot_dir, path, name, counter):
214	klukas	1.9	output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
215			canvas.SaveAs(output_file_name)
216	klukas	1.8	if options.ext == "pdf":
217	klukas	1.9	numbered_pdf_name = "%.3i.pdf" % counter
218			shutil.copy(output_file_name, numbered_pdf_name)
219	klukas	1.8	report_progress(counter, 1)
220
221
222			def report_progress(counter, divisor):
223			if counter % divisor == 0:
224			print "\r%i plots written to %s" % (counter, options.output),
225			sys.stdout.flush()
226
227			def merge_pdf():
228			print "Writing merged pdf..."
229			os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
230			"-dAutoRotatePages=/All "
231			"-sOutputFile=%s.pdf " % options.output +
232			"[0-9][0-9][0-9].pdf")
233			os.system("rm [0-9]*.pdf")
234
235
236	klukas	1.10	def display_page_number(page_number):
237			page_text = ROOT.TText()
238			page_text.SetTextSize(0.03)
239			page_text.SetTextAlign(33)
240			page_text.DrawTextNDC(0.97, 0.985, "%i" % page_number)
241
242
243			def display_overflow(stack, hist):
244			nbins = hist.GetNbinsX()
245			x = 0.5 * (hist.GetBinLowEdge(nbins) +
246			hist.GetBinLowEdge(nbins + 1))
247			y = stack.GetMinimum("nostack")
248			display_bin_text(x, y, nbins, "Overflow")
249
250
251			def display_underflow(stack, hist):
252			nbins = hist.GetNbinsX()
253			x = 0.5 * (hist.GetBinLowEdge(1) +
254			hist.GetBinLowEdge(2))
255			y = stack.GetMinimum("nostack")
256			display_bin_text(x, y, nbins, "Underflow")
257
258
259			def display_bin_text(x, y, nbins, text):
260			bin_text = ROOT.TText()
261			bin_text.SetTextSize(min(1. / nbins, 0.04))
262			bin_text.SetTextAlign(12)
263			bin_text.SetTextAngle(90)
264			bin_text.SetTextColor(13)
265			bin_text.SetTextFont(42)
266			bin_text.DrawText(x, y, text)
267	klukas	1.1
268
269	klukas	1.10	#### Run main program
270	klukas	1.1	if __name__ == "__main__":
271	klukas	1.8	if options.timing:
272	klukas	1.9	import profile
273			profile.run('main()', 'fooprof')
274	klukas	1.8	import pstats
275			p = pstats.Stats('fooprof')
276			p.sort_stats('cumulative').print_stats(15)
277			else:
278			sys.exit(main())
279	klukas	1.1