VHbb/python/evaluateMVA.py

#!/usr/bin/env python
import sys
import os
import ROOT 
from ROOT import TFile
from array import array
from math import sqrt
from copy import copy
#suppres the EvalInstace conversion warning bug
import warnings
warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
from optparse import OptionParser
from BetterConfigParser import BetterConfigParser
from samplesclass import sample
from mvainfos import mvainfo
import pickle
from progbar import progbar
from printcolor import printc

#CONFIGURE

#load config
#os.mkdir(path+'/sys')
argv = sys.argv
parser = OptionParser()
parser.add_option("-U", "--update", dest="update", default=0,
                      help="update infofile")
parser.add_option("-D", "--discr", dest="discr", default="",
                      help="discriminators to be added")
parser.add_option("-P", "--path", dest="path", default="",
                      help="path to samples")
parser.add_option("-S", "--samples", dest="names", default="",
                      help="samples you want to run on")
parser.add_option("-C", "--config", dest="config", default=[], action="append",
                      help="configuration file")
(opts, args) = parser.parse_args(argv)
if opts.config =="":
        opts.config = "config"
config = BetterConfigParser()
#config.read('./config7TeV_ZZ')
config.read(opts.config)
anaTag = config.get("Analysis","tag")

#get locations:
Wdir=config.get('Directories','Wdir')
MVASubdir=config.get('Directories','MVAdir')
samplesinfo=config.get('Directories','samplesinfo')

#systematics
systematics=config.get('systematics','systematics')
systematics=systematics.split(' ')

#TreeVar Array
#MVA_Vars={}
#for systematic in systematics:
#    MVA_Vars[systematic]=config.get('treeVars',systematic)
#    MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')

######################
#Evaluate multi: Must Have same treeVars!!!

Apath=opts.path
infofile = open(samplesinfo,'r')
info = pickle.load(infofile)
infofile.close()
arglist=opts.discr #RTight_blavla,bsbsb

namelistIN=opts.names
namelist=namelistIN.split(',')

doinfo=bool(int(opts.update))

MVAlist=arglist.split(',')
MVAdir=config.get('Directories','vhbbpath')

#CONFIG
#factory
factoryname=config.get('factory','factoryname')
#MVA
#MVAnames=[]
#for MVA in MVAlist:
#    print MVA
#    MVAnames.append(config.get(MVA,'MVAname'))
#print Wdir+'/weights/'+factoryname+'_'+MVAname+'.info'
#MVAinfofiles=[]
MVAinfos=[]
for MVAname in MVAlist:
    MVAinfofile = open(MVAdir+'/data/'+factoryname+'_'+MVAname+'.info','r')
    MVAinfos.append(pickle.load(MVAinfofile))
    MVAinfofile.close()
    
treeVarSet=MVAinfos[0].varset
#variables
#TreeVar Array
MVA_Vars={}
for systematic in systematics:
    MVA_Vars[systematic]=config.get(treeVarSet,systematic)
    MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
#Spectators:
#spectators=config.get(treeVarSet,'spectators')
#spectators=spectators.split(' ')
#progbar quatsch
longe=40
#Workdir
workdir=ROOT.gDirectory.GetPath()
#os.mkdir(Apath+'/MVAout')

#Book TMVA readers: MVAlist=["MMCC_bla","CC5050_bla"]
readers=[]
for MVA in MVAlist:
    readers.append(ROOT.TMVA.Reader("!Color:!Silent"))

#define variables and specatators
MVA_var_buffer = []
MVA_var_buffer4 = []
for i in range(len( MVA_Vars['Nominal'])):
    MVA_var_buffer.append(array( 'f', [ 0 ] ))
    for reader in readers:
        reader.AddVariable( MVA_Vars['Nominal'][i],MVA_var_buffer[i])
#MVA_spectator_buffer = []
#for i in range(len(spectators)):
#    MVA_spectator_buffer.append(array( 'f', [ 0 ] ))
#    for reader in readers:
#        reader.AddSpectator(spectators[i],MVA_spectator_buffer[i])
#Load raeder
for i in range(0,len(readers)):
    readers[i].BookMVA(MVAinfos[i].MVAname,MVAdir+'/data/'+MVAinfos[i].getweightfile())
#--> Now the MVA is booked

#Apply samples
infofile = open(samplesinfo,'r')
Ainfo = pickle.load(infofile)
infofile.close()

#eval
for job in Ainfo:
    if eval(job.active):
        if job.name in namelist:
            #get trees:
            input = TFile.Open(Apath+'/'+job.getpath(),'read')
            outfile = TFile.Open(Apath+'/'+MVASubdir+job.prefix+job.identifier+'.root','recreate')
            input.cd()
            obj = ROOT.TObject
            for key in ROOT.gDirectory.GetListOfKeys():
                input.cd()
                obj = key.ReadObj()
                #print obj.GetName()
                if obj.GetName() == job.tree:
                    continue
                outfile.cd()
                #print key.GetName()
                obj.Write(key.GetName())
            tree = input.Get(job.tree)
            nEntries = tree.GetEntries()
            outfile.cd()
            newtree = tree.CloneTree(0)

            #MCs:
            if job.type != 'DATA':
                MVA_formulas={}
                MVA_formulas4={}
                for systematic in systematics: 
                    #print '\t\t - ' + systematic
                    MVA_formulas[systematic]=[]
                    MVA_formulas4[systematic]=[]
                    #create TTreeFormulas
                    for j in range(len( MVA_Vars['Nominal'])):
                        MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),MVA_Vars[systematic][j],tree))
                        MVA_formulas4[systematic].append(ROOT.TTreeFormula("MVA_formula4%s_%s"%(j,systematic),MVA_Vars['Nominal'][j]+'+('+MVA_Vars[systematic][j]+'-'+MVA_Vars['Nominal'][j]+')*4',tree))#HERE change
                outfile.cd()
                #Setup Branches
                MVAbranches=[]
                MVAbranches4=[]
                for i in range(0,len(readers)):
                    MVAbranches.append(array('f',[0]*9))
                    MVAbranches4.append(array('f',[0]*9))
                    newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down/F')
                    newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches4[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down/F')
                print '\n--> ' + job.name +':'
                #progbar setup
                if nEntries >= longe:
                    step=int(nEntries/longe)
                    long=longe
                else:
                    long=nEntries
                    step = 1
                bar=progbar(long)
                #Fill event by event:
                for entry in range(0,nEntries):
                    if entry % step == 0:
                        bar.move()
                    #load entry
                    tree.GetEntry(entry)
                    for systematic in systematics:
                        for j in range(len( MVA_Vars['Nominal'])):
                            MVA_var_buffer[j][0] = MVA_formulas[systematic][j].EvalInstance()
                            
                        for j in range(0,len(readers)):
                            MVAbranches[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
                            
                        for j in range(len( MVA_Vars['Nominal'])):
                            MVA_var_buffer[j][0] = MVA_formulas4[systematic][j].EvalInstance()
                            
                        for j in range(0,len(readers)):
                            MVAbranches4[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
                    #Fill:
                    newtree.Fill()
                newtree.AutoSave()
                outfile.Close()
                
            #DATA:
            if job.type == 'DATA':
                #MVA Formulas
                MVA_formulas_Nominal = []
                #create TTreeFormulas
                for j in range(len( MVA_Vars['Nominal'])):
                    MVA_formulas_Nominal.append(ROOT.TTreeFormula("MVA_formula%s_Nominal"%j, MVA_Vars['Nominal'][j],tree))
                outfile.cd()
                MVAbranches=[]
                for i in range(0,len(readers)):
                    MVAbranches.append(array('f',[0]))
                    newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F') 
                    newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches[i],'nominal/F') 
                #progbar           
                print '\n--> ' + job.name +':'
                if nEntries >= longe:
                    step=int(nEntries/longe)
                    long=longe
                else:
                    long=nEntries
                    step = 1
                bar=progbar(long)
                #Fill event by event:
                for entry in range(0,nEntries):
                    if entry % step == 0:
                        bar.move()
                    #load entry
                    tree.GetEntry(entry)
                    #nominal:
                    for j in range(len( MVA_Vars['Nominal'])):
                            MVA_var_buffer[j][0] = MVA_formulas_Nominal[j].EvalInstance()
                            
                    for j in range(0,len(readers)):
                        MVAbranches[j][0]= readers[j].EvaluateMVA(MVAinfos[j].MVAname)
                    newtree.Fill()
                newtree.AutoSave()
                outfile.Close()

print '\n'

#Update Info:
if doinfo:
    for job in Ainfo:        
        for MVAinfo in MVAinfos:
            job.addcomment('Added MVA %s'%MVAinfo.MVAname)
        job.addpath(MVAdir)
    infofile = open(samplesinfo,'w')
    pickle.dump(Ainfo,infofile)
    infofile.close()


Revision:	1.15
Committed:	Thu Oct 11 16:53:25 2012 UTC (12 years, 7 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	hcpPreApp
Changes since 1.14:	+4 -3 lines
Log Message:	Only one place for samples info
#	User	Rev	Content
1	peller	1.1	#!/usr/bin/env python
2			import sys
3			import os
4			import ROOT
5			from ROOT import TFile
6			from array import array
7			from math import sqrt
8			from copy import copy
9			#suppres the EvalInstace conversion warning bug
10			import warnings
11			warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
12	nmohr	1.9	from optparse import OptionParser
13	nmohr	1.6	from BetterConfigParser import BetterConfigParser
14	peller	1.5	from samplesclass import sample
15	peller	1.1	from mvainfos import mvainfo
16			import pickle
17			from progbar import progbar
18			from printcolor import printc
19
20			#CONFIGURE
21
22			#load config
23	nmohr	1.9	#os.mkdir(path+'/sys')
24	nmohr	1.10	argv = sys.argv
25	nmohr	1.9	parser = OptionParser()
26	nmohr	1.10	parser.add_option("-U", "--update", dest="update", default=0,
27			help="update infofile")
28			parser.add_option("-D", "--discr", dest="discr", default="",
29			help="discriminators to be added")
30			parser.add_option("-P", "--path", dest="path", default="",
31			help="path to samples")
32			parser.add_option("-S", "--samples", dest="names", default="",
33			help="samples you want to run on")
34	nmohr	1.11	parser.add_option("-C", "--config", dest="config", default=[], action="append",
35	nmohr	1.10	help="configuration file")
36	nmohr	1.9	(opts, args) = parser.parse_args(argv)
37	nmohr	1.10	if opts.config =="":
38	nmohr	1.9	opts.config = "config"
39	nmohr	1.6	config = BetterConfigParser()
40	peller	1.14	#config.read('./config7TeV_ZZ')
41	nmohr	1.9	config.read(opts.config)
42			anaTag = config.get("Analysis","tag")
43	peller	1.1
44			#get locations:
45			Wdir=config.get('Directories','Wdir')
46	nmohr	1.13	MVASubdir=config.get('Directories','MVAdir')
47	nmohr	1.15	samplesinfo=config.get('Directories','samplesinfo')
48	peller	1.1
49			#systematics
50			systematics=config.get('systematics','systematics')
51			systematics=systematics.split(' ')
52
53			#TreeVar Array
54	nmohr	1.7	#MVA_Vars={}
55			#for systematic in systematics:
56			# MVA_Vars[systematic]=config.get('treeVars',systematic)
57			# MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
58	peller	1.1
59			######################
60			#Evaluate multi: Must Have same treeVars!!!
61
62	nmohr	1.10	Apath=opts.path
63	nmohr	1.15	infofile = open(samplesinfo,'r')
64	nmohr	1.9	info = pickle.load(infofile)
65			infofile.close()
66	nmohr	1.10	arglist=opts.discr #RTight_blavla,bsbsb
67	peller	1.1
68	nmohr	1.10	namelistIN=opts.names
69	peller	1.4	namelist=namelistIN.split(',')
70
71	nmohr	1.10	doinfo=bool(int(opts.update))
72	peller	1.1
73			MVAlist=arglist.split(',')
74	nmohr	1.12	MVAdir=config.get('Directories','vhbbpath')
75	peller	1.1
76			#CONFIG
77			#factory
78			factoryname=config.get('factory','factoryname')
79			#MVA
80			#MVAnames=[]
81			#for MVA in MVAlist:
82			# print MVA
83			# MVAnames.append(config.get(MVA,'MVAname'))
84			#print Wdir+'/weights/'+factoryname+'_'+MVAname+'.info'
85			#MVAinfofiles=[]
86			MVAinfos=[]
87			for MVAname in MVAlist:
88	nmohr	1.12	MVAinfofile = open(MVAdir+'/data/'+factoryname+'_'+MVAname+'.info','r')
89	peller	1.1	MVAinfos.append(pickle.load(MVAinfofile))
90			MVAinfofile.close()
91
92			treeVarSet=MVAinfos[0].varset
93			#variables
94			#TreeVar Array
95			MVA_Vars={}
96			for systematic in systematics:
97			MVA_Vars[systematic]=config.get(treeVarSet,systematic)
98			MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
99			#Spectators:
100			#spectators=config.get(treeVarSet,'spectators')
101			#spectators=spectators.split(' ')
102			#progbar quatsch
103			longe=40
104			#Workdir
105			workdir=ROOT.gDirectory.GetPath()
106			#os.mkdir(Apath+'/MVAout')
107
108			#Book TMVA readers: MVAlist=["MMCC_bla","CC5050_bla"]
109			readers=[]
110			for MVA in MVAlist:
111			readers.append(ROOT.TMVA.Reader("!Color:!Silent"))
112
113			#define variables and specatators
114			MVA_var_buffer = []
115	peller	1.5	MVA_var_buffer4 = []
116	peller	1.1	for i in range(len( MVA_Vars['Nominal'])):
117			MVA_var_buffer.append(array( 'f', [ 0 ] ))
118			for reader in readers:
119			reader.AddVariable( MVA_Vars['Nominal'][i],MVA_var_buffer[i])
120			#MVA_spectator_buffer = []
121			#for i in range(len(spectators)):
122			# MVA_spectator_buffer.append(array( 'f', [ 0 ] ))
123			# for reader in readers:
124			# reader.AddSpectator(spectators[i],MVA_spectator_buffer[i])
125			#Load raeder
126			for i in range(0,len(readers)):
127	nmohr	1.12	readers[i].BookMVA(MVAinfos[i].MVAname,MVAdir+'/data/'+MVAinfos[i].getweightfile())
128	peller	1.1	#--> Now the MVA is booked
129
130			#Apply samples
131	nmohr	1.15	infofile = open(samplesinfo,'r')
132	peller	1.1	Ainfo = pickle.load(infofile)
133			infofile.close()
134
135			#eval
136	peller	1.4	for job in Ainfo:
137	peller	1.8	if eval(job.active):
138			if job.name in namelist:
139			#get trees:
140	nmohr	1.11	input = TFile.Open(Apath+'/'+job.getpath(),'read')
141	nmohr	1.13	outfile = TFile.Open(Apath+'/'+MVASubdir+job.prefix+job.identifier+'.root','recreate')
142	peller	1.4	input.cd()
143	peller	1.8	obj = ROOT.TObject
144			for key in ROOT.gDirectory.GetListOfKeys():
145			input.cd()
146			obj = key.ReadObj()
147			#print obj.GetName()
148			if obj.GetName() == job.tree:
149			continue
150			outfile.cd()
151			#print key.GetName()
152			obj.Write(key.GetName())
153			tree = input.Get(job.tree)
154			nEntries = tree.GetEntries()
155	peller	1.4	outfile.cd()
156	peller	1.8	newtree = tree.CloneTree(0)
157	peller	1.4
158	peller	1.8	#MCs:
159			if job.type != 'DATA':
160			MVA_formulas={}
161			MVA_formulas4={}
162			for systematic in systematics:
163			#print '\t\t - ' + systematic
164			MVA_formulas[systematic]=[]
165			MVA_formulas4[systematic]=[]
166			#create TTreeFormulas
167			for j in range(len( MVA_Vars['Nominal'])):
168			MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),MVA_Vars[systematic][j],tree))
169			MVA_formulas4[systematic].append(ROOT.TTreeFormula("MVA_formula4%s_%s"%(j,systematic),MVA_Vars['Nominal'][j]+'+('+MVA_Vars[systematic][j]+'-'+MVA_Vars['Nominal'][j]+')*4',tree))#HERE change
170			outfile.cd()
171			#Setup Branches
172			MVAbranches=[]
173			MVAbranches4=[]
174			for i in range(0,len(readers)):
175			MVAbranches.append(array('f',[0]*9))
176			MVAbranches4.append(array('f',[0]*9))
177			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down/F')
178			newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches4[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down/F')
179			print '\n--> ' + job.name +':'
180			#progbar setup
181			if nEntries >= longe:
182			step=int(nEntries/longe)
183			long=longe
184			else:
185			long=nEntries
186			step = 1
187			bar=progbar(long)
188			#Fill event by event:
189			for entry in range(0,nEntries):
190			if entry % step == 0:
191			bar.move()
192			#load entry
193			tree.GetEntry(entry)
194			for systematic in systematics:
195			for j in range(len( MVA_Vars['Nominal'])):
196			MVA_var_buffer[j][0] = MVA_formulas[systematic][j].EvalInstance()
197
198			for j in range(0,len(readers)):
199			MVAbranches[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
200
201			for j in range(len( MVA_Vars['Nominal'])):
202			MVA_var_buffer[j][0] = MVA_formulas4[systematic][j].EvalInstance()
203
204			for j in range(0,len(readers)):
205			MVAbranches4[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
206			#Fill:
207			newtree.Fill()
208			newtree.AutoSave()
209			outfile.Close()
210
211			#DATA:
212			if job.type == 'DATA':
213			#MVA Formulas
214			MVA_formulas_Nominal = []
215	peller	1.4	#create TTreeFormulas
216			for j in range(len( MVA_Vars['Nominal'])):
217	peller	1.8	MVA_formulas_Nominal.append(ROOT.TTreeFormula("MVA_formula%s_Nominal"%j, MVA_Vars['Nominal'][j],tree))
218			outfile.cd()
219			MVAbranches=[]
220			for i in range(0,len(readers)):
221			MVAbranches.append(array('f',[0]))
222			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F')
223			newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches[i],'nominal/F')
224			#progbar
225			print '\n--> ' + job.name +':'
226			if nEntries >= longe:
227			step=int(nEntries/longe)
228			long=longe
229			else:
230			long=nEntries
231			step = 1
232			bar=progbar(long)
233			#Fill event by event:
234			for entry in range(0,nEntries):
235			if entry % step == 0:
236			bar.move()
237			#load entry
238			tree.GetEntry(entry)
239			#nominal:
240	peller	1.4	for j in range(len( MVA_Vars['Nominal'])):
241	peller	1.8	MVA_var_buffer[j][0] = MVA_formulas_Nominal[j].EvalInstance()
242
243	peller	1.4	for j in range(0,len(readers)):
244	peller	1.8	MVAbranches[j][0]= readers[j].EvaluateMVA(MVAinfos[j].MVAname)
245			newtree.Fill()
246			newtree.AutoSave()
247			outfile.Close()
248	peller	1.1
249	peller	1.4	print '\n'
250	peller	1.1
251			#Update Info:
252			if doinfo:
253			for job in Ainfo:
254			for MVAinfo in MVAinfos:
255			job.addcomment('Added MVA %s'%MVAinfo.MVAname)
256	peller	1.4	job.addpath(MVAdir)
257	nmohr	1.15	infofile = open(samplesinfo,'w')
258	peller	1.1	pickle.dump(Ainfo,infofile)
259			infofile.close()
260
261