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.1 by klukas, Wed Nov 18 21:55:17 2009 UTC vs.
Revision 1.10 by klukas, Fri Jan 8 20:36:24 2010 UTC

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

Diff Legend

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