EWP-BplusToKstMuMu-AngAna/Code/Selection/ComparisonTool/Plots.py

from ROOT import gROOT, gDirectory, gStyle, TChain, TTree, TAxis, TH1D, TLine, TGaxis
from Utils import expressionList, variableTag
from Globals import *

def getHistHeight(hist1,hist2):
    minY = 1e3
    maxY = 0.0
    for b in xrange(hist1.GetXaxis().GetNbins()): #hist1 and hist2 have same nBins
        maxTmp = max(hist1.GetBinContent(b+1),hist2.GetBinContent(b+1))
        if (maxY < maxTmp): maxY = maxTmp
        minTmp = min(hist1.GetBinContent(b+1),hist2.GetBinContent(b+1))
        if (minY > minTmp): minY = minTmp
    if (minY < 0.0): minY = 1.1*minY    
    else: minY = 0.001
 
    return minY, maxY*1.1    
 
def defineXAxisLabels(variable1,variable2):

    xTag = variable1.replace("TMath::","")
    for expr in expressionList():
        xTag = variable1.replace(expr.lower(),expr)

    if (variable1==variable2):
        return xTag
    else:     
        xTag2 = variable2.replace("TMath::","")
        for expr in expressionList():
            xTag2 = variable2.replace(expr.lower(),expr)
        return xTag + " vs " + xTag2 #TODO


def makePlot(hist1, hist2, histRatio, pad1, pad2, xAxisTitle,  Min, Max, lowerPlotHeight = 0.6):
   
    textsize = gStyle.GetTextSize();

    # modifiy pad for the ratio plot
    pad1.Clear()
    pad1.SetBorderSize  (0)
    pad1.SetBottomMargin(1e-6)
    pad1.SetTopMargin(pad1.GetTopMargin() / ( 1.0 - lowerPlotHeight) )
    pad1.cd()

    #Set XAxis of the ratio plot #TODO
    TGaxis().SetExponentOffset(1e+9,1e+9,"y") #offset = pad size * 1e+7
    histRatio.GetXaxis().SetNoExponent() #<-- spoils MaxDigits settings, so don't use it on other axis
    histRatio.GetXaxis().SetLabelSize(0.0) #don't print labels
    histRatio.GetXaxis().SetRangeUser(Min, Max)
    histRatio.GetXaxis().SetTickLength(histRatio.GetXaxis().GetTickLength()/(1-lowerPlotHeight))
    histRatio.GetXaxis().SetNdivisions(gStyle.GetNdivisions("X"))

    #Set YAxis of the ratio plot
    histRatio.GetYaxis().SetTitle("ratio red/black")
    histRatio.GetYaxis().SetTitleSize(textsize/(1-lowerPlotHeight))
    histRatio.GetYaxis().SetLabelSize(textsize/(1-lowerPlotHeight))
    histRatio.GetYaxis().SetTickLength(histRatio.GetYaxis().GetTickLength())
    histRatio.GetYaxis().SetNdivisions(gStyle.GetNdivisions("Y"))
    histRatio.GetYaxis().SetTitleOffset(1.06*(1-lowerPlotHeight))
    histRatio.GetYaxis().SetRangeUser(0.001, 3)
    histRatio.Draw("E1")

    #Add a line at ratio = 1
    unityline2 = TLine(Min, 0.8, Max, 1.0)
    unityline2.SetLineStyle(2)
    unityline2.SetLineColor(4)
    unityline2.Draw()


    #Clear pad for the comparison plot
    pad2.Clear()
    pad2.SetBorderSize(0)
    pad2.SetTopMargin(1e-6)
    pad2.SetBottomMargin(pad2.GetBottomMargin() / lowerPlotHeight )
    pad2.cd()

    #Set XAxis of the comparison plot
    TGaxis().SetExponentOffset(1e+9,1e+9,"x") #offset = pad size * 1e+7
    hist1.GetXaxis().SetTitleOffset(1.05)
    hist1.GetXaxis().SetTitleSize  (textsize/lowerPlotHeight)
    hist1.GetXaxis().SetLabelSize  (textsize/lowerPlotHeight)
    hist1.GetXaxis().SetTickLength (histRatio.GetXaxis().GetTickLength()/(lowerPlotHeight/(1-lowerPlotHeight)))
    hist1.GetXaxis().SetNdivisions (gStyle.GetNdivisions("X"))
    hist1.GetXaxis().SetTitle(xAxisTitle) 
    hist1.GetXaxis().SetRangeUser(Min, Max)


    #Set YAxis of the comparison plot
    minY, maxY = getHistHeight(hist1,hist2)
    hist1.GetYaxis().SetRangeUser(minY,maxY) 
    hist1.GetYaxis().SetTitle("weighted events")
    hist1.GetYaxis().SetLabelSize  (textsize/lowerPlotHeight)
    hist1.GetYaxis().SetTitleSize  (textsize/lowerPlotHeight)
    hist1.GetYaxis().SetTickLength(histRatio.GetYaxis().GetTickLength())
    hist1.GetYaxis().SetNdivisions (gStyle.GetNdivisions("Y"))
    hist1.GetYaxis().SetTitleOffset(1.06*lowerPlotHeight)
    hist1.Draw("E1")

    hist2.SetLineColor(2)
    hist2.SetMarkerColor(2)
    hist2.Draw("E1SAME")

    return

def defineLegendTags(optionsDict1, optionsDict2):
    tag1 = "" #initialize empty strings for tags
    tag2 = ""
    
    for key in optionsDict1:
        if (optionsDict1[key] != optionsDict2[key]):
            if (key == 'year' or key == 'magnet'):
                tag1 = tag1 + " " + str(optionsDict1[key])
                tag2 = tag2 + " " + str(optionsDict2[key]) 
            elif (key == 'Run'):
                tag1 = tag1 + " Run" + optionsDict1[key] 
                tag2 = tag2 + " Tun" + optionsDict2[key] 
            else:
                if(optionsDict1[key]): tag1 = tag1 + " " + key
                if(optionsDict2[key]): tag2 = tag2 + " " + key 

    return tag1, tag2            


  

def getPlotFileName(variable1, optionsDict1, cut1, variable2, optionsDict2, cut2):

    #First define dataset tag
    if ( optionsDict1["year"] != optionsDict2["year"] or optionsDict1["magnet"] != optionsDict1["magnet"]):
        plotName = str(optionsDict1["year"]) + str(optionsDict1["magnet"])+ "_" + str(optionsDict2["year"]) + str(optionsDict2["magnet"])
    else:
        plotName = str(optionsDict1["year"]) + str(optionsDict1["magnet"])

    #Define name tag
    if (variable1==variable2):
        plotName = plotName + "_" + variableTag(variable1) 
    else:     
        plotName = plotName + "_" + variableTag(variable2) + "_" + variableTag(variable1) 

    #Define sample tag
    if (optionsDict1["MC"]):   plotName1 = "MC"     
    else:                      plotName1 = "data"       
    if (optionsDict1["Preselected"]): 
        plotName1 = plotName1 + "_Presel"
        if (optionsDict1["BDTed"]):
            plotName1 = plotName1 + "_BDTed"
    else:  plotName1 = plotName1 + "_Strip"   
    if (KshortChannel()):  
        if (optionsDict1["KshortDecaysInVelo"]): plotName1 = plotName1 + "_LL"
        else:                                    plotName1 = plotName1 + "_DD"

    if (optionsDict2["MC"]):    plotName2 = "MC"     
    else:                       plotName2 = "data"       
    if (optionsDict2["Preselected"]):
        plotName2 = plotName2 + "_Presel"
        if (optionsDict2["BDTed"]):
            plotName2 = plotName2 + "_BDTed"
    else:  plotName2 = plotName2 + "_Strip" 
    if (KshortChannel()):
        if (optionsDict2["KshortDecaysInVelo"]): plotName2 = plotName2 + "_LL"
        else:                                    plotName2 = plotName2 + "_DD"

    #Define MC/data tags
    if (optionsDict1["MC"]):  
        for keyMC in ['TM','ReferenceChannel', 'PHSP','bWeighted', 'b2Dweighted']: #TODO possibly add branch name
            if(optionsDict1[keyMC]): plotName1 = plotName1 + "_" + keyMC
    else:
        if(optionsDict1["sWeighted"]): plotName1 = plotName1 + "_sWeighted"

    if (optionsDict2["MC"]):  
        for keyMC in ['TM','ReferenceChannel', 'PHSP','bWeighted', 'b2Dweighted']:
            if(optionsDict2[keyMC]): plotName2 = plotName2 + "_" + keyMC
    else:
        if(optionsDict2["sWeighted"]): plotName2 = plotName2 + "_sWeighted"
    
    # if (optionsDict1["UseLowQ2Range"]) plotName1 = plotName1 + "_LowQ2"
    # else:                              plotName1 = plotName1 + "_HighQ2"
    # if (optionsDict2["UseLowQ2Range"]) plotName2 = plotName2 + "_LowQ2"
    # else:                              plotName2 = plotName2 + "_HighQ2"

    #Strinp the extra cut on TM
    if (optionsDict1["TM"] and ("TMedBKGCAT=1" in cut1)): cut1 = cut1.replace("TMedBKGCAT=1","",-1) 
    if (optionsDict2["TM"] and ("TMedBKGCAT=1" in cut2)): cut2 = cut2.replace("TMedBKGCAT=1","",-1) 

    plotName1 = (plotName1 + "_"+variableTag(cut1)) if (cut1 != "") else plotName1
    plotName2 = (plotName2 + "_"+variableTag(cut2)) if (cut2 != "") else plotName2  
    
    plotName = plotName + "_" + plotName1 + "_" + plotName2  
    plotName = plotName.replace(" ","")

    if (KshortChannel()):
        return "/auto/data/dgerick/B2Kstmumu/compare_plots/" + plotName + ".eps"
    else:
        return "/home/lhcb/kopecna/B2KstarMuMu/data/compare_plots/" + plotName + ".eps"