TMVATraining.C

Go to the documentation of this file.

#include "TMVA/Factory.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TRegexp.h"
#include "TEventList.h"
#include "TLeaf.h"

#include <iostream>
#include <map>
#include <vector>
#include <algorithm>
#include <utility>


int gettype(TTree *t, TString varname)
{
        if (t->GetBranch(varname)==0) return -1;
        
        TString leaftype = t->GetLeaf(varname)->GetTypeName();
        
        if (leaftype=="Float_t")         return 0; 
        else if (leaftype=="Int_t")      return 1;  
        else if (leaftype=="Bool_t") return 2;
        
        return -1;
}

// ---------------------------------------------------------------

int SplitString(TString s, TString delim, TString *toks, int maxtoks)
{
        TObjArray *tok = s.Tokenize(delim);
        int N = tok->GetEntries();      
        for (int i=0;i<N;++i) 
                if (i<maxtoks) 
                {
                        toks[i] = ((TObjString*)tok->At(i))->String();
                        toks[i].ReplaceAll("\t","");
                        toks[i] = toks[i].Strip(TString::kBoth);
                }
        return N;
}

// ---------------------------------------------------------------

TString getFromCut(TString vars)
{
        TString toks[50];
        int n=SplitString(vars, "&&", toks, 50);
        TRegexp rvar("[_a-zA-Z][_a-zA-Z0-9]*");
        
        TString res=" ";
        
        for (int i=0;i<n;++i)
        {
                TString v = toks[i](rvar);
                if (v!="")
                {
                        if (v=="tag" || res.Contains(" "+v+" ")) continue;
                        res+=v+" ";
                }
        }
        
        res = res.Strip(TString::kBoth);
        
        return res;
}

// ---------------------------------------------------------------

int TMVATraining(TString fname, TString vars, TString precut="", TString treename="")
{
        if (vars.Contains("&&")) vars = getFromCut(vars);
        cout <<"Vars : "<<vars<<endl;

        TString sigcut = "tag&&mode%1000!=900";
        TString bkgcut = "tag&&mode%1000==900";
        
        if (precut!="")
        {
                cout <<"Precut : "<<precut<<endl;
                sigcut += "&&" + precut;
                bkgcut += "&&" + precut;
        }
        
        TRegexp rntp("n[0-9][0-9][0-9]");
        TRegexp rmod("M[0-9][0-9][0-9]");
        if (treename=="") treename = fname(rntp);
        TString smode = fname(rmod);

        TString code = smode+treename(1,3);
        
        TFile *f = TFile::Open(fname);
        TTree *t =(TTree*) f->Get(treename);
        
        TFile* outputFile = TFile::Open( "TMVA_"+code+".root", "RECREATE" );
        
        TMVA::Factory *factory = new TMVA::Factory( code, outputFile, "Silent:!V:Transformations=I;N;D");

        TString toks[30];
        int N = SplitString(vars," ",toks,30);
        
        for (int i=0;i<N;++i)
        {
                int btype = gettype(t, toks[i]);
                if (btype==0 || btype==1)
                        factory->AddVariable(toks[i], btype==0?'F':'I');
        }   
        
        factory->AddTree( t, "Signal",     1.0, sigcut.Data());
        factory->AddTree( t, "Background", 1.0, bkgcut.Data());
        
        int nsig = t->GetEntries(sigcut);
        int nbkg = t->GetEntries(bkgcut);
        
        factory->PrepareTrainingAndTestTree( "", int(nsig*0.8), int(nbkg*0.8), int(nsig*0.19), int(nbkg*0.19)); 

//      factory->BookMethod( TMVA::Types::kLikelihood, "Likelihood","!V:NAvEvtPerBin=50" );
        //factory->BookMethod( TMVA::Types::kMLP, "MLP", "!V:NCycles=50:HiddenLayers=10,10:TestRate=5" );

        // new verion 4.2.0
//      factory->BookMethod( TMVA::Types::kBDT, "BDT", "!V:nTrees=400:BoostType=AdaBoost:nCuts=10:MaxDepth=3:UseFisherCuts:DoPreselection" );
        // old verion 4.1.3
        factory->BookMethod( TMVA::Types::kBDT, "BDT", "!V:nTrees=400:BoostType=AdaBoost:nCuts=10:NNodesMax=10" );
        
        factory->TrainAllMethods();  
        //factory->TestAllMethods();
        //factory->EvaluateAllMethods();
        
        outputFile->Close();
        delete factory;
  return 0;
}