| 1 |
import ROOT
|
| 2 |
from array import array
|
| 3 |
from mvainfos import mvainfo
|
| 4 |
|
| 5 |
class MvaEvaluator:
|
| 6 |
def __init__(self, config, MVAinfo):
|
| 7 |
self.varset = MVAinfo.varset
|
| 8 |
#Define reader
|
| 9 |
self.reader = ROOT.TMVA.Reader("!Color:!Silent")
|
| 10 |
MVAdir=config.get('Directories','vhbbpath')
|
| 11 |
self.systematics=config.get('systematics','systematics').split(' ')
|
| 12 |
self.MVA_Vars={}
|
| 13 |
self.MVAname = MVAinfo.MVAname
|
| 14 |
for systematic in self.systematics:
|
| 15 |
self.MVA_Vars[systematic]=config.get(self.varset,systematic)
|
| 16 |
self.MVA_Vars[systematic]=self.MVA_Vars[systematic].split(' ')
|
| 17 |
#define variables and specatators
|
| 18 |
self.MVA_var_buffer = []
|
| 19 |
for i in range(len( self.MVA_Vars['Nominal'])):
|
| 20 |
self.MVA_var_buffer.append(array( 'f', [ 0 ] ))
|
| 21 |
self.reader.AddVariable( self.MVA_Vars['Nominal'][i],self.MVA_var_buffer[i])
|
| 22 |
self.reader.BookMVA(MVAinfo.MVAname,MVAdir+'/data/'+MVAinfo.getweightfile())
|
| 23 |
#--> Now the MVA is booked
|
| 24 |
|
| 25 |
def setVariables(self,tree,job):
|
| 26 |
#Set formulas for all vars
|
| 27 |
self.MVA_formulas={}
|
| 28 |
for systematic in self.systematics:
|
| 29 |
if job.type == 'DATA' and not systematic == 'Nominal':
|
| 30 |
continue
|
| 31 |
self.MVA_formulas[systematic]=[]
|
| 32 |
for j in range(len( self.MVA_Vars['Nominal'])):
|
| 33 |
self.MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),self.MVA_Vars[systematic][j],tree))
|
| 34 |
|
| 35 |
def evaluate(self,MVA_values,job):
|
| 36 |
#Evaluate all vars and fill the branches
|
| 37 |
for systematic in self.systematics:
|
| 38 |
for j in range(len( self.MVA_Vars['Nominal'])):
|
| 39 |
if job.type == 'DATA' and not systematic == 'Nominal':
|
| 40 |
continue
|
| 41 |
self.MVA_var_buffer[j][0] = self.MVA_formulas[systematic][j].EvalInstance()
|
| 42 |
MVA_values[self.systematics.index(systematic)] = self.reader.EvaluateMVA(self.MVAname)
|
