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.6 by klukas, Fri Nov 20 20:16:30 2009 UTC vs.
Revision 1.13 by klukas, Wed Feb 3 19:09:56 2010 UTC

#	Line 1 | Line 1
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	<	## For more information, use the -h option:
6	<	##     ./overlayHists.py -h
8	>	######## Import python libraries #############################################
9
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
23	–	rgbvals = [[82, 124, 219],
24	–	[145, 83, 207],
25	–	[231, 139, 77],
26	–	[114, 173, 117],
27	–	[67, 77, 83]]
28	–
29	–	## Import python libraries
10		import sys
11		import optparse
12	+	import shutil
13		import os
14		import re
15
16	<	## Import ROOT in batch mode
17	<	if '-h' not in sys.argv:
18	<	sys.argv.append('-b')
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	<	if os.path.exists('rootlogon.C'): ROOT.gROOT.Macro('rootlogon.C')
21	<	sys.argv.remove('-b')
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	<	colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
26	<	c1 = ROOT.TCanvas()
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
45	–	## Parse options
46	–	parser = optparse.OptionParser(usage=usage)
47	–	parser.add_option('-n', '--normalize', action="store_true", default=False,
48	–	help="area normalize all histograms")
49	–	parser.add_option('-e', '--ext', default="pdf",
50	–	help="choose an output extension; default is pdf")
51	–	parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
52	–	help="name of output directory; default is 'overlaidHists'")
53	–	parser.add_option('-m', '--match', default="", metavar="REGEX",
54	–	help="only make plots for paths containing the specified "
55	–	"regular expression (use '.*' for wildcard)")
56	–	options, arguments = parser.parse_args()
57	–	plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
58	–	regex = re.compile(options.match)
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:file_name.find(".root")]
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
#	Line 69 | Line 68 | class RootFile:
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
74	–	def main():
75	–	files = [RootFile(filename) for filename in arguments]
76	–	if len(files) == 0:
77	–	parser.print_help()
78	–	sys.exit(0)
79	–	process_directory("", files)
80	–	print
81	–	if options.ext == "pdf":
82	–	print "Writing merged pdf..."
83	–	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	–
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	<	dir_to_make = "%s/%s" % (plot_dir, path)
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()
#	Line 100 | Line 94 | def process_directory(path, files):
94		new_path = "%s/%s" % (path, obj.GetName())
95		if obj.IsA().InheritsFrom("TDirectory"):
96		process_directory(new_path, files)
97	<	if (regex.search(new_path) and
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	<	counter = next_counter()
108	<	name = obj.GetName()
109	<	hist = files[0].file.GetDirectory(path).Get(name)
110	<	title = hist.GetTitle()
111	<	x_title = hist.GetXaxis().GetTitle()
112	<	y_title = hist.GetYaxis().GetTitle()
113	<	if "Norm" in name or options.normalize:
114	<	y_title = "Fraction of Events in Bin"
115	<	hist.Draw()
116	<	stack = ROOT.THStack("st%.3i" % int(counter), title)
117	<	legend_height = 0.04 * len(files) + 0.02
118	<	legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
119	<	c1.SetLogx("Logx" in name)
120	<	c1.SetLogy("Logy" in name)
121	<	for i, file in enumerate(files):
122	<	hist = file.file.GetDirectory(path).Get(name)
123	<	if not hist: continue
124	<	hist.Draw()
125	<	hist.SetTitle(file.name)
126	<	color = colors[i % len(colors)]
127	<	hist.SetLineColor(color)
128	<	## hist.SetMarkerColor(color)
129	<	## hist.SetMarkerStyle(i + 1)
130	<	if "Norm" in name or options.normalize:
131	<	integral = hist.Integral()
132	<	hist.Scale(1. / integral)
133	<	stack.Add(hist)
134	<	legend.AddEntry(hist)
135	<	stack.Draw("nostack p H")
136	<	stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
137	<	if "Eff" in name:
138	<	stack.Draw("nostack e p")
139	<	stack.SetMaximum(1.)
140	<	stack.SetMinimum(0.)
141	<	legend.Draw()
142	<	if options.ext == "pdf":
143	<	c1.SaveAs("%.3i.pdf" % counter)
144	<	c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext))
145	<	print "\r%i plots written to %s" % (counter, options.output),
146	<	sys.stdout.flush()
147	<
102	>	process_hist(path, new_path, files, obj)
103
104
105	<	def counter_generator():
106	<	k = 0
107	<	while True:
108	<	k += 1
109	<	yield k
110	<	next_counter = counter_generator().next
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	<	sys.exit(main())
366	>	main()
367

Diff Legend

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