Update 'Utils3.py'

Renata Kopecná 2022-02-01 17:26:51 +01:00
parent 1a61301eab
commit 93e005aaba

@ -5,152 +5,86 @@ Helper functions originally created for [[the comparison tool|Comparison-tool]].
## Classes
### 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)
Class used for pretty parser error printouts
## Global variables
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"
]
expressionList():
return ["Sqrt"
,"Abs"
,"Sin"
,"Cos"
,"Exp"
,"Log"
]
### intVarsList
A list of trree branches that are saved as integers
## Functions
### def isInt(variable=""):
A function that returns true for variables that are integers
### def treeName(MC=False, TM=False, Preselected=False)
Returns the correct tree name for given sample.
### def getOptionsDictionary(year, Run, magnet, MC, TM, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, weightBranch, KshortDecaysInVelo, UseLowQ2Range)
Returns a dicitonary with all possible saved options
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 expressionList()
* **Return**
* A list of used expressions, used for pasring any input formula
### def isInt()
* **Parameters**
* variable=""
* **Return**
* True is the variable is stored as an integer, false otherwise
### def checkYearSample(optionsDict)
Checks whether the selected data sample defined by optionsDict makes sense. If yes, returns True, otherwise False
### def treeName()
* **Parameters**
* MC=False, TM=False, Preselected=False
* **Return**
* The correct tree name for given sample.
### def getOptionsDictionary(y)
* **Parameters**
* year
* Run
* magnet
* MC
* TM
* ReferenceChannel
* PHSP
* Preselected
* BDTed
* sWeighted
* bWeighted
* b2Dweighted
* weightBranch
* KshortDecaysInVelo
* UseLowQ2Range
* **Reurn**
* A dicitonary with all possible saved options in
### def checkYearSample()
* **Parameters**
* optionsDict
* **Return**
* Checks whether the selected data sample defined by optionsDict makes sense. If yes, returns True, otherwise False
### def checkYearSample(year,MC,ReferenceChannel,PHSP,KshortChannel)
Checks whether the selected data sample defined by the options makes sense. If yes, returns True, otherwise False
### def checkYearSample()
* **Parameters**
* year,MC,ReferenceChannel,PHSP,KshortChannel
* **Return**
* Checks whether the selected data sample defined by the options makes sense. If yes, returns True, otherwise False
### def getTreePath(optionsDict, verbose)
Calls [[getPathForPython|getPathForPython]].
getPathForPython usage:
getPathForPython command year Run magnet Preselected MC ReferenceChannel PHSP KshortDecayInVelo Selection UseLowQ2Range
It has to have everything to keep it easy
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 getTreeList(year, Run, MC,TM,Reference,PHSP, verbose): #WIP
#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
yearArr = yearRunArray(year,Run, MC,Reference,PHSP)
optionsDict = {
'MC' : int(MC),
'TM' : int(TM),
'Ref' : int(Reference),
'PHSP' : int(PHSP),
}
#if (verbose): print(yearArr)
filePath = []
for yr in yearArr.T:
if (verbose): print ("Opening year", yr)
s_yr = str(yr[0])
s_run = str(yr[1])
options = "BDToutput "+s_yr+" " + s_run+ " both 1 %(MC)r %(Ref)r %(PHSP)r 0 0 0" %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.append(subprocess.check_output(command + options, shell=True,universal_newlines=True).split('\n')[1])
if (verbose): print ("Files read: ", filePath)
return filePath
### def getTreePath()
* **Parameters**
* optionsDict, verbose
* **Return**
* Returns the path to the tree based on the optionsDict. Calls [[getPathForPython|getPathForPython]].
### def addToTChain()
* **Parameters**
* tree, optionsDict, verbose
* **Return**
### def addToTChain(tree, optionsDict, verbose):
filePath, filePath2 = getTreePath(optionsDict=optionsDict, verbose=verbose)
#Get first file name
@ -168,161 +102,98 @@ def getTreeList(year, Run, MC,TM,Reference,PHSP, verbose): #WIP
tree.Add(filePath)
### def addTMathTags(varName):
### def addTMathTags()
* **Parameters**
* varName
* **Return**
*
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 list_of_DTF_vars ()
### def replace_variables_to_DTF(variable = ""):
* **Return**
* A list of variables that can have a DTF suffix
### def replace_variables_to_DTF()
Pretty much appends \"_DTF\" to everything
* **Parameters**
* variable = ""
* **Return**
* the updated name of the variable (with \"_DTF\" appended)
#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 getListOfUsedVariables(
Parses the desired cut on the given variable to be readable by python
* **Parameters**
* variable, cut
* **Return**
* List of branches that need to be loaded
*
### def evaluateCut()
* **Parameters**
* cut,variablesDict
* **Return**
* parses and evaluates the cut using eval()
### 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 makeVariablePythonFriendly()
Removes ROOT tags and lowers all the letters
* **Parameters**
* variable
* **Return**
* variable parsed to be readable by python
### def evaluateVariable(variable,variablesDict):
### def evaluateVariable()
* **Parameters**
* variable,variablesDict
* **Return**
*
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 evaluateWeight() #TODO
* **Parameters**
* variablesDict,optionsDict
* **Return**
* weight for the given optionsDict
### 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 datasetTag()
* **Parameters**
* optionsDict
* **Return**
* Tag (name) for the desired dataset
### def weightTag(optionsDict):
tag = ""
if (optionsDict['sWeighted']): tag = "_sWeighted"
if (optionsDict['bWeighted']): tag = "_1DWeight"
if (optionsDict['b2Dweighted']): tag = "_2DWeight"
return tag
### def weightTag()
* **Parameters**
* optionsDict
* **Return**
* weight tag for the given optionsDict
*
### def KshortDecaysInVeloTag()
* **Parameters**
* KshortDecaysInVelo=False
* **Return**
* Empty string if K+pi0, otherwise '\_LL' if KshortDecaysInVelo else '\_DD'
### 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 variableTag()
* **Parameters**
* variable=""
* **Return**
* parsed and cleanedup variable name for saving the file. It is not perfext but works somewhat
### def histName(variable, optionsDict):
### def histName()
* **Parameters**
* variable, optionsDict
* **Return**
*
name = variableTag(variable)
name = name.replace("TMath::","") #remove functions
for expr in (expressionList()): #For whatever reason I have to but the brackets there
@ -333,79 +204,25 @@ def getTreeList(year, Run, MC,TM,Reference,PHSP, verbose): #WIP
return name
### def stopWatch(value):
'''From seconds to Days;Hours:Minutes;Seconds'''
### def stopWatch()
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)
Prints the time passed between `value` and now.
* **Parameters**
* value
print (Days,";",Hours,":",Minutes,";",Seconds)
### def getTreeWithPairingBranch()
Adds the branch with the corresponding pairing function to the tree. Pairing function is a unique number calculated from the event and run numbers. This is calculated using the [Cantor pairing function](https://en.wikipedia.org/wiki/Pairing_function#Cantor_pairing_function)
* **Parameters**
* treeFile,tree,run,test
* **Return**
* The new tree file and the new tree
### def checkMCyear()
* **Parameters**
* year, ReferenceChannel, PHSP
* **Return**
### def getTreeWithPairingBranch(treeFile,tree,run,test):
print ("Creating the pairing function...")
#Load needed branches
eventNumber = array('L',[0])
runNumber = array('L',[0])
#Activate and read branches
tree.SetBranchStatus('*',0)
tree.SetBranchStatus('eventNumber',1)
tree.SetBranchStatus('runNumber',1)
tree.SetBranchAddress('eventNumber',eventNumber)
tree.SetBranchAddress('runNumber',runNumber)
#Add new branch
pairingNumber = array('L',[0])
b_pairingNumber = tree.Branch("pairingNumber", pairingNumber, 'pairingNumber/l')
#Loop over events and save
pairingNumber[0] = 10
for evt in range (tree.GetEntries()):
tree.GetEntry(evt) #pi(a,b)= 1/2(a+b)(a+b+1)+b
if (run ==2): rn = runNumber[0] - 100000
else: rn = runNumber[0]
if (eventNumber[0]>6100000000): #ONLY FOR DATA 2017!!!
rn = 1 #Because the high numbers come from only one run
en = eventNumber[0]-6100000000
else:
en = eventNumber[0]
pairingNumber[0] = int((en+rn)*(en+rn+1)/2+rn)
if (test): print (en,rn,pairingNumber[0])
b_pairingNumber.Fill()
#Get tree path
treePath = treeFile.GetPath()
treePath = treePath.replace(".root:/",".root")
if (test): print ("treePath:", treePath)
#Get tree name
treeName = tree.GetName()
if (test): print ("treeName:", treeName)
#Just for safety reasons, close the file and open it again
print ("Writing into file",treePath)
treeFile.Write("",TFile.kOverwrite)
treeFile.Close()
print ("Open file",treePath)
treeFileNew = TFile.Open(treePath,"UPDATE")
treeNew = treeFileNew.Get(treeName)
return treeFileNew,treeNew
### def checkMCyear(year, ReferenceChannel, PHSP):
if ((not KshortChannel) and (not ReferenceChannel) and (not PHSP) and year == 2015): return 2016
else: return year