VHbb/python/addingSamples.py

#!/usr/bin/env python

import sys,hashlib,ROOT
ROOT.gROOT.SetBatch(True)
from array import array
from optparse import OptionParser

def getTotal( bin, fileList ):
    total = 0.
    for i in range(0,len(fileList)):
        total = total + fileList[i][2][bin]
    return total


argv = sys.argv
parser = OptionParser()

parser.add_option("-C", "--config", dest="config", default=[], action="append",
                              help="configuration defining bins")
parser.add_option("-A", "--apply", dest="apply", default=False, action="store_true",
                              help="configuration defining the")
parser.add_option("-W", "--weight", dest="weights", default=True, action="store_false",
                              help="Calculate the weights")
parser.add_option("-V", "--validate", dest="validate", default=True, action="store_false",
                              help="Make validation plot the weights")
(opts, args) = parser.parse_args(argv)
if opts.config =="":
    opts.config = "config"

from myutils import BetterConfigParser

print opts.config
config = BetterConfigParser()
config.read(opts.config)
prefix = config.get('LHEWeights','prefix')
newpostfix = config.get('LHEWeights','newpostfix')
inclusive = config.get('LHEWeights','inclusive')
files = eval(config.get('LHEWeights','fileList'))
lheBin = eval(config.get('LHEWeights','lheBin'))
fileList = []
for file in files:
    new_entry = [prefix+file+'.root',files[file],[]]
    if file == inclusive:
        fileList.insert(0, new_entry)
    else:
        fileList.append(new_entry)
print fileList
print lheBin

def get_weights(fileList,lheBin):    
    for file in fileList:
        #file.append( 1 )
        #continue
        print file
        infile = ROOT.TFile.Open(file[0],"READ")
        tree = infile.Get('tree')
        for bin in lheBin:
            file[2].append( float(tree.GetEntries(bin)) )
        count = infile.Get('CountWithPU')
        file.append( count.GetBinContent(1) )
        infile.Close()
        del tree

    CountIncl = fileList[0][3]
    print fileList

    #total -> total numer of events in each lheBin
    print 'Calculating total'
    total = []
    weight= []
    for bin in range(0, len(lheBin) ):
        #weight.append(1.)
        #continue
        total.append(getTotal(bin, fileList))
        #print lheBin[bin]
        #to better stich we need the highest stat (that should correspond to the binned sample relative to the bin)
        fileList.sort( key=lambda file: file[2][bin], reverse=True )
        #print 'After sorting'
        #print fileList
        if total[bin] > 0.:
            #the first is always the one with the highest N in the bin: 
            weight.append( (fileList[0][1]/fileList[0][3]) * (CountIncl/2950.0) * fileList[0][2][bin]/total[bin] )
        else:
            weight.append(1.)

    weight_map = {}
    for bin in range(0, len(lheBin) ):
        print '%s: %.2f' %(lheBin[bin], weight[bin])
        weight_map[lheBin[bin]] = weight[bin]
    return weight_map


def apply_weights(fileList,weight_map,inclusive,newpostfix):    
    #now add the branch with the weight normalized to the inclusive
    for file in fileList:
        outfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),'RECREATE')
        infile = ROOT.TFile.Open(file[0],"READ")
        histoInfile = ROOT.TFile.Open(inclusive+'.root',"READ")
        histoInfile.cd()
        obj = ROOT.TObject
        for key in ROOT.gDirectory.GetListOfKeys():
            histoInfile.cd()
            obj = key.ReadObj()
            #print obj.GetName()
            if obj.GetName() == 'tree':
                continue
            outfile.cd()
            #print key.GetName()
            obj.Write(key.GetName())

        infile.cd()
        tree = infile.Get('tree')
        outfile.cd()
        newtree = tree.CloneTree(0)
        lheWeight = array('f',[0])
        newtree.Branch('lheWeight',lheWeight,'lheWeight/F')

        nEntries = tree.GetEntries()
        theBinForms = {}
        for bin in weight_map:
            theBinForms[bin] = ROOT.TTreeFormula("Bin_formula_%s"%(bin),bin,tree)
        for entry in range(0,nEntries):
            tree.GetEntry(entry)
            match_bin = None
            for bin in weight_map:
                if theBinForms[bin].EvalInstance() > 0.:
                    match_bin = bin
            if match_bin:
                lheWeight[0] = weight_map[match_bin]
            else:
                lheWeight[0] = 1.
            newtree.Fill()
                   
        newtree.AutoSave()
        outfile.Write()
        outfile.Close()
        infile.Close()
        del tree
        del newtree

def do_validation(fileList,inclusive,newpostfix):
    histo = ROOT.TH1F("lheV_pt","lheV_pt",300,0,300)
    histo1 = ROOT.TH1F("lheV_ptInc","lheV_ptInc",300,0,300)
    histo.SetDirectory(0)
    histo1.SetDirectory(0)

    for file in fileList: 
        h_name=str(hashlib.sha224(file[0]).hexdigest())
        print file[0]
        print h_name
        tfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),"read")
        tree = tfile.Get("tree")
        h_tmp = ROOT.TH1F(h_name,h_name,300,0,300)
        #ROOT.gDirectory.ls()
        tree.Draw("lheV_pt>>"+str(h_name),"(lheWeight)*(lheV_pt < 300.)","goff")
        if inclusive in file[0]:
            h_tmp1 = ROOT.TH1F(h_name+'Inc',h_name+'Inc',300,0,300)
            tree.Draw("lheV_pt>>"+str(h_name+'Inc'),"(lheV_pt < 300.)","goff")
            histo1.Add(h_tmp1)
        histo.Add(h_tmp)
        tfile.Close()
        del tree


    canvas = ROOT.TCanvas("lheV_pt","lheV_pt",600,600)
    ROOT.gPad.SetLogy()
    histo.Draw()
    histo1.Draw('SAME')
    print histo.Integral(100,300)
    print histo1.Integral(100,300)
    histo1.SetLineColor(ROOT.kRed)
    canvas.Print("validation_lheV_pt.pdf","pdf")

if opts.weights:
    weight_map = get_weights(fileList,lheBin)
    config.set('LHEWeights', 'weights_per_bin', '%s' %weight_map)
    f = open('8TeVconfig/lhe_weights', 'w')
    for section in config.sections():
        if not section == 'LHEWeights':
            config.remove_section(section)
    config.write(f)
    f.close()
elif opts.apply:
    weight_map = eval(config.get('LHEWeights', 'weights_per_bin'))
if opts.apply:
    apply_weights(fileList,weight_map,prefix+inclusive,newpostfix)
if opts.validate:
    do_validation(fileList,inclusive,newpostfix)
Revision:	1.7
Committed:	Tue Mar 5 17:02:59 2013 UTC (12 years, 2 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
Changes since 1.6:	+5 -2 lines
Log Message:	Fix to apply weights
#	User	Rev	Content
1	nmohr	1.6	#!/usr/bin/env python
2	bortigno	1.1
3	nmohr	1.6	import sys,hashlib,ROOT
4	nmohr	1.4	ROOT.gROOT.SetBatch(True)
5	bortigno	1.1	from array import array
6	nmohr	1.4	from optparse import OptionParser
7	bortigno	1.1
8			def getTotal( bin, fileList ):
9			total = 0.
10			for i in range(0,len(fileList)):
11			total = total + fileList[i][2][bin]
12			return total
13
14	nmohr	1.4
15			argv = sys.argv
16			parser = OptionParser()
17
18			parser.add_option("-C", "--config", dest="config", default=[], action="append",
19	nmohr	1.5	help="configuration defining bins")
20			parser.add_option("-A", "--apply", dest="apply", default=False, action="store_true",
21			help="configuration defining the")
22			parser.add_option("-W", "--weight", dest="weights", default=True, action="store_false",
23			help="Calculate the weights")
24			parser.add_option("-V", "--validate", dest="validate", default=True, action="store_false",
25			help="Make validation plot the weights")
26	nmohr	1.4	(opts, args) = parser.parse_args(argv)
27			if opts.config =="":
28			opts.config = "config"
29
30			from myutils import BetterConfigParser
31
32			print opts.config
33			config = BetterConfigParser()
34			config.read(opts.config)
35			prefix = config.get('LHEWeights','prefix')
36	nmohr	1.5	newpostfix = config.get('LHEWeights','newpostfix')
37	nmohr	1.4	inclusive = config.get('LHEWeights','inclusive')
38			files = eval(config.get('LHEWeights','fileList'))
39			lheBin = eval(config.get('LHEWeights','lheBin'))
40			fileList = []
41			for file in files:
42	nmohr	1.5	new_entry = [prefix+file+'.root',files[file],[]]
43	nmohr	1.4	if file == inclusive:
44			fileList.insert(0, new_entry)
45			else:
46			fileList.append(new_entry)
47	bortigno	1.1	print fileList
48	nmohr	1.4	print lheBin
49	bortigno	1.1
50	nmohr	1.4	def get_weights(fileList,lheBin):
51			for file in fileList:
52			#file.append( 1 )
53			#continue
54			print file
55	nmohr	1.5	infile = ROOT.TFile.Open(file[0],"READ")
56	nmohr	1.4	tree = infile.Get('tree')
57			for bin in lheBin:
58			file[2].append( float(tree.GetEntries(bin)) )
59			count = infile.Get('CountWithPU')
60			file.append( count.GetBinContent(1) )
61			infile.Close()
62			del tree
63	bortigno	1.1
64	nmohr	1.4	CountIncl = fileList[0][3]
65			print fileList
66	bortigno	1.1
67	nmohr	1.4	#total -> total numer of events in each lheBin
68			print 'Calculating total'
69			total = []
70			weight= []
71			for bin in range(0, len(lheBin) ):
72			#weight.append(1.)
73			#continue
74			total.append(getTotal(bin, fileList))
75			#print lheBin[bin]
76	bortigno	1.1	#to better stich we need the highest stat (that should correspond to the binned sample relative to the bin)
77	nmohr	1.4	fileList.sort( key=lambda file: file[2][bin], reverse=True )
78			#print 'After sorting'
79			#print fileList
80			if total[bin] > 0.:
81			#the first is always the one with the highest N in the bin:
82			weight.append( (fileList[0][1]/fileList[0][3]) * (CountIncl/2950.0) * fileList[0][2][bin]/total[bin] )
83			else:
84			weight.append(1.)
85	bortigno	1.1
86	nmohr	1.4	weight_map = {}
87			for bin in range(0, len(lheBin) ):
88			print '%s: %.2f' %(lheBin[bin], weight[bin])
89			weight_map[lheBin[bin]] = weight[bin]
90			return weight_map
91	bortigno	1.1
92
93	nmohr	1.5	def apply_weights(fileList,weight_map,inclusive,newpostfix):
94	nmohr	1.4	#now add the branch with the weight normalized to the inclusive
95			for file in fileList:
96	nmohr	1.5	outfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),'RECREATE')
97			infile = ROOT.TFile.Open(file[0],"READ")
98	nmohr	1.7	histoInfile = ROOT.TFile.Open(inclusive+'.root',"READ")
99	nmohr	1.4	histoInfile.cd()
100			obj = ROOT.TObject
101			for key in ROOT.gDirectory.GetListOfKeys():
102			histoInfile.cd()
103			obj = key.ReadObj()
104			#print obj.GetName()
105			if obj.GetName() == 'tree':
106			continue
107			outfile.cd()
108			#print key.GetName()
109			obj.Write(key.GetName())
110	bortigno	1.1
111	nmohr	1.4	infile.cd()
112			tree = infile.Get('tree')
113			outfile.cd()
114			newtree = tree.CloneTree(0)
115			lheWeight = array('f',[0])
116			newtree.Branch('lheWeight',lheWeight,'lheWeight/F')
117
118			nEntries = tree.GetEntries()
119			theBinForms = {}
120			for bin in weight_map:
121			theBinForms[bin] = ROOT.TTreeFormula("Bin_formula_%s"%(bin),bin,tree)
122			for entry in range(0,nEntries):
123			tree.GetEntry(entry)
124			match_bin = None
125			for bin in weight_map:
126			if theBinForms[bin].EvalInstance() > 0.:
127			match_bin = bin
128			if match_bin:
129			lheWeight[0] = weight_map[match_bin]
130			else:
131			lheWeight[0] = 1.
132			newtree.Fill()
133	bortigno	1.1
134	nmohr	1.4	newtree.AutoSave()
135			outfile.Write()
136			outfile.Close()
137			infile.Close()
138			del tree
139			del newtree
140
141	nmohr	1.5	def do_validation(fileList,inclusive,newpostfix):
142	nmohr	1.4	histo = ROOT.TH1F("lheV_pt","lheV_pt",300,0,300)
143			histo1 = ROOT.TH1F("lheV_ptInc","lheV_ptInc",300,0,300)
144	nmohr	1.6	histo.SetDirectory(0)
145			histo1.SetDirectory(0)
146	nmohr	1.4
147			for file in fileList:
148	nmohr	1.6	h_name=str(hashlib.sha224(file[0]).hexdigest())
149			print file[0]
150	nmohr	1.4	print h_name
151	nmohr	1.5	tfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),"read")
152	nmohr	1.4	tree = tfile.Get("tree")
153			h_tmp = ROOT.TH1F(h_name,h_name,300,0,300)
154	nmohr	1.6	#ROOT.gDirectory.ls()
155	nmohr	1.4	tree.Draw("lheV_pt>>"+str(h_name),"(lheWeight)*(lheV_pt < 300.)","goff")
156			if inclusive in file[0]:
157			h_tmp1 = ROOT.TH1F(h_name+'Inc',h_name+'Inc',300,0,300)
158			tree.Draw("lheV_pt>>"+str(h_name+'Inc'),"(lheV_pt < 300.)","goff")
159			histo1.Add(h_tmp1)
160	nmohr	1.5	histo.Add(h_tmp)
161			tfile.Close()
162			del tree
163	nmohr	1.4
164
165			canvas = ROOT.TCanvas("lheV_pt","lheV_pt",600,600)
166			ROOT.gPad.SetLogy()
167			histo.Draw()
168			histo1.Draw('SAME')
169	nmohr	1.6	print histo.Integral(100,300)
170			print histo1.Integral(100,300)
171	nmohr	1.4	histo1.SetLineColor(ROOT.kRed)
172			canvas.Print("validation_lheV_pt.pdf","pdf")
173
174	nmohr	1.5	if opts.weights:
175			weight_map = get_weights(fileList,lheBin)
176			config.set('LHEWeights', 'weights_per_bin', '%s' %weight_map)
177			f = open('8TeVconfig/lhe_weights', 'w')
178	nmohr	1.7	for section in config.sections():
179			if not section == 'LHEWeights':
180			config.remove_section(section)
181	nmohr	1.5	config.write(f)
182			f.close()
183	nmohr	1.6	elif opts.apply:
184	nmohr	1.7	weight_map = eval(config.get('LHEWeights', 'weights_per_bin'))
185	nmohr	1.5	if opts.apply:
186			apply_weights(fileList,weight_map,prefix+inclusive,newpostfix)
187			if opts.validate:
188			do_validation(fileList,inclusive,newpostfix)