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')


#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

#for axample
#0 5 0
#and
#5 -1 1

start=int(sys.argv[3])
stop=int(sys.argv[4])
doinfo=bool(int(sys.argv[5]))

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[start:stop]:
    #get trees:
    input = TFile.Open(job.getpath(),'read')
    outfile = TFile.Open(job.path+'/MVAout/'+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()


#Update Info:
if doinfo:
    for job in Ainfo:        
        for MVAinfo in MVAinfos:
            job.addcomment('Added MVA %s'%MVAinfo.MVAname)
        job.addpath('/MVAout')
    infofile = open(Apath+'/MVAout/samples.info','w')
    pickle.dump(Ainfo,infofile)
    infofile.close()


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