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]
    xSecIncl = fileList[0][1]
    print fileList
    print 'Inclusive cross-section %.2f pb' %xSecIncl

    #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/xSecIncl) * 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.8
Committed:	Thu Mar 14 17:09:31 2013 UTC (12 years, 2 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	lhcp_UnblindFix, hcp_Unblind, lhcp_11April, LHCP_PreAppFixAfterFreeze, LHCP_PreAppFreeze, HEAD
Changes since 1.7:	+3 -1 lines
Log Message:	Remove hardcoded cross-section
#	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	nmohr	1.8	xSecIncl = fileList[0][1]
66	nmohr	1.4	print fileList
67	nmohr	1.8	print 'Inclusive cross-section %.2f pb' %xSecIncl
68	bortigno	1.1
69	nmohr	1.4	#total -> total numer of events in each lheBin
70			print 'Calculating total'
71			total = []
72			weight= []
73			for bin in range(0, len(lheBin) ):
74			#weight.append(1.)
75			#continue
76			total.append(getTotal(bin, fileList))
77			#print lheBin[bin]
78	bortigno	1.1	#to better stich we need the highest stat (that should correspond to the binned sample relative to the bin)
79	nmohr	1.4	fileList.sort( key=lambda file: file[2][bin], reverse=True )
80			#print 'After sorting'
81			#print fileList
82			if total[bin] > 0.:
83			#the first is always the one with the highest N in the bin:
84	nmohr	1.8	weight.append( (fileList[0][1]/fileList[0][3]) * (CountIncl/xSecIncl) * fileList[0][2][bin]/total[bin] )
85	nmohr	1.4	else:
86			weight.append(1.)
87	bortigno	1.1
88	nmohr	1.4	weight_map = {}
89			for bin in range(0, len(lheBin) ):
90			print '%s: %.2f' %(lheBin[bin], weight[bin])
91			weight_map[lheBin[bin]] = weight[bin]
92			return weight_map
93	bortigno	1.1
94
95	nmohr	1.5	def apply_weights(fileList,weight_map,inclusive,newpostfix):
96	nmohr	1.4	#now add the branch with the weight normalized to the inclusive
97			for file in fileList:
98	nmohr	1.5	outfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),'RECREATE')
99			infile = ROOT.TFile.Open(file[0],"READ")
100	nmohr	1.7	histoInfile = ROOT.TFile.Open(inclusive+'.root',"READ")
101	nmohr	1.4	histoInfile.cd()
102			obj = ROOT.TObject
103			for key in ROOT.gDirectory.GetListOfKeys():
104			histoInfile.cd()
105			obj = key.ReadObj()
106			#print obj.GetName()
107			if obj.GetName() == 'tree':
108			continue
109			outfile.cd()
110			#print key.GetName()
111			obj.Write(key.GetName())
112	bortigno	1.1
113	nmohr	1.4	infile.cd()
114			tree = infile.Get('tree')
115			outfile.cd()
116			newtree = tree.CloneTree(0)
117			lheWeight = array('f',[0])
118			newtree.Branch('lheWeight',lheWeight,'lheWeight/F')
119
120			nEntries = tree.GetEntries()
121			theBinForms = {}
122			for bin in weight_map:
123			theBinForms[bin] = ROOT.TTreeFormula("Bin_formula_%s"%(bin),bin,tree)
124			for entry in range(0,nEntries):
125			tree.GetEntry(entry)
126			match_bin = None
127			for bin in weight_map:
128			if theBinForms[bin].EvalInstance() > 0.:
129			match_bin = bin
130			if match_bin:
131			lheWeight[0] = weight_map[match_bin]
132			else:
133			lheWeight[0] = 1.
134			newtree.Fill()
135	bortigno	1.1
136	nmohr	1.4	newtree.AutoSave()
137			outfile.Write()
138			outfile.Close()
139			infile.Close()
140			del tree
141			del newtree
142
143	nmohr	1.5	def do_validation(fileList,inclusive,newpostfix):
144	nmohr	1.4	histo = ROOT.TH1F("lheV_pt","lheV_pt",300,0,300)
145			histo1 = ROOT.TH1F("lheV_ptInc","lheV_ptInc",300,0,300)
146	nmohr	1.6	histo.SetDirectory(0)
147			histo1.SetDirectory(0)
148	nmohr	1.4
149			for file in fileList:
150	nmohr	1.6	h_name=str(hashlib.sha224(file[0]).hexdigest())
151			print file[0]
152	nmohr	1.4	print h_name
153	nmohr	1.5	tfile = ROOT.TFile.Open(file[0].replace('.root',newpostfix),"read")
154	nmohr	1.4	tree = tfile.Get("tree")
155			h_tmp = ROOT.TH1F(h_name,h_name,300,0,300)
156	nmohr	1.6	#ROOT.gDirectory.ls()
157	nmohr	1.4	tree.Draw("lheV_pt>>"+str(h_name),"(lheWeight)*(lheV_pt < 300.)","goff")
158			if inclusive in file[0]:
159			h_tmp1 = ROOT.TH1F(h_name+'Inc',h_name+'Inc',300,0,300)
160			tree.Draw("lheV_pt>>"+str(h_name+'Inc'),"(lheV_pt < 300.)","goff")
161			histo1.Add(h_tmp1)
162	nmohr	1.5	histo.Add(h_tmp)
163			tfile.Close()
164			del tree
165	nmohr	1.4
166
167			canvas = ROOT.TCanvas("lheV_pt","lheV_pt",600,600)
168			ROOT.gPad.SetLogy()
169			histo.Draw()
170			histo1.Draw('SAME')
171	nmohr	1.6	print histo.Integral(100,300)
172			print histo1.Integral(100,300)
173	nmohr	1.4	histo1.SetLineColor(ROOT.kRed)
174			canvas.Print("validation_lheV_pt.pdf","pdf")
175
176	nmohr	1.5	if opts.weights:
177			weight_map = get_weights(fileList,lheBin)
178			config.set('LHEWeights', 'weights_per_bin', '%s' %weight_map)
179			f = open('8TeVconfig/lhe_weights', 'w')
180	nmohr	1.7	for section in config.sections():
181			if not section == 'LHEWeights':
182			config.remove_section(section)
183	nmohr	1.5	config.write(f)
184			f.close()
185	nmohr	1.6	elif opts.apply:
186	nmohr	1.7	weight_map = eval(config.get('LHEWeights', 'weights_per_bin'))
187	nmohr	1.5	if opts.apply:
188			apply_weights(fileList,weight_map,prefix+inclusive,newpostfix)
189			if opts.validate:
190			do_validation(fileList,inclusive,newpostfix)