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