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 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]
    c1 = 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('-e', '--ext', default="pdf", 
                  help="choose an output extension; default is pdf")
parser.add_option('-o', '--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:file_name.find(".root")]
        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"
            hist.Draw()
            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)
            c1.SetLogx("Logx" in name)
            c1.SetLogy("Logy" in name)
            for i, file in enumerate(files):
                hist = file.file.GetDirectory(path).Get(name)
                if not hist: continue
                hist.Draw()
                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 "Norm" in name or options.normalize:
                    integral = hist.Integral()
                    hist.Scale(1. / integral)
                stack.Add(hist)
                legend.AddEntry(hist)
            stack.Draw("nostack p H")
            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.)
            legend.Draw()
            save_plot(stack, plot_dir, path, name, counter)
            

def save_plot(stack, plot_dir, path, name, counter):
    if options.ext == "pdf":
        c1.SaveAs("%.3i.pdf" % counter)
    c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext))
    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 pstats
        p = pstats.Stats('fooprof')
        p.sort_stats('cumulative').print_stats(15)
    else:
        sys.exit(main())

Revision:	1.8
Committed:	Tue Nov 24 22:18:26 2009 UTC (15 years, 5 months ago) by klukas
Content type:	text/x-python
Branch:	MAIN
Changes since 1.7:	+38 -13 lines
Log Message:	Added option for timing summary; broke up into smaller functions elucidate timing
#	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			import os
35			import re
36
37			## Import ROOT in batch mode
38	klukas	1.7	if '-h' not in sys.argv and len(sys.argv) > 1:
39	klukas	1.1	sys.argv.append('-b')
40			import ROOT
41	klukas	1.7	if os.path.exists('rootlogon.C'):
42			ROOT.gROOT.Macro('rootlogon.C')
43			else:
44			os.system('echo -e "{\n}\n" >> rootlogon.C')
45			ROOT.gROOT.Macro('rootlogon.C')
46			os.remove('rootlogon.C')
47	klukas	1.1	sys.argv.remove('-b')
48			ROOT.gErrorIgnoreLevel = ROOT.kWarning
49	klukas	1.4	colors = [ROOT.TColor.GetColor(rgb[0], rgb[1], rgb[2]) for rgb in rgbvals]
50	klukas	1.1	c1 = ROOT.TCanvas()
51
52			## Parse options
53			parser = optparse.OptionParser(usage=usage)
54			parser.add_option('-n', '--normalize', action="store_true", default=False,
55			help="area normalize all histograms")
56	klukas	1.7	parser.add_option('-m', '--markers', action="store_true", default=False,
57			help="add markers to histograms")
58	klukas	1.1	parser.add_option('-e', '--ext', default="pdf",
59	klukas	1.3	help="choose an output extension; default is pdf")
60	klukas	1.1	parser.add_option('-o', '--output', default="overlaidHists", metavar="NAME",
61	klukas	1.3	help="name of output directory; default is 'overlaidHists'")
62	klukas	1.7	parser.add_option('--match', default="", metavar="REGEX",
63	klukas	1.3	help="only make plots for paths containing the specified "
64			"regular expression (use '.*' for wildcard)")
65	klukas	1.8	parser.add_option('--timing', action="store_true", default=False,
66			help="output timing information")
67	klukas	1.1	options, arguments = parser.parse_args()
68			plot_dir = "%s/%s" % (os.path.abspath('.'), options.output)
69			regex = re.compile(options.match)
70
71
72	klukas	1.4
73	klukas	1.1	class RootFile:
74			def __init__(self, file_name):
75			self.name = file_name[0:file_name.find(".root")]
76			self.file = ROOT.TFile(file_name, "read")
77			if self.file.IsZombie():
78			print "Error opening %s, exiting..." % file_name
79			sys.exit(1)
80			def Get(self, object_name):
81			return self.file.Get(object_name)
82
83
84
85			def main():
86	klukas	1.4	files = [RootFile(filename) for filename in arguments]
87	klukas	1.2	if len(files) == 0:
88			parser.print_help()
89			sys.exit(0)
90	klukas	1.1	process_directory("", files)
91	klukas	1.8	print ""
92	klukas	1.1	if options.ext == "pdf":
93	klukas	1.8	merge_pdf()
94	klukas	1.1
95
96
97			def process_directory(path, files):
98			dir_to_make = "%s/%s" % (plot_dir, path)
99			if not os.path.exists(dir_to_make):
100			os.mkdir(dir_to_make)
101			keys = files[0].file.GetDirectory(path).GetListOfKeys()
102			key = keys[0]
103			while key:
104			obj = key.ReadObj()
105			key = keys.After(key)
106			new_path = "%s/%s" % (path, obj.GetName())
107			if obj.IsA().InheritsFrom("TDirectory"):
108			process_directory(new_path, files)
109			if (regex.search(new_path) and
110			obj.IsA().InheritsFrom("TH1") and
111			not obj.IsA().InheritsFrom("TH2") and
112			not obj.IsA().InheritsFrom("TH3")):
113			counter = next_counter()
114			name = obj.GetName()
115			hist = files[0].file.GetDirectory(path).Get(name)
116			title = hist.GetTitle()
117			x_title = hist.GetXaxis().GetTitle()
118			y_title = hist.GetYaxis().GetTitle()
119			if "Norm" in name or options.normalize:
120			y_title = "Fraction of Events in Bin"
121			hist.Draw()
122			stack = ROOT.THStack("st%.3i" % int(counter), title)
123	klukas	1.5	legend_height = 0.04 * len(files) + 0.02
124			legend = ROOT.TLegend(0.65, 0.89 - legend_height, 0.87, 0.89)
125	klukas	1.1	c1.SetLogx("Logx" in name)
126			c1.SetLogy("Logy" in name)
127			for i, file in enumerate(files):
128			hist = file.file.GetDirectory(path).Get(name)
129			if not hist: continue
130			hist.Draw()
131			hist.SetTitle(file.name)
132			color = colors[i % len(colors)]
133			hist.SetLineColor(color)
134	klukas	1.7	if options.markers:
135			hist.SetMarkerColor(color)
136			hist.SetMarkerStyle(marker_styles[i])
137			else:
138			hist.SetMarkerSize(0)
139	klukas	1.1	if "Norm" in name or options.normalize:
140			integral = hist.Integral()
141	klukas	1.4	hist.Scale(1. / integral)
142	klukas	1.1	stack.Add(hist)
143			legend.AddEntry(hist)
144			stack.Draw("nostack p H")
145			stack.SetTitle("%s;%s;%s" % (title, x_title, y_title))
146			if "Eff" in name:
147			stack.Draw("nostack e p")
148			stack.SetMaximum(1.)
149			stack.SetMinimum(0.)
150			legend.Draw()
151	klukas	1.8	save_plot(stack, plot_dir, path, name, counter)
152	klukas	1.1
153
154
155	klukas	1.8	def save_plot(stack, plot_dir, path, name, counter):
156			if options.ext == "pdf":
157			c1.SaveAs("%.3i.pdf" % counter)
158			c1.SaveAs("%s/%s/%s.%s" % (plot_dir, path, name, options.ext))
159			report_progress(counter, 1)
160
161
162
163			def report_progress(counter, divisor):
164			if counter % divisor == 0:
165			print "\r%i plots written to %s" % (counter, options.output),
166			sys.stdout.flush()
167
168
169
170			def merge_pdf():
171			print "Writing merged pdf..."
172			os.system("gs -q -dBATCH -dNOPAUSE -sDEVICE=pdfwrite "
173			"-dAutoRotatePages=/All "
174			"-sOutputFile=%s.pdf " % options.output +
175			"[0-9][0-9][0-9].pdf")
176			os.system("rm [0-9]*.pdf")
177
178
179
180	klukas	1.1	def counter_generator():
181			k = 0
182			while True:
183			k += 1
184			yield k
185			next_counter = counter_generator().next
186
187
188
189			if __name__ == "__main__":
190	klukas	1.8	if options.timing:
191			import cProfile
192			cProfile.run('main()', 'fooprof')
193			import pstats
194			p = pstats.Stats('fooprof')
195			p.sort_stats('cumulative').print_stats(15)
196			else:
197			sys.exit(main())
198	klukas	1.1