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[5:]
parser = OptionParser()
parser.add_option("-C", "--config", dest="config", default=[], action="append",
                      help="configuration defining the plots to make")
(opts, args) = parser.parse_args(argv)
if opts.config ==[]:
        opts.config = "config"
print opts.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=sys.argv[1]
infofile = open(Apath+'/samples.info','r')
info = pickle.load(infofile)
infofile.close()
arglist=sys.argv[2] #RTight_blavla,bsbsb

namelistIN=sys.argv[3]
namelist=namelistIN.split(',')

doinfo=bool(int(sys.argv[4]))

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