from ROOT import gROOT, gDirectory, gStyle, TChain, TTree, TAxis, TH1D import subprocess import time import numpy as np import re #for splitting with multiple deliminers from array import array import numexpr import argparse import sys import os import re from Globals import * ''' Tools used in compareUlitmate.py ''' #Define a list of variables that are considered integers intVarsList = ["nTracks" ,"nLongTracks" ,"nVeloTracks" ,"nTTracks" ,"nDownstreamTracks" ,"nUpstreamTracks" ,"nCandidate" ,"totCandidates" ,"nSPDHits" ,"nPV" ,"B_plus_ID" ,"B_plus_TRUEID" ,"K_star_plus_ID" ,"K_plus_ID" ,"pi_zero_resolved_ID" ,"mu_plus_ID" ,"mu_minus_ID" ,"Polarity" ,"nDiMuonMassBin" ,"TMedBKGCAT" ] def expressionList(): return ["Sqrt" ,"Abs" ,"Sin" ,"Cos" ,"Exp" ,"Log" ] #Define a function that returns true for variables that are integers def isInt(variable=""): if variable in intVarsList: return True else: return False def treeName(MC=False, TM=False, Preselected=False): if (TM or (MC and Preselected)): #TODO return "DecayTreeTruthMatched" elif (not Preselected): return "b2KstKs0pimumu_Tuple" if KshortChannel() else "b2KstKpi0mumuResolvedTuple/DecayTree" else: return "DecayTree" def getOptionsDictionary(year, Run, magnet, MC, TM, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, weightBranch, KshortDecaysInVelo, UseLowQ2Range): dictionaryTmp = { 'year': year, 'Run' : Run, 'magnet' : magnet, 'MC' : int(MC), 'TM' : int(TM), 'ReferenceChannel' : int(ReferenceChannel), 'PHSP' : int(PHSP), 'Preselected' : int(Preselected), 'BDTed' : BDTed, 'sWeighted' : sWeighted, 'bWeighted' : bWeighted, 'b2Dweighted' : b2Dweighted, 'weightBranch' : weightBranch, 'KshortDecaysInVelo' : int(KshortDecaysInVelo), 'UseLowQ2Range' : UseLowQ2Range } return dictionaryTmp def checkYearSample(optionsDict): #If wrong year/MC, skip if optionsDict['year'] not in yearList(): return False, "Wrong Year!" if (not optionsDict['MC']): return True, "" else: if (KshortChannel()): if (optionsDict['PHSP']): return True, "" elif (optionsDict['ReferenceChannel']): if (optionsDict['year'] < 2013 or optionsDict['year'] == 2016): return True, "" #Referenche channel available only for Run I else: return False, "RefChannel only available for Run I!" else: return True, "" else: if (optionsDict['PHSP']): if (optionsDict['year'] < 2018): return True, "" else: return False, "" elif (optionsDict['ReferenceChannel']): if (optionsDict['year'] < 2017): return True, "" #Referenche channel available only for Run I else: return False, "RefChannel only available for < 2017!" else: if (optionsDict['year'] > 2017): return False, "No MC for 2018!" else: return True, "" def getTreePath(optionsDict, verbose): #getPathForPython usage: # getPathForPython command year Run magnet Preselected MC ReferenceChannel PHSP KshortDecayInVelo Selection UseLowQ2Range #It has to have everything to keep it easy #Get first file name filePath2 = "" command = "/home/lhcb/kopecna/B2KstarMuMu/code/ewp-Bplus2Kstmumu-AngAna/CodeForTests/CompareUltimate/getPathForPython " #TODO maybe replace by a current path if(optionsDict['BDTed']): options = "BDToutput %(year)i %(Run)i %(magnet)s 0 %(MC)r %(ReferenceChannel)r %(PHSP)r %(KshortDecaysInVelo)r 0 %(UseLowQ2Range)r" %optionsDict if verbose: print "[DEBUG]\t\tRunning a c++ script to get the name of the file to open: ", command + options #Read second line from the c++ script (first one is the Hello...) filePath = subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1] elif (optionsDict['sWeighted'] or optionsDict['bWeighted'] or optionsDict['b2Dweighted']): options = "BDTinput %(year)i %(Run)i %(magnet)s 0 %(MC)r %(ReferenceChannel)r %(PHSP)r %(KshortDecaysInVelo)r 0 %(UseLowQ2Range)r" %optionsDict if verbose: print "[DEBUG]\t\tRunning a c++ script to get the name of the file to open: ", command + options #Read second line from the c++ script (first one is the Hello...) filePath = subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1] else: if (optionsDict['magnet']=="both"): #down options = "input %(year)i %(Run)i down %(Preselected)r %(MC)r %(ReferenceChannel)r %(PHSP)r %(KshortDecaysInVelo)r 0 %(UseLowQ2Range)r" %optionsDict if verbose: print "[DEBUG]\t\tRunning a c++ script to get the name of the file to open: ", command + options filePath = subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1] #up options = "input %(year)i %(Run)i up %(Preselected)r %(MC)r %(ReferenceChannel)r %(PHSP)r %(KshortDecaysInVelo)r 0 %(UseLowQ2Range)r" %optionsDict if verbose: print "[DEBUG]\t\tRunning a c++ script to get the name of the file to open: ", command + options filePath2 = subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1] else: options = "input %(year)i %(Run)i %(magnet)s %(Preselected)r %(MC)r %(ReferenceChannel)r %(PHSP)r %(KshortDecaysInVelo)r 0 %(UseLowQ2Range)r" %optionsDict if verbose: print "[DEBUG]\t\tRunning a c++ script to get the name of the file to open: ", command + options filePath = subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1] return filePath, filePath2 def addToTChain(tree, optionsDict, verbose): filePath, filePath2 = getTreePath(optionsDict=optionsDict, verbose=verbose) #Get first file name if(optionsDict['BDTed'] or optionsDict['sWeighted'] or optionsDict['bWeighted'] or optionsDict['b2Dweighted']): print "[INFO]\tOpening", filePath tree.Add(filePath) else: if (optionsDict['magnet']=="both"): print "[INFO]\tOpening", filePath tree.Add(filePath) print "[INFO]\tOpening", filePath2 tree.Add(filePath2) else: print "[INFO]\tOpening", filePath tree.Add(filePath) def addTMathTags(varName): for expr in expressionList(): varName = varName.replace(expr,"TMath::"+expr) return varName def list_of_DTF_vars (): return [ "_PT", "_PX", "_PY", "_PZ", "_PE", "_P" "_ETA" , "_MM", "B_plus_M" , "K_star_plus_M" , "pi_zero_resolved_plus_M" ] def replace_variables_to_DTF(variable = ""): #Append "_DTF" to everything in the list above by for var in list_of_DTF_vars(): if (variable.find(var) != -1): positions = [i.start() for i in re.finditer(var, variable)] #print positions for i in range(len(positions)): pos=positions[i] variable = variable[:pos+len(var)] + "_DTF" + variable[pos+len(var):] positions = [x+4 for x in positions] #Remove false DTF variales, such as MinIP (searching for _M) removeList = re.findall('DTF[A-Z]', variable) #print removeList for remove in removeList: positions = [i.start() for i in re.finditer(remove, variable)] for i in range(len(positions)): pos=positions[i] variable = variable[:pos-1] + variable[pos+3:] positions = [x-4 for x in positions] return variable def getListOfUsedVariables(variable, cut): tmp = str(variable) + "|" + str(cut) #| is there as a separator tmp = tmp.replace("gamma1","gammaa") tmp = tmp.replace("gamma2","gammab") #Protect from removing numbers tmp = re.sub(r'[0-9]+', '', tmp) #remove numbers tmp = tmp.replace("gammaa","gamma1") tmp = tmp.replace("gammab","gamma2") #Protect from removing numbers tmp = tmp.replace("CHI","CHI2") #check CHI2 numbers tmp = tmp.replace("DOCA","DOCA1") #check CHI2 numbers tmp = tmp.replace("Q","Q2") #check Q2 numbers #TODO tmp = tmp.replace(".","") #remove decimals left from numbers tmp = tmp.replace("-","|") #remove minus from negative numbers, keep an operator tmp = tmp.replace("TMath::","") #remove functions for expr in expressionList(): tmp = tmp.replace(expr,"") #Remove expressions tmp = tmp.replace(expr.lower(),"") #remove expressions without capital tmp = tmp.replace("<=","<") tmp = tmp.replace(">=",">") tmp = tmp.replace("==","=") tmp = tmp.replace("&&","&") tmp = tmp.replace("||","|") tmp = tmp.replace("!","") #tmp = re.sub(r"\s+", "", tmp) #Removes whitespaces, s matches all whitespaces branchList = re.split('[+ - * / ( ) < > = & | ]', tmp) branchList = filter(None, branchList) #removes empty strings branchList = np.unique(branchList) #Removes duplicates return branchList def evaluateCut(cut,variablesDict): #TODO: should be easy to let it be able take operations into account if (cut == ""): return True tmp = str(cut) for varName,varValue in variablesDict.items(): tmp = tmp.replace(str(varName),str(varValue[0])) tmp = tmp.replace ("||"," or ") tmp = tmp.replace ("&&"," and ") tmp = tmp.replace ("="," == ") tmp = tmp.replace ("! =="," != ") #in case #print "[DEBUG]\t\t", tmp return eval(tmp) def makeVariablePythonFriendly(variable): tmp = variable tmp = tmp.replace("TMath::","") #remove ROOT tags tmp = tmp.lower() #replaces e.g. Log by log and so on return tmp def evaluateVariable(variable,variablesDict): tmp = variable for varName,varValue in variablesDict.items(): tmp = tmp.replace(str(varName),str(varValue[0])) tmp = makeVariablePythonFriendly(tmp) return numexpr.evaluate(tmp).item() def evaluateWeight(variablesDict,optionsDict): #TODO if (optionsDict["sWeighted"]): return variablesDict["N_Bplus_sw"][0] elif (optionsDict["bWeighted"]): return variablesDict["weight_nLongTracks"][0] elif (optionsDict["b2Dweighted"]): return variablesDict["weight2D_nLongTracks"][0] else: return 1.0 def datasetTag(optionsDict): tag = "" if (optionsDict['MC']): if (optionsDict['ReferenceChannel']): tag = "RefChannel" elif (optionsDict['PHSP']): tag = "PHSP" else: tag = "MC" if (optionsDict['TM']): tag = tag + "_TM" else: tag = "data" if not (optionsDict['Preselected']): tag = tag + "_Strip" return tag def weightTag(optionsDict): tag = "" if (optionsDict['sWeighted']): tag = "_sWeighted" if (optionsDict['bWeighted']): tag = "_1DWeight" if (optionsDict['b2Dweighted']): tag = "_2DWeight" return tag def KshortDecaysInVeloTag(KshortDecaysInVelo=False): return "" if not KshortChannel() else ("_LL" if KshortDecaysInVelo else "_DD" ) def variableTag(variable=""): name = variable.replace("TMath::","") name = name.replace("(","") name = name.replace(")","") name = name.replace("|","") name = name.replace("/","_over_") name = name.replace("*","_x_") return name def histName(variable, optionsDict): name = variableTag(variable) name = name.replace("TMath::","") #remove functions for expr in (expressionList()): #For whatever reason I have to but the brackets there name = name.replace(expr.lower(),expr) #Add capitals to functions name = "Hist_" + str(optionsDict['year']) + "_" + datasetTag(optionsDict) name = name + KshortDecaysInVeloTag(optionsDict) + weightTag(optionsDict) return name def stopWatch(value): '''From seconds to Days;Hours:Minutes;Seconds''' valueD = (((value/365)/24)/60) Days = int (valueD) valueH = (valueD-Days)*365 Hours = int(valueH) valueM = (valueH - Hours)*24 Minutes = int(valueM) valueS = (valueM - Minutes)*60 Seconds = int(valueS) print Days,";",Hours,":",Minutes,";",Seconds class ShowArgumentsParser(argparse.ArgumentParser): def error(self, message): sys.stderr.write('error: %s\n\n' %message) parser.print_usage(sys.stderr) sys.stderr.write('\n'+self.description) sys.exit(2) def checkMCyear(year, ReferenceChannel, PHSP): if ((not KshortChannel()) and (not ReferenceChannel) and (not PHSP) and year == 2015): return 2016 else: return year