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="",
                      help="configuration file")
(opts, args) = parser.parse_args(argv)
if opts.config =="":
        opts.config = "config"
config = BetterConfigParser()
config.read(opts.config)
anaTag = config.get("Analysis","tag")

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

MVAdir=config.get('Directories','MVAdir')

#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(Apath+'/samples.info','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(',')

#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(Wdir+'/weights/'+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,MVAinfos[i].getweightfile())
#--> Now the MVA is booked

#Apply samples
infofile = open(Apath+'/samples.info','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(job.getpath(),'read')
            outfile = TFile.Open(job.path+'/'+MVAdir+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(Apath+MVAdir+'/samples.info','w')
    pickle.dump(Ainfo,infofile)
    infofile.close()


Revision:	1.10
Committed:	Thu Aug 9 13:52:04 2012 UTC (12 years, 9 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
Changes since 1.9:	+16 -9 lines
Log Message:	Fixes for 7 and 8 TeV
#	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			parser.add_option("-C", "--config", dest="config", default="",
35			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	nmohr	1.9	config.read(opts.config)
41			anaTag = config.get("Analysis","tag")
42	peller	1.1
43			#get locations:
44			Wdir=config.get('Directories','Wdir')
45
46	peller	1.4	MVAdir=config.get('Directories','MVAdir')
47	peller	1.1
48			#systematics
49			systematics=config.get('systematics','systematics')
50			systematics=systematics.split(' ')
51
52			#TreeVar Array
53	nmohr	1.7	#MVA_Vars={}
54			#for systematic in systematics:
55			# MVA_Vars[systematic]=config.get('treeVars',systematic)
56			# MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
57	peller	1.1
58			######################
59			#Evaluate multi: Must Have same treeVars!!!
60
61	nmohr	1.10	Apath=opts.path
62	nmohr	1.9	infofile = open(Apath+'/samples.info','r')
63			info = pickle.load(infofile)
64			infofile.close()
65	nmohr	1.10	arglist=opts.discr #RTight_blavla,bsbsb
66	peller	1.1
67	nmohr	1.10	namelistIN=opts.names
68	peller	1.4	namelist=namelistIN.split(',')
69
70	nmohr	1.10	doinfo=bool(int(opts.update))
71	peller	1.1
72			MVAlist=arglist.split(',')
73
74			#CONFIG
75			#factory
76			factoryname=config.get('factory','factoryname')
77			#MVA
78			#MVAnames=[]
79			#for MVA in MVAlist:
80			# print MVA
81			# MVAnames.append(config.get(MVA,'MVAname'))
82			#print Wdir+'/weights/'+factoryname+'_'+MVAname+'.info'
83			#MVAinfofiles=[]
84			MVAinfos=[]
85			for MVAname in MVAlist:
86			MVAinfofile = open(Wdir+'/weights/'+factoryname+'_'+MVAname+'.info','r')
87			MVAinfos.append(pickle.load(MVAinfofile))
88			MVAinfofile.close()
89
90			treeVarSet=MVAinfos[0].varset
91			#variables
92			#TreeVar Array
93			MVA_Vars={}
94			for systematic in systematics:
95			MVA_Vars[systematic]=config.get(treeVarSet,systematic)
96			MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
97			#Spectators:
98			#spectators=config.get(treeVarSet,'spectators')
99			#spectators=spectators.split(' ')
100			#progbar quatsch
101			longe=40
102			#Workdir
103			workdir=ROOT.gDirectory.GetPath()
104			#os.mkdir(Apath+'/MVAout')
105
106			#Book TMVA readers: MVAlist=["MMCC_bla","CC5050_bla"]
107			readers=[]
108			for MVA in MVAlist:
109			readers.append(ROOT.TMVA.Reader("!Color:!Silent"))
110
111			#define variables and specatators
112			MVA_var_buffer = []
113	peller	1.5	MVA_var_buffer4 = []
114	peller	1.1	for i in range(len( MVA_Vars['Nominal'])):
115			MVA_var_buffer.append(array( 'f', [ 0 ] ))
116			for reader in readers:
117			reader.AddVariable( MVA_Vars['Nominal'][i],MVA_var_buffer[i])
118			#MVA_spectator_buffer = []
119			#for i in range(len(spectators)):
120			# MVA_spectator_buffer.append(array( 'f', [ 0 ] ))
121			# for reader in readers:
122			# reader.AddSpectator(spectators[i],MVA_spectator_buffer[i])
123			#Load raeder
124			for i in range(0,len(readers)):
125			readers[i].BookMVA(MVAinfos[i].MVAname,MVAinfos[i].getweightfile())
126			#--> Now the MVA is booked
127
128			#Apply samples
129			infofile = open(Apath+'/samples.info','r')
130			Ainfo = pickle.load(infofile)
131			infofile.close()
132
133			#eval
134	peller	1.4	for job in Ainfo:
135	peller	1.8	if eval(job.active):
136			if job.name in namelist:
137			#get trees:
138			input = TFile.Open(job.getpath(),'read')
139			outfile = TFile.Open(job.path+'/'+MVAdir+job.prefix+job.identifier+'.root','recreate')
140	peller	1.4	input.cd()
141	peller	1.8	obj = ROOT.TObject
142			for key in ROOT.gDirectory.GetListOfKeys():
143			input.cd()
144			obj = key.ReadObj()
145			#print obj.GetName()
146			if obj.GetName() == job.tree:
147			continue
148			outfile.cd()
149			#print key.GetName()
150			obj.Write(key.GetName())
151			tree = input.Get(job.tree)
152			nEntries = tree.GetEntries()
153	peller	1.4	outfile.cd()
154	peller	1.8	newtree = tree.CloneTree(0)
155	peller	1.4
156	peller	1.8	#MCs:
157			if job.type != 'DATA':
158			MVA_formulas={}
159			MVA_formulas4={}
160			for systematic in systematics:
161			#print '\t\t - ' + systematic
162			MVA_formulas[systematic]=[]
163			MVA_formulas4[systematic]=[]
164			#create TTreeFormulas
165			for j in range(len( MVA_Vars['Nominal'])):
166			MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),MVA_Vars[systematic][j],tree))
167			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
168			outfile.cd()
169			#Setup Branches
170			MVAbranches=[]
171			MVAbranches4=[]
172			for i in range(0,len(readers)):
173			MVAbranches.append(array('f',[0]*9))
174			MVAbranches4.append(array('f',[0]*9))
175			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')
176			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')
177			print '\n--> ' + job.name +':'
178			#progbar setup
179			if nEntries >= longe:
180			step=int(nEntries/longe)
181			long=longe
182			else:
183			long=nEntries
184			step = 1
185			bar=progbar(long)
186			#Fill event by event:
187			for entry in range(0,nEntries):
188			if entry % step == 0:
189			bar.move()
190			#load entry
191			tree.GetEntry(entry)
192			for systematic in systematics:
193			for j in range(len( MVA_Vars['Nominal'])):
194			MVA_var_buffer[j][0] = MVA_formulas[systematic][j].EvalInstance()
195
196			for j in range(0,len(readers)):
197			MVAbranches[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
198
199			for j in range(len( MVA_Vars['Nominal'])):
200			MVA_var_buffer[j][0] = MVA_formulas4[systematic][j].EvalInstance()
201
202			for j in range(0,len(readers)):
203			MVAbranches4[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
204			#Fill:
205			newtree.Fill()
206			newtree.AutoSave()
207			outfile.Close()
208
209			#DATA:
210			if job.type == 'DATA':
211			#MVA Formulas
212			MVA_formulas_Nominal = []
213	peller	1.4	#create TTreeFormulas
214			for j in range(len( MVA_Vars['Nominal'])):
215	peller	1.8	MVA_formulas_Nominal.append(ROOT.TTreeFormula("MVA_formula%s_Nominal"%j, MVA_Vars['Nominal'][j],tree))
216			outfile.cd()
217			MVAbranches=[]
218			for i in range(0,len(readers)):
219			MVAbranches.append(array('f',[0]))
220			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F')
221			newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches[i],'nominal/F')
222			#progbar
223			print '\n--> ' + job.name +':'
224			if nEntries >= longe:
225			step=int(nEntries/longe)
226			long=longe
227			else:
228			long=nEntries
229			step = 1
230			bar=progbar(long)
231			#Fill event by event:
232			for entry in range(0,nEntries):
233			if entry % step == 0:
234			bar.move()
235			#load entry
236			tree.GetEntry(entry)
237			#nominal:
238	peller	1.4	for j in range(len( MVA_Vars['Nominal'])):
239	peller	1.8	MVA_var_buffer[j][0] = MVA_formulas_Nominal[j].EvalInstance()
240
241	peller	1.4	for j in range(0,len(readers)):
242	peller	1.8	MVAbranches[j][0]= readers[j].EvaluateMVA(MVAinfos[j].MVAname)
243			newtree.Fill()
244			newtree.AutoSave()
245			outfile.Close()
246	peller	1.1
247	peller	1.4	print '\n'
248	peller	1.1
249			#Update Info:
250			if doinfo:
251			for job in Ainfo:
252			for MVAinfo in MVAinfos:
253			job.addcomment('Added MVA %s'%MVAinfo.MVAname)
254	peller	1.4	job.addpath(MVAdir)
255			infofile = open(Apath+MVAdir+'/samples.info','w')
256	peller	1.1	pickle.dump(Ainfo,infofile)
257			infofile.close()
258
259