PhD-Kopecna-Renata/pythonScripts/prettyPlots.py

271 lines
8.9 KiB
Python
Raw Normal View History

from ROOT import gROOT, TFile, gStyle
from ROOT import kFALSE, kTRUE
from ROOT import RooPlot, RooHist
from ROOT import TPaveText, TGraph, TCanvas, TLatex, TF1, TH1D, gPad
from ROOT import kBlack, kBlue, kRed
from math import exp,log
import collections #To select subsets of dictionaries
verbose = True
def isMatched(name):
return ("matched" in name)
def isFull(name):
return ("full" in name)
def rangeMass(method):
if (method == "Long"): return [2625.0,3575.0]
if (method == "Velo"): return [2925.0,3275.0]
if (method == "T"): return [2625.0,3575.0]
def binEdges(name):
nameTmp = name.replace(" ","") #Remove spaces, they are useless
#split makes a list, first part we throw away
#Then we split the string of binLow<VAR<binHigh
nameList = (nameTmp.split(',')[1]).split('<')
return (nameList[0],nameList[2]) #We return the first and the third one
def methodTex(method):
method_tex = TLatex()
method_tex.SetNDC(True)
method_tex.SetTextFont(132)
method_tex.SetTextSize(0.06)
method_tex.SetTextAlign(13)
method_tex.DrawLatex(0.205, 0.905, method+" method")
return
def mainTag():
tex = TLatex()
tex.SetNDC(True)
tex.SetTextFont(132)
tex.SetTextSize(0.06)
tex.SetTextAlign(33)
tex.DrawLatex(0.875, 0.905, "This thesis")
def matchTag(match,full): #This order is crucial
if (match): return "Matched"
if (full): return "Full"
else: return "Failed"
def matchTagDraw(match, full):
tex = TLatex()
tex.SetNDC(True)
tex.SetTextFont(132)
tex.SetTextSize(0.06)
tex.SetTextAlign(13)
tex.DrawLatex(0.205, 0.805, matchTag(match,full))
def ConvertRooHistToTH1D(hist):
nBins = hist.GetN()
newHist = TH1D(hist.GetName(),hist.GetTitle(),nBins, hist.GetXaxis().GetXmin(),hist.GetXaxis().GetXmax())
for bin in range(nBins):
#print(bin,hist.GetPointY(bin))
newHist.SetBinContent(bin+1,hist.GetPointY(bin))
print(newHist.GetEntries())
return newHist
def parsePaveText(text, param):
parStr = text.GetLineWith(param).GetTitle()
parStr.replace(" "," ") #Replace double spaces by one to avoid trouble
parList = parStr.split(" ")
return (float(parList[-3]),float(parList[-1])) #the list is then [mess, mess, something, val,#pm,err]
def getCanvas(method):
canvas = TCanvas("c_"+method, "c_"+method, 10,10,650,600)
canvas.SetBottomMargin(0.15)
canvas.SetRightMargin(0.05)
canvas.SetLeftMargin(0.16)
canvas.SetTopMargin(0.0675)
return canvas
def designRooHist(hist,method):
#Set range
hist.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
#Set line design
hist.SetLineWidth(2)
hist.SetLineColor(kBlack)
#Set title
hist.SetTitle("")
#Set X axis
hist.GetXaxis().SetTitle("m_{#mu^{+}#mu^{-}} [MeV]")
hist.GetXaxis().SetNdivisions(505)
hist.GetXaxis().SetTitleFont(132)
hist.GetXaxis().SetTitleSize(0.06)
hist.GetXaxis().SetTitleOffset(1.1)
hist.GetXaxis().SetLabelFont(132)
hist.GetXaxis().SetLabelSize(0.05)
#Set Y axis
hist.GetYaxis().SetTitle("Entries (a.u.)")
hist.GetYaxis().SetNdivisions(505)
hist.GetYaxis().SetTitleFont(132)
hist.GetYaxis().SetTitleSize(0.06)
hist.GetYaxis().SetTitleOffset(1.3)
hist.GetYaxis().SetLabelFont(132)
hist.GetYaxis().SetLabelSize(0.05)
return
def getHistLimits(hist):
xmin = hist.GetXaxis().GetXmin()
xmax = hist.GetXaxis().GetXmax()
ymin = hist.GetBinContent(1) #Not technically ymin, but an y at xmin
ymax = hist.GetBinContent(100) #Not technically ymin, but an y at xmax
return xmin, xmax, ymin, ymax
def designSignal(hist, method):
hist.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
hist.SetLineWidth(2)
hist.SetLineStyle(7)
hist.SetLineColor(kRed+2)
return
def getPlotName(outputPath,
var, bin, method, matchTag):
return outputPath+"/"+var+"_bin"+str(bin) +"_"+method+ "_" + matchTag
def plotYield(histRoo, bkg_entries, bkg_exp,
method, match, full,
plotName ):
#create canvases
canvas = getCanvas(method)
#Draw the hist
h_all = ConvertRooHistToTH1D(histRoo)
designRooHist(frDictList[match]['hist'],method)
hist.GetYaxis().SetRangeUser(0,h_all.GetMaximum()*1.25)
histRoo.Draw()
#Get the hist values at edges and get the background histogram
xmin, xmax, ymin, ymax = getHistLimits(h_all)
bkg = TF1("bkg","expo", xmin, xmax)
a = log(
(ymin-ymax)/(exp(xmin*bkg_exp)-exp(xmax*bkg_exp))
)
bkg.SetParameters(a,bkg_exp)
print (a,bkg_exp)
#Fill the bkg histogram
h_bkg = TH1D("h_bkg","h_bkg",100,xmin,xmax)
h_bkg.FillRandom("bkg",int(bkg_entries))
h_bkg.GetXaxis().SetRangeUser(rangeMass(method)[0],rangeMass(method)[1])
#h_bkg.Draw("SAMEAL")
#Get signal
h_all.Add(h_bkg,-1.0) #Substract the h_bkg from h_all
designSignal(h_all, method)
h_all.Draw("SAMEAL")
#Draw main tag
mainTag()
#Draw the method
methodTex(method)
#Draw matched/failed
matchTagDraw(match,full)
canvas.SetTitle("")
canvas.SaveAs(plotName+".root","root")
canvas.Print(plotName+".eps")
return
def Plot(simVer = "Sim09h", year = "2018_25ns", data = True,
method = "Long", var = "ETA", simFit = True,
outputPath = "./prettyPlots"
):
gROOT.SetBatch(kTRUE)
gStyle.SetTextFont(132)
dataTag = "Data" if data else "MC"
folderName = "./results/"+year+"_WG/"
if (not data): folderName = fileName + "/" +simVer + "/"
fileName = "trackEff_" + dataTag + "_" + var + "_" + method + "_method.root"
canvas_list = ["matched", "full", "fail"]
inputfile = TFile(folderName+fileName)
print ("Openning", folderName+fileName)
frameDictList = []
#Get the frames
for item in inputfile.GetListOfKeys():
#if any(cvs in item.GetName() for cvs in canvas_list):
if (item.GetClassName() != "RooPlot"): continue
if (verbose): print("Got item " + item.GetName() + " of class " + item.GetClassName())
frame = RooPlot()
frame = inputfile.Get(item.GetName())
name = frame.GetTitle()
hist = RooHist()
graph = TGraph()
text = TPaveText()
hist = frame.getObject(0)
graph = frame.getObject(1)
if ( (simFit) and isMatched(name)): text = frame.getObject(2)
frameDict ={"frame": frame,
"matched": isMatched(name),
"bins": (0,0) if (var=="Full") else binEdges(name),
"hist": hist,
"graph": graph,
"text": text
}
frameDictList.append(frameDict)
#Sort into dictionaries based on the same bin
result = collections.defaultdict(list)
for frameDict in frameDictList:
result[frameDict['bins']].append(frameDict)
sortedList = list(result.values())
#Check there are only two entries per list
for frDictList in sortedList:
if (len(frDictList)!=2):
print("ERROR: somehow there are more/less than two frames per bins, check.")
return
#Now loop over the bins and plot each function+data
#Matched and Failed has to be treated seperately as the failed ones don't have TPaveText
for i,frDictList in enumerate (sortedList):
if (simFit):
#This is needed only cause full sample doesn't have a pavetext in failed sample... sigh
b_M = parsePaveText(frDictList[0]['text'],"tau")[0]
b_F = parsePaveText(frDictList[0]['text'],"tauF")[0]
bkg_entries_M = parsePaveText(frDictList[0]['text'],"background_yield")[0]
bkg_eff = parsePaveText(frDictList[0]['text'],"efficiency_bkg")[0]
bkg_entries_F = bkg_entries_M/bkg_eff - bkg_entries_M #eff = M/(M+F)
plotName = getPlotName(outputPath,var,i,method,matchTag(True,False))
plotYield(frDictList[0]['hist'],bkg_entries_M,b_M,method,True,False,plotName)
plotName = getPlotName(outputPath,var,i,method,matchTag(False,False))
plotYield(frDictList[1]['hist'],bkg_entries_F,b_F,method,False,False,plotName)
else:
for matched in [True,False]:
b = parsePaveText(frDictList[matched]['text'],"tau")[0]
bkg_entries = parsePaveText(frDictList[matched]['text'],"background_yield")[0]
plotName = getPlotName(outputPath,var,i,method,matchTag(False,True))
plotYield(frDictList[matched]['hist'],method,bkg_entries,b,matched,True,outputPath)
return
#Plot()
Plot(method="Long",var="Full")
#Plot(method="Velo",var="Full")
#Plot(method="T",var="Full")
#Plot(method="Long",var="Full",SimFit=False)
#Plot(method="Velo",var="Full",match=False)
#Plot(method="T",var="Full",match=False)