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 *

parser = OptionParser()

parser.add_option("-l", "--localConfig", dest="localConfig",
                  help="local configuration file")

parser.add_option("-c", "--condorDir", dest="condorDir",
                  help="condor output directory")

parser.add_option("-n", "--normalize", action="store_true", dest="normalize", default=False,
                  help="normalize total background MC yield to the data")

(options, args) = parser.parse_args()

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

if options.condorDir:
    condor_dir = "condor/%s" % options.condorDir
else: #get most recent condor submission directory
   dir_list = []
   for directory in os.listdir("./condor/"):
       if directory.find("condor_") is not -1:
           dir_list.append(directory)
   if len(dir_list) is 0:
       sys.exit("Cannot find last condor working directory")
   dir_list.sort(reverse=True)
   condor_dir = "condor/%s" % dir_list[0]


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