VHbb/python/evaluateMVA.py

#!/usr/bin/env python
from __future__ import print_function
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
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)

if opts.config =="":
        opts.config = "config"

#Import after configure to get help message
from myutils import BetterConfigParser, progbar, printc, mvainfo, ParseInfo

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

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

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

info = ParseInfo(samplesinfo,INpath)

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

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

samples = info.get_samples(namelist)
for job in samples:
    #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=long(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')
Revision:	1.22
Committed:	Fri Jan 25 16:18:00 2013 UTC (12 years, 3 months ago) by nmohr
Content type:	text/x-python
Branch:	MAIN
Changes since 1.21:	+66 -92 lines
Log Message:	Restructuring, still to be validated, workspace writing missing
#	User	Rev	Content
1	peller	1.1	#!/usr/bin/env python
2	nmohr	1.22	from __future__ import print_function
3	peller	1.1	import sys
4			import os
5			import ROOT
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	peller	1.1	import pickle
14	peller	1.21
15	peller	1.1
16			#CONFIGURE
17	nmohr	1.18	ROOT.gROOT.SetBatch(True)
18	nmohr	1.22	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	nmohr	1.22	if opts.config =="":
38			opts.config = "config"
39	peller	1.21
40	nmohr	1.22	#Import after configure to get help message
41			from myutils import BetterConfigParser, progbar, printc, mvainfo, ParseInfo
42	nmohr	1.18
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.15	samplesinfo=config.get('Directories','samplesinfo')
51	peller	1.1
52			#systematics
53	peller	1.17	INpath = config.get('Directories','MVAin')
54			OUTpath = config.get('Directories','MVAout')
55
56	nmohr	1.22	info = ParseInfo(samplesinfo,INpath)
57	peller	1.20
58	nmohr	1.10	arglist=opts.discr #RTight_blavla,bsbsb
59	peller	1.1
60	nmohr	1.10	namelistIN=opts.names
61	peller	1.4	namelist=namelistIN.split(',')
62
63	peller	1.20	#doinfo=bool(int(opts.update))
64	peller	1.1
65			MVAlist=arglist.split(',')
66	nmohr	1.18
67	peller	1.1	#CONFIG
68			#factory
69			factoryname=config.get('factory','factoryname')
70	bortigno	1.19
71			#load the namespace
72			VHbbNameSpace=config.get('VHbbNameSpace','library')
73			ROOT.gSystem.Load(VHbbNameSpace)
74
75	peller	1.1	#MVA
76			MVAinfos=[]
77	nmohr	1.18	MVAdir=config.get('Directories','vhbbpath')
78	peller	1.1	for MVAname in MVAlist:
79	nmohr	1.12	MVAinfofile = open(MVAdir+'/data/'+factoryname+'_'+MVAname+'.info','r')
80	peller	1.1	MVAinfos.append(pickle.load(MVAinfofile))
81			MVAinfofile.close()
82
83			longe=40
84			#Workdir
85			workdir=ROOT.gDirectory.GetPath()
86
87	nmohr	1.18	class MvaEvaluater:
88			def __init__(self, config, MVAinfo):
89			self.varset = MVAinfo.varset
90			#Define reader
91			self.reader = ROOT.TMVA.Reader("!Color:!Silent")
92			MVAdir=config.get('Directories','vhbbpath')
93			self.systematics=config.get('systematics','systematics').split(' ')
94			self.MVA_Vars={}
95			self.MVAname = MVAinfo.MVAname
96			for systematic in self.systematics:
97			self.MVA_Vars[systematic]=config.get(self.varset,systematic)
98			self.MVA_Vars[systematic]=self.MVA_Vars[systematic].split(' ')
99			#define variables and specatators
100			self.MVA_var_buffer = []
101			for i in range(len( self.MVA_Vars['Nominal'])):
102			self.MVA_var_buffer.append(array( 'f', [ 0 ] ))
103			self.reader.AddVariable( self.MVA_Vars['Nominal'][i],self.MVA_var_buffer[i])
104			self.reader.BookMVA(MVAinfo.MVAname,MVAdir+'/data/'+MVAinfo.getweightfile())
105			#--> Now the MVA is booked
106
107			def setBranches(self,tree,job):
108			#Set formulas for all vars
109			self.MVA_formulas={}
110			for systematic in self.systematics:
111			if job.type == 'DATA' and not systematic == 'Nominal': continue
112			self.MVA_formulas[systematic]=[]
113			for j in range(len( self.MVA_Vars['Nominal'])):
114			self.MVA_formulas[systematic].append(ROOT.TTreeFormula("MVA_formula%s_%s"%(j,systematic),self.MVA_Vars[systematic][j],tree))
115
116			def evaluate(self,MVAbranches,job):
117			#Evaluate all vars and fill the branches
118			for systematic in self.systematics:
119			for j in range(len( self.MVA_Vars['Nominal'])):
120			if job.type == 'DATA' and not systematic == 'Nominal': continue
121			self.MVA_var_buffer[j][0] = self.MVA_formulas[systematic][j].EvalInstance()
122			MVAbranches[self.systematics.index(systematic)] = self.reader.EvaluateMVA(self.MVAname)
123
124
125			theMVAs = []
126			for mva in MVAinfos:
127			theMVAs.append(MvaEvaluater(config,mva))
128
129
130	peller	1.1	#eval
131	nmohr	1.22
132			samples = info.get_samples(namelist)
133			for job in samples:
134			#get trees:
135			print(INpath+'/'+job.prefix+job.identifier+'.root')
136			input = ROOT.TFile.Open(INpath+'/'+job.prefix+job.identifier+'.root','read')
137			print(OUTpath+'/'+job.prefix+job.identifier+'.root')
138			outfile = ROOT.TFile.Open(OUTpath+'/'+job.prefix+job.identifier+'.root','recreate')
139			input.cd()
140			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			outfile.cd()
153			newtree = tree.CloneTree(0)
154	nmohr	1.18
155	nmohr	1.22	#Set branch adress for all vars
156			for i in range(0,len(theMVAs)):
157			theMVAs[i].setBranches(tree,job)
158			outfile.cd()
159			#Setup Branches
160			MVAbranches=[]
161			for i in range(0,len(theMVAs)):
162			if job.type == 'Data':
163			MVAbranches.append(array('f',[0]))
164			newtree.Branch(MVAinfos[i].MVAname,MVAbranches[i],'nominal/F')
165			else:
166			MVAbranches.append(array('f',[0]*11))
167			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')
168			MVA_formulas_Nominal = []
169			print('\n--> ' + job.name +':')
170			#progbar setup
171			if nEntries >= longe:
172			step=long(nEntries/longe)
173			long=longe
174			else:
175			long=nEntries
176			step = 1
177			bar=progbar(long)
178			#Fill event by event:
179			for entry in range(0,nEntries):
180			if entry % step == 0:
181			bar.move()
182			#load entry
183			tree.GetEntry(entry)
184	peller	1.8
185	nmohr	1.22	for i in range(0,len(theMVAs)):
186			theMVAs[i].evaluate(MVAbranches[i],job)
187			#Fill:
188			newtree.Fill()
189			newtree.AutoSave()
190			outfile.Close()
191	nmohr	1.18
192	nmohr	1.22	print('\n')