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()
            
            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.Draw("HIST")
            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].Draw("HIST")
            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].Draw("E")
            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):
            Stack.GetXaxis().SetTitle(xAxisLabel)
        elif(numSignalSamples is not 0):
            SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
        elif(numDataSamples is not 0):
            DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)

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

        LumiLabel.Draw()
        if options.normalize:
            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.7
Committed:	Tue Feb 5 10:50:45 2013 UTC (12 years, 2 months ago) by lantonel
Content type:	text/x-python
Branch:	MAIN
Changes since 1.6:	+3 -22 lines
Log Message:	switched to using the processingUtilities features
#	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
134			if( types[sample] == "bgMC"):
135	lantonel	1.3
136			numBgMCSamples += 1
137	lantonel	1.1
138			Histogram.SetFillStyle(1001)
139			Histogram.SetFillColor(colors[sample])
140			Histogram.SetLineColor(1)
141			Histogram.SetLineStyle(1)
142			Histogram.SetLineWidth(1)
143	lantonel	1.4
144			backgroundIntegral += Histogram.Integral()
145
146			BgMCLegend.AddEntry(Histogram,labels[sample],"F")
147			BgMCHistograms.append(Histogram)
148	lantonel	1.1
149			elif( types[sample] == "signalMC"):
150	lantonel	1.3
151			numSignalSamples += 1
152	lantonel	1.1
153			Histogram.SetFillStyle(0)
154			Histogram.SetLineColor(colors[sample])
155			Histogram.SetLineStyle(1)
156			Histogram.SetLineWidth(2)
157
158			SignalMCLegend.AddEntry(Histogram,labels[sample],"L")
159			SignalMCHistograms.append(Histogram)
160
161			elif( types[sample] == "data"):
162	lantonel	1.3
163			numDataSamples += 1
164	lantonel	1.1
165			Histogram.SetFillStyle(0)
166			Histogram.SetLineColor(colors[sample])
167			Histogram.SetLineStyle(1)
168			Histogram.SetLineWidth(2)
169	lantonel	1.4
170			dataIntegral += Histogram.Integral()
171	lantonel	1.1
172			BgMCLegend.AddEntry(Histogram,labels[sample],"LEP")
173			DataHistograms.append(Histogram)
174
175	lantonel	1.5	if dataIntegral > 0 and backgroundIntegral > 0:
176	lantonel	1.4	scaleFactor = dataIntegral/backgroundIntegral
177			for bgMCHist in BgMCHistograms:
178			if options.normalize:
179			bgMCHist.Scale(scaleFactor)
180			Stack.Add(bgMCHist)
181
182	lantonel	1.1
183			stackMax = Stack.GetMaximum()
184			finalMax = stackMax
185			for signalMCHist in SignalMCHistograms:
186			if(signalMCHist.GetMaximum() > finalMax):
187			finalMax = signalMCHist.GetMaximum()
188			for dataHist in DataHistograms:
189			if(dataHist.GetMaximum() > finalMax):
190			finalMax = dataHist.GetMaximum()
191	lantonel	1.4
192
193			if len(DataHistograms) is 1:
194			dataIntegral += DataHistograms[0].Integral()
195	lantonel	1.1
196	lantonel	1.3	outputFile.cd(rootDirectory+"/"+channel)
197
198			if(numBgMCSamples is not 0):
199			Stack.Draw("HIST")
200			Stack.SetMaximum(1.1*finalMax)
201			for signalMCHist in SignalMCHistograms:
202			signalMCHist.Draw("HIST SAME")
203			for dataHist in DataHistograms:
204			dataHist.Draw("E SAME")
205
206			elif(numSignalSamples is not 0):
207			SignalMCHistograms[0].Draw("HIST")
208			SignalMCHistograms[0].SetMaximum(1.1*finalMax)
209			for signalMCHist in SignalMCHistograms:
210			if(signalMCHist is not SignalMCHistograms[0]):
211			signalMCHist.Draw("HIST SAME")
212			for dataHist in DataHistograms:
213			dataHist.Draw("E SAME")
214			elif(numDataSamples is not 0):
215			DataHistograms[0].Draw("E")
216			DataHistograms[0].SetMaximum(1.1*finalMax)
217			for dataHist in DataHistograms:
218			if(dataHist is not DataHistograms[0]):
219			dataHist.Draw("E SAME")
220
221	lantonel	1.4
222
223	lantonel	1.3
224			if(numBgMCSamples is not 0):
225			Stack.GetXaxis().SetTitle(xAxisLabel)
226			elif(numSignalSamples is not 0):
227			SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
228			elif(numDataSamples is not 0):
229			DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
230
231			if(numBgMCSamples is not 0 or numDataSamples is not 0):
232			BgMCLegend.Draw()
233			if(numSignalSamples is not 0):
234			SignalMCLegend.Draw()
235	lantonel	1.1
236			LumiLabel.Draw()
237	lantonel	1.4	if options.normalize:
238			NormLabel = TPaveLabel(0.1,0.75,0.35,0.85,"MC scaled to data","NDC")
239			NormLabel.SetBorderSize(0)
240			NormLabel.SetFillColor(0)
241			NormLabel.SetFillStyle(0)
242			NormLabel.Draw()
243
244
245	lantonel	1.1	Canvas.Write()
246
247	ahart	1.6	for dataset in processed_datasets:
248			dataset_file = "%s/%s.root_tmp" % (condor_dir,dataset)
249			os.remove(dataset_file)
250	lantonel	1.1
251			outputFile.Close()