ViewVC Help
View File | Revision Log | Show Annotations | Root Listing
root/cvsroot/UserCode/RootMacros/overlayHists.py
Revision: 1.4
Committed: Thu Nov 19 19:42:07 2009 UTC (15 years, 5 months ago) by klukas
Content type: text/x-python
Branch: MAIN
Changes since 1.3: +9 -14 lines
Log Message:
Various small improvements for readability

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 klukas 1.4 rgbvals = [[82, 124, 219],
24     [145, 83, 207],
25     [231, 139, 77],
26     [114, 173, 117],
27     [67, 77, 83]]
28 klukas 1.1
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 klukas 1.4 colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
43 klukas 1.1 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 klukas 1.3 help="choose an output extension; default is pdf")
51 klukas 1.1 parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
52 klukas 1.3 help="name of output directory; default is 'overlaidHists'")
53 klukas 1.1 parser.add_option('-m', '--match', default="", metavar="REGEX",
54 klukas 1.3 help="only make plots for paths containing the specified "
55     "regular expression (use '.*' for wildcard)")
56 klukas 1.1 options, arguments = parser.parse_args()
57     plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
58     regex = re.compile(options.match)
59    
60    
61 klukas 1.4
62 klukas 1.1 class RootFile:
63     def __init__(self, file_name):
64     self.name = file_name[0:file_name.find(".root")]
65     self.file = ROOT.TFile(file_name, "read")
66     if self.file.IsZombie():
67     print "Error opening %s, exiting..." % file_name
68     sys.exit(1)
69     def Get(self, object_name):
70     return self.file.Get(object_name)
71    
72    
73    
74     def main():
75 klukas 1.4 files = [RootFile(filename) for filename in arguments]
76 klukas 1.2 if len(files) == 0:
77     parser.print_help()
78     sys.exit(0)
79 klukas 1.1 process_directory("", files)
80     if options.ext == "pdf":
81     os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
82     "-dAutoRotatePages=/All "
83     "-sOutputFile=%s.pdf " % options.output +
84     "[0-9][0-9][0-9].pdf")
85     os.system("rm [0-9]*.pdf")
86     print "Wrote %i plots to %s" % (next_counter() - 1, options.output)
87    
88    
89    
90     def process_directory(path, files):
91     dir_to_make = "%s/%s" % (plot_dir, path)
92     if not os.path.exists(dir_to_make):
93     os.mkdir(dir_to_make)
94     keys = files[0].file.GetDirectory(path).GetListOfKeys()
95     key = keys[0]
96     while key:
97     obj = key.ReadObj()
98     key = keys.After(key)
99     new_path = "%s/%s" % (path, obj.GetName())
100     if obj.IsA().InheritsFrom("TDirectory"):
101     process_directory(new_path, files)
102     if (regex.search(new_path) and
103     obj.IsA().InheritsFrom("TH1") and
104     not obj.IsA().InheritsFrom("TH2") and
105     not obj.IsA().InheritsFrom("TH3")):
106     counter = next_counter()
107     name = obj.GetName()
108     hist = files[0].file.GetDirectory(path).Get(name)
109     title = hist.GetTitle()
110     x_title = hist.GetXaxis().GetTitle()
111     y_title = hist.GetYaxis().GetTitle()
112     if "Norm" in name or options.normalize:
113     y_title = "Fraction of Events in Bin"
114     hist.Draw()
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 klukas 1.4 hist.Scale(1. / integral)
131 klukas 1.1 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     def counter_generator():
147     k = 0
148     while True:
149     k += 1
150     yield k
151     next_counter = counter_generator().next
152    
153    
154    
155     if __name__ == "__main__":
156     sys.exit(main())
157