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

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

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