187 lines
6.8 KiB
Python
187 lines
6.8 KiB
Python
#Core code for removal of multiple cnadidates from the sample
|
|
#Renata Kopecna
|
|
|
|
from ROOT import TTree, TChain, TFile, TVector3, TObject
|
|
import numpy as np
|
|
import itertools
|
|
#import matplotlib.pyplot as plt
|
|
from array import array
|
|
from collections import Counter
|
|
import time
|
|
|
|
import sys
|
|
sys.path.insert(0,'/home/lhcb/kopecna/B2KstarMuMu_clean/Code/Selection/ComparisonTool')
|
|
#import Utils
|
|
from Utils3 import getTreePath, getOptionsDictionary,treeName, stopWatch, getTreeWithPairingBranch
|
|
|
|
test = False
|
|
#Read Inclusive sample (horribly hacked in)
|
|
IsInc = False
|
|
|
|
def isMultipleCandidate(year, Run, Data, MC, TM, TMpid, gammaTM, ReferenceChannel, PHSP, KshortDecaysInVelo,verbose):
|
|
|
|
#read file
|
|
sWeighted = False
|
|
bWeighted = False
|
|
b2Dweighted = False
|
|
|
|
if ((TM + TMpid) > 1): raise Exception("Select only one of TM or TMpid!")
|
|
|
|
optionsDictionary = getOptionsDictionary(int(year),Run,"both",MC,TM,ReferenceChannel,PHSP,
|
|
True, True, sWeighted,bWeighted,b2Dweighted,"",KshortDecaysInVelo,False)
|
|
|
|
#Add tree to TChain (only one, since it has to be BDTed file)
|
|
treePath, treePath2 = getTreePath(optionsDictionary,verbose) #treePath2 is not used
|
|
if (IsInc): treePath = treePath.replace("KplusPi0Resolved","Inclusive",-1)
|
|
if (verbose): print("Opening " +treePath)
|
|
treeFileTmp = TFile.Open(treePath,"UPDATE")
|
|
treeTmp = treeFileTmp.Get(treeName(MC,TM,True))
|
|
|
|
if (not MC and TM):
|
|
print ("[WARN]\t\tTM cannot be set to True for data! setting TM to False")
|
|
TM = False
|
|
|
|
if (test) : nEvts = 23 #testing
|
|
else: nEvts = treeTmp.GetEntries()
|
|
|
|
#Check if the pairing function branch is in the tree
|
|
#If yes, continue, if not, create the branch:
|
|
if not("pairingNumber" in treeTmp.GetListOfBranches()):
|
|
treeFile, tree = getTreeWithPairingBranch(treeFileTmp,treeTmp,Run,test)
|
|
else: treeFile, tree = treeFileTmp,treeTmp
|
|
|
|
|
|
#Load needed branches
|
|
pairingNumber = array('L',[0])
|
|
MLPresponse = array('d',[0])
|
|
TMed = array('i',[0])
|
|
TM_gammas = array('i',[0])
|
|
|
|
#Activate and read branches
|
|
tree.SetBranchStatus('*',0)
|
|
tree.SetBranchStatus('pairingNumber',1)
|
|
tree.SetBranchStatus('MLPresponse',1)
|
|
if (TM): tree.SetBranchStatus('TMedBKGCAT',1)
|
|
elif (TMpid):
|
|
tree.SetBranchStatus('TMed',1)
|
|
tree.SetBranchStatus('TM_gammas',1)
|
|
#elif (pi0TMoff): tree.SetBranchStatus('TMed_noPi0',1)
|
|
tree.SetBranchAddress('pairingNumber',pairingNumber)
|
|
tree.SetBranchAddress('MLPresponse',MLPresponse)
|
|
if (TM): tree.SetBranchAddress('TMedBKGCAT',TMed)
|
|
elif (TMpid):
|
|
tree.SetBranchAddress('TMed',TMed)
|
|
tree.SetBranchAddress('TM_gammas',TM_gammas)
|
|
#elif (pi0TMoff): tree.SetBranchAddress('TMed_noPi0',TMed)
|
|
|
|
|
|
#Load events into a 3D array: [row,column] (1 row = 1 event)
|
|
evtNumberArray = np.zeros((nEvts,3), dtype=float)
|
|
for evt in range (nEvts):
|
|
tree.GetEntry(evt)
|
|
evtNumberArray[evt,0] = pairingNumber[0]
|
|
evtNumberArray[evt,1] = MLPresponse[0]
|
|
if (TM): evtNumberArray[evt,2] = TMed[0]
|
|
if (TMpid):
|
|
if (gammaTM): evtNumberArray[evt,2] = TMed[0] and TM_gammas[0]<4
|
|
else: evtNumberArray[evt,2] = TMed[0] and TM_gammas[0]<6
|
|
|
|
|
|
if (test): print (evtNumberArray)
|
|
|
|
print ("Loaded MLP+event number array.")
|
|
|
|
#Get all MLP values
|
|
MLPvalues = np.sort(evtNumberArray[:,1])
|
|
if (test): print (MLPvalues)
|
|
|
|
#Select only TMed events
|
|
if (TM or TMpid):
|
|
evtNumberArray = evtNumberArray[np.where(evtNumberArray[:,2] == 1)]
|
|
if(test): print (evtNumberArray)
|
|
|
|
|
|
#Run the first round to remove lonely events first (huge speed gain)
|
|
dictCountMLP = {}
|
|
unique, counts = np.unique(evtNumberArray[:,0],return_counts = True)
|
|
dictEvtCount = dict(zip(unique, counts))
|
|
for key in dictEvtCount.keys():
|
|
key_int = int(key)
|
|
if (dictEvtCount[key_int] == 1):
|
|
dictCountMLP[key_int]=0.0
|
|
evtNumberArray = evtNumberArray[np.where(evtNumberArray[:,0] != key_int)]
|
|
if(test): print (evtNumberArray)
|
|
|
|
#Loop over all possible MLP values
|
|
count = 0
|
|
start = time.time()
|
|
for MLPcut in MLPvalues:
|
|
if (verbose and count%1000 == 0): print ("Reading event", count)
|
|
#Select only events with MLP values than the previous smallest value
|
|
tmpArr = evtNumberArray[(evtNumberArray[:,1]>=MLPcut)]
|
|
if (len(tmpArr)==0): break
|
|
unique, counts = np.unique(tmpArr[:,0],return_counts = True)
|
|
dictEvtCount = dict(zip(unique, counts))
|
|
if (test): print ("dictEvtCount\n",dictEvtCount)
|
|
for key in dictEvtCount.keys():
|
|
key_int = int(key)
|
|
if key_int not in dictCountMLP:
|
|
if dictEvtCount[key_int] == 1:
|
|
dictCountMLP[key_int]=MLPcut
|
|
count += 1
|
|
end = time.time()
|
|
print ("Needed ", stopWatch(end-start), "to finish.")
|
|
if (test): print (dictCountMLP)
|
|
|
|
print ("Multiple candidates dictionary created")
|
|
|
|
tree.SetBranchStatus('*',1)
|
|
|
|
#Add new branch
|
|
IsAloneAt = array('d',[0])
|
|
if (MC):
|
|
if (TM): b_IsAloneAt = tree.Branch("IsAloneAt", IsAloneAt, 'IsAloneAt/D')
|
|
elif (TMpid):
|
|
if (gammaTM): b_IsAloneAt = tree.Branch("IsAloneAt_TMed", IsAloneAt, 'IsAloneAt_TMed/D')
|
|
else: b_IsAloneAt = tree.Branch("IsAloneAt_TMed_rndGamma", IsAloneAt, 'IsAloneAt_TMed_rndGamma/D')
|
|
else: b_IsAloneAt = tree.Branch("IsAloneAtNotTM", IsAloneAt, 'IsAloneAtNotTM/D')
|
|
else:
|
|
b_IsAloneAt = tree.Branch("IsAloneAt", IsAloneAt, 'IsAloneAt/D')
|
|
|
|
|
|
if (TM or TMpid):
|
|
for evt in range (nEvts):
|
|
tree.GetEntry(evt)
|
|
if (test): print (pairingNumber[0],dictCountMLP.get(pairingNumber[0]))
|
|
if (not TMed[0]): IsAloneAt[0] = 2.0
|
|
elif (TMpid and (gammaTM and TM_gammas[0]>3)): IsAloneAt[0] = 2.0
|
|
elif (TMpid and (not gammaTM and TM_gammas[0]==6)): IsAloneAt[0] = 2.0
|
|
else:
|
|
if (pairingNumber[0] in dictCountMLP.keys()): #check makes it slower, but oh well
|
|
IsAloneAt[0] = dictCountMLP.get(pairingNumber[0])
|
|
else:
|
|
IsAloneAt[0] = -1.0
|
|
b_IsAloneAt.Fill()
|
|
|
|
else:
|
|
for evt in range (nEvts):
|
|
tree.GetEntry(evt)
|
|
if (test): print (pairingNumber[0],dictCountMLP.get(pairingNumber[0]))
|
|
if (pairingNumber[0] in dictCountMLP.keys()): #check makes it slower, but oh well
|
|
IsAloneAt[0] = dictCountMLP.get(pairingNumber[0])
|
|
else:
|
|
IsAloneAt[0] = -1.0
|
|
b_IsAloneAt.Fill()
|
|
|
|
print ("Writing into file...")
|
|
treeFile.Write("",TFile.kOverwrite)
|
|
treeFile.Close()
|
|
|
|
print ("All done, exiting!")
|
|
return
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
removeMultiple()
|