ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/RootMacros/overlayHists.py
Revision: 1.2
Committed: Thu Nov 19 15:21:39 2009 UTC (15 years, 5 months ago) by klukas
Content type: text/x-python
Branch: MAIN
Changes since 1.1: +3 -0 lines
Log Message:
Now displays help if given no arguments rather than crashing

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     function: overlays histograms from several files with identical structure
12    
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]]
26    
27     ## Import python libraries
28     import sys
29     import optparse
30     import os
31     import re
32    
33     ## Import ROOT in batch mode
34     if '-h' not in sys.argv:
35     sys.argv.append('-b')
36     import ROOT
37     if os.path.exists('rootlogon.C'): ROOT.gROOT.Macro('rootlogon.C')
38     sys.argv.remove('-b')
39     ROOT.gErrorIgnoreLevel = ROOT.kWarning
40     colors = []
41     for rgb in rgbcolors:
42     colors.append(ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]))
43     c1 = ROOT.TCanvas()
44    
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="specify the type (extension) of the output files")
51     parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
52     help="specify the name of the output file/directory")
53     parser.add_option('-m', '--match', default="", metavar="REGEX",
54     help="only make plots for paths containing REGEX")
55     options, arguments = parser.parse_args()
56     plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
57     regex = re.compile(options.match)
58    
59    
60     class RootFile:
61     def __init__(self, file_name):
62     self.name = file_name[0:file_name.find(".root")]
63     self.file = ROOT.TFile(file_name, "read")
64     if self.file.IsZombie():
65     print "Error opening %s, exiting..." % file_name
66     sys.exit(1)
67     def Get(self, object_name):
68     return self.file.Get(object_name)
69    
70    
71    
72     def main():
73     files = []
74     for filename in arguments: files.append(RootFile(filename))
75 klukas 1.2 if len(files) == 0:
76     parser.print_help()
77     sys.exit(0)
78 klukas 1.1 process_directory("", files)
79     if options.ext == "pdf":
80     os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
81     "-dAutoRotatePages=/All "
82     "-sOutputFile=%s.pdf " % options.output +
83     "[0-9][0-9][0-9].pdf")
84     os.system("rm [0-9]*.pdf")
85     print "Wrote %i plots to %s" % (next_counter() - 1, options.output)
86    
87    
88    
89     def process_directory(path, files):
90     dir_to_make = "%s/%s" % (plot_dir, path)
91     if not os.path.exists(dir_to_make):
92     os.mkdir(dir_to_make)
93     keys = files[0].file.GetDirectory(path).GetListOfKeys()
94     key = keys[0]
95     while key:
96     obj = key.ReadObj()
97     key = keys.After(key)
98     new_path = "%s/%s" % (path, obj.GetName())
99     if obj.IsA().InheritsFrom("TDirectory"):
100     process_directory(new_path, files)
101     if (regex.search(new_path) and
102     obj.IsA().InheritsFrom("TH1") and
103     not obj.IsA().InheritsFrom("TH2") and
104     not obj.IsA().InheritsFrom("TH3")):
105     counter = next_counter()
106     name = obj.GetName()
107     hist = files[0].file.GetDirectory(path).Get(name)
108     title = hist.GetTitle()
109     x_title = hist.GetXaxis().GetTitle()
110     y_title = hist.GetYaxis().GetTitle()
111     if "Norm" in name or options.normalize:
112     y_title = "Fraction of Events in Bin"
113     hist.Draw()
114     hists = []
115     stack = ROOT.THStack("st%.3i" % int(counter), title)
116     legend = ROOT.TLegend(0.65, 0.77, 0.87, 0.89)
117     c1.SetLogx("Logx" in name)
118     c1.SetLogy("Logy" in name)
119     for i, file in enumerate(files):
120     hist = file.file.GetDirectory(path).Get(name)
121     if not hist: continue
122     hist.Draw()
123     hist.SetTitle(file.name)
124     color = colors[i % len(colors)]
125     hist.SetLineColor(color)
126     hist.SetMarkerColor(color)
127     hist.SetMarkerStyle(i + 1)
128     if "Norm" in name or options.normalize:
129     integral = hist.Integral()
130     hist.Scale(1 / integral)
131     stack.Add(hist)
132     legend.AddEntry(hist)
133     stack.Draw("nostack p H")
134     stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
135     if "Eff" in name:
136     stack.Draw("nostack e p")
137     stack.SetMaximum(1.)
138     stack.SetMinimum(0.)
139     legend.Draw()
140     if options.ext == "pdf":
141     c1.SaveAs("%.3i.pdf" % counter)
142     c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext))
143    
144    
145    
146    
147     def counter_generator():
148     k = 0
149     while True:
150     k += 1
151     yield k
152     next_counter = counter_generator().next
153    
154    
155    
156    
157     if __name__ == "__main__":
158     sys.exit(main())
159