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 BetterConfigParser import BetterConfigParser
from samplesclass import sample
from mvainfos import mvainfo
import pickle
from progbar import progbar
from printcolor import printc

#CONFIGURE

#load config
config = BetterConfigParser()
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 = []
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 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.7
Committed:	Tue Jun 19 22:55:47 2012 UTC (12 years, 10 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	ICHEP8TeV, ichep8TeV
Changes since 1.6:	+4 -4 lines
Log Message:	Not used
#	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.6	from BetterConfigParser import BetterConfigParser
13	peller	1.5	from samplesclass import sample
14	peller	1.1	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	nmohr	1.6	config = BetterConfigParser()
23	peller	1.1	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	nmohr	1.7	#MVA_Vars={}
36			#for systematic in systematics:
37			# MVA_Vars[systematic]=config.get('treeVars',systematic)
38			# MVA_Vars[systematic]=MVA_Vars[systematic].split(' ')
39	peller	1.1
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	peller	1.5	MVA_var_buffer4 = []
93	peller	1.1	for i in range(len( MVA_Vars['Nominal'])):
94			MVA_var_buffer.append(array( 'f', [ 0 ] ))
95			for reader in readers:
96			reader.AddVariable( MVA_Vars['Nominal'][i],MVA_var_buffer[i])
97			#MVA_spectator_buffer = []
98			#for i in range(len(spectators)):
99			# MVA_spectator_buffer.append(array( 'f', [ 0 ] ))
100			# for reader in readers:
101			# reader.AddSpectator(spectators[i],MVA_spectator_buffer[i])
102			#Load raeder
103			for i in range(0,len(readers)):
104			readers[i].BookMVA(MVAinfos[i].MVAname,MVAinfos[i].getweightfile())
105			#--> Now the MVA is booked
106
107			#Apply samples
108			infofile = open(Apath+'/samples.info','r')
109			Ainfo = pickle.load(infofile)
110			infofile.close()
111
112			#eval
113	peller	1.4	for job in Ainfo:
114
115			if job.name in namelist:
116			#get trees:
117			input = TFile.Open(job.getpath(),'read')
118			outfile = TFile.Open(job.path+'/'+MVAdir+job.prefix+job.identifier+'.root','recreate')
119	peller	1.1	input.cd()
120	peller	1.4	obj = ROOT.TObject
121			for key in ROOT.gDirectory.GetListOfKeys():
122			input.cd()
123			obj = key.ReadObj()
124			#print obj.GetName()
125			if obj.GetName() == job.tree:
126			continue
127			outfile.cd()
128			#print key.GetName()
129			obj.Write(key.GetName())
130			tree = input.Get(job.tree)
131			nEntries = tree.GetEntries()
132	peller	1.1	outfile.cd()
133	peller	1.4	newtree = tree.CloneTree(0)
134
135			#MCs:
136			if job.type != 'DATA':
137			MVA_formulas={}
138	peller	1.5	MVA_formulas4={}
139	peller	1.4	for systematic in systematics:
140			#print '\t\t - ' + systematic
141			MVA_formulas[systematic]=[]
142	peller	1.5	MVA_formulas4[systematic]=[]
143	peller	1.4	#create TTreeFormulas
144			for j in range(len( MVA_Vars['Nominal'])):
145			MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),MVA_Vars[systematic][j],tree))
146	peller	1.5	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
147	peller	1.4	outfile.cd()
148			#Setup Branches
149			MVAbranches=[]
150	peller	1.5	MVAbranches4=[]
151	peller	1.4	for i in range(0,len(readers)):
152			MVAbranches.append(array('f',[0]*9))
153	peller	1.5	MVAbranches4.append(array('f',[0]*9))
154	peller	1.4	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')
155	peller	1.5	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')
156	peller	1.4	print '\n--> ' + job.name +':'
157			#progbar setup
158			if nEntries >= longe:
159			step=int(nEntries/longe)
160			long=longe
161			else:
162			long=nEntries
163			step = 1
164			bar=progbar(long)
165			#Fill event by event:
166			for entry in range(0,nEntries):
167			if entry % step == 0:
168			bar.move()
169			#load entry
170			tree.GetEntry(entry)
171			for systematic in systematics:
172			for j in range(len( MVA_Vars['Nominal'])):
173			MVA_var_buffer[j][0] = MVA_formulas[systematic][j].EvalInstance()
174
175			for j in range(0,len(readers)):
176			MVAbranches[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
177	peller	1.5
178			for j in range(len( MVA_Vars['Nominal'])):
179			MVA_var_buffer[j][0] = MVA_formulas4[systematic][j].EvalInstance()
180
181			for j in range(0,len(readers)):
182			MVAbranches4[j][systematics.index(systematic)] = readers[j].EvaluateMVA(MVAinfos[j].MVAname)
183	peller	1.4	#Fill:
184			newtree.Fill()
185			newtree.AutoSave()
186			outfile.Close()
187
188			#DATA:
189			if job.type == 'DATA':
190			#MVA Formulas
191			MVA_formulas_Nominal = []
192	peller	1.1	#create TTreeFormulas
193			for j in range(len( MVA_Vars['Nominal'])):
194	peller	1.4	MVA_formulas_Nominal.append(ROOT.TTreeFormula("MVA_formula%s_Nominal"%j, MVA_Vars['Nominal'][j],tree))
195			outfile.cd()
196			MVAbranches=[]
197			for i in range(0,len(readers)):
198			MVAbranches.append(array('f',[0]))
199			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F')
200	peller	1.5	newtree.Branch(MVAinfos[i].MVAname+'_4',MVAbranches[i],'nominal/F')
201	peller	1.4	#progbar
202			print '\n--> ' + job.name +':'
203			if nEntries >= longe:
204			step=int(nEntries/longe)
205			long=longe
206			else:
207			long=nEntries
208			step = 1
209			bar=progbar(long)
210			#Fill event by event:
211			for entry in range(0,nEntries):
212			if entry % step == 0:
213			bar.move()
214			#load entry
215			tree.GetEntry(entry)
216			#nominal:
217	peller	1.1	for j in range(len( MVA_Vars['Nominal'])):
218	peller	1.4	MVA_var_buffer[j][0] = MVA_formulas_Nominal[j].EvalInstance()
219
220	peller	1.1	for j in range(0,len(readers)):
221	peller	1.4	MVAbranches[j][0]= readers[j].EvaluateMVA(MVAinfos[j].MVAname)
222			newtree.Fill()
223			newtree.AutoSave()
224			outfile.Close()
225	peller	1.1
226	peller	1.4	print '\n'
227	peller	1.1
228			#Update Info:
229			if doinfo:
230			for job in Ainfo:
231			for MVAinfo in MVAinfos:
232			job.addcomment('Added MVA %s'%MVAinfo.MVAname)
233	peller	1.4	job.addpath(MVAdir)
234			infofile = open(Apath+MVAdir+'/samples.info','w')
235	peller	1.1	pickle.dump(Ainfo,infofile)
236			infofile.close()
237
238