ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/RootMacros/overlayHists.py
Revision: 1.1
Committed: Wed Nov 18 21:55:17 2009 UTC (15 years, 5 months ago) by klukas
Content type: text/x-python
Branch: MAIN
Log Message:
Full-featured utility to overlay corresponding histograms from several files

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     process_directory("", files)
76     if options.ext == "pdf":
77     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    
84    
85    
86     def process_directory(path, files):
87     dir_to_make = "%s/%s" % (plot_dir, path)
88     if not os.path.exists(dir_to_make):
89     os.mkdir(dir_to_make)
90     keys = files[0].file.GetDirectory(path).GetListOfKeys()
91     key = keys[0]
92     while key:
93     obj = key.ReadObj()
94     key = keys.After(key)
95     new_path = "%s/%s" % (path, obj.GetName())
96     if obj.IsA().InheritsFrom("TDirectory"):
97     process_directory(new_path, files)
98     if (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()
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    
151    
152    
153    
154     if __name__ == "__main__":
155     sys.exit(main())
156