VHbb/python/evaluateMVA.py

#!/usr/bin/env python
import sys
import os
import ROOT 
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
import pickle
from myutils import parse_info

#CONFIGURE
ROOT.gROOT.SetBatch(True)
print 'hello'
#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("-I", "--inpath", dest="inpath", default="",
#                      help="path to samples")
#parser.add_option("-O", "--outpath", dest="outpath", default="",
#                      help="path where to store output 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)

#from samplesclass import sample
from mvainfos import mvainfo
from progbar import progbar
from printcolor import printc

if opts.config =="":
        opts.config = "config"
config = BetterConfigParser()
#config.read('./config7TeV_ZZ')
config.read(opts.config)
anaTag = config.get("Analysis","tag")

#get locations:
Wdir=config.get('Directories','Wdir')
MVASubdir=config.get('Directories','MVAdir')
samplesinfo=config.get('Directories','samplesinfo')

#systematics


INpath = config.get('Directories','MVAin')
OUTpath = config.get('Directories','MVAout')

info = parse_info(samplesinfo,INpath)

#infofile = open(samplesinfo,'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')

#load the namespace
VHbbNameSpace=config.get('VHbbNameSpace','library')
ROOT.gSystem.Load(VHbbNameSpace)

#MVA
MVAinfos=[]
MVAdir=config.get('Directories','vhbbpath')
for MVAname in MVAlist:
    MVAinfofile = open(MVAdir+'/data/'+factoryname+'_'+MVAname+'.info','r')
    MVAinfos.append(pickle.load(MVAinfofile))
    MVAinfofile.close()
    
longe=40
#Workdir
workdir=ROOT.gDirectory.GetPath()


#Apply samples
#infofile = open(samplesinfo,'r')
#Ainfo = pickle.load(infofile)
#infofile.close()


class MvaEvaluater:
    def __init__(self, config, MVAinfo):
        self.varset = MVAinfo.varset
        #Define reader
        self.reader = ROOT.TMVA.Reader("!Color:!Silent")
        MVAdir=config.get('Directories','vhbbpath')
        self.systematics=config.get('systematics','systematics').split(' ')
        self.MVA_Vars={}
        self.MVAname = MVAinfo.MVAname
        for systematic in self.systematics:
            self.MVA_Vars[systematic]=config.get(self.varset,systematic)
            self.MVA_Vars[systematic]=self.MVA_Vars[systematic].split(' ')
        #define variables and specatators
        self.MVA_var_buffer = []
        for i in range(len( self.MVA_Vars['Nominal'])):
            self.MVA_var_buffer.append(array( 'f', [ 0 ] ))
            self.reader.AddVariable( self.MVA_Vars['Nominal'][i],self.MVA_var_buffer[i])
        self.reader.BookMVA(MVAinfo.MVAname,MVAdir+'/data/'+MVAinfo.getweightfile())
        #--> Now the MVA is booked

    def setBranches(self,tree,job):
        #Set formulas for all vars
        self.MVA_formulas={}
        for systematic in self.systematics: 
            if job.type == 'DATA' and not systematic == 'Nominal': continue
            self.MVA_formulas[systematic]=[]
            for j in range(len( self.MVA_Vars['Nominal'])):
                self.MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),self.MVA_Vars[systematic][j],tree))

    def evaluate(self,MVAbranches,job):
        #Evaluate all vars and fill the branches
        for systematic in self.systematics:
            for j in range(len( self.MVA_Vars['Nominal'])):
                if job.type == 'DATA' and not systematic == 'Nominal': continue
                self.MVA_var_buffer[j][0] = self.MVA_formulas[systematic][j].EvalInstance()                
            MVAbranches[self.systematics.index(systematic)] = self.reader.EvaluateMVA(self.MVAname)


theMVAs = []
for mva in MVAinfos:
    theMVAs.append(MvaEvaluater(config,mva))


#eval
for job in info:
    if eval(job.active):
        if job.name in namelist:
            #get trees:
            print INpath+'/'+job.prefix+job.identifier+'.root'
            input = ROOT.TFile.Open(INpath+'/'+job.prefix+job.identifier+'.root','read')
            print OUTpath+'/'+job.prefix+job.identifier+'.root'
            outfile = ROOT.TFile.Open(OUTpath+'/'+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)
            

            #Set branch adress for all vars
            for i in range(0,len(theMVAs)):
                theMVAs[i].setBranches(tree,job)
            outfile.cd()
            #Setup Branches
            MVAbranches=[]
            for i in range(0,len(theMVAs)):
                if job.type == 'Data':
                    MVAbranches.append(array('f',[0]))
                    newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F') 
                else:
                    MVAbranches.append(array('f',[0]*11))
                    newtree.Branch(theMVAs[i].MVAname,MVAbranches[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down:beff1_up:beff1_down/F')
                MVA_formulas_Nominal = []
            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 i in range(0,len(theMVAs)):
                    theMVAs[i].evaluate(MVAbranches[i],job)
                #Fill:
                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(samplesinfo,'w')
#    pickle.dump(Ainfo,infofile)
#    infofile.close()
Revision:	1.20
Committed:	Wed Jan 16 07:46:03 2013 UTC (12 years, 3 months ago) by peller
Content type:	text/x-python
Branch:	MAIN
Changes since 1.19:	+20 -19 lines
Log Message:	changed samples.info to parser
#	User	Rev	Content
1	peller	1.1	#!/usr/bin/env python
2			import sys
3			import os
4			import ROOT
5			from array import array
6			from math import sqrt
7			from copy import copy
8			#suppres the EvalInstace conversion warning bug
9			import warnings
10			warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
11	nmohr	1.9	from optparse import OptionParser
12	nmohr	1.6	from BetterConfigParser import BetterConfigParser
13	peller	1.1	import pickle
14	peller	1.20	from myutils import parse_info
15	peller	1.1
16			#CONFIGURE
17	nmohr	1.18	ROOT.gROOT.SetBatch(True)
18	peller	1.17	print 'hello'
19	peller	1.1	#load config
20	nmohr	1.9	#os.mkdir(path+'/sys')
21	nmohr	1.10	argv = sys.argv
22	nmohr	1.9	parser = OptionParser()
23	nmohr	1.10	parser.add_option("-U", "--update", dest="update", default=0,
24			help="update infofile")
25			parser.add_option("-D", "--discr", dest="discr", default="",
26			help="discriminators to be added")
27	peller	1.17	#parser.add_option("-I", "--inpath", dest="inpath", default="",
28			# help="path to samples")
29			#parser.add_option("-O", "--outpath", dest="outpath", default="",
30			# help="path where to store output samples")
31	nmohr	1.10	parser.add_option("-S", "--samples", dest="names", default="",
32			help="samples you want to run on")
33	nmohr	1.11	parser.add_option("-C", "--config", dest="config", default=[], action="append",
34	nmohr	1.10	help="configuration file")
35	nmohr	1.9	(opts, args) = parser.parse_args(argv)
36	nmohr	1.18
37	peller	1.20	#from samplesclass import sample
38	nmohr	1.18	from mvainfos import mvainfo
39			from progbar import progbar
40			from printcolor import printc
41
42	nmohr	1.10	if opts.config =="":
43	nmohr	1.9	opts.config = "config"
44	nmohr	1.6	config = BetterConfigParser()
45	peller	1.14	#config.read('./config7TeV_ZZ')
46	nmohr	1.9	config.read(opts.config)
47			anaTag = config.get("Analysis","tag")
48	peller	1.1
49			#get locations:
50			Wdir=config.get('Directories','Wdir')
51	nmohr	1.13	MVASubdir=config.get('Directories','MVAdir')
52	nmohr	1.15	samplesinfo=config.get('Directories','samplesinfo')
53	peller	1.1
54			#systematics
55
56
57	peller	1.17	INpath = config.get('Directories','MVAin')
58			OUTpath = config.get('Directories','MVAout')
59
60	peller	1.20	info = parse_info(samplesinfo,INpath)
61
62			#infofile = open(samplesinfo,'r')
63			#info = pickle.load(infofile)
64			#infofile.close()
65	nmohr	1.10	arglist=opts.discr #RTight_blavla,bsbsb
66	peller	1.1
67	nmohr	1.10	namelistIN=opts.names
68	peller	1.4	namelist=namelistIN.split(',')
69
70	peller	1.20	#doinfo=bool(int(opts.update))
71	peller	1.1
72			MVAlist=arglist.split(',')
73	nmohr	1.18
74	peller	1.1
75			#CONFIG
76			#factory
77			factoryname=config.get('factory','factoryname')
78	bortigno	1.19
79			#load the namespace
80			VHbbNameSpace=config.get('VHbbNameSpace','library')
81			ROOT.gSystem.Load(VHbbNameSpace)
82
83	peller	1.1	#MVA
84			MVAinfos=[]
85	nmohr	1.18	MVAdir=config.get('Directories','vhbbpath')
86	peller	1.1	for MVAname in MVAlist:
87	nmohr	1.12	MVAinfofile = open(MVAdir+'/data/'+factoryname+'_'+MVAname+'.info','r')
88	peller	1.1	MVAinfos.append(pickle.load(MVAinfofile))
89			MVAinfofile.close()
90
91			longe=40
92			#Workdir
93			workdir=ROOT.gDirectory.GetPath()
94
95	bortigno	1.19
96	peller	1.1	#Apply samples
97	peller	1.20	#infofile = open(samplesinfo,'r')
98			#Ainfo = pickle.load(infofile)
99			#infofile.close()
100	peller	1.1
101	nmohr	1.18
102			class MvaEvaluater:
103			def __init__(self, config, MVAinfo):
104			self.varset = MVAinfo.varset
105			#Define reader
106			self.reader = ROOT.TMVA.Reader("!Color:!Silent")
107			MVAdir=config.get('Directories','vhbbpath')
108			self.systematics=config.get('systematics','systematics').split(' ')
109			self.MVA_Vars={}
110			self.MVAname = MVAinfo.MVAname
111			for systematic in self.systematics:
112			self.MVA_Vars[systematic]=config.get(self.varset,systematic)
113			self.MVA_Vars[systematic]=self.MVA_Vars[systematic].split(' ')
114			#define variables and specatators
115			self.MVA_var_buffer = []
116			for i in range(len( self.MVA_Vars['Nominal'])):
117			self.MVA_var_buffer.append(array( 'f', [ 0 ] ))
118			self.reader.AddVariable( self.MVA_Vars['Nominal'][i],self.MVA_var_buffer[i])
119			self.reader.BookMVA(MVAinfo.MVAname,MVAdir+'/data/'+MVAinfo.getweightfile())
120			#--> Now the MVA is booked
121
122			def setBranches(self,tree,job):
123			#Set formulas for all vars
124			self.MVA_formulas={}
125			for systematic in self.systematics:
126			if job.type == 'DATA' and not systematic == 'Nominal': continue
127			self.MVA_formulas[systematic]=[]
128			for j in range(len( self.MVA_Vars['Nominal'])):
129			self.MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),self.MVA_Vars[systematic][j],tree))
130
131			def evaluate(self,MVAbranches,job):
132			#Evaluate all vars and fill the branches
133			for systematic in self.systematics:
134			for j in range(len( self.MVA_Vars['Nominal'])):
135			if job.type == 'DATA' and not systematic == 'Nominal': continue
136			self.MVA_var_buffer[j][0] = self.MVA_formulas[systematic][j].EvalInstance()
137			MVAbranches[self.systematics.index(systematic)] = self.reader.EvaluateMVA(self.MVAname)
138
139
140			theMVAs = []
141			for mva in MVAinfos:
142			theMVAs.append(MvaEvaluater(config,mva))
143
144
145	peller	1.1	#eval
146	peller	1.20	for job in info:
147	peller	1.8	if eval(job.active):
148			if job.name in namelist:
149			#get trees:
150	peller	1.17	print INpath+'/'+job.prefix+job.identifier+'.root'
151	nmohr	1.18	input = ROOT.TFile.Open(INpath+'/'+job.prefix+job.identifier+'.root','read')
152	peller	1.17	print OUTpath+'/'+job.prefix+job.identifier+'.root'
153	nmohr	1.18	outfile = ROOT.TFile.Open(OUTpath+'/'+job.prefix+job.identifier+'.root','recreate')
154	peller	1.4	input.cd()
155	peller	1.8	obj = ROOT.TObject
156			for key in ROOT.gDirectory.GetListOfKeys():
157			input.cd()
158			obj = key.ReadObj()
159			#print obj.GetName()
160			if obj.GetName() == job.tree:
161			continue
162			outfile.cd()
163			#print key.GetName()
164			obj.Write(key.GetName())
165			tree = input.Get(job.tree)
166			nEntries = tree.GetEntries()
167	peller	1.4	outfile.cd()
168	peller	1.8	newtree = tree.CloneTree(0)
169	nmohr	1.18
170	peller	1.4
171	nmohr	1.18	#Set branch adress for all vars
172			for i in range(0,len(theMVAs)):
173			theMVAs[i].setBranches(tree,job)
174			outfile.cd()
175			#Setup Branches
176			MVAbranches=[]
177			for i in range(0,len(theMVAs)):
178			if job.type == 'Data':
179	peller	1.8	MVAbranches.append(array('f',[0]))
180			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F')
181			else:
182	nmohr	1.18	MVAbranches.append(array('f',[0]*11))
183			newtree.Branch(theMVAs[i].MVAname,MVAbranches[i],'nominal:JER_up:JER_down:JES_up:JES_down:beff_up:beff_down:bmis_up:bmis_down:beff1_up:beff1_down/F')
184			MVA_formulas_Nominal = []
185			print '\n--> ' + job.name +':'
186			#progbar setup
187			if nEntries >= longe:
188			step=int(nEntries/longe)
189			long=longe
190			else:
191			long=nEntries
192			step = 1
193			bar=progbar(long)
194			#Fill event by event:
195			for entry in range(0,nEntries):
196			if entry % step == 0:
197			bar.move()
198			#load entry
199			tree.GetEntry(entry)
200	peller	1.8
201	nmohr	1.18	for i in range(0,len(theMVAs)):
202			theMVAs[i].evaluate(MVAbranches[i],job)
203			#Fill:
204			newtree.Fill()
205			newtree.AutoSave()
206			outfile.Close()
207
208	peller	1.4	print '\n'
209	peller	1.1
210			#Update Info:
211	peller	1.20	#if doinfo:
212			# for job in Ainfo:
213			# for MVAinfo in MVAinfos:
214			# job.addcomment('Added MVA %s'%MVAinfo.MVAname)
215			# job.addpath(MVAdir)
216			# infofile = open(samplesinfo,'w')
217			# pickle.dump(Ainfo,infofile)
218			# infofile.close()