ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/RootMacros/overlayHists.py
Revision: 1.3
Committed: Thu Nov 19 19:32:52 2009 UTC (15 years, 5 months ago) by klukas
Content type: text/x-python
Branch: MAIN
Changes since 1.2: +7 -4 lines
Log Message:
Improved help documentation

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     and also merging all images into a single file (if output is pdf)
14 klukas 1.1
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]]
28    
29     ## Import python libraries
30     import sys
31     import optparse
32     import os
33     import re
34    
35     ## Import ROOT in batch mode
36     if '-h' not in sys.argv:
37     sys.argv.append('-b')
38     import ROOT
39     if os.path.exists('rootlogon.C'): ROOT.gROOT.Macro('rootlogon.C')
40     sys.argv.remove('-b')
41     ROOT.gErrorIgnoreLevel = ROOT.kWarning
42     colors = []
43     for rgb in rgbcolors:
44     colors.append(ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]))
45     c1 = ROOT.TCanvas()
46    
47     ## Parse options
48     parser = optparse.OptionParser(usage=usage)
49     parser.add_option('-n', '--normalize', action="store_true", default=False,
50     help="area normalize all histograms")
51     parser.add_option('-e', '--ext', default="pdf",
52 klukas 1.3 help="choose an output extension; default is pdf")
53 klukas 1.1 parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
54 klukas 1.3 help="name of output directory; default is 'overlaidHists'")
55 klukas 1.1 parser.add_option('-m', '--match', default="", metavar="REGEX",
56 klukas 1.3 help="only make plots for paths containing the specified "
57     "regular expression (use '.*' for wildcard)")
58 klukas 1.1 options, arguments = parser.parse_args()
59     plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
60     regex = re.compile(options.match)
61    
62    
63     class RootFile:
64     def __init__(self, file_name):
65     self.name = file_name[0:file_name.find(".root")]
66     self.file = ROOT.TFile(file_name, "read")
67     if self.file.IsZombie():
68     print "Error opening %s, exiting..." % file_name
69     sys.exit(1)
70     def Get(self, object_name):
71     return self.file.Get(object_name)
72    
73    
74    
75     def main():
76     files = []
77     for filename in arguments: files.append(RootFile(filename))
78 klukas 1.2 if len(files) == 0:
79     parser.print_help()
80     sys.exit(0)
81 klukas 1.1 process_directory("", files)
82     if options.ext == "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     print "Wrote %i plots to %s" % (next_counter() - 1, options.output)
89    
90    
91    
92     def process_directory(path, files):
93     dir_to_make = "%s/%s" % (plot_dir, path)
94     if not os.path.exists(dir_to_make):
95     os.mkdir(dir_to_make)
96     keys = files[0].file.GetDirectory(path).GetListOfKeys()
97     key = keys[0]
98     while key:
99     obj = key.ReadObj()
100     key = keys.After(key)
101     new_path = "%s/%s" % (path, obj.GetName())
102     if obj.IsA().InheritsFrom("TDirectory"):
103     process_directory(new_path, files)
104     if (regex.search(new_path) and
105     obj.IsA().InheritsFrom("TH1") and
106     not obj.IsA().InheritsFrom("TH2") and
107     not obj.IsA().InheritsFrom("TH3")):
108     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
156    
157    
158    
159    
160     if __name__ == "__main__":
161     sys.exit(main())
162