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 ConfigParser import SafeConfigParser
from samplesinfo import sample
from mvainfos import mvainfo
import pickle
from progbar import progbar
from printcolor import printc

#CONFIGURE

#load config
config = SafeConfigParser()
config.read('./config')

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