addNVtxWeight

#!/usr/bin/env python

import argparse
import sys
import os
from array import array

"""
Setup argument parser
"""

parser = argparse.ArgumentParser(description="This program creates a MC Tag-and-Probe (TnP) tree with an additional weight branch from given MC and data TnP input files. The weigths are calculated from the ratio of number of vertices in data relative to MC (weight = numVtxData/numVtxMC). Because the full input MC files are copied to the output and a new branch is added to the output file, this process can take a while.")
parser.add_argument("filenamesInputData", help="Path to the input Tag-And-Probe ROOT Data files")
parser.add_argument("filenamesInputMC", help="Path to the input Tag-And-Probe ROOT MC files")
parser.add_argument("filenameOutputMC", help="Path to the output Tag-And-Probe ROOT MC file with added weight branch")
parser.add_argument("-d", "--directory", default="tpTree", help="Directory in the input ROOT file which contains the Tag-And-Probe tree")
parser.add_argument("-t", "--tree", default="fitter_tree", help="Name of the tree holding the variables")
parser.add_argument("-b", "--branch", default="tag_nVertices", help="Branch name with number of vertices")
parser.add_argument("-c", "--cut", default="", help="Cut string which is applied on number of vertices branch")
parser.add_argument("-hf", "--histogramFilename", default="control_nVtx.png", help="Output filename of control histogram with number of vertices")
parser.add_argument("-hr", "--histogramRange", default="100,-0.5,99.5", help="Range of control histogram given as \"bins,min,max\"")
parser.add_argument("-v", "--verbosity", default=True, help="Set verbosity to [0, 1]")
args = parser.parse_args()

"""
Read input files
"""

from ROOT import * # import this here, otherwise it overwrites the argparse stuff
gROOT.SetBatch(True) # set ROOT to batch mode, this suppresses printing canvases

# Get path to tree
treePath = os.path.join(args.directory,args.tree)
if args.verbosity==1:
    print('Used path to tree in files:')
    print('---------------------------')
    print(treePath)
    print('')

# Setup tree chains
chainData = TChain(treePath)
if args.verbosity==1:
    print('Input files data:')
    print('-----------------')
for filename in args.filenamesInputData.split(' '):
    if args.verbosity==1:
        print(filename)
    chainData.AddFile(filename)
if args.verbosity==1:
    print('')

chainMC = TChain(treePath)
if args.verbosity==1:
    print('Input files MC:')
    print('---------------')
for filename in args.filenamesInputMC.split(' '):
    if args.verbosity==1:
        print(filename)
    chainMC.AddFile(filename)
if args.verbosity==1:
    print('')

"""
Make histograms of number of vertices and create control plot
"""

if args.verbosity==1:
    print('Make control histogram:')
    print('-----------------------')
c1 = TCanvas("c1", "c1")
c1.Divide(2,1)
c1.cd(1)
chainData.Draw(args.branch+'>>hData('+args.histogramRange+')', args.cut)
c1.cd(2)
chainMC.Draw(args.branch+'>>hMC('+args.histogramRange+')', args.cut)
c1.Print(args.histogramFilename)
if args.verbosity==1:
    print('')

"""
Add weight branch to MC tree chain
"""

# Calculate weights
hData = gROOT.FindObject('hData')
hMC = gROOT.FindObject('hMC')
hData.Scale(1.0/hData.Integral())
hMC.Scale(1.0/hMC.Integral())

weights = [1.0]*(hData.GetNbinsX()+1)
for i in range(1, len(weights)):
    nMC = hMC.GetBinContent(i)
    nData = hData.GetBinContent(i)
    if nMC > 0:
        weights[i-1] = nData/nMC
    else:
        weights[i-1] = 1.0

# Create new file with directory
fileOutputMC = TFile.Open(args.filenameOutputMC, 'recreate')
fileOutputMC.mkdir(args.directory).cd()

# Clone tree from input MC
progressbarWidth = 40
if args.verbosity==1:
    print('Adding weight column:')
    print('---------------------')
#    sys.stdout.write('Progress: [{}]'.format('-'*progressbarWidth))
    sys.stdout.flush() # this forces to print the stdout buffer
    sys.stdout.write('\b'*(progressbarWidth+1)) # return to start of line, after '['


treeOutput = chainMC.CloneTree(0)
weight = array('f', [1.0])
treeOutput.Branch('weight', weight, 'weight/F')
numEvents = chainMC.GetEntries()

for i in range(numEvents):
    chainMC.GetEntry(i)
    weight[0] = weights[int(getattr(chainMC, args.branch))] # NOTE we need this struct because otherwise PyROOT somehow drops objects...
    treeOutput.Fill()
    if args.verbosity==1:
        if i%int(numEvents/(progressbarWidth-1))==0:
            sys.stdout.write('+')
            sys.stdout.flush()

if args.verbosity==1:
    sys.stdout.write('\n')
    print('')

# Write file, close file and print filename
fileOutputMC.Write()
fileOutputMC.Close()
if args.verbosity==1:
    print('Output file MC:')
    print('---------------')
    print(args.filenameOutputMC)
    print('')