1 |
#!/usr/bin/env python
|
2 |
from samplesclass import sample
|
3 |
from printcolor import printc
|
4 |
import pickle
|
5 |
import ROOT
|
6 |
from ROOT import TFile, TTree
|
7 |
import ROOT
|
8 |
from array import array
|
9 |
from ConfigParser import SafeConfigParser
|
10 |
import sys
|
11 |
from mvainfos import mvainfo
|
12 |
|
13 |
#warnings.filterwarnings( action='ignore', category=RuntimeWarning, message='creating converter.*' )
|
14 |
|
15 |
|
16 |
#usage: ./train run gui
|
17 |
|
18 |
|
19 |
#CONFIGURE
|
20 |
|
21 |
#load config
|
22 |
config = SafeConfigParser()
|
23 |
config.read('./config')
|
24 |
|
25 |
#get locations:
|
26 |
Wdir=config.get('Directories','Wdir')
|
27 |
|
28 |
#systematics
|
29 |
systematics=config.get('systematics','systematics')
|
30 |
systematics=systematics.split(' ')
|
31 |
|
32 |
weightF=config.get('Weights','weightF')
|
33 |
|
34 |
def getTree(job,cut):
|
35 |
Tree = ROOT.TChain(job.tree)
|
36 |
Tree.Add(job.getpath())
|
37 |
#Tree.SetDirectory(0)
|
38 |
CuttedTree=Tree.CopyTree(cut)
|
39 |
CuttedTree.SetDirectory(0)
|
40 |
print '\t--> read in %s'%job.name
|
41 |
return CuttedTree
|
42 |
|
43 |
def getScale(job):
|
44 |
input = TFile.Open(job.getpath())
|
45 |
CountWithPU = input.Get("CountWithPU")
|
46 |
CountWithPU2011B = input.Get("CountWithPU2011B")
|
47 |
#print lumi*xsecs[i]/hist.GetBinContent(1)
|
48 |
return float(job.lumi)*float(job.xsec)*float(job.sf)/(0.46502*CountWithPU.GetBinContent(1)+0.53498*CountWithPU2011B.GetBinContent(1))*2/float(job.split)
|
49 |
|
50 |
run=sys.argv[1]
|
51 |
gui=sys.argv[2]
|
52 |
|
53 |
|
54 |
#CONFIG
|
55 |
#factory
|
56 |
factoryname=config.get('factory','factoryname')
|
57 |
factorysettings=config.get('factory','factorysettings')
|
58 |
#MVA
|
59 |
MVAtype=config.get(run,'MVAtype')
|
60 |
MVAname=run
|
61 |
MVAsettings=config.get(run,'MVAsettings')
|
62 |
fnameOutput = Wdir +'/weights/'+factoryname+'_'+MVAname+'.root'
|
63 |
#locations
|
64 |
path=config.get(run,'path')
|
65 |
|
66 |
TCutname=config.get(run, 'treeCut')
|
67 |
TCut=config.get('Cuts',TCutname)
|
68 |
#print TCut
|
69 |
|
70 |
#signals
|
71 |
signals=config.get(run,'signals')
|
72 |
signals=signals.split(' ')
|
73 |
#backgrounds
|
74 |
backgrounds=config.get(run,'backgrounds')
|
75 |
backgrounds=backgrounds.split(' ')
|
76 |
|
77 |
treeVarSet=config.get(run,'treeVarSet')
|
78 |
|
79 |
#variables
|
80 |
#TreeVar Array
|
81 |
MVA_Vars={}
|
82 |
MVA_Vars['Nominal']=config.get(treeVarSet,'Nominal')
|
83 |
MVA_Vars['Nominal']=MVA_Vars['Nominal'].split(' ')
|
84 |
#Spectators:
|
85 |
#spectators=config.get(treeVarSet,'spectators')
|
86 |
#spectators=spectators.split(' ')
|
87 |
|
88 |
#TRAINING samples
|
89 |
infofile = open(path+'/samples.info','r')
|
90 |
info = pickle.load(infofile)
|
91 |
infofile.close()
|
92 |
|
93 |
#Workdir
|
94 |
workdir=ROOT.gDirectory.GetPath()
|
95 |
|
96 |
|
97 |
TrainCut='%s && EventForTraining==1'%TCut
|
98 |
EvalCut='%s && EventForTraining==0'%TCut
|
99 |
|
100 |
#load TRAIN trees
|
101 |
Tbackgrounds = []
|
102 |
TbScales = []
|
103 |
Tsignals = []
|
104 |
TsScales = []
|
105 |
|
106 |
for job in info:
|
107 |
if job.name in signals:
|
108 |
Tsignal = getTree(job,TrainCut)
|
109 |
ROOT.gDirectory.Cd(workdir)
|
110 |
TsScale = getScale(job)
|
111 |
Tsignals.append(Tsignal)
|
112 |
TsScales.append(TsScale)
|
113 |
|
114 |
if job.name in backgrounds:
|
115 |
Tbackground = getTree(job,TrainCut)
|
116 |
ROOT.gDirectory.Cd(workdir)
|
117 |
TbScale = getScale(job)
|
118 |
Tbackgrounds.append(Tbackground)
|
119 |
TbScales.append(TbScale)
|
120 |
|
121 |
#load EVALUATE trees
|
122 |
Ebackgrounds = []
|
123 |
EbScales = []
|
124 |
Esignals = []
|
125 |
EsScales = []
|
126 |
|
127 |
for job in info:
|
128 |
if job.name in signals:
|
129 |
Esignal = getTree(job,EvalCut)
|
130 |
ROOT.gDirectory.Cd(workdir)
|
131 |
EsScale = getScale(job)
|
132 |
Esignals.append(Esignal)
|
133 |
EsScales.append(EsScale)
|
134 |
|
135 |
if job.name in backgrounds:
|
136 |
Ebackground = getTree(job,EvalCut)
|
137 |
ROOT.gDirectory.Cd(workdir)
|
138 |
EbScale = getScale(job)
|
139 |
Ebackgrounds.append(Ebackground)
|
140 |
EbScales.append(EbScale)
|
141 |
|
142 |
output = ROOT.TFile.Open(fnameOutput, "RECREATE")
|
143 |
factory = ROOT.TMVA.Factory(factoryname, output, factorysettings)
|
144 |
|
145 |
#set input trees
|
146 |
for i in range(len(Tsignals)):
|
147 |
|
148 |
factory.AddSignalTree(Tsignals[i], TsScales[i], ROOT.TMVA.Types.kTraining)
|
149 |
factory.AddSignalTree(Esignals[i], EsScales[i], ROOT.TMVA.Types.kTesting)
|
150 |
|
151 |
for i in range(len(Tbackgrounds)):
|
152 |
if (Tbackgrounds[i].GetEntries()>0):
|
153 |
factory.AddBackgroundTree(Tbackgrounds[i], TbScales[i], ROOT.TMVA.Types.kTraining)
|
154 |
|
155 |
if (Ebackgrounds[i].GetEntries()>0):
|
156 |
factory.AddBackgroundTree(Ebackgrounds[i], EbScales[i], ROOT.TMVA.Types.kTesting)
|
157 |
|
158 |
|
159 |
for var in MVA_Vars['Nominal']:
|
160 |
factory.AddVariable(var,'D') # add the variables
|
161 |
#for var in spectators:
|
162 |
# factory.AddSpectator(var,'D') #add specators
|
163 |
|
164 |
#Execute TMVA
|
165 |
factory.SetSignalWeightExpression(weightF)
|
166 |
factory.Verbose()
|
167 |
factory.BookMethod(MVAtype,MVAname,MVAsettings)
|
168 |
factory.TrainAllMethods()
|
169 |
factory.TestAllMethods()
|
170 |
factory.EvaluateAllMethods()
|
171 |
output.Write()
|
172 |
|
173 |
#WRITE INFOFILE
|
174 |
infofile = open(Wdir+'/weights/'+factoryname+'_'+MVAname+'.info','w')
|
175 |
info=mvainfo(MVAname)
|
176 |
info.factoryname=factoryname
|
177 |
info.factorysettings=factorysettings
|
178 |
info.MVAtype=MVAtype
|
179 |
info.MVAsettings=MVAsettings
|
180 |
info.weightfilepath=Wdir+'/weights'
|
181 |
info.path=path
|
182 |
info.varset=treeVarSet
|
183 |
info.vars=MVA_Vars['Nominal']
|
184 |
#info.spectators=spectators
|
185 |
pickle.dump(info,infofile)
|
186 |
infofile.close()
|
187 |
|
188 |
# open the TMVA Gui
|
189 |
if gui == 'gui':
|
190 |
ROOT.gROOT.ProcessLine( ".L TMVAGui.C")
|
191 |
ROOT.gROOT.ProcessLine( "TMVAGui(\"%s\")" % fnameOutput )
|
192 |
ROOT.gApplication.Run()
|
193 |
|
194 |
|