Configuration/scripts/makePlots.py

#!/usr/bin/env python
import sys
import os
import re
from array import *
from decimal import *
from optparse import OptionParser
from OSUT3Analysis.Configuration.configurationOptions import *
from OSUT3Analysis.Configuration.processingUtilities import *

parser = OptionParser()
parser = set_commandline_arguments(parser)

(options, args) = parser.parse_args()

if options.localConfig:
    sys.path.append(os.getcwd())
    exec("from " + options.localConfig.rstrip('.py') + " import *")

condor_dir = set_condor_output_dir(options)


from ROOT import TFile, gROOT, gStyle, gDirectory, TStyle, THStack, TH1F, TCanvas, TString, TLegend, TArrow, THStack, TIter, TKey, TPaveLabel

gROOT.SetBatch()
gStyle.SetOptStat(0)
gStyle.SetCanvasBorderMode(0)
gStyle.SetPadBorderMode(0)
gStyle.SetPadColor(0)
gStyle.SetCanvasColor(0)
gStyle.SetTextFont(42)
gROOT.ForceStyle()
outputFile = TFile(condor_dir + "/stacked_histograms.root", "RECREATE")

channels = []
processed_datasets = []

#### check which input datasets have valid output files
for sample in datasets:
    fileName = condor_dir + "/" + sample + ".root"
    if not os.path.exists(fileName):
        continue
    testFile = TFile(fileName)
    if not (testFile.IsZombie()):
        processed_datasets.append(sample)

if len(processed_datasets) is 0:
    sys.exit("No datasets have been processed")

#### open first input file and re-make its directory structure in the output file
testFile = TFile(condor_dir + "/" + processed_datasets[0] + ".root")
testFile.cd()
for key in testFile.GetListOfKeys():
    if (key.GetClassName() != "TDirectoryFile"):
        continue
    outputFile.cd()
    outputFile.mkdir(key.GetName())
    rootDirectory = key.GetName()
    
    testFile.cd(key.GetName())
    for key2 in gDirectory.GetListOfKeys():
        if (key2.GetClassName() != "TDirectoryFile"):
            continue
        outputFile.cd(key.GetName())
        gDirectory.mkdir(key2.GetName())
        channels.append(key2.GetName())


weight = intLumi / 10000.0
for dataset in processed_datasets:
    dataset_file = "%s/%s.root" % (condor_dir,dataset)
    os.system("mergeTFileServiceHistograms -i %s -o %s -w %g" % (dataset_file, dataset_file + "_tmp", weight))

for channel in channels: # loop over final states, which each have their own directory

    testFile.cd(rootDirectory+"/"+channel)
    histograms = []
    for key in gDirectory.GetListOfKeys():
        if re.match ('TH1', key.GetClassName()):
            histograms.append(key.GetName())
    
    for histogramName in histograms: # loop over histograms in the current directory

        numBgMCSamples = 0
        numDataSamples = 0
        numSignalSamples = 0

        Stack = THStack("stack",histogramName)
        
        if(intLumi < 1000.):
            LumiText = "L_{int} = " + str(intLumi) + " pb^{-1}"
        else:
            getcontext().prec = 2 
            LumiInFb = intLumi/1000.
            LumiText = "L_{int} = " + str(LumiInFb) + " fb^{-1}"

        LumiLabel = TPaveLabel(0.1,0.8,0.34,0.9,LumiText,"NDC")
        LumiLabel.SetBorderSize(0)
        LumiLabel.SetFillColor(0)
        LumiLabel.SetFillStyle(0)
        
        BgMCLegend = TLegend(0.70,0.65,0.99,0.89, "Data & Bkgd. MC")
        BgMCLegend.SetBorderSize(0)
        BgMCLegend.SetFillColor(0)
        BgMCLegend.SetFillStyle(0)
        SignalMCLegend = TLegend(0.45,0.65,0.70,0.89,"Signal MC")
        SignalMCLegend.SetBorderSize(0)
        SignalMCLegend.SetFillColor(0)
        SignalMCLegend.SetFillStyle(0)

        
        outputFile.cd(rootDirectory+"/"+channel)
        Canvas = TCanvas(histogramName)
        BgMCHistograms = []
        SignalMCHistograms = []
        DataHistograms = []

        backgroundIntegral = 0
        dataIntegral = 0
        scaleFactor = 1

        
        for sample in processed_datasets: # loop over different samples as listed in configurationOptions.py
            dataset_file = "%s/%s.root_tmp" % (condor_dir,sample)
            inputFile = TFile(dataset_file)
            if(inputFile.IsZombie()):
                continue
            Histogram = inputFile.Get(rootDirectory+"/"+channel+"/"+histogramName).Clone()
            Histogram.SetDirectory(0)
            inputFile.Close()
            xAxisLabel = Histogram.GetXaxis().GetTitle()
            histoTitle = Histogram.GetTitle()
            
            if( types[sample] == "bgMC"):

                numBgMCSamples += 1
                
                Histogram.SetFillStyle(1001)
                Histogram.SetFillColor(colors[sample])
                Histogram.SetLineColor(1)
                Histogram.SetLineStyle(1)
                Histogram.SetLineWidth(1)

                backgroundIntegral += Histogram.Integral()
            
                BgMCLegend.AddEntry(Histogram,labels[sample],"F")
                BgMCHistograms.append(Histogram)
                
            elif( types[sample] == "signalMC"):

                numSignalSamples += 1
                
                Histogram.SetFillStyle(0)
                Histogram.SetLineColor(colors[sample])
                Histogram.SetLineStyle(1)
                Histogram.SetLineWidth(2)
                
                SignalMCLegend.AddEntry(Histogram,labels[sample],"L")
                SignalMCHistograms.append(Histogram)

            elif( types[sample] == "data"):

                numDataSamples += 1
                
                Histogram.SetFillStyle(0)
                Histogram.SetLineColor(colors[sample])
                Histogram.SetLineStyle(1)
                Histogram.SetLineWidth(2)

                dataIntegral += Histogram.Integral()
                
                BgMCLegend.AddEntry(Histogram,labels[sample],"LEP")
                DataHistograms.append(Histogram)

        if dataIntegral > 0 and backgroundIntegral > 0:
            scaleFactor = dataIntegral/backgroundIntegral
        for bgMCHist in BgMCHistograms:
            if options.normalize:
                bgMCHist.Scale(scaleFactor)
            Stack.Add(bgMCHist)

                
        stackMax = Stack.GetMaximum()
        finalMax = stackMax
        for signalMCHist in SignalMCHistograms:
            if(signalMCHist.GetMaximum() > finalMax):
                finalMax = signalMCHist.GetMaximum()
        for dataHist in DataHistograms:
            if(dataHist.GetMaximum() > finalMax):
                finalMax = dataHist.GetMaximum()


        if len(DataHistograms) is 1:
            dataIntegral += DataHistograms[0].Integral()
 
        outputFile.cd(rootDirectory+"/"+channel)        

        if(numBgMCSamples is not 0):
            Stack.SetTitle(histoTitle)
            Stack.Draw("HIST")
            Stack.GetXaxis().SetTitle(xAxisLabel)
            Stack.SetMaximum(1.1*finalMax)
            for signalMCHist in SignalMCHistograms:
                signalMCHist.Draw("HIST SAME")
            for dataHist in DataHistograms:
                dataHist.Draw("E SAME")
                                
        elif(numSignalSamples is not 0):
            SignalMCHistograms[0].SetTitle(histoTitle)
            SignalMCHistograms[0].Draw("HIST")
            SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
            SignalMCHistograms[0].SetMaximum(1.1*finalMax)
            for signalMCHist in SignalMCHistograms:
                if(signalMCHist is not SignalMCHistograms[0]):
                    signalMCHist.Draw("HIST SAME")
            for dataHist in DataHistograms:
                dataHist.Draw("E SAME")
        elif(numDataSamples is not 0):
            DataHistograms[0].SetTitle(histoTitle)
            DataHistograms[0].Draw("E")
            DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
            DataHistograms[0].SetMaximum(1.1*finalMax)
            for dataHist in DataHistograms:
                if(dataHist is not DataHistograms[0]):
                    dataHist.Draw("E SAME")

            
        if(numBgMCSamples is not 0 or numDataSamples is not 0):
            BgMCLegend.Draw()
        if(numSignalSamples is not 0):
            SignalMCLegend.Draw()

        LumiLabel.Draw()
        if options.normalize and numBgMCSamples > 0 and numDataSamples > 0:
            NormLabel = TPaveLabel(0.1,0.75,0.35,0.85,"MC scaled to data","NDC")
            NormLabel.SetBorderSize(0)
            NormLabel.SetFillColor(0)
            NormLabel.SetFillStyle(0)
            NormLabel.Draw()


        Canvas.Write()

for dataset in processed_datasets:
    dataset_file = "%s/%s.root_tmp" % (condor_dir,dataset)
    os.remove(dataset_file)

outputFile.Close()         
Revision:	1.8
Committed:	Tue Feb 5 13:52:03 2013 UTC (12 years, 2 months ago) by lantonel
Content type:	text/x-python
Branch:	MAIN
Changes since 1.7:	+9 -11 lines
Log Message:	fixed the setting of the canvas title
#	User	Rev	Content
1	lantonel	1.1	#!/usr/bin/env python
2			import sys
3			import os
4	ahart	1.6	import re
5	lantonel	1.1	from array import *
6	lantonel	1.3	from decimal import *
7	lantonel	1.1	from optparse import OptionParser
8	lantonel	1.2	from OSUT3Analysis.Configuration.configurationOptions import *
9	lantonel	1.7	from OSUT3Analysis.Configuration.processingUtilities import *
10	lantonel	1.1
11			parser = OptionParser()
12	lantonel	1.7	parser = set_commandline_arguments(parser)
13	lantonel	1.4
14	lantonel	1.1	(options, args) = parser.parse_args()
15
16			if options.localConfig:
17			sys.path.append(os.getcwd())
18			exec("from " + options.localConfig.rstrip('.py') + " import *")
19
20	lantonel	1.7	condor_dir = set_condor_output_dir(options)
21	lantonel	1.1
22
23			from ROOT import TFile, gROOT, gStyle, gDirectory, TStyle, THStack, TH1F, TCanvas, TString, TLegend, TArrow, THStack, TIter, TKey, TPaveLabel
24
25			gROOT.SetBatch()
26			gStyle.SetOptStat(0)
27			gStyle.SetCanvasBorderMode(0)
28			gStyle.SetPadBorderMode(0)
29			gStyle.SetPadColor(0)
30			gStyle.SetCanvasColor(0)
31			gStyle.SetTextFont(42)
32	lantonel	1.3	gROOT.ForceStyle()
33	lantonel	1.1	outputFile = TFile(condor_dir + "/stacked_histograms.root", "RECREATE")
34
35			channels = []
36			processed_datasets = []
37
38			#### check which input datasets have valid output files
39			for sample in datasets:
40			fileName = condor_dir + "/" + sample + ".root"
41			if not os.path.exists(fileName):
42			continue
43			testFile = TFile(fileName)
44			if not (testFile.IsZombie()):
45			processed_datasets.append(sample)
46
47			if len(processed_datasets) is 0:
48			sys.exit("No datasets have been processed")
49
50			#### open first input file and re-make its directory structure in the output file
51			testFile = TFile(condor_dir + "/" + processed_datasets[0] + ".root")
52			testFile.cd()
53			for key in testFile.GetListOfKeys():
54			if (key.GetClassName() != "TDirectoryFile"):
55			continue
56			outputFile.cd()
57			outputFile.mkdir(key.GetName())
58			rootDirectory = key.GetName()
59
60			testFile.cd(key.GetName())
61			for key2 in gDirectory.GetListOfKeys():
62			if (key2.GetClassName() != "TDirectoryFile"):
63			continue
64			outputFile.cd(key.GetName())
65			gDirectory.mkdir(key2.GetName())
66			channels.append(key2.GetName())
67
68
69	ahart	1.6	weight = intLumi / 10000.0
70			for dataset in processed_datasets:
71			dataset_file = "%s/%s.root" % (condor_dir,dataset)
72			os.system("mergeTFileServiceHistograms -i %s -o %s -w %g" % (dataset_file, dataset_file + "_tmp", weight))
73	lantonel	1.1
74			for channel in channels: # loop over final states, which each have their own directory
75
76			testFile.cd(rootDirectory+"/"+channel)
77			histograms = []
78			for key in gDirectory.GetListOfKeys():
79	ahart	1.6	if re.match ('TH1', key.GetClassName()):
80			histograms.append(key.GetName())
81	lantonel	1.1
82			for histogramName in histograms: # loop over histograms in the current directory
83
84	lantonel	1.3	numBgMCSamples = 0
85			numDataSamples = 0
86			numSignalSamples = 0
87
88	lantonel	1.1	Stack = THStack("stack",histogramName)
89
90			if(intLumi < 1000.):
91			LumiText = "L_{int} = " + str(intLumi) + " pb^{-1}"
92			else:
93	lantonel	1.3	getcontext().prec = 2
94			LumiInFb = intLumi/1000.
95	lantonel	1.1	LumiText = "L_{int} = " + str(LumiInFb) + " fb^{-1}"
96
97	lantonel	1.3	LumiLabel = TPaveLabel(0.1,0.8,0.34,0.9,LumiText,"NDC")
98			LumiLabel.SetBorderSize(0)
99			LumiLabel.SetFillColor(0)
100			LumiLabel.SetFillStyle(0)
101	lantonel	1.1
102			BgMCLegend = TLegend(0.70,0.65,0.99,0.89, "Data & Bkgd. MC")
103	lantonel	1.3	BgMCLegend.SetBorderSize(0)
104			BgMCLegend.SetFillColor(0)
105			BgMCLegend.SetFillStyle(0)
106	lantonel	1.1	SignalMCLegend = TLegend(0.45,0.65,0.70,0.89,"Signal MC")
107	lantonel	1.3	SignalMCLegend.SetBorderSize(0)
108			SignalMCLegend.SetFillColor(0)
109			SignalMCLegend.SetFillStyle(0)
110	lantonel	1.1
111
112
113			outputFile.cd(rootDirectory+"/"+channel)
114			Canvas = TCanvas(histogramName)
115	lantonel	1.4	BgMCHistograms = []
116	lantonel	1.1	SignalMCHistograms = []
117			DataHistograms = []
118	lantonel	1.4
119			backgroundIntegral = 0
120			dataIntegral = 0
121			scaleFactor = 1
122
123	lantonel	1.1
124			for sample in processed_datasets: # loop over different samples as listed in configurationOptions.py
125	ahart	1.6	dataset_file = "%s/%s.root_tmp" % (condor_dir,sample)
126			inputFile = TFile(dataset_file)
127	lantonel	1.1	if(inputFile.IsZombie()):
128			continue
129			Histogram = inputFile.Get(rootDirectory+"/"+channel+"/"+histogramName).Clone()
130			Histogram.SetDirectory(0)
131	ahart	1.6	inputFile.Close()
132	lantonel	1.1	xAxisLabel = Histogram.GetXaxis().GetTitle()
133	lantonel	1.8	histoTitle = Histogram.GetTitle()
134	lantonel	1.1
135			if( types[sample] == "bgMC"):
136	lantonel	1.3
137			numBgMCSamples += 1
138	lantonel	1.1
139			Histogram.SetFillStyle(1001)
140			Histogram.SetFillColor(colors[sample])
141			Histogram.SetLineColor(1)
142			Histogram.SetLineStyle(1)
143			Histogram.SetLineWidth(1)
144	lantonel	1.4
145			backgroundIntegral += Histogram.Integral()
146
147			BgMCLegend.AddEntry(Histogram,labels[sample],"F")
148			BgMCHistograms.append(Histogram)
149	lantonel	1.1
150			elif( types[sample] == "signalMC"):
151	lantonel	1.3
152			numSignalSamples += 1
153	lantonel	1.1
154			Histogram.SetFillStyle(0)
155			Histogram.SetLineColor(colors[sample])
156			Histogram.SetLineStyle(1)
157			Histogram.SetLineWidth(2)
158
159			SignalMCLegend.AddEntry(Histogram,labels[sample],"L")
160			SignalMCHistograms.append(Histogram)
161
162			elif( types[sample] == "data"):
163	lantonel	1.3
164			numDataSamples += 1
165	lantonel	1.1
166			Histogram.SetFillStyle(0)
167			Histogram.SetLineColor(colors[sample])
168			Histogram.SetLineStyle(1)
169			Histogram.SetLineWidth(2)
170	lantonel	1.4
171			dataIntegral += Histogram.Integral()
172	lantonel	1.1
173			BgMCLegend.AddEntry(Histogram,labels[sample],"LEP")
174			DataHistograms.append(Histogram)
175
176	lantonel	1.5	if dataIntegral > 0 and backgroundIntegral > 0:
177	lantonel	1.4	scaleFactor = dataIntegral/backgroundIntegral
178			for bgMCHist in BgMCHistograms:
179			if options.normalize:
180			bgMCHist.Scale(scaleFactor)
181			Stack.Add(bgMCHist)
182
183	lantonel	1.1
184			stackMax = Stack.GetMaximum()
185			finalMax = stackMax
186			for signalMCHist in SignalMCHistograms:
187			if(signalMCHist.GetMaximum() > finalMax):
188			finalMax = signalMCHist.GetMaximum()
189			for dataHist in DataHistograms:
190			if(dataHist.GetMaximum() > finalMax):
191			finalMax = dataHist.GetMaximum()
192	lantonel	1.4
193
194			if len(DataHistograms) is 1:
195			dataIntegral += DataHistograms[0].Integral()
196	lantonel	1.1
197	lantonel	1.3	outputFile.cd(rootDirectory+"/"+channel)
198
199			if(numBgMCSamples is not 0):
200	lantonel	1.8	Stack.SetTitle(histoTitle)
201	lantonel	1.3	Stack.Draw("HIST")
202	lantonel	1.8	Stack.GetXaxis().SetTitle(xAxisLabel)
203	lantonel	1.3	Stack.SetMaximum(1.1*finalMax)
204			for signalMCHist in SignalMCHistograms:
205			signalMCHist.Draw("HIST SAME")
206			for dataHist in DataHistograms:
207			dataHist.Draw("E SAME")
208
209			elif(numSignalSamples is not 0):
210	lantonel	1.8	SignalMCHistograms[0].SetTitle(histoTitle)
211	lantonel	1.3	SignalMCHistograms[0].Draw("HIST")
212	lantonel	1.8	SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
213	lantonel	1.3	SignalMCHistograms[0].SetMaximum(1.1*finalMax)
214			for signalMCHist in SignalMCHistograms:
215			if(signalMCHist is not SignalMCHistograms[0]):
216			signalMCHist.Draw("HIST SAME")
217			for dataHist in DataHistograms:
218			dataHist.Draw("E SAME")
219			elif(numDataSamples is not 0):
220	lantonel	1.8	DataHistograms[0].SetTitle(histoTitle)
221	lantonel	1.3	DataHistograms[0].Draw("E")
222	lantonel	1.8	DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
223	lantonel	1.3	DataHistograms[0].SetMaximum(1.1*finalMax)
224			for dataHist in DataHistograms:
225			if(dataHist is not DataHistograms[0]):
226			dataHist.Draw("E SAME")
227
228	lantonel	1.8
229	lantonel	1.3	if(numBgMCSamples is not 0 or numDataSamples is not 0):
230			BgMCLegend.Draw()
231			if(numSignalSamples is not 0):
232			SignalMCLegend.Draw()
233	lantonel	1.1
234			LumiLabel.Draw()
235	lantonel	1.8	if options.normalize and numBgMCSamples > 0 and numDataSamples > 0:
236	lantonel	1.4	NormLabel = TPaveLabel(0.1,0.75,0.35,0.85,"MC scaled to data","NDC")
237			NormLabel.SetBorderSize(0)
238			NormLabel.SetFillColor(0)
239			NormLabel.SetFillStyle(0)
240			NormLabel.Draw()
241
242
243	lantonel	1.1	Canvas.Write()
244
245	ahart	1.6	for dataset in processed_datasets:
246			dataset_file = "%s/%s.root_tmp" % (condor_dir,dataset)
247			os.remove(dataset_file)
248	lantonel	1.1
249			outputFile.Close()