#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()