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 *


### parse the command-line options

parser = OptionParser()
parser = set_commandline_arguments(parser)
parser.add_option("--ylog", action="store_true", dest="setLogY", default=False,          
                  help="Set logarithmic scale on vertical axis on all plots")    
parser.add_option("--ymin", dest="setYMin", 
                  help="Minimum of y axis")      
parser.add_option("--ymax", dest="setYMax", 
                  help="Maximum of y axis")      

(arguments, args) = parser.parse_args()


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

#### 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, THStack, TIter, TKey, TPaveLabel, gPad


### setting ROOT options so our plots will look awesome and everyone will love us

gROOT.SetBatch()
gStyle.SetOptStat(0)
gStyle.SetCanvasBorderMode(0)
gStyle.SetPadBorderMode(0)
gStyle.SetPadColor(0)
gStyle.SetCanvasColor(0)
gStyle.SetTextFont(42)
gStyle.SetCanvasDefH(600)
gStyle.SetCanvasDefW(600)
gStyle.SetCanvasDefX(0)
gStyle.SetCanvasDefY(0)
gStyle.SetPadTopMargin(0.05)
gStyle.SetPadBottomMargin(0.13)
gStyle.SetPadLeftMargin(0.15)
gStyle.SetPadRightMargin(0.05)
gStyle.SetTitleColor(1, "XYZ")
gStyle.SetTitleFont(42, "XYZ")
gStyle.SetTitleSize(0.05, "XYZ")
gStyle.SetTitleXOffset(0.95)
gStyle.SetTitleYOffset(1.25)
gStyle.SetTextAlign(12)
gStyle.SetLabelColor(1, "XYZ")
gStyle.SetLabelFont(42, "XYZ")
gStyle.SetLabelOffset(0.007, "XYZ")
gStyle.SetLabelSize(0.05, "XYZ")
gStyle.SetAxisColor(1, "XYZ")
gStyle.SetStripDecimals(True)
gStyle.SetTickLength(0.03, "XYZ")
gStyle.SetNdivisions(510, "XYZ")
gStyle.SetPadTickX(1)
gStyle.SetPadTickY(1)
gROOT.ForceStyle()


#set the text for the luminosity label
if(intLumi < 1000.):
    LumiText = "L_{int} = " + str(intLumi) + " pb^{-1}"
    LumiText = "L_{int} = " + str.format('{0:.1f}', LumiInPb) + " pb^{-1}"
else:
    LumiInFb = intLumi/1000.
    LumiText = "L_{int} = " + str.format('{0:.1f}', LumiInFb) + " fb^{-1}"

    
#bestest place for lumi. label, in top left corner
topLeft_x_left    = 0.1375839
topLeft_x_right   = 0.4580537
topLeft_y_bottom  = 0.8479021
topLeft_y_top     = 0.9475524
topLeft_y_offset  = 0.035
                                                                        
                                                                            
##########################################################################################################################################
##########################################################################################################################################
##########################################################################################################################################

# some fancy-ass code from Andrzej Zuranski to merge bins in the ratio plot until the error goes below some threshold
def ratioHistogram( dataHist, mcHist, relErrMax=0.10):

    def groupR(group):
        Data,MC = [float(sum(hist.GetBinContent(i) for i in group)) for hist in [dataHist,mcHist]]
        return (Data-MC)/MC if MC else 0
    
    def groupErr(group):
        Data,MC = [float(sum(hist.GetBinContent(i) for i in group)) for hist in [dataHist,mcHist]]
        dataErr2,mcErr2 = [sum(hist.GetBinError(i)**2 for i in group) for hist in [dataHist,mcHist]]
        return abs(math.sqrt( (dataErr2+mcErr2)/(Data-MC)**2 + mcErr2/MC**2 ) * (Data-MC)/MC) if Data and MC else 0

    def regroup(groups):
        err,iG = max( (groupErr(g),groups.index(g)) for g in groups )
        if err < relErrMax or len(groups)<3 : return groups
        iH = max( [iG-1,iG+1], key = lambda i: groupErr(groups[i]) if 0<=i<len(groups) else -1 )
        iLo,iHi = sorted([iG,iH])
        return regroup(groups[:iLo] + [groups[iLo]+groups[iHi]] + groups[iHi+1:])

    #don't rebin the histograms of the number of a given object (except for the pileup ones)
    if ((dataHist.GetName().find("num") is not -1 and dataHist.GetName().find("Primaryvertexs") is -1) or
        dataHist.GetName().find("CutFlow")  is not -1 or
        dataHist.GetName().find("GenMatch") is not -1): 
        ratio = dataHist.Clone()
        ratio.Add(mcHist,-1)
        ratio.Divide(mcHist)
        ratio.SetTitle("")
    else:
        groups = regroup( [(i,) for i in range(1,1+dataHist.GetNbinsX())] )
        ratio = TH1F("ratio","",len(groups), array('d', [dataHist.GetBinLowEdge(min(g)) for g in groups ] + [dataHist.GetXaxis().GetBinUpEdge(dataHist.GetNbinsX())]) )
        for i,g in enumerate(groups) :
            ratio.SetBinContent(i+1,groupR(g))
            ratio.SetBinError(i+1,groupErr(g))

    ratio.GetYaxis().SetTitle("#frac{Data-MC}{MC}")
    ratio.SetLineColor(1)
    ratio.SetLineWidth(2)
    return ratio

##########################################################################################################################################
##########################################################################################################################################
##########################################################################################################################################


def MakeOneDHist(pathToDir,histogramName): 

    
    numBgMCSamples = 0
    numDataSamples = 0
    numSignalSamples = 0
    
    Stack = THStack("stack",histogramName)

    LumiLabel = TPaveLabel(topLeft_x_left,topLeft_y_bottom,topLeft_x_right,topLeft_y_top,LumiText,"NDC")
    LumiLabel.SetBorderSize(0)
    LumiLabel.SetFillColor(0)
    LumiLabel.SetFillStyle(0)
    
    NormLabel = TPaveLabel()
    NormLabel.SetDrawOption("NDC")
    NormLabel.SetX1NDC(topLeft_x_left)
    NormLabel.SetX2NDC(topLeft_x_right)
    
    NormLabel.SetBorderSize(0)
    NormLabel.SetFillColor(0)
    NormLabel.SetFillStyle(0)
    
    NormText = ""
    if arguments.normalizeToUnitArea:
        NormText = "Scaled to unit area"
    elif arguments.normalizeToData:
        NormText = "MC scaled to data"
        NormLabel.SetLabel(NormText)
        
        
    BgMCLegend = TLegend()
    BgTitle = BgMCLegend.AddEntry(0, "Data & Bkgd. MC", "H")
    BgTitle.SetTextAlign(22)
    BgTitle.SetTextFont(62)
    BgMCLegend.SetBorderSize(0)
    BgMCLegend.SetFillColor(0)
    BgMCLegend.SetFillStyle(0)
    SignalMCLegend = TLegend()
    SignalTitle = SignalMCLegend.AddEntry(0, "Signal MC", "H")
    SignalTitle.SetTextAlign(22)
    SignalTitle.SetTextFont(62)
    SignalMCLegend.SetBorderSize(0)
    SignalMCLegend.SetFillColor(0)
    SignalMCLegend.SetFillStyle(0)
    
    outputFile.cd(pathToDir)
    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(pathToDir+"/"+histogramName).Clone()
        Histogram.SetDirectory(0)
        inputFile.Close()
        if arguments.rebinFactor:
            RebinFactor = int(arguments.rebinFactor)
            #don't rebin histograms which will have less than 5 bins or any gen-matching histograms
            if Histogram.GetNbinsX() >= RebinFactor*5 and Histogram.GetTitle().find("GenMatch") is -1:
                Histogram.Rebin(RebinFactor)

        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)
                    
    #scaling histograms as per user's specifications
    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
    legendIndex = 0
    for Histogram in DataHistograms:
        BgMCLegend.AddEntry(Histogram,DataLegendEntries[legendIndex],"LEP")
        legendIndex = legendIndex+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")
        for Histogram in BgMCHistograms:
            if Histogram is not BgMCHistograms[0]:
                ErrorHisto.Add(Histogram)


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


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


    ### finding the maximum value of anything going on the canvas, so we 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.15*finalMax
    if arguments.setYMax:  
        finalMax = float(arguments.setYMax)


    ### Drawing histograms to canvas

    outputFile.cd(pathToDir)        
        
    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()
        if arguments.setLogY:
            gPad.SetLogy()
        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: # the first thing to draw to the canvas is a bgMC sample
        
        yAxisMin = 0.0001
        if arguments.setYMin:
            yAxisMin = float(arguments.setYMin)
        if not arguments.noStack: # draw unstacked background samples
            Stack.SetTitle(histoTitle)
            Stack.Draw("HIST")
            Stack.GetXaxis().SetTitle(xAxisLabel)
            Stack.SetMaximum(finalMax)
            Stack.SetMinimum(yAxisMin)
            if makeRatioPlots or makeDiffPlots:
                Stack.GetHistogram().GetXaxis().SetLabelSize(0)
            #draw shaded error bands
            ErrorHisto.Draw("A E2 SAME")
                
        else: #draw the stacked backgrounds
            BgMCHistograms[0].SetTitle(histoTitle)
            BgMCHistograms[0].Draw("HIST")
            BgMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
            BgMCHistograms[0].SetMaximum(finalMax)
            BgMCHistograms[0].SetMinimum(yAxisMin)            
            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: # the first thing to draw to the canvas is a signalMC sample 
        SignalMCHistograms[0].SetTitle(histoTitle)
        SignalMCHistograms[0].Draw("HIST")
        SignalMCHistograms[0].GetXaxis().SetTitle(xAxisLabel)
        SignalMCHistograms[0].SetMaximum(finalMax)
        SignalMCHistograms[0].SetMinimum(yAxisMin)

        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): # the first thing to draw to the canvas is a data sample 
        DataHistograms[0].SetTitle(histoTitle)
        DataHistograms[0].Draw("E")
        DataHistograms[0].GetXaxis().SetTitle(xAxisLabel)
        DataHistograms[0].SetMaximum(finalMax)
        DataHistograms[0].SetMinimum(yAxisMin)
        for dataHist in DataHistograms:
            if(dataHist is not DataHistograms[0]):
                dataHist.Draw("A E SAME")


    #legend coordinates, empirically determined :-)
    x_left = 0.6761745
    x_right = 0.9328859
    x_width = x_right - x_left
    y_max = 0.9335664
    entry_height = 0.05

    if(numBgMCSamples is not 0 or numDataSamples is not 0): #then draw the data & bgMC legend

        numExtraEntries = 1 # count the legend title
        BgMCLegend.SetX1NDC(x_left)
        if numBgMCSamples > 0:
            numExtraEntries = numExtraEntries + 1 # count the stat. errors entry
            
        BgMCLegend.SetY1NDC(y_max-entry_height*(numExtraEntries+numBgMCSamples+numDataSamples))
        BgMCLegend.SetX2NDC(x_right)
        BgMCLegend.SetY2NDC(y_max)
        BgMCLegend.Draw()

        if(numSignalSamples is not 0): #then draw the signalMC legend to the left of the other one
            SignalMCLegend.SetX1NDC(x_left-x_width)
            SignalMCLegend.SetY1NDC(y_max-entry_height*(1+numSignalSamples)) # add one for the title
            SignalMCLegend.SetX2NDC(x_left)
            SignalMCLegend.SetY2NDC(y_max)
            SignalMCLegend.Draw()

    elif numSignalSamples is not 0: #draw the signalMC legend in the upper right corner
        SignalMCLegend.SetX1NDC(x_left)
        SignalMCLegend.SetY1NDC(y_max-entry_height*(1+numSignalSamples)) # add one for the title
        SignalMCLegend.SetX2NDC(x_right)
        SignalMCLegend.SetY2NDC(y_max)
        SignalMCLegend.Draw()


    if not arguments.normalizeToUnitArea or numDataSamples > 0: #don't draw the lumi label if there's no data and it's scaled to unit area
        LumiLabel.Draw()
        if arguments.normalizeToUnitArea or arguments.normalizeToData:
            #move the normalization label down before drawing if we drew the lumi. label
            NormLabel.SetY1NDC(topLeft_y_bottom-topLeft_y_offset)
            NormLabel.SetY2NDC(topLeft_y_top-topLeft_y_offset)
            NormLabel.Draw()

    elif arguments.normalizeToUnitArea or arguments.normalizeToData:
        NormLabel.SetY1NDC(topLeft_y_bottom)
        NormLabel.SetY2NDC(topLeft_y_top)
        NormLabel.Draw()


    if makeRatioPlots or makeDiffPlots:
        Canvas.cd(2)
        BgSum = Stack.GetStack().Last()
        if makeRatioPlots:
            Comparison = ratioHistogram(DataHistograms[0],BgSum)
        elif makeDiffPlots:
            Comparison = DataHistograms[0].Clone("diff")
            Comparison.Add(BgSum,-1)
            Comparison.SetTitle("")
            Comparison.GetYaxis().SetTitle("Data-MC")
        Comparison.GetXaxis().SetTitle(xAxisLabel)
        Comparison.GetYaxis().CenterTitle()
        Comparison.GetYaxis().SetTitleSize(0.1)
        Comparison.GetYaxis().SetTitleOffset(0.4)
        Comparison.GetXaxis().SetTitleSize(0.15)
        Comparison.GetYaxis().SetLabelSize(0.1)
        Comparison.GetXaxis().SetLabelSize(0.15)
        if makeRatioPlots:
            RatioYRange = 1.15
            if arguments.ratioYRange:
                RatioYRange = float(arguments.ratioYRange)
            Comparison.GetYaxis().SetRangeUser(-1*RatioYRange, RatioYRange)
        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 arguments.savePDFs:
        pathToDirString = pathToDir.replace(' ','_').replace('<','lt').replace('>','gt').replace('(','').replace(')','').replace('=','eq')
        Canvas.SaveAs(condor_dir+"/stacked_histograms_pdfs/"+pathToDirString+"/"+histogramName+".pdf")


##########################################################################################################################################
##########################################################################################################################################
##########################################################################################################################################

def MakeTwoDHist(pathToDir,histogramName):
            numBgMCSamples = 0
            numDataSamples = 0
            numSignalSamples = 0

            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(pathToDir)
            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(pathToDir+"/"+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(pathToDir)        

            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()


##########################################################################################################################################
##########################################################################################################################################
##########################################################################################################################################

processed_datasets = []

condor_dir = set_condor_output_dir(arguments)

#### 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")


#### make output file
outputFileName = "stacked_histograms.root"
if arguments.outputFileName:
    outputFileName = arguments.outputFileName

outputFile = TFile(condor_dir + "/" + outputFileName, "RECREATE")


#### use the first input file as a template and make stacked versions of all its histograms
inputFile = TFile(condor_dir + "/" + processed_datasets[0] + ".root")
inputFile.cd()
outputFile.cd()

if arguments.savePDFs:
    os.system("rm -r %s/stacked_histograms_pdfs" % (condor_dir))
    os.system("mkdir %s/stacked_histograms_pdfs" % (condor_dir))
    

#get root directory in the first layer, generally "OSUAnalysis"
for key in inputFile.GetListOfKeys():
    if (key.GetClassName() != "TDirectoryFile"):
        continue
    rootDirectory = key.GetName()
    outputFile.mkdir(rootDirectory)
    if arguments.savePDFs:
        os.system("mkdir %s/stacked_histograms_pdfs/%s" % (condor_dir,rootDirectory.replace(' ','_').replace('<','lt').replace('>','gt').replace('(','').replace(')','').replace('=','eq')))

    #cd to root directory and look for histograms
    inputFile.cd(rootDirectory)
    for key2 in gDirectory.GetListOfKeys():

        if re.match ('TH1', key2.GetClassName()): # found a 1-D histogram
            MakeOneDHist(rootDirectory,key2.GetName())
        elif re.match ('TH2', key2.GetClassName()) and arguments.draw2DPlots: # found a 2-D histogram        
            MakeTwoDHist(rootDirectory,key2.GetName())

        elif (key2.GetClassName() == "TDirectoryFile"): # found a directory, cd there and look for histograms
            level2Directory = rootDirectory+"/"+key2.GetName()

            #make a corresponding directory in the output file
            outputFile.cd(rootDirectory)
            gDirectory.mkdir(key2.GetName())
            if arguments.savePDFs:
                os.system("mkdir %s/stacked_histograms_pdfs/%s" % (condor_dir,level2Directory.replace(' ','_').replace('<','lt').replace('>','gt').replace('(','').replace(')','').replace('=','eq')))

            #####################################################
            ###  This layer is typically the "channels" layer ###
            #####################################################

            inputFile.cd(level2Directory)
            for key3 in gDirectory.GetListOfKeys():
                if re.match ('TH1', key3.GetClassName()): # found a 1-D histogram
                    MakeOneDHist(level2Directory,key3.GetName())
                elif re.match ('TH2', key3.GetClassName()) and arguments.draw2DPlots: # found a 2-D histogram        
                    MakeTwoDHist(level2Directory,key3.GetName())

                elif (key3.GetClassName() == "TDirectoryFile"): # found a directory, cd there and look for histograms
                    level3Directory = level2Directory+"/"+key3.GetName()
                    
                    #make a corresponding directory in the output file
                    outputFile.cd(level2Directory)
                    gDirectory.mkdir(key3.GetName())
                    if arguments.savePDFs:
                        os.system("mkdir %s/stacked_histograms_pdfs/%s" % (condor_dir,level3Directory.replace(' ','_').replace('<','lt').replace('>','gt').replace('(','').replace(')','').replace('=','eq')))

                    #################################################
                    ###  This layer is typically the "cuts" layer ###
                    #################################################

                    inputFile.cd(level3Directory)
                    for key3 in gDirectory.GetListOfKeys():
                        if re.match ('TH1', key3.GetClassName()): # found a 1-D histogram
                            MakeOneDHist(level3Directory,key3.GetName())
                        elif re.match ('TH2', key3.GetClassName()) and arguments.draw2DPlots: # found a 2-D histogram        
                            MakeTwoDHist(level3Directory,key3.GetName())


outputFile.Close()         
Revision:	1.44
Committed:	Fri Jun 21 13:06:43 2013 UTC (11 years, 10 months ago) by wulsin
Content type:	text/x-python
Branch:	MAIN
Changes since 1.43:	+19 -6 lines
Log Message:	Add options --ylog, --ymin, --ymax, to specify y-axis range