from ROOT import gROOT, gDirectory, gStyle, TChain, TTree, TAxis, TH1D, TLegend, TCanvas, TPad, TLine #TODO check what needs to be loaded where import subprocess import time import numpy as np import re #for splitting with multiple deliminers from array import array import numexpr from Utils import * from Plots import * from LHCbStyle import * from Globals import * ''' Tool for comparing two variables from two datasets. Yes, writting it in Python3 would be smarter, but default is Python2 on all servers so far ''' Min = 0.0 Max = 0.0 #TODO stripping comparisons def compareUltimate( variable1 = "Abs(K_star_plus_PT-K_plus_PT)/10", variable2 = "Sqrt(B_plus_PT)", year1 = 2011, year2 = 2012, cut1 = "", cut2 = "", magnet1 = "both", MC1 = False, TM1 = False, ReferenceChannel1 = False, PHSP1 = False, Preselected1 = True, BDTed1 = False, magnet2 = "both", MC2 = True, TM2 = False, ReferenceChannel2 = False, PHSP2 = False, Preselected2 = True, BDTed2 = False, sWeighted1 = True, bWeighted1 = False, b2Dweighted1 = False, weightBranch1 = firstMCweight(), sWeighted2 = False, bWeighted2 = False, b2Dweighted2 = True, weightBranch2 = firstMCweight(), bPrint = False, KshortDecaysInVelo1 = False, KshortDecaysInVelo2 = False, verbose = True): #TODO: add run option #if((plotLog1 and plotLogOneMinus1) or (plotLog2 and plotLogOneMinus2)): # raise Exception("Only Log( var ) or Log( 1 - var ) option possible, but both flags were set!") #No need to set it in here, will be taken care of by definning variable #The only requiremnt here that for cuts, no operations are allowed for now (no sqrt(x) < 0.5 or so) gROOT.SetBatch(1) #Set LHCb style for plots LHCbStyle() #TODO add check for valid particles #Force TM for weighted MC if (b2Dweighted1 or bWeighted1): TM1 = True if (b2Dweighted2 or bWeighted2): TM2 = True if (MC1): year1 = checkMCyear(year1,ReferenceChannel1, PHSP1) if (MC2): year2 = checkMCyear(year2,ReferenceChannel2, PHSP2) if (magnet1 == "down"): cut1 = ("Polarity=65535") if (cut1=="") else (cut1 + "&&Polarity=65535") #Somehow, it overflows to 2^16-1 if (magnet1 == "up"): cut1 = ("Polarity=1") if (cut1=="") else (cut1 + "&&Polarity=1") if (magnet2 == "down"): cut2 = ("Polarity=65535") if (cut2=="") else (cut2 + "&&Polarity=65535")#Somehow, it overflows to 2^16-1 if (magnet2 == "up"): cut2 = ("Polarity=1") if (cut2=="") else (cut2 + "&&Polarity=1") if (TM1): cut1 = ("TMedBKGCAT=1") if (cut1=="") else (cut1 + "&& TMedBKGCAT=1") if (TM2): cut2 = ("TMedBKGCAT=1") if (cut2=="") else (cut2 + "&& TMedBKGCAT=1") #raise an exception if MC and data is selected if ( (not MC1 and (TM1 or ReferenceChannel1 or PHSP1)) or (not MC2 and (TM2 or ReferenceChannel1 or PHSP1)) ) : raise Exception('Data cannot be TruthMatched nor RefChan nor PHSP! Check your options!') #Set DTF if required if (UseDTF()): variable1 = replace_variables_to_DTF(variable1) variable2 = replace_variables_to_DTF(variable2) if (verbose): print "[DEBUG]\t\tUsing variables ", variable1, variable2 #Get Run Run1 = 1 if (year1 < 2013) else 2 Run2 = 1 if (year2 < 2013) else 2 #Set up a dictionary with dataset information optionsDictionary1 = getOptionsDictionary(year1, Run1, magnet1, MC1, TM1, ReferenceChannel1, PHSP1, Preselected1, BDTed1, sWeighted1, bWeighted1, b2Dweighted1, weightBranch1,KshortDecaysInVelo1, False) optionsDictionary2 = getOptionsDictionary(year2, Run2, magnet2, MC2, TM2, ReferenceChannel2, PHSP2, Preselected2, BDTed2, sWeighted2, bWeighted2, b2Dweighted2, weightBranch2,KshortDecaysInVelo2, False) #Check for year+sample if (not checkYearSample(optionsDictionary1)[0]): print checkYearSample(optionsDictionary1)[1] return 0 if (not checkYearSample(optionsDictionary2)[0]): print checkYearSample(optionsDictionary2)[1] return 0 #In case of the need of build, uncomment #import os #os.system("g++ getPathForPython.cc `root-config --libs --cflags` -o getPathForPython") #Read the trees tree1 = TChain(treeName(MC1,TM1,Preselected1)) tree2 = TChain(treeName(MC2,TM2,Preselected2)) addToTChain(tree1, optionsDictionary1, verbose) addToTChain(tree2, optionsDictionary2, verbose) #Load the branches into basic histograms histName1 = histName(variable1,optionsDictionary1)+"1" histName2 = histName(variable2,optionsDictionary2)+"2" drawVar1 = addTMathTags(variable1) drawVar2 = addTMathTags(variable2) if verbose: print "[DEBUG]\t\tVariable names are: ", histName1, histName2 print "[DEBUG]\t\tDrawing variables: ", drawVar1 , drawVar2 #Create default histograms directly from the tree tree1.Draw("%s >> %s" % (drawVar1,histName1) ) hist1Drawn = gDirectory.Get(histName1) if verbose: print "[DEBUG]\t\tDrawing into ", hist1Drawn.GetName() tree2.Draw("%s >> %s" % (drawVar2,histName2) ) hist2Drawn = gDirectory.Get(histName2) if verbose: print "[DEBUG]\t\tDrawing into ", hist2Drawn.GetName() #Check if histograms are empty if (hist1Drawn.GetEntries()==0): raise Exception('Created histogram1 is empty! Check your cuts and variable!') if (hist2Drawn.GetEntries()==0): raise Exception('Created histogram2 is empty! Check your cuts and variable!') #Check if the variable is integer isVar1Int = isInt(variable1) isVar2Int = isInt(variable2) hist1InfoList = [] hist2InfoList = [] dictVarPlotLimits = dictVarPlotLimitsKshort(KshortDecaysInVelo1 or KshortDecaysInVelo2) if KshortChannel() else dictVarPlotLimitsKplus() if (isVar1Int): hist1InfoList = dictVarPlotLimitsInt()[variable1] else: # Check for non-empty bins at the edges and set xaxis range accordingly nBins = hist1Drawn.GetXaxis().GetNbins() firstBin = 0 lastBin = nBins+1 # check if variable needs manually adjusted range if (variable1 in dictVarPlotLimits.keys()): if ("Min" in dictVarPlotLimits[variable1].keys()): # If min is defined, adjust min1 = dictVarPlotLimits[variable1]["Min"] else: while (hist1Drawn.GetBinContent(firstBin) == 0): # If not, loop over to get first bin firstBin = firstBin+1 min1 = hist1Drawn.GetXaxis().GetBinLowEdge(firstBin) if ("Max" in dictVarPlotLimits[variable1].keys()): # If max is defined, adjust max1 = dictVarPlotLimits[variable1]["Max"] else: while (hist1Drawn.GetBinContent(lastBin) == 0): # If not, loop over to get the last bin lastBin = lastBin-1 max1 = hist1Drawn.GetXaxis().GetBinUpEdge(lastBin) else: # if variable doesn't need manual adjustement, loop over to get edge bins while (hist1Drawn.GetBinContent(firstBin) == 0): firstBin = firstBin+1 min1 = hist1Drawn.GetXaxis().GetBinLowEdge(firstBin) while (hist1Drawn.GetBinContent(lastBin) == 0): lastBin = lastBin-1 max1 = hist1Drawn.GetXaxis().GetBinUpEdge(lastBin) hist1InfoList = [nBins/2, min1, max1] if (isVar2Int): hist2InfoList = dictVarPlotLimitsInt()[variable2] else: # Check for non-empty bins at the edges and set xaxis range accordingly nBins = hist2Drawn.GetXaxis().GetNbins() firstBin = 0 lastBin = nBins+1 # check if variable needs manually adjusted range if (variable2 in dictVarPlotLimits.keys()): if ("Min" in dictVarPlotLimits[variable2].keys()): # If min is defined, adjust min2 = dictVarPlotLimits[variable2]["Min"] else: while (hist2Drawn.GetBinContent(firstBin) == 0): # If not, loop over to get first bin firstBin = firstBin+2 min2 = hist2Drawn.GetXaxis().GetBinLowEdge(firstBin) if ("Max" in dictVarPlotLimits[variable2].keys()): # If max is defined, adjust max2 = dictVarPlotLimits[variable2]["Max"] else: while (hist2Drawn.GetBinContent(lastBin) == 0): # If not, loop over to get the last bin lastBin = lastBin-2 max2 = hist2Drawn.GetXaxis().GetBinUpEdge(lastBin) else: # if variable doesn't need manual adjustement, loop over to get edge bins while (hist2Drawn.GetBinContent(firstBin) == 0): firstBin = firstBin+2 min2 = hist2Drawn.GetXaxis().GetBinLowEdge(firstBin) while (hist2Drawn.GetBinContent(lastBin) == 0): lastBin = lastBin-1 max2 = hist2Drawn.GetXaxis().GetBinUpEdge(lastBin) hist2InfoList = [nBins/2, min2, max2] if verbose: print "[DEBUG]\t\tVariable %s with %i bins from %f to %f." %(variable1,hist1InfoList[0],hist1InfoList[1],hist1InfoList[2]) print "[DEBUG]\t\tVariable %s with %i bins from %f to %f." % (variable2,hist2InfoList[0],hist2InfoList[1],hist2InfoList[2]) #Set xranges and nBins to be equal for both variables #nBins are necessary for proper histogram division histInfoList = [ max(hist1InfoList[0], hist2InfoList[0]), min(hist1InfoList[1], hist2InfoList[1]), max(hist1InfoList[2], hist2InfoList[2]) ] #TODO have a function for Xaxis name and for tags #parse variables and cuts in a way that all needed branches are loaded variables1List = [] variables1Dict = {} variables2List = [] variables2Dict = {} #Create new histograms hist1 = TH1D("hist1","hist1",histInfoList[0],histInfoList[1],histInfoList[2]) hist2 = TH1D("hist2","hist2",histInfoList[0],histInfoList[1],histInfoList[2]) hist1.Sumw2() hist2.Sumw2() #Load variables from only branches actually used tree1.SetBranchStatus("*",0) tree2.SetBranchStatus("*",0) listOfVars1 = getListOfUsedVariables(variable1,cut1) listOfVars2 = getListOfUsedVariables(variable2,cut2) #add weight branches if (sWeighted1): listOfVars1 = np.append(listOfVars1,"N_Bplus_sw") if (sWeighted2): listOfVars2 = np.append(listOfVars2,"N_Bplus_sw") if (bWeighted1 and TM1): listOfVars1 = np.append(listOfVars1,"weight_nLongTracks") #TODO: check the branch name elif (b2Dweighted1 and TM1): listOfVars1 = np.append(listOfVars1,"weight2D_nLongTracks") if (bWeighted2 and TM2): listOfVars2 = np.append(listOfVars2,"weight_nLongTracks") elif (b2Dweighted2 and TM2): listOfVars2 = np.append(listOfVars2,"weight2D_nLongTracks") if verbose: print "[DEBUG]\t\tActivating branches in tree1: ", listOfVars1 print "[DEBUG]\t\tActivating branches in tree2: ", listOfVars2 #Assign variables with either a double or integer #WeAllHateRoot for var in listOfVars1: tree1.SetBranchStatus(var,1) if(isInt(var)): variables1Dict[var] = array('i',[0]) else: variables1Dict[var] = array('d',[0]) tree1.SetBranchAddress(var,variables1Dict[var]) for var in listOfVars2: tree2.SetBranchStatus(var,1) if(isInt(var)): variables2Dict[var] = array('i',[0]) else: variables2Dict[var] = array('d',[0]) tree2.SetBranchAddress(var,variables2Dict[var]) if verbose: print "[DEBUG]\t\tDefining variables for branches in tree1: ", variables1Dict print "[DEBUG]\t\tDefining variables for branches in tree2: ", variables2Dict if verbose: print "[DEBUG]\t\tUsing cuts1: ", cut1 print "[DEBUG]\t\tUsing cuts2: ", cut2 for evt in xrange(tree1.GetEntries()): if ( (evt+1) % 10000 < 0.01): print "[INFO]\tReding event ", evt+1, " from tree1" tree1.GetEntry(evt) if not evaluateCut(cut1,variables1Dict): continue hist1.Fill(evaluateVariable(variable1,variables1Dict),evaluateWeight(variables1Dict,optionsDictionary1)) for evt in xrange(tree2.GetEntries()): if ( (evt+1) % 10000 < 0.01): print "[INFO]\tReding event ", evt+1, " from tree2" tree2.GetEntry(evt) if not evaluateCut(cut2,variables2Dict): continue hist2.Fill(evaluateVariable(variable2,variables2Dict),evaluateWeight(variables2Dict,optionsDictionary2)) #Clone histograms to make a ratio #TODO check if it's really necessary hist1Clone = hist1.Clone('hist1Clone') hist2Clone = hist2.Clone('hist2Clone') #Normalize histograms if (hist1Clone.Integral() == 0): raise Exception('Filled histogram1 is empty! Check your options!') if (hist2Clone.Integral() == 0): raise Exception('Filled histogram2 is empty! Check your options!') hist1Clone.Scale(1.0/hist1Clone.Integral()) hist2Clone.Scale(1.0/hist2Clone.Integral()) histRatio = hist2Clone.Clone('histRatio') histRatio.Divide(hist1Clone) histRatio.Draw() #Plotting lowerPlotHeight = 0.6; canvas = TCanvas('canvas','ratio and superposition') pad1 = TPad("pad1", "ratio",0.0,lowerPlotHeight,1.0,1.0,0); pad2 = TPad("pad2", "superposition",0.0,0.0,1.0,lowerPlotHeight,0); pad1.Draw(); pad2.Draw(); xAxisLabel = defineXAxisLabels(variable1,variable2) makePlot(hist1Clone, hist2Clone, histRatio, pad1, pad2, xAxisLabel, histInfoList[1], histInfoList[2], lowerPlotHeight) #TODO put this into plotting function, where it god-knows-why doesn't work pad1.cd() unityline = TLine(histInfoList[1], 1.0, histInfoList[2], 1.0) unityline.SetLineStyle(2) unityline.SetLineColor(4) unityline.Draw() #Legend definition and drawing LegendXposition = 0.4; LegendYposition = 0.62; pad2.cd() leg = TLegend(LegendXposition,LegendYposition,LegendXposition + 0.25, LegendYposition + 0.28); tag1, tag2 = defineLegendTags(optionsDictionary1,optionsDictionary2) pad2.cd() if (len(variable1) > 30): leg.AddEntry(hist1Clone,"#splitline{%s}{#splitline{%s}{%s}}" % (variable1, tag1, cut1), "lpe"); else: leg.AddEntry(hist1Clone,"#splitline{%s%s}{%s}" % (variable1, tag1, cut1), "lpe"); if (len(variable2) > 30): leg.AddEntry(hist2Clone,"#splitline{%s}{#splitline{%s}{%s}}" % (variable2, tag2, cut2), "lpe"); else: leg.AddEntry(hist2Clone,"#splitline{%s%s}{%s}" % (variable2, tag2, cut2), "lpe"); leg.SetTextSize(0.04/lowerPlotHeight); leg.Draw() canvas.Print(getPlotFileName(variable1, optionsDictionary1, cut1, variable2, optionsDictionary2, cut2)) return if __name__ == '__main__': compareUltimate()