EWP-BplusToKstMuMu-AngAna/Code/Selection/RemoveMultipleCandidatesSource.py

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