#!/usr/bin/env python
# coding: utf-8

# <img src="http://oproject.org/img/ROOT.png" height="30%" width="30%">
# <img src="http://oproject.org/img/tmvalogo.png" height="30%" width="30%">
# 
# # TMVA Variable Importance

# In[ ]:


from ROOT import TFile,TMVA, TTree, TCut
get_ipython().run_line_magic('jsroot', 'on')


# ## Declare DataLoader

# In[ ]:


TMVA.Tools.Instance()

inputFile = TFile.Open("https://raw.githubusercontent.com/iml-wg/tmvatutorials/master/inputdata.root")

loader = TMVA.DataLoader("dataset")

loader.AddVariable("var1")
loader.AddVariable("var2")
loader.AddVariable("var3")
loader.AddVariable("var4")
loader.AddVariable("var5 := var1*var2")


# ## Setting up dataset

# In[ ]:


tsignal = TTree()
tbackground = TTree()
inputFile.GetObject("Sig", tsignal)
inputFile.GetObject("Bkg", tbackground)

mycuts = TCut()
mycutb = TCut()
   
loader.AddSignalTree(tsignal, 1)           #signal weight = 1
loader.AddBackgroundTree (tbackground, 1)  #background weight = 1 

loader.PrepareTrainingAndTestTree(mycuts, mycutb,
                                  "nTrain_Signal=1000:nTrain_Background=1000:nTest_Signal=500:nTest_Background=500:SplitMode=Random:NormMode=NumEvents:!V" )


# ## Variable importance

# In[ ]:


vi = TMVA.VariableImportance(loader)

vi.BookMethod(TMVA.Types.kBDT, "BDT",
"!V:NTrees=5:MinNodeSize=2.5%:MaxDepth=2:BoostType=AdaBoost:AdaBoostBeta=0.5:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20" )

vi.SetType(TMVA.kShort)

vi.Evaluate()


# In[ ]:


results=vi.GetResults()
results.Print()
results.Draw()