Configuration/scripts/makePlots.py

#!/usr/bin/env python
import sys
import os
import re
from math import *
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)
(arguments, args) = parser.parse_args()

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


outputFileName = "stacked_histograms.root"
if arguments.outputFileName:
    outputFileName = arguments.outputFileName

condor_dir = set_condor_output_dir(arguments)


#### deal with conflicting arguments
if arguments.normalizeToData and arguments.normalizeToUnitArea:
    print "Conflicting normalizations requsted, will normalize to unit area"
    arguments.normalizeToData = False
if arguments.normalizeToData and arguments.noStack:
    print "You have asked to scale non-stacked backgrounds to data.  This is a very strange request.  Will normalize to unit area instead"
    arguments.normalizeToData = False
    arguments.normalizeToUnitArea = True
if arguments.makeRatioPlots and arguments.makeDiffPlots:
    print "You have requested both ratio and difference plots.  Will make just ratio plots instead"
    arguments.makeRatioPlots = False

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

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 + "/" + outputFileName, "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 testFile.IsZombie() or not testFile.GetNkeys():
        continue
    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)
##     fin = TFile (dataset_file)
##     flags = fin.Get ("flags")
##     noWeights = flags and flags.GetBinContent (1)
##     fin.Close ()

##     if types[dataset] != "data" and not noWeights:
##         os.system("mergeTFileServiceHistograms -i %s -o %s -w %g" % (dataset_file, dataset_file + "_tmp", weight))
##     else:
##         os.system("mergeTFileServiceHistograms -i %s -o %s -w %g" % (dataset_file, dataset_file + "_tmp", 1.0))


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

    testFile.cd(rootDirectory+"/"+channel)

    for key in gDirectory.GetListOfKeys(): # loop over histograms in the current directory
        histogramName = key.GetName()
        
        if re.match ('TH1', key.GetClassName()): # plot a 1-D histogram
                        
            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()
            BgTitle = BgMCLegend.AddEntry(0, "Data & Bkgd. MC", "H")
            BgTitle.SetTextAlign(21)
            BgTitle.SetTextFont(62)
            BgMCLegend.SetBorderSize(0)
            BgMCLegend.SetFillColor(0)
            BgMCLegend.SetFillStyle(0)
            SignalMCLegend = TLegend()
            SignalTitle = SignalMCLegend.AddEntry(0, "Signal MC", "H")
            SignalTitle.SetTextAlign(21)
            SignalTitle.SetTextFont(62)
            SignalMCLegend.SetBorderSize(0)
            SignalMCLegend.SetFillColor(0)
            SignalMCLegend.SetFillStyle(0)
            
            outputFile.cd(rootDirectory+"/"+channel)
            Canvas = TCanvas(histogramName)
            BgMCHistograms = []
            BgMCLegendEntries = []
            SignalMCHistograms = []
            SignalMCLegendEntries = []
            DataHistograms = []
            DataLegendEntries = []
            
            
            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" % (condor_dir,sample)
                inputFile = TFile(dataset_file)
                Histogram = inputFile.Get(rootDirectory+"/"+channel+"/"+histogramName).Clone()
                Histogram.SetDirectory(0)
                if arguments.rebinFactor:
                    RebinFactor = int(arguments.rebinFactor)
                    if Histogram.GetNbinsX() >= RebinFactor*10:
                        Histogram.Rebin(RebinFactor)
                inputFile.Close()
                xAxisLabel = Histogram.GetXaxis().GetTitle()
                histoTitle = Histogram.GetTitle()

                legLabel = labels[sample]
                if (arguments.printYields):
                    yieldHist = Histogram.Integral()
                    legLabel = legLabel + " (%.1f)" % yieldHist

                if( types[sample] == "bgMC"):
                    
                    numBgMCSamples += 1
                    backgroundIntegral += Histogram.Integral()

                    Histogram.SetLineStyle(1)
                    if(arguments.noStack):
                        Histogram.SetFillStyle(0)
                        Histogram.SetLineColor(colors[sample])
                        Histogram.SetLineWidth(2)
                    else:
                        Histogram.SetFillStyle(1001)
                        Histogram.SetFillColor(colors[sample])
                        Histogram.SetLineColor(1)
                        Histogram.SetLineWidth(1)

                    BgMCLegendEntries.append(legLabel) 
                    BgMCHistograms.append(Histogram)

                    
                elif( types[sample] == "signalMC"):
                    
                    numSignalSamples += 1

                    Histogram.SetFillStyle(0)
                    Histogram.SetLineColor(colors[sample])
                    Histogram.SetLineStyle(1)
                    Histogram.SetLineWidth(2)
                    if(arguments.normalizeToUnitArea and Histogram.Integral() > 0):
                        Histogram.Scale(1./Histogram.Integral())

                    SignalMCLegendEntries.append(legLabel)
                    SignalMCHistograms.append(Histogram)
                    
                elif( types[sample] == "data"):

                    numDataSamples += 1
                    dataIntegral += Histogram.Integral()

                    Histogram.SetFillStyle(0)
                    Histogram.SetLineColor(colors[sample])
                    Histogram.SetLineStyle(1)
                    Histogram.SetLineWidth(2)
                    if(arguments.normalizeToUnitArea and Histogram.Integral() > 0):
                        Histogram.Scale(1./Histogram.Integral())

                    DataLegendEntries.append(legLabel)
                    DataHistograms.append(Histogram)
                    
            if dataIntegral > 0 and backgroundIntegral > 0:
                scaleFactor = dataIntegral/backgroundIntegral
            for bgMCHist in BgMCHistograms:
                if arguments.normalizeToData:
                    bgMCHist.Scale(scaleFactor)
                if arguments.normalizeToUnitArea and not arguments.noStack and backgroundIntegral > 0:
                    bgMCHist.Scale(1./backgroundIntegral)
                elif arguments.normalizeToUnitArea and arguments.noStack and bgMCHist.Integral() > 0:
                    bgMCHist.Scale(1./bgMCHist.Integral())
                if not arguments.noStack:
                    Stack.Add(bgMCHist)

                
            ### formatting data histograms and adding to legend
            counter = 0
            for Histogram in DataHistograms:
                    BgMCLegend.AddEntry(Histogram,DataLegendEntries[counter],"LEP").SetTextFont (42)
                    counter = counter+1


            ### creating the histogram to represent the statistical errors on the stack
            if numBgMCSamples is not 0 and not arguments.noStack: 
                ErrorHisto = BgMCHistograms[0].Clone("errors")
                ErrorHisto.SetFillStyle(3001)
                ErrorHisto.SetFillColor(13)
                ErrorHisto.SetLineWidth(0)
                BgMCLegend.AddEntry(ErrorHisto,"Stat. Errors","F").SetTextFont (42)
                for Histogram in BgMCHistograms:
                    if Histogram is not BgMCHistograms[0]:
                        ErrorHisto.Add(Histogram)


            ### formatting bgMC histograms and adding to legend
            counter = numBgMCSamples-1
            for Histogram in reversed(BgMCHistograms):
                    if(arguments.noStack):
                        BgMCLegend.AddEntry(Histogram,BgMCLegendEntries[counter],"L").SetTextFont (42)
                    else:
                        BgMCLegend.AddEntry(Histogram,BgMCLegendEntries[counter],"F").SetTextFont (42)
                    counter = counter-1

            ### formatting signalMC histograms and adding to legend
            counter = 0
            for Histogram in SignalMCHistograms:
                    SignalMCLegend.AddEntry(Histogram,SignalMCLegendEntries[counter],"L").SetTextFont (42)
                    counter = counter+1


            ### finding the maximum value of anything going on the canvas, and know how to set the y-axis
            finalMax = 0
            if numBgMCSamples is not 0 and not arguments.noStack:
                finalMax = ErrorHisto.GetMaximum() + ErrorHisto.GetBinError(ErrorHisto.GetMaximumBin())
            else:
                for bgMCHist in BgMCHistograms:
                    if(bgMCHist.GetMaximum() > finalMax):
                        finalMax = bgMCHist.GetMaximum()
            for signalMCHist in SignalMCHistograms:
                if(signalMCHist.GetMaximum() > finalMax):
                    finalMax = signalMCHist.GetMaximum()
            for dataHist in DataHistograms:
                if(dataHist.GetMaximum() > finalMax):
                    finalMax = dataHist.GetMaximum() + dataHist.GetBinError(dataHist.GetMaximumBin())
            finalMax = 1.1*finalMax


            ### Drawing histograms to canvas


            outputFile.cd(rootDirectory+"/"+channel)        

            makeRatioPlots = arguments.makeRatioPlots
            makeDiffPlots = arguments.makeDiffPlots

            if numBgMCSamples is 0 or numDataSamples is not 1:
                makeRatioPlots = False
                makeDiffPlots = False
            if makeRatioPlots or makeDiffPlots:
                Canvas.SetFillStyle(0)
                Canvas.Divide(1,2)
                Canvas.cd(1)
                gPad.SetPad(0.01,0.25,0.99,0.99)
                gPad.SetMargin(0.1,0.05,0.02,0.07)
                gPad.SetFillStyle(0)
                gPad.Update()
                gPad.Draw()
                Canvas.cd(2)
                gPad.SetPad(0.01,0.01,0.99,0.25)
                #format: gPad.SetMargin(l,r,b,t)
                gPad.SetMargin(0.1,0.05,0.4,0.02)
                gPad.SetFillStyle(0)
                gPad.SetGridy(1)
                gPad.Update()
                gPad.Draw()

                Canvas.cd(1)

            if numBgMCSamples is not 0:

                if not arguments.noStack:
                    Stack.SetTitle(histoTitle)
                    Stack.Draw("HIST")
                    Stack.GetXaxis().SetTitle(xAxisLabel)
                    Stack.SetMaximum(finalMax)
                    Stack.SetMinimum(0.0001)
                    if makeRatioPlots or makeDiffPlots:
                        Stack.GetHistogram().GetXaxis().SetLabelSize(0)
                    #draw shaded error bands
                    ErrorHisto.Draw("A E2 SAME")

                    
                else:
                    BgMCHistograms[0].SetTitle(histoTitle)
                    BgMCHistograms[0].Draw("HIST")
                    BgMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                    BgMCHistograms[0].SetMaximum(finalMax)
                    BgMCHistograms[0].SetMinimum(0.0001)
                    for bgMCHist in BgMCHistograms:
                        bgMCHist.Draw("A HIST SAME")
                for signalMCHist in SignalMCHistograms:
                    signalMCHist.Draw("A HIST SAME")
                for dataHist in DataHistograms:
                    dataHist.Draw("A E SAME")
                                
            elif numSignalSamples is not 0:
                SignalMCHistograms[0].SetTitle(histoTitle)
                SignalMCHistograms[0].Draw("HIST")
                SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                SignalMCHistograms[0].SetMaximum(finalMax)
                SignalMCHistograms[0].SetMinimum(0.0001)
                for signalMCHist in SignalMCHistograms:
                    if(signalMCHist is not SignalMCHistograms[0]):
                        signalMCHist.Draw("A HIST SAME")
                for dataHist in DataHistograms:
                    dataHist.Draw("A E SAME")

            elif(numDataSamples is not 0):
                DataHistograms[0].SetTitle(histoTitle)
                DataHistograms[0].Draw("E")
                DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                DataHistograms[0].SetMaximum(finalMax)
                DataHistograms[0].SetMinimum(0.0001)
                for dataHist in DataHistograms:
                    if(dataHist is not DataHistograms[0]):
                        dataHist.Draw("A E SAME")

            
            if(numBgMCSamples is not 0 or numDataSamples is not 0):
                BgMCLegend.SetX1NDC(0.75)
                if numBgMCSamples is 0:
                    BgMCLegend.SetY1NDC(0.9-0.05*(1+numDataSamples))
                else:
                    BgMCLegend.SetY1NDC(0.9-0.05*(2+numBgMCSamples+numDataSamples))
                BgMCLegend.SetX2NDC(0.9)
                BgMCLegend.SetY2NDC(0.9)
                BgMCLegend.Draw()
                if(numSignalSamples is not 0):
                    SignalMCLegend.SetX1NDC(0.6)
                    SignalMCLegend.SetY1NDC(0.9-0.05*(1+numSignalSamples))
                    SignalMCLegend.SetX2NDC(0.75)
                    SignalMCLegend.SetY2NDC(0.9)
                    SignalMCLegend.Draw()
            elif numSignalSamples is not 0:
                SignalMCLegend.SetX1NDC(0.75)
                SignalMCLegend.SetY1NDC(0.9-0.05*(1+numSignalSamples))
                SignalMCLegend.SetX2NDC(0.9)
                SignalMCLegend.SetY2NDC(0.9)
                SignalMCLegend.Draw()

            if not arguments.normalizeToUnitArea or numDataSamples > 0:
                LumiLabel.Draw()
            if arguments.normalizeToData and numBgMCSamples > 0 and numDataSamples > 0:
                if numSignalSamples is 0:
                    NormLabel = TPaveLabel(0.6,0.85,0.75,0.9,"MC scaled to data","NDC")
                else:
                    NormLabel = TPaveLabel(0.45,0.85,0.6,0.9,"MC scaled to data","NDC")
                NormLabel.SetBorderSize(0)
                NormLabel.SetFillColor(0)
                NormLabel.SetFillStyle(0)
                NormLabel.Draw()
            elif arguments.normalizeToUnitArea:
                if (numBgMCSamples is not 0 or numDataSamples is not 0) and numSignalSamples is 0:
                    NormLabel = TPaveLabel(0.6,0.85,0.75,0.9,"Scaled to unit area","NDC")
                else:
                    NormLabel = TPaveLabel(0.45,0.85,0.6,0.9,"Scaled to unit area","NDC")
                NormLabel.SetBorderSize(0)
                NormLabel.SetFillColor(0)
                NormLabel.SetFillStyle(0)
                NormLabel.Draw()


            if makeRatioPlots or makeDiffPlots:
                Canvas.cd(2)
                BgSum = Stack.GetStack().Last()
                Comparison = DataHistograms[0].Clone()
                Comparison.Add(BgSum,-1)
                if not makeDiffPlots:        
                      Comparison.Divide(BgSum)
                Comparison.SetTitle("")
                Comparison.GetXaxis().SetTitle(xAxisLabel)
                if makeRatioPlots:
                    Comparison.GetYaxis().SetTitle("#frac{Data-MC}{MC}")
                elif makeDiffPlots:
                    Comparison.GetYaxis().SetTitle("Data-MC")
                Comparison.GetYaxis().CenterTitle()
                Comparison.GetYaxis().SetTitleSize(0.1)
                Comparison.GetYaxis().SetTitleOffset(0.35)                
                Comparison.GetXaxis().SetTitleSize(0.15)
                Comparison.GetYaxis().SetLabelSize(0.1)
                Comparison.GetXaxis().SetLabelSize(0.15)
                if makeRatioPlots:
                    Comparison.GetYaxis().SetRangeUser(-1.15,1.15)
                elif makeDiffPlots:
                    YMax = Comparison.GetMaximum()
                    YMin = Comparison.GetMinimum()
                    if YMax <= 0 and YMin <= 0:
                        Comparison.GetYaxis().SetRangeUser(-1.2*YMin,0)
                    elif YMax >= 0 and YMin >= 0:
                        Comparison.GetYaxis().SetRangeUser(0,1.2*YMax)
                    else: #axis crosses y=0
                        if abs(YMax) > abs(YMin):
                            Comparison.GetYaxis().SetRangeUser(-1.2*YMax,1.2*YMax)
                        else:
                            Comparison.GetYaxis().SetRangeUser(-1.2*YMin,1.2*YMin)
                            
                Comparison.GetYaxis().SetNdivisions(205)
                Comparison.Draw()
            Canvas.Write()

        if re.match ('TH2', key.GetClassName()) and arguments.draw2DPlots: # plot a 2-D histogram        

            numBgMCSamples = 0
            numDataSamples = 0
            numSignalSamples = 0

            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.76,0.65,0.99,0.9)
            BgMCLegend.AddEntry (0, "Data & Bkgd. MC", "H").SetTextFont (62)
            BgMCLegend.SetBorderSize(0)
            BgMCLegend.SetFillColor(0)
            BgMCLegend.SetFillStyle(0)
            SignalMCLegend = TLegend(0.76,0.135,0.99,0.377)
            SignalMCLegend.AddEntry (0, "Signal MC", "H").SetTextFont (62)
            SignalMCLegend.SetBorderSize(0)
            SignalMCLegend.SetFillColor(0)
            SignalMCLegend.SetFillStyle(0)

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

            for sample in processed_datasets: # loop over different samples as listed in configurationOptions.py
                dataset_file = "%s/%s.root" % (condor_dir,sample)
                inputFile = TFile(dataset_file)
                Histogram = inputFile.Get(rootDirectory+"/"+channel+"/"+histogramName).Clone()
                Histogram.SetDirectory(0)
                RebinFactor = int(arguments.rebinFactor)
                if arguments.rebinFactor and Histogram.GetNbinsX() >= RebinFactor*10 and Histogram.GetNbinsY() >= RebinFactor*10:
                     Histogram.Rebin2D(RebinFactor)
                inputFile.Close()
                xAxisLabel = Histogram.GetXaxis().GetTitle()
                yAxisLabel = Histogram.GetYaxis().GetTitle()
                histoTitle = Histogram.GetTitle()
            
                if( types[sample] == "bgMC"):

                    numBgMCSamples += 1
                    Histogram.SetMarkerColor(colors[sample])
                    Histogram.SetFillColor(colors[sample])                    
                    BgMCLegend.AddEntry(Histogram,labels[sample],"F").SetTextFont (42)
                    BgMCHistograms.append(Histogram)
                
                elif( types[sample] == "signalMC"):

                    numSignalSamples += 1
                    Histogram.SetMarkerColor(colors[sample])
                    Histogram.SetFillColor(colors[sample])
                    SignalMCLegend.AddEntry(Histogram,labels[sample],"F").SetTextFont (42)
                    SignalMCHistograms.append(Histogram)

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

                    numDataSamples += 1
                    Histogram.SetMarkerColor(colors[sample])
                    Histogram.SetFillColor(colors[sample])
                    BgMCLegend.AddEntry(Histogram,labels[sample],"F").SetTextFont (42)
                    DataHistograms.append(Histogram)


            outputFile.cd(rootDirectory+"/"+channel)        

            if(numBgMCSamples is not 0):
                BgMCHistograms[0].SetTitle(histoTitle)
                BgMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                BgMCHistograms[0].GetYaxis().SetTitle(yAxisLabel)
                BgMCHistograms[0].Draw()            
                for signalMCHist in SignalMCHistograms:
                    signalMCHist.Draw("SAME")
                for dataHist in DataHistograms:
                    dataHist.Draw("SAME")
                                
            elif(numSignalSamples is not 0):
                SignalMCHistograms[0].SetTitle(histoTitle)
                SignalMCHistograms[0].Draw()
                SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                SignalMCHistograms[0].GetYaxis().SetTitle(yAxisLabel)            
                for signalMCHist in SignalMCHistograms:
                    if(signalMCHist is not SignalMCHistograms[0]):
                        signalMCHist.Draw("SAME")
                for dataHist in DataHistograms:
                    dataHist.Draw("SAME")

            elif(numDataSamples is not 0):
                DataHistograms[0].SetTitle(histoTitle)
                DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
                DataHistograms[0].GetYaxis().SetTitle(yAxisLabel)            
                DataHistograms[0].Draw()
                for dataHist in DataHistograms:
                    if(dataHist is not DataHistograms[0]):
                        dataHist.Draw("SAME")

            
            if(numBgMCSamples is not 0 or numDataSamples is not 0):
                BgMCLegend.Draw()
            if(numSignalSamples is not 0):
                SignalMCLegend.Draw()
            if not arguments.normalizeToUnitArea or numDataSamples > 0:                
                LumiLabel.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.34
Committed:	Tue Apr 30 09:34:47 2013 UTC (12 years ago) by jbrinson
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	V02-01-00, V01-01-00, V01-00-00, V02-00-00, V00-01-00
Changes since 1.33:	+4 -6 lines
Log Message:	Fixed bug in adding yields to legends