ViewVC Help

View File | Revision Log | Show Annotations | Root Listing

root/cvsroot/UserCode/RootMacros/overlayHists.py

Comparing UserCode/RootMacros/overlayHists.py (file contents):
Revision 1.3 by klukas, Thu Nov 19 19:32:52 2009 UTC vs.
Revision 1.12 by klukas, Fri Jan 15 04:20:33 2010 UTC

#	Line 1 | Line 1
1		#!/usr/bin/env python
2
3		## Created by Jeff Klukas (klukas@wisc.edu), November 2009
4	+	## Updated January 2010
5
6		## For more information, use the -h option:
7		##     ./overlayHists.py -h
#	Line 10 | Line 11 | usage="""usage: %prog [options] file1.ro
11
12		function: overlays corresponding histograms from several files, dumping the
13		images into an identical directory structure in the local directory
14	<	and also merging all images into a single file (if output is pdf)
14	>	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
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
23	–	rgbcolors = [[82, 124, 219],
24	–	[145, 83, 207],
25	–	[231, 139, 77],
26	–	[114, 173, 117],
27	–	[67, 77, 83]]
17
18	<	## Import python libraries
18	>	#### Define colors and styles
19	>	plot_colors_rgb = [[82, 124, 219],
20	>	[212,58,143],
21	>	[231, 139, 77],
22	>	[145, 83, 207],
23	>	[114, 173, 117],
24	>	[67, 77, 83]]
25	>	marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]
26	>
27	>
28	>	#### Import python libraries
29		import sys
30		import optparse
31	+	import shutil
32		import os
33		import re
34
35	<	## Import ROOT in batch mode
36	<	if '-h' not in sys.argv:
35	>
36	>	#### Import ROOT in batch mode
37	>	if '-h' not in sys.argv and len(sys.argv) > 1:
38		sys.argv.append('-b')
39		import ROOT
40	<	if os.path.exists('rootlogon.C'): ROOT.gROOT.Macro('rootlogon.C')
40	>	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		sys.argv.remove('-b')
47		ROOT.gErrorIgnoreLevel = ROOT.kWarning
48	<	colors = []
49	<	for rgb in rgbcolors:
44	<	colors.append(ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]))
45	<	c1 = ROOT.TCanvas()
48	>	colors = [ROOT.TColor.GetColor(r,g,b) for r,g,b in plot_colors_rgb]
49	>	canvas = ROOT.TCanvas()
50
51	<	## Parse options
51	>
52	>	#### Parse options
53		parser = optparse.OptionParser(usage=usage)
49	–	parser.add_option('-n', '--normalize', action="store_true", default=False,
50	–	help="area normalize all histograms")
54		parser.add_option('-e', '--ext', default="pdf",
55		help="choose an output extension; default is pdf")
56	<	parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
56	>	parser.add_option('-o', '--opt', default="",
57	>	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	>	parser.add_option('--output', default="overlaidHists", metavar="NAME",
64		help="name of output directory; default is 'overlaidHists'")
65	<	parser.add_option('-m', '--match', default="", metavar="REGEX",
65	>	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	>	parser.add_option('--match', default="", metavar="REGEX",
70		help="only make plots for paths containing the specified "
71		"regular expression (use '.*' for wildcard)")
72	+	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	+	"always have that option turned on."
80	+	)
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	+	group1.add_option('--ratio', action="store_true", default=False,
94	+	help="'Ratio': display a ratio plot below the normal plot")
95	+	parser.add_option_group(group1)
96		options, arguments = parser.parse_args()
97		plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
98		regex = re.compile(options.match)
99
100
101	+	#### Define classes and utility functions
102		class RootFile:
103	+	"""A wrapper for TFiles, allowing quick access to the name and Get."""
104		def __init__(self, file_name):
105	<	self.name = file_name[0:file_name.find(".root")]
105	>	self.name = file_name[0:-5]
106		self.file = ROOT.TFile(file_name, "read")
107		if self.file.IsZombie():
108		print "Error opening %s, exiting..." % file_name
#	Line 71 | Line 111 | class RootFile:
111		return self.file.Get(object_name)
112
113
114	+	def counter_generator():
115	+	"""Incremement the counter used to number plots."""
116	+	k = 0
117	+	while True:
118	+	k += 1
119	+	yield k
120	+	next_counter = counter_generator().next
121	+
122
123	+	#### Define the main program functions
124		def main():
125	<	files = []
126	<	for filename in arguments: files.append(RootFile(filename))
125	>	"""Initialize files and enter into process_directory loop"""
126	>	files = [RootFile(filename) for filename in arguments]
127		if len(files) == 0:
128		parser.print_help()
129		sys.exit(0)
130		process_directory("", files)
131	+	print ""
132		if options.ext == "pdf":
133	<	os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
84	<	"-dAutoRotatePages=/All "
85	<	"-sOutputFile=%s.pdf " % options.output +
86	<	"[0-9][0-9][0-9].pdf")
87	<	os.system("rm [0-9]*.pdf")
88	<	print "Wrote %i plots to %s" % (next_counter() - 1, options.output)
89	<
133	>	merge_pdf()
134
135
136		def process_directory(path, files):
137	+	"""Loop through all histograms in the directory and plot them."""
138		dir_to_make = "%s/%s" % (plot_dir, path)
139		if not os.path.exists(dir_to_make):
140		os.mkdir(dir_to_make)
#	Line 101 | Line 146 | def process_directory(path, files):
146		new_path = "%s/%s" % (path, obj.GetName())
147		if obj.IsA().InheritsFrom("TDirectory"):
148		process_directory(new_path, files)
149	+	#### If obj is a desired histogram, process it
150		if (regex.search(new_path) and
151		obj.IsA().InheritsFrom("TH1") and
152		not obj.IsA().InheritsFrom("TH2") and
153		not obj.IsA().InheritsFrom("TH3")):
154	<	counter = next_counter()
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 "Norm" in name or options.normalize:
115	<	y_title = "Fraction of Events in Bin"
116	<	hist.Draw()
117	<	hists = []
118	<	stack = ROOT.THStack("st%.3i" % int(counter), title)
119	<	legend = ROOT.TLegend(0.65, 0.77, 0.87, 0.89)
120	<	c1.SetLogx("Logx" in name)
121	<	c1.SetLogy("Logy" in name)
122	<	for i, file in enumerate(files):
123	<	hist = file.file.GetDirectory(path).Get(name)
124	<	if not hist: continue
125	<	hist.Draw()
126	<	hist.SetTitle(file.name)
127	<	color = colors[i % len(colors)]
128	<	hist.SetLineColor(color)
129	<	hist.SetMarkerColor(color)
130	<	hist.SetMarkerStyle(i + 1)
131	<	if "Norm" in name or options.normalize:
132	<	integral = hist.Integral()
133	<	hist.Scale(1 / integral)
134	<	stack.Add(hist)
135	<	legend.AddEntry(hist)
136	<	stack.Draw("nostack p H")
137	<	stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
138	<	if "Eff" in name:
139	<	stack.Draw("nostack e p")
140	<	stack.SetMaximum(1.)
141	<	stack.SetMinimum(0.)
142	<	legend.Draw()
143	<	if options.ext == "pdf":
144	<	c1.SaveAs("%.3i.pdf" % counter)
145	<	c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext))
146	<
147	<
148	<
149	<
150	<	def counter_generator():
151	<	k = 0
152	<	while True:
153	<	k += 1
154	<	yield k
155	<	next_counter = counter_generator().next
154	>	process_hist(path, new_path, files, obj)
155
156
157	+	#### This is where all the plotting actually happens
158	+	def process_hist(path, new_path, files, obj):
159	+	"""Overlay the plots and apply options."""
160	+	counter = next_counter() # used for page numbers
161	+	name = obj.GetName()
162	+	hist = files[0].file.GetDirectory(path).Get(name)
163	+	title = hist.GetTitle()
164	+	x_title = hist.GetXaxis().GetTitle()
165	+	y_title = hist.GetYaxis().GetTitle()
166	+	if options.normalize or (options.special and "Norm" in name):
167	+	y_title = "Fraction of Events in Bin"
168	+	hists = []
169	+	#### Apply options to hist from each file
170	+	for i, file in enumerate(files):
171	+	hist = file.file.GetDirectory(path).Get(name)
172	+	if not hist: continue
173	+	hist.SetTitle(file.name)
174	+	color = colors[i % len(colors)]
175	+	hist.SetLineColor(color)
176	+	if options.markers:
177	+	hist.SetMarkerColor(color)
178	+	hist.SetMarkerStyle(marker_styles[i])
179	+	else:
180	+	hist.SetMarkerSize(0)
181	+	if options.overflow or (options.special and "Overflow" in name):
182	+	nbins = hist.GetNbinsX()
183	+	overflow = hist.GetBinContent(nbins + 1)
184	+	hist.AddBinContent(nbins, overflow)
185	+	if options.underflow or (options.special and "Underflow" in name):
186	+	underflow = hist.GetBinContent(0)
187	+	hist.AddBinContent(1, underflow)
188	+	if options.normalize or (options.special and "Norm" in name):
189	+	integral = hist.Integral()
190	+	if integral: hist.Scale(1. / integral)
191	+	hists.append(hist)
192	+	#### Combine hists in a THStack and draw
193	+	pads = [canvas]
194	+	stack = ROOT.THStack("st%.3i" % int(counter), title)
195	+	legend_height = 0.04 * len(files) + 0.02
196	+	legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
197	+	for hist in hists:
198	+	stack.Add(hist)
199	+	legend.AddEntry(hist)
200	+	stack.Draw()
201	+	stack.GetXaxis().SetTitle(x_title)
202	+	stack.GetYaxis().SetTitle(y_title)
203	+	if options.ratio or (options.special and "Ratio" in name):
204	+	pads, stack, stack_ratio = add_ratio_plot(hists, stack, counter)
205	+	pads[1].cd()
206	+	stack_ratio.Draw("nostack")
207	+	pads[0].cd()
208	+	pads[0].SetLogx(options.logx or (options.special and "Logx" in name))
209	+	pads[0].SetLogy(options.logy or (options.special and "Logy" in name))
210	+	stack.Draw("nostack p H" + options.opt)
211	+	if options.numbering:
212	+	display_page_number(counter)
213	+	if options.efficiency or (options.special and "Eff" in name):
214	+	stack.Draw("nostack e p" + options.opt)
215	+	stack.SetMaximum(1.)
216	+	stack.SetMinimum(0.)
217	+	if options.overflow or (options.special and "Overflow" in name):
218	+	display_overflow(stack, hist)
219	+	if options.underflow or (options.special and "Underflow" in name):
220	+	display_underflow(stack, hist)
221	+	legend.Draw()
222	+	save_plot(stack, plot_dir, path, name, counter)
223	+
224	+
225	+	#### Define the extra functions used by the main routines
226	+	def save_plot(stack, plot_dir, path, name, counter):
227	+	output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
228	+	canvas.SaveAs(output_file_name)
229	+	if options.ext == "pdf":
230	+	numbered_pdf_name = "%.3i.pdf" % counter
231	+	shutil.copy(output_file_name, numbered_pdf_name)
232	+	report_progress(counter, 1)
233	+
234	+
235	+	def report_progress(counter, divisor):
236	+	if counter % divisor == 0:
237	+	print "\r%i plots written to %s" % (counter, options.output),
238	+	sys.stdout.flush()
239	+
240	+	def merge_pdf():
241	+	print "Writing merged pdf..."
242	+	os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
243	+	"-dAutoRotatePages=/All "
244	+	"-sOutputFile=%s.pdf " % options.output +
245	+	"[0-9][0-9][0-9].pdf")
246	+	os.system("rm [0-9]*.pdf")
247	+
248	+
249	+	def display_page_number(page_number):
250	+	page_text = ROOT.TText()
251	+	page_text.SetTextSize(0.03)
252	+	page_text.SetTextAlign(33)
253	+	page_text.DrawTextNDC(0.97, 0.985, "%i" % page_number)
254	+
255	+
256	+	def display_overflow(stack, hist):
257	+	nbins = hist.GetNbinsX()
258	+	x = 0.5 * (hist.GetBinLowEdge(nbins) +
259	+	hist.GetBinLowEdge(nbins + 1))
260	+	y = stack.GetMinimum("nostack")
261	+	display_bin_text(x, y, nbins, "Overflow")
262	+
263	+
264	+	def display_underflow(stack, hist):
265	+	nbins = hist.GetNbinsX()
266	+	x = 0.5 * (hist.GetBinLowEdge(1) +
267	+	hist.GetBinLowEdge(2))
268	+	y = stack.GetMinimum("nostack")
269	+	display_bin_text(x, y, nbins, "Underflow")
270	+
271	+
272	+	def display_bin_text(x, y, nbins, text):
273	+	bin_text = ROOT.TText()
274	+	bin_text.SetTextSize(min(1. / nbins, 0.04))
275	+	bin_text.SetTextAlign(12)
276	+	bin_text.SetTextAngle(90)
277	+	bin_text.SetTextColor(13)
278	+	bin_text.SetTextFont(42)
279	+	bin_text.DrawText(x, y, text)
280	+
281	+
282	+	def add_ratio_plot(hists, stack, counter):
283	+	"""Divide canvas into two parts, and plot the ratio on the bottom."""
284	+	## Both pads are set to the full canvas size to maintain font sizes
285	+	## Fill style 4000 used to ensure pad transparency because of this
286	+	div = 0.3 # portion of canvas to use for ratio plot
287	+	margins = [ROOT.gStyle.GetPadTopMargin(), ROOT.gStyle.GetPadBottomMargin()]
288	+	useable_height = 1 - (margins[0] + margins[1])
289	+	canvas.Clear()
290	+	pad = ROOT.TPad("mainPad", "mainPad", 0., 0., 1., 1.)
291	+	pad.SetFillStyle(4000)
292	+	pad.Draw()
293	+	pad.SetBottomMargin(margins[1] + div * useable_height)
294	+	pad_ratio = ROOT.TPad("ratioPad", "ratioPad", 0., 0., 1., 1.);
295	+	pad_ratio.SetFillStyle(4000)
296	+	pad_ratio.Draw()
297	+	pad_ratio.SetTopMargin(margins[0] + (1 - div) * useable_height)
298	+	pad.cd()
299	+	stack.Draw()
300	+	stack_ratio = ROOT.THStack("stRatio%.3i" % int(counter),
301	+	";%s;Ratio" % stack.GetXaxis().GetTitle())
302	+	for hist in hists[1:]:
303	+	ratio_hist = hist.Clone()
304	+	ratio_hist.Divide(hists[0])
305	+	stack_ratio.Add(ratio_hist)
306	+	stack_ratio.Draw()
307	+	stack_ratio.GetYaxis().SetNdivisions(507) # Avoids crowded labels
308	+	stack.GetXaxis().SetBinLabel(1, "") # Don't show numbers below top plot
309	+	stack.GetXaxis().SetTitle("")
310	+	if stack.GetYaxis().GetTitle() == "":
311	+	stack.GetYaxis().SetTitle("Content")
312	+	# Avoid overlap of y-axis numbers by supressing zero
313	+	if stack.GetMinimum() / stack.GetMaximum() < 0.25:
314	+	stack.SetMinimum(stack.GetMaximum() / 10000)
315	+	return [pad, pad_ratio], stack, stack_ratio
316
317
318	+	#### Run main program
319		if __name__ == "__main__":
320	<	sys.exit(main())
320	>	if options.timing:
321	>	import profile
322	>	profile.run('main()', 'fooprof')
323	>	import pstats
324	>	p = pstats.Stats('fooprof')
325	>	p.sort_stats('cumulative').print_stats(15)
326	>	else:
327	>	sys.exit(main())
328

Diff Legend

–	Removed lines
+	Added lines
<	Changed lines
>	Changed lines