ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/RootMacros/overlayHists.py
Revision: 1.10
Committed: Fri Jan 8 20:36:24 2010 UTC (15 years, 3 months ago) by klukas
Content type: text/x-python
Branch: MAIN
Changes since 1.9: +145 -104 lines
Log Message: 
Broke the main code into more manageable functions; standardized the special plotting options

File Contents

# User Rev Content
1 klukas 1.1 #!/usr/bin/env python
2    
3     ## Created by Jeff Klukas (klukas@wisc.edu), November 2009
4    
5     ## For more information, use the -h option:
6     ## ./overlayHists.py -h
7    
8     ## Define usage string for help option
9     usage="""usage: %prog [options] file1.root file2.root file3.root ...
10    
11 klukas 1.3 function: overlays corresponding histograms from several files, dumping the
12     images into an identical directory structure in the local directory
13 klukas 1.10 and also merging all images into a single file (if output is pdf)"""
14 klukas 1.1
15    
16 klukas 1.10 #### Define colors and styles
17 klukas 1.4 rgbvals = [[82, 124, 219],
18 klukas 1.7 [212,58,143],
19     [231, 139, 77],
20 klukas 1.4 [145, 83, 207],
21     [114, 173, 117],
22     [67, 77, 83]]
23 klukas 1.7 marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]
24 klukas 1.1
25 klukas 1.10
26     #### Import python libraries
27 klukas 1.1 import sys
28     import optparse
29 klukas 1.9 import shutil
30 klukas 1.1 import os
31     import re
32    
33 klukas 1.10
34     #### Import ROOT in batch mode
35 klukas 1.7 if '-h' not in sys.argv and len(sys.argv) > 1:
36 klukas 1.1 sys.argv.append('-b')
37     import ROOT
38 klukas 1.7 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 klukas 1.1 sys.argv.remove('-b')
45     ROOT.gErrorIgnoreLevel = ROOT.kWarning
46 klukas 1.4 colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
47 klukas 1.9 canvas = ROOT.TCanvas()
48 klukas 1.1
49 klukas 1.10
50     #### Parse options
51 klukas 1.1 parser = optparse.OptionParser(usage=usage)
52     parser.add_option('-e', '--ext', default="pdf",
53 klukas 1.3 help="choose an output extension; default is pdf")
54 klukas 1.9 parser.add_option('-o', '--opt', default="",
55 klukas 1.10 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 klukas 1.9 parser.add_option('--output', default="overlaidHists", metavar="NAME",
62 klukas 1.3 help="name of output directory; default is 'overlaidHists'")
63 klukas 1.10 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 klukas 1.7 parser.add_option('--match', default="", metavar="REGEX",
68 klukas 1.3 help="only make plots for paths containing the specified "
69     "regular expression (use '.*' for wildcard)")
70 klukas 1.10 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 klukas 1.1 options, arguments = parser.parse_args()
93     plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
94     regex = re.compile(options.match)
95    
96    
97 klukas 1.10 #### Define classes and utility functions
98 klukas 1.1 class RootFile:
99 klukas 1.10 """A wrapper for TFiles, allowing quick access to the name and Get."""
100 klukas 1.1 def __init__(self, file_name):
101 klukas 1.9 self.name = file_name[0:-5]
102 klukas 1.1 self.file = ROOT.TFile(file_name, "read")
103     if self.file.IsZombie():
104     print "Error opening %s, exiting..." % file_name
105     sys.exit(1)
106     def Get(self, object_name):
107     return self.file.Get(object_name)
108    
109    
110 klukas 1.10 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 klukas 1.1
119 klukas 1.10 #### Define the main program functions
120 klukas 1.1 def main():
121 klukas 1.10 """Initialize files and enter into process_directory loop"""
122 klukas 1.4 files = [RootFile(filename) for filename in arguments]
123 klukas 1.2 if len(files) == 0:
124     parser.print_help()
125     sys.exit(0)
126 klukas 1.1 process_directory("", files)
127 klukas 1.8 print ""
128 klukas 1.1 if options.ext == "pdf":
129 klukas 1.8 merge_pdf()
130 klukas 1.1
131    
132     def process_directory(path, files):
133 klukas 1.10 """Loop through all histograms in the directory and plot them."""
134 klukas 1.1 dir_to_make = "%s/%s" % (plot_dir, path)
135     if not os.path.exists(dir_to_make):
136     os.mkdir(dir_to_make)
137     keys = files[0].file.GetDirectory(path).GetListOfKeys()
138     key = keys[0]
139     while key:
140     obj = key.ReadObj()
141     key = keys.After(key)
142     new_path = "%s/%s" % (path, obj.GetName())
143     if obj.IsA().InheritsFrom("TDirectory"):
144     process_directory(new_path, files)
145 klukas 1.10 #### If obj is a desired histogram, process it
146 klukas 1.1 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 klukas 1.10 process_hist(path, new_path, files, obj)
151    
152 klukas 1.1
153 klukas 1.10 #### 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 klukas 1.1
209 klukas 1.10
210     #### Define the extra functions used by the main routines
211 klukas 1.8 def save_plot(stack, plot_dir, path, name, counter):
212 klukas 1.9 output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
213     canvas.SaveAs(output_file_name)
214 klukas 1.8 if options.ext == "pdf":
215 klukas 1.9 numbered_pdf_name = "%.3i.pdf" % counter
216     shutil.copy(output_file_name, numbered_pdf_name)
217 klukas 1.8 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 klukas 1.10 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 klukas 1.1
266    
267 klukas 1.10 #### Run main program
268 klukas 1.1 if __name__ == "__main__":
269 klukas 1.8 if options.timing:
270 klukas 1.9 import profile
271     profile.run('main()', 'fooprof')
272 klukas 1.8 import pstats
273     p = pstats.Stats('fooprof')
274     p.sort_stats('cumulative').print_stats(15)
275     else:
276     sys.exit(main())
277 klukas 1.1