UserCode/RootMacros/overlayHists.py

#!/usr/bin/env python

## Created by Jeff Klukas (klukas@wisc.edu), November 2009

## For more information, use the -h option:
##     ./overlayHists.py -h

## Define usage string for help option
usage="""usage: %prog [options] file1.root file2.root file3.root ...

function: overlays corresponding histograms from several files, dumping the
          images into an identical directory structure in the local directory
          and also merging all images into a single file (if output is pdf)

naming: histograms whose names contain certain key terms will be handled
        specially.  Use this to your advantage!
  'Eff' : y-axis will be scaled from 0 to 1
  'Norm': plot will be area normalized
  'Logx': x-axis will be on log scale
  'Logy': y-axis will be on log scale"""

## Define colors and styles
rgbvals = [[82, 124, 219],
           [212,58,143],
           [231, 139, 77],
           [145, 83, 207],
           [114, 173, 117],
           [67, 77, 83]]
marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]

## Import python libraries
import sys
import optparse
import shutil
import os
import re

## Import ROOT in batch mode
if '-h' not in sys.argv and len(sys.argv) > 1:
    sys.argv.append('-b')
    import ROOT
    if os.path.exists('rootlogon.C'):
        ROOT.gROOT.Macro('rootlogon.C')
    else:
        os.system('echo -e "{\n}\n" >> rootlogon.C')
        ROOT.gROOT.Macro('rootlogon.C')
        os.remove('rootlogon.C')
    sys.argv.remove('-b')
    ROOT.gErrorIgnoreLevel = ROOT.kWarning
    colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
    canvas = ROOT.TCanvas()

## Parse options
parser = optparse.OptionParser(usage=usage)
parser.add_option('-n', '--normalize', action="store_true", default=False,
                  help="area normalize all histograms")
parser.add_option('-m', '--markers', action="store_true", default=False,
                  help="add markers to histograms")
parser.add_option('--logx', action="store_true", default=False,
                  help="Force log scale for X axis on all histograms")
parser.add_option('--logy', action="store_true", default=False,
                  help="Force log scale for Y axis on all histograms")
parser.add_option('-e', '--ext', default="pdf", 
                  help="choose an output extension; default is pdf")
parser.add_option('-o', '--opt', default="",
                  help="argument to be passed to THStack::Draw")
parser.add_option('--output', default="overlaidHists", metavar="NAME",
                  help="name of output directory; default is 'overlaidHists'")
parser.add_option('--match', default="", metavar="REGEX",
                  help="only make plots for paths containing the specified "
                  "regular expression (use '.*' for wildcard)")
parser.add_option('--timing', action="store_true", default=False,
                  help="output timing information")
options, arguments = parser.parse_args()
plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
regex = re.compile(options.match)


class RootFile:
    def __init__(self, file_name):
        self.name = file_name[0:-5]
        self.file = ROOT.TFile(file_name, "read")
        if self.file.IsZombie():
            print "Error opening %s, exiting..." % file_name
            sys.exit(1)
    def Get(self, object_name):
        return self.file.Get(object_name)


def main():
    files = [RootFile(filename) for filename in arguments]
    if len(files) == 0:
        parser.print_help()
        sys.exit(0)
    process_directory("", files)
    print ""
    if options.ext == "pdf":
        merge_pdf()


def process_directory(path, files):
    dir_to_make = "%s/%s" % (plot_dir, path)
    if not os.path.exists(dir_to_make):
        os.mkdir(dir_to_make)
    keys = files[0].file.GetDirectory(path).GetListOfKeys()
    key = keys[0]
    while key:
        obj = key.ReadObj()
        key = keys.After(key)
        new_path = "%s/%s" % (path, obj.GetName())
        if obj.IsA().InheritsFrom("TDirectory"):
            process_directory(new_path, files)
        if (regex.search(new_path) and
            obj.IsA().InheritsFrom("TH1") and
            not obj.IsA().InheritsFrom("TH2") and
            not obj.IsA().InheritsFrom("TH3")):
            counter = next_counter()
            name = obj.GetName()
            hist = files[0].file.GetDirectory(path).Get(name)
            title = hist.GetTitle()
            x_title = hist.GetXaxis().GetTitle()
            y_title = hist.GetYaxis().GetTitle()
            if "Norm" in name or options.normalize:
                y_title = "Fraction of Events in Bin"
            stack = ROOT.THStack("st%.3i" % int(counter), title)
            legend_height = 0.04 * len(files) + 0.02
            legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
            canvas.SetLogx("Logx" in name or options.logx)
            canvas.SetLogy("Logy" in name or options.logy)
            for i, file in enumerate(files):
                hist = file.file.GetDirectory(path).Get(name)
                if not hist: continue
                hist.SetTitle(file.name)
                color = colors[i % len(colors)]
                hist.SetLineColor(color)
                if options.markers:
                    hist.SetMarkerColor(color)
                    hist.SetMarkerStyle(marker_styles[i])
                else:
                    hist.SetMarkerSize(0)
                if "Overflow" in name:
                    nbins = hist.GetNbinsX()
                    overflow = hist.GetBinContent(nbins + 1)
                    hist.AddBinContent(nbins, overflow)
                if "Underflow" in name:
                    underflow = hist.GetBinContent(0)
                    hist.AddBinContent(1, underflow)
                if "Norm" in name or options.normalize:
                    integral = hist.Integral()
                    hist.Scale(1. / integral)
                stack.Add(hist)
                legend.AddEntry(hist)
            stack.Draw("nostack p H" + options.opt)
            stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
            if "Eff" in name:
                stack.Draw("nostack e p")
                stack.SetMaximum(1.)
                stack.SetMinimum(0.)
            if "Overflow" in name:
                nbins = hist.GetNbinsX()
                x = 0.5 * (hist.GetBinLowEdge(nbins) +
                           hist.GetBinLowEdge(nbins + 1))
                y = stack.GetMinimum("nostack")
                overflow_text = ROOT.TText(x, y, "Overflow")
                overflow_text.SetTextSize(min(1. / nbins, 0.04))
                overflow_text.SetTextAlign(12)
                overflow_text.SetTextAngle(90)
                overflow_text.SetTextColor(13)
                overflow_text.SetTextFont(42)
                overflow_text.Draw()
            if "Underflow" in name:
                nbins = hist.GetNbinsX()
                x = hist.GetBinLowEdge(1)
                y = hist.GetMinimum()
                underflow_text = ROOT.TText(x, y, "Underflow")
                underflow_text.SetTextSize(1. / nbins)
                underflow_text.SetTextAlign(13)
                underflow_text.SetTextAngle(90)
                underflow_text.SetTextColor(13)
                underflow_text.Draw()
            legend.Draw()
            save_plot(stack, plot_dir, path, name, counter)
            

def save_plot(stack, plot_dir, path, name, counter):
    output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
    canvas.SaveAs(output_file_name)
    if options.ext == "pdf":
        numbered_pdf_name = "%.3i.pdf" % counter
        shutil.copy(output_file_name, numbered_pdf_name)
    report_progress(counter, 1)


def report_progress(counter, divisor):
    if counter % divisor == 0:
        print "\r%i plots written to %s" % (counter, options.output),
        sys.stdout.flush()


def merge_pdf():
    print "Writing merged pdf..."
    os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
              "-dAutoRotatePages=/All "
              "-sOutputFile=%s.pdf " % options.output +
              "[0-9][0-9][0-9].pdf")
    os.system("rm [0-9]*.pdf")


def counter_generator():
    k = 0
    while True:
        k += 1
        yield k
next_counter = counter_generator().next


if __name__ == "__main__":
    if options.timing:
##         import cProfile
##         cProfile.run('main()', 'fooprof')
        import profile
        profile.run('main()', 'fooprof')
        import pstats
        p = pstats.Stats('fooprof')
        p.sort_stats('cumulative').print_stats(15)
    else:
        sys.exit(main())

Revision:	1.9
Committed:	Wed Jan 6 21:18:09 2010 UTC (15 years, 3 months ago) by klukas
Content type:	text/x-python
Branch:	MAIN
Changes since 1.8:	+50 -12 lines
Log Message:	Added logx, logy, opt, overflow, and underflow capabilities
#	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	klukas	1.7	## Define colors and styles
23	klukas	1.4	rgbvals = [[82, 124, 219],
24	klukas	1.7	[212,58,143],
25			[231, 139, 77],
26	klukas	1.4	[145, 83, 207],
27			[114, 173, 117],
28			[67, 77, 83]]
29	klukas	1.7	marker_styles = [3, 4, 5, 25, 26, 27, 28, 30]
30	klukas	1.1
31			## Import python libraries
32			import sys
33			import optparse
34	klukas	1.9	import shutil
35	klukas	1.1	import os
36			import re
37
38			## Import ROOT in batch mode
39	klukas	1.7	if '-h' not in sys.argv and len(sys.argv) > 1:
40	klukas	1.1	sys.argv.append('-b')
41			import ROOT
42	klukas	1.7	if os.path.exists('rootlogon.C'):
43			ROOT.gROOT.Macro('rootlogon.C')
44			else:
45			os.system('echo -e "{\n}\n" >> rootlogon.C')
46			ROOT.gROOT.Macro('rootlogon.C')
47			os.remove('rootlogon.C')
48	klukas	1.1	sys.argv.remove('-b')
49			ROOT.gErrorIgnoreLevel = ROOT.kWarning
50	klukas	1.4	colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
51	klukas	1.9	canvas = ROOT.TCanvas()
52	klukas	1.1
53			## Parse options
54			parser = optparse.OptionParser(usage=usage)
55			parser.add_option('-n', '--normalize', action="store_true", default=False,
56			help="area normalize all histograms")
57	klukas	1.7	parser.add_option('-m', '--markers', action="store_true", default=False,
58			help="add markers to histograms")
59	klukas	1.9	parser.add_option('--logx', action="store_true", default=False,
60			help="Force log scale for X axis on all histograms")
61			parser.add_option('--logy', action="store_true", default=False,
62			help="Force log scale for Y axis on all histograms")
63	klukas	1.1	parser.add_option('-e', '--ext', default="pdf",
64	klukas	1.3	help="choose an output extension; default is pdf")
65	klukas	1.9	parser.add_option('-o', '--opt', default="",
66			help="argument to be passed to THStack::Draw")
67			parser.add_option('--output', default="overlaidHists", metavar="NAME",
68	klukas	1.3	help="name of output directory; default is 'overlaidHists'")
69	klukas	1.7	parser.add_option('--match', default="", metavar="REGEX",
70	klukas	1.3	help="only make plots for paths containing the specified "
71			"regular expression (use '.*' for wildcard)")
72	klukas	1.8	parser.add_option('--timing', action="store_true", default=False,
73			help="output timing information")
74	klukas	1.1	options, arguments = parser.parse_args()
75			plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
76			regex = re.compile(options.match)
77
78
79	klukas	1.4
80	klukas	1.1	class RootFile:
81			def __init__(self, file_name):
82	klukas	1.9	self.name = file_name[0:-5]
83	klukas	1.1	self.file = ROOT.TFile(file_name, "read")
84			if self.file.IsZombie():
85			print "Error opening %s, exiting..." % file_name
86			sys.exit(1)
87			def Get(self, object_name):
88			return self.file.Get(object_name)
89
90
91
92			def main():
93	klukas	1.4	files = [RootFile(filename) for filename in arguments]
94	klukas	1.2	if len(files) == 0:
95			parser.print_help()
96			sys.exit(0)
97	klukas	1.1	process_directory("", files)
98	klukas	1.8	print ""
99	klukas	1.1	if options.ext == "pdf":
100	klukas	1.8	merge_pdf()
101	klukas	1.1
102
103
104			def process_directory(path, files):
105			dir_to_make = "%s/%s" % (plot_dir, path)
106			if not os.path.exists(dir_to_make):
107			os.mkdir(dir_to_make)
108			keys = files[0].file.GetDirectory(path).GetListOfKeys()
109			key = keys[0]
110			while key:
111			obj = key.ReadObj()
112			key = keys.After(key)
113			new_path = "%s/%s" % (path, obj.GetName())
114			if obj.IsA().InheritsFrom("TDirectory"):
115			process_directory(new_path, files)
116			if (regex.search(new_path) and
117			obj.IsA().InheritsFrom("TH1") and
118			not obj.IsA().InheritsFrom("TH2") and
119			not obj.IsA().InheritsFrom("TH3")):
120			counter = next_counter()
121			name = obj.GetName()
122			hist = files[0].file.GetDirectory(path).Get(name)
123			title = hist.GetTitle()
124			x_title = hist.GetXaxis().GetTitle()
125			y_title = hist.GetYaxis().GetTitle()
126			if "Norm" in name or options.normalize:
127			y_title = "Fraction of Events in Bin"
128			stack = ROOT.THStack("st%.3i" % int(counter), title)
129	klukas	1.5	legend_height = 0.04 * len(files) + 0.02
130			legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
131	klukas	1.9	canvas.SetLogx("Logx" in name or options.logx)
132			canvas.SetLogy("Logy" in name or options.logy)
133	klukas	1.1	for i, file in enumerate(files):
134			hist = file.file.GetDirectory(path).Get(name)
135			if not hist: continue
136			hist.SetTitle(file.name)
137			color = colors[i % len(colors)]
138			hist.SetLineColor(color)
139	klukas	1.7	if options.markers:
140			hist.SetMarkerColor(color)
141			hist.SetMarkerStyle(marker_styles[i])
142			else:
143			hist.SetMarkerSize(0)
144	klukas	1.9	if "Overflow" in name:
145			nbins = hist.GetNbinsX()
146			overflow = hist.GetBinContent(nbins + 1)
147			hist.AddBinContent(nbins, overflow)
148			if "Underflow" in name:
149			underflow = hist.GetBinContent(0)
150			hist.AddBinContent(1, underflow)
151	klukas	1.1	if "Norm" in name or options.normalize:
152			integral = hist.Integral()
153	klukas	1.4	hist.Scale(1. / integral)
154	klukas	1.1	stack.Add(hist)
155			legend.AddEntry(hist)
156	klukas	1.9	stack.Draw("nostack p H" + options.opt)
157	klukas	1.1	stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
158			if "Eff" in name:
159			stack.Draw("nostack e p")
160			stack.SetMaximum(1.)
161			stack.SetMinimum(0.)
162	klukas	1.9	if "Overflow" in name:
163			nbins = hist.GetNbinsX()
164			x = 0.5 * (hist.GetBinLowEdge(nbins) +
165			hist.GetBinLowEdge(nbins + 1))
166			y = stack.GetMinimum("nostack")
167			overflow_text = ROOT.TText(x, y, "Overflow")
168			overflow_text.SetTextSize(min(1. / nbins, 0.04))
169			overflow_text.SetTextAlign(12)
170			overflow_text.SetTextAngle(90)
171			overflow_text.SetTextColor(13)
172			overflow_text.SetTextFont(42)
173			overflow_text.Draw()
174			if "Underflow" in name:
175			nbins = hist.GetNbinsX()
176			x = hist.GetBinLowEdge(1)
177			y = hist.GetMinimum()
178			underflow_text = ROOT.TText(x, y, "Underflow")
179			underflow_text.SetTextSize(1. / nbins)
180			underflow_text.SetTextAlign(13)
181			underflow_text.SetTextAngle(90)
182			underflow_text.SetTextColor(13)
183			underflow_text.Draw()
184	klukas	1.1	legend.Draw()
185	klukas	1.8	save_plot(stack, plot_dir, path, name, counter)
186	klukas	1.1
187
188
189	klukas	1.8	def save_plot(stack, plot_dir, path, name, counter):
190	klukas	1.9	output_file_name = "%s/%s/%s.%s" % (plot_dir, path, name, options.ext)
191			canvas.SaveAs(output_file_name)
192	klukas	1.8	if options.ext == "pdf":
193	klukas	1.9	numbered_pdf_name = "%.3i.pdf" % counter
194			shutil.copy(output_file_name, numbered_pdf_name)
195	klukas	1.8	report_progress(counter, 1)
196
197
198
199			def report_progress(counter, divisor):
200			if counter % divisor == 0:
201			print "\r%i plots written to %s" % (counter, options.output),
202			sys.stdout.flush()
203
204
205
206			def merge_pdf():
207			print "Writing merged pdf..."
208			os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
209			"-dAutoRotatePages=/All "
210			"-sOutputFile=%s.pdf " % options.output +
211			"[0-9][0-9][0-9].pdf")
212			os.system("rm [0-9]*.pdf")
213
214
215
216	klukas	1.1	def counter_generator():
217			k = 0
218			while True:
219			k += 1
220			yield k
221			next_counter = counter_generator().next
222
223
224
225			if __name__ == "__main__":
226	klukas	1.8	if options.timing:
227	klukas	1.9	## import cProfile
228			## cProfile.run('main()', 'fooprof')
229			import profile
230			profile.run('main()', 'fooprof')
231	klukas	1.8	import pstats
232			p = pstats.Stats('fooprof')
233			p.sort_stats('cumulative').print_stats(15)
234			else:
235			sys.exit(main())
236	klukas	1.1