Suggested speed improvements for the MCResampling tool

diff

Index: Utils.py
===================================================================
--- Utils.py	(revision 183973)
+++ Utils.py	(working copy)
@@ -3,11 +3,9 @@
 
 
 def __AddVar__(self, var, name, res):
-    has_branche = res.has_key(name)
-            
-    if has_branche:
+    try:
         res[name][0] = float(var)
-    else:
+    except KeyError:
         comp = array('f', [0])
         self.Branch(name, comp, name+'/F')
         comp[0]= float(var)
@@ -20,32 +18,28 @@
 
 
 def __AddVarInt__(self, var, name, res):
-        has_branche = res.has_key(name)
-
-        if has_branche:
-            res[name][0] = int(var)
-        else:
-            comp = array('i', [0])
-            self.Branch(name, comp, name+'/I')
-            comp[0]= int(var)
-            aux = {name: comp}
-            res.update(aux)
-        return res
+    try:
+        res[name][0] = int(var)
+    except KeyError:
+        comp = array('i', [0])
+        self.Branch(name, comp, name+'/I')
+        comp[0]= int(var)
+        aux = {name: comp}
+        res.update(aux)
+    return res
         
 TTree.AddVarInt = __AddVarInt__
 
 def __AddVarLong__(self, var, name, res):
-        has_branche = res.has_key(name)
-
-        if has_branche:
-            res[name][0] = int(var)
-        else:
-            comp = array('l', [0])
-            self.Branch(name, comp, name+'/l')
-            comp[0]= int(var)
-            aux = {name: comp}
-            res.update(aux)
-        return res
+    try:
+        res[name][0] = int(var)
+    except KeyError:
+        comp = array('l', [0])
+        self.Branch(name, comp, name+'/l')
+        comp[0]= int(var)
+        aux = {name: comp}
+        res.update(aux)
+    return res
         
 TTree.AddVarLong = __AddVarLong__
 
Index: MakePIDCorrection.py
===================================================================
--- MakePIDCorrection.py	(revision 183973)
+++ MakePIDCorrection.py	(working copy)
@@ -15,6 +15,8 @@
 import os.path
 
 
+
+
 def getRandomFromTNSparse(sparse,tmp,nDim,dim):
     """
     Generates the radom values for all dimensions simultaniously and saves them into tmp
@@ -42,7 +44,7 @@
         sys.exit(2)
 
 
-if '__main__' == __name__:
+def main():
     start()
     print ""
     parser = ShowArgumentsParserMCreweight(
@@ -170,7 +172,8 @@
 
     tmpResults_dict = {}
 
-    treeN = TTree('DecayTree', 'DecayTree')
+    fileN = TFile(opts.mcFile.replace(".root","_PID-corrected.root"),"RECREATE")
+    treeN = treeToCorrect.CloneTree(0)
     res = {}
 
     for iE in xrange(nEntries):
@@ -194,20 +197,29 @@
             nDim = len(dim)
             tmp = array.array('d', std.vector(float)(nDim))
             test = getRandomFromTNSparse(tmp_histos[0],tmp,nDim,dim)
-            if(test==0):
-                continue
-            for i in xrange(nDim):
-                tmpParticleResults_dict[axis_dict[i+3]] = tmp[i]
+            ##### The following lines slightly change the error handling from older versions. 
+            # Before, whenever test came out as test==0,
+            # it would use the values in tmpResults_dit[p] from the previous 
+            # iteration. Thus basically resample a random value
+            # It now will simply add a -1000
+            if test==0 :
+                for i in xrange(nDim):
+                    tmpParticleResults_dict[axis_dict[i+3]] = -1000
+            else: 
+                for i in xrange(nDim):
+                    tmpParticleResults_dict[axis_dict[i+3]] = tmp[i]
             tmpResults_dict[p] = tmpParticleResults_dict
 
-        treeN.SetDirectory(0)
-        CopyTree(treeToCorrect,treeN,res)
         for part in YourPart_dict.keys():
             for pid in pidVars:
                 treeN.AddVar(tmpResults_dict[part][pid],part+"_"+pid+"corr",res)
         treeN.Fill()
-    fileN = TFile(opts.mcFile.replace(".root","_PID-corrected.root"),"RECREATE")
     treeN.Write()
     fileN.Save()
     fileN.Close()
     print "Done, have a nice day!"
+
+
+
+if '__main__' == __name__:
+    main()

MakePIDCorrection.py

#################################################################################
###
###
###
#################################################################################

from PIDPerfScripts.StartScreen import *

from ROOT import *
from Utils import *

import array
import argparse
import sys
import os.path




def getRandomFromTNSparse(sparse,tmp,nDim,dim):
    """
    Generates the radom values for all dimensions simultaniously and saves them into tmp
    Returns 1 for success and 0 if the bins were empty
    """

    for i in xrange(len(dim)):
        tmp_hist = sparse.Projection(dim[i])
        nBins = tmp_hist.GetNbinsX()
        inte = tmp_hist.ComputeIntegral()
        if(nBins!=0 and inte!=0):
            tmp_val = tmp_hist.GetRandom()
            tmp[i] = tmp_val
        else:
            tmp_hist.IsA().Destructor(tmp_hist)
            return 0
        tmp_hist.IsA().Destructor(tmp_hist)
    return 1

class ShowArgumentsParserMCreweight(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 main():
    start()
    print ""
    parser = ShowArgumentsParserMCreweight(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        prog=os.path.basename(sys.argv[0]),
        description=("""Make PID correction for a given:
        a) MC file <mcFile> ("Complete path to the file")
        b) Path to the TTree inside the MC file <mcFilePathToTree>
        c) particle type <partName> (\"K\", \"P\", \"Pi\",\"Mu\", "Muplus" and so on)
        d) PID variable to correct for each particle, <pidCut>

For a full list of arguments, do: 'python {0} -h'
""").format(os.path.basename(sys.argv[0]))
        )

    ## add the positional arguments
    parser.add_argument('mcFile', metavar='<mcFile>',
                        help="Sets the MC file to correct the PID variables")

    parser.add_argument('mcFilePathToTree',metavar='<mcFilePathToTree>',
                        help="Sets the internal path to the TTree of the MC file")

    parser.add_argument('partName', metavar='<partName>',
                        help="Sets the particle type")

    parser.add_argument('pidVars', metavar='<pidVars>',
                        help="Sets the PID variables to store")
    #------------------------------------------------------------------------------
    parser.add_argument('pidLibrary', metavar='<pidLibrary>',
                        help="Sets the library file with the PID distributions")

    opts = parser.parse_args()

    pid_dict = {'DLLK':'PIDK','DLLp':'PIDp','DLLmu':'PIDmu','DLLpi':'PIDpi','ProbNNK':'ProbNNK','ProbNNp':'ProbNNp','ProbNNmu':'ProbNNmu','ProbNNpi':'ProbNNpi'}



    # load the the file to correct
    fileToCorrect = TFile(opts.mcFile)
    treeToCorrect = fileToCorrect.Get(opts.mcFilePathToTree)


    # load the pid library file
    fileLibrary = TFile(opts.pidLibrary)

    # set the particles to correct
    YourPart = opts.partName

    if YourPart.startswith("["):
        if not YourPart.endswith("]"):
            parser.error("Invalid Particles string %s" %YourPart)
        YourPart = YourPart[1:-1].split(',')
    elif YourPart.startswith("("):
        if not YourPart.endswith(")"):
            parser.error("Invalid Particles string %s" %YourPart)
        YourPart = YourPart[1:-1].split(',')
    else:
        YourPart = (YourPart,)
    if len(YourPart)>0:
        YourPart_dict={}

    for ypt in YourPart:
        tmps=ypt.split(':')
        YourPart_dict[tmps[0]] = tmps[1]

    # set pid variables to store
    pidVars = opts.pidVars

    if pidVars.startswith("["):
        if not pidVars.endswith("]"):
            parser.error("Invalid pid variables string %s" %pidVars)
        pidVars = pidVars[1:-1].split(',')
    elif pidVars.startswith("("):
        if not pidVars.endswith(")"):
            parser.error("Invalid pid variables string %s" %pidVars)
        pidVars = pidVars[1:-1].split(',')
    else:
        pidVars = (pidVars,)
    ct=0

    pidVars_dict={}
    for v in pidVars:
        pidVars_dict[ct]=v
        ct=ct+1


    #Start of PID corection
    #----------------------
    #load the libraries for each particle
    
    print "YourPart_dict : "+str(YourPart_dict)
    print "fileLibrary : "+str(fileLibrary)     
    print str(fileLibrary.Get("lab/K"))

    part_dict = {}
    for i in YourPart_dict:
        tmp = []
        tmp.append(fileLibrary.Get(YourPart_dict[i]).Get("histo_PID"))
        tmp.append(fileLibrary.Get(YourPart_dict[i]).Get("H1D_P"))
        tmp.append(fileLibrary.Get(YourPart_dict[i]).Get("H1D_ETA"))
        tmp.append(fileLibrary.Get(YourPart_dict[i]).Get("H1D_nTracks"))
        part_dict[i] = tmp

    #get dictionary for the particle to axis
    listOfAxes = tmp[0].GetListOfAxes()
    axis_dict = {}
    for j in xrange(3,listOfAxes.GetEntries()):
        axis_dict[j] = pid_dict[listOfAxes[j].GetName()]


    ### start to loop over entries in the TTree
    nEntries = treeToCorrect.GetEntries()
    nDim = len(axis_dict)

    dim = array.array("i",std.vector(int)(nDim))
    for j in xrange(nDim):
        dim[j] = j+3

    mom_dict = {}
    eta_dict = {}
    for p in YourPart_dict.keys():
        mom = (array.array('d',[0]))
        eta = (array.array('d',[0]))
        treeToCorrect.SetBranchAddress( p+"_"+"P", mom )
        treeToCorrect.SetBranchAddress( p +"_"+"ETA", eta )
        mom_dict[p] = mom
        eta_dict[p] = eta
        del mom
        del eta

    tmpResults_dict = {}

    fileN = TFile(opts.mcFile.replace(".root","_PID-corrected.root"),"RECREATE")
    # treeN = TTree('DecayTree', 'DecayTree')
    treeN = treeToCorrect.CloneTree(0)
    res = {}

    for iE in xrange(nEntries):
        treeToCorrect.GetEntry(iE)
        if(iE%100)==0:
            print nEntries-iE
        for p in YourPart_dict.keys():
            tmpParticleResults_dict = {}

            tmp_histos = part_dict[p]
            #find correct bin
            bin_P = tmp_histos[1].FindBin(mom_dict[p][0])
            bin_ETA = tmp_histos[2].FindBin(eta_dict[p][0])
            bin_nTrack = tmp_histos[3].FindBin(treeToCorrect.nTracks)


            #set projection range
            tmp_histos[0].GetAxis(0).SetRange(bin_P,bin_P)
            tmp_histos[0].GetAxis(1).SetRange(bin_ETA,bin_ETA)
            tmp_histos[0].GetAxis(2).SetRange(bin_nTrack,bin_nTrack)
            nDim = len(dim)
            tmp = array.array('d', std.vector(float)(nDim))
            test = getRandomFromTNSparse(tmp_histos[0],tmp,nDim,dim)
            if(test==0):
                continue
            for i in xrange(nDim):
                tmpParticleResults_dict[axis_dict[i+3]] = tmp[i]
            tmpResults_dict[p] = tmpParticleResults_dict

        # treeN.SetDirectory(0)
        # CopyTree(treeToCorrect,treeN,res)
        for part in YourPart_dict.keys():
            for pid in pidVars:
                treeN.AddVar(tmpResults_dict[part][pid],part+"_"+pid+"corr",res)
        treeN.Fill()
    treeN.Write()
    fileN.Save()
    fileN.Close()
    print "Done, have a nice day!"



if '__main__' == __name__:
    main()

Utils.py

from ROOT import *
from array import array


def __AddVar__(self, var, name, res):
    try:
        res[name][0] = float(var)
    except KeyError:
        comp = array('f', [0])
        self.Branch(name, comp, name+'/F')
        comp[0]= float(var)
        aux = {name: comp}
        res.update(aux)
    return res
    

TTree.AddVar = __AddVar__


def __AddVarInt__(self, var, name, res):
    try:
        res[name][0] = int(var)
    except KeyError:
        comp = array('i', [0])
        self.Branch(name, comp, name+'/I')
        comp[0]= int(var)
        aux = {name: comp}
        res.update(aux)
    return res
        
TTree.AddVarInt = __AddVarInt__

def __AddVarLong__(self, var, name, res):
    try:
        res[name][0] = int(var)
    except KeyError:
        comp = array('l', [0])
        self.Branch(name, comp, name+'/l')
        comp[0]= int(var)
        aux = {name: comp}
        res.update(aux)
    return res
        
TTree.AddVarLong = __AddVarLong__


def CopyTree(tree, newtree, res_tuple):

    branches = tree.GetListOfBranches()
    var_list = []
    for b in branches:
        var_list.append(b.GetName())

    for n in var_list:
        if type(tree.__getattr__(n))==type(float()):
            res_tuple = newtree.AddVar(tree.__getattr__(n), n, res_tuple)
        elif  type(tree.__getattr__(n))==type(int()):
            res_tuple = newtree.AddVarInt(tree.__getattr__(n), n, res_tuple)


def CopyBranches(tree, newtree, res_tuple, list_branches):

    for n in list_branches:
        if type(tree.__getattr__(n))==type(float()):
            res_tuple = newtree.AddVar(tree.__getattr__(n), n, res_tuple)
        elif  type(tree.__getattr__(n))==type(int()):
            res_tuple = newtree.AddVarInt(tree.__getattr__(n), n, res_tuple)
        elif type(tree.__getattr__(n))==type(long()):
            res_tuple = newtree.AddVarLong(tree.__getattr__(n), n, res_tuple)