TrainMVA.C

Go to the documentation of this file.

#include "TFile.h"
#include "TMVA/Factory.h"
#include "TTree.h"

#include <iostream>
#include <map>
#include <string>
#include <vector>

#include "MVAPIDResult.h"

void BuildTree(std::string inFile, std::string outFile)
{
  TFile* fIn = new TFile(inFile.c_str());
  TTree* tr = (TTree*)fIn->Get("pid/MVAPID");
  std::vector<anab::MVAPIDResult>* mvares = 0;
  tr->SetBranchAddress("MVAResult", &mvares);
  TFile* fOut = new TFile(outFile.c_str(), "RECREATE");
  TTree* mvaTree = new TTree("mvaTree", "mvaTree");

  float evalRatio, concentration, coreHaloRatio, conicalness;
  float dEdxStart, dEdxEnd, dEdxEndRatio;
  float length;
  int isTrack, isStoppingReco;

  mvaTree->Branch("evalRatio", &evalRatio, "evalRatio/F");
  mvaTree->Branch("concentration", &concentration, "concentration/F");
  mvaTree->Branch("coreHaloRatio", &coreHaloRatio, "coreHaloRatio/F");
  mvaTree->Branch("conicalness", &conicalness, "conicalness/F");
  mvaTree->Branch("dEdxStart", &dEdxStart, "dEdxStart/F");
  mvaTree->Branch("dEdxEnd", &dEdxEnd, "dEdxEnd/F");
  mvaTree->Branch("dEdxEndRatio", &dEdxEndRatio, "dEdxEndRatio/F");
  mvaTree->Branch("length", &length, "length/F");
  mvaTree->Branch("isTrack", &isTrack, "isTrack/I");
  mvaTree->Branch("isStoppingReco", &isStoppingReco, "isStoppingReco/I");

  for (int iEntry = 0; iEntry < tr->GetEntries(); ++iEntry) {
    tr->GetEntry(iEntry);
    if (!mvares->size()) continue;

    anab::MVAPIDResult* biggestTrack = &((*mvares)[0]);
    for (unsigned int iRes = 0; iRes != mvares->size(); ++iRes) {
      if ((((*mvares)[iRes]).nSpacePoints) > (biggestTrack->nSpacePoints)) {
        biggestTrack = &((*mvares)[iRes]);
      }
    }

    evalRatio = biggestTrack->evalRatio;
    concentration = biggestTrack->concentration;
    coreHaloRatio = biggestTrack->coreHaloRatio;
    conicalness = biggestTrack->conicalness;
    dEdxStart = biggestTrack->dEdxStart;
    dEdxEnd = biggestTrack->dEdxEnd;
    dEdxEndRatio = biggestTrack->dEdxEndRatio;
    length = biggestTrack->length;
    isTrack = biggestTrack->isTrack;
    isStoppingReco = biggestTrack->isStoppingReco;

    mvaTree->Fill();
  }

  fIn->Close();
  fOut->Write();
  fOut->Close();
}

void TrainMVA(std::vector<std::string> signalFiles,
              std::vector<std::string> backgroundFiles,
              std::string outputFile,
              std::string jobName)
{

  TFile* fOut = new TFile(outputFile.c_str(), "RECREATE");
  TMVA::Factory* factory = new TMVA::Factory(jobName.c_str(), fOut, "");

  std::vector<TTree*> sigTrees;

  for (std::vector<std::string>::iterator fIter = signalFiles.begin(); fIter != signalFiles.end();
       ++fIter) {
    TFile* fIn = new TFile(fIter->c_str());
    factory->AddSignalTree((TTree*)fIn->Get("mvaTree"));
  }

  for (std::vector<std::string>::iterator fIter = backgroundFiles.begin();
       fIter != backgroundFiles.end();
       ++fIter) {
    TFile* fIn = new TFile(fIter->c_str());
    factory->AddBackgroundTree((TTree*)fIn->Get("mvaTree"));
  }

  factory->AddVariable("evalRatio", 'F');
  factory->AddVariable("concentration", 'F');
  factory->AddVariable("coreHaloRatio", 'F');
  factory->AddVariable("conicalness", 'F');
  factory->AddVariable("dEdxStart", 'F');
  factory->AddVariable("dEdxEnd", 'F');
  factory->AddVariable("dEdxEndRatio", 'F');

  factory->BookMethod(TMVA::Types::kTMlpANN, "ANN", "");
  factory->BookMethod(TMVA::Types::kBDT, "BDT", "");
  factory->TrainAllMethods();
  factory->TestAllMethods();
  factory->EvaluateAllMethods();
  fOut->Write();
  fOut->Close();
}

void PrintRes(std::string inFile)
{

  TFile* fIn = new TFile(inFile.c_str());
  TTree* tr = (TTree*)fIn->Get("pid/ANAB");
  std::vector<anab::MVAPIDResult>* mvares = 0;
  tr->SetBranchAddress("MVAResult", &mvares);

  for (int iEntry = 0; iEntry < tr->GetEntries(); ++iEntry) {
    tr->GetEntry(iEntry);
    if (!mvares->size()) continue;

    anab::MVAPIDResult* biggestTrack = &((*mvares)[0]);
    for (unsigned int iRes = 0; iRes != mvares->size(); ++iRes) {
      if ((((*mvares)[iRes]).nSpacePoints) > (biggestTrack->nSpacePoints)) {
        biggestTrack = &((*mvares)[iRes]);
      }
    }

    std::cout << biggestTrack->mvaOutput.at(string("ANN")) << std::endl;
  }

  fIn->Close();
}
/*
  std::vector<std::string> methods;
  methods.push_back("TMlp_ANN");
  methods.push_back("BDT");
  //methods.push_back("Cuts");
  //methods.push_back("kNN");
  //methods.push_back("Likelihood");
  gStyle->SetOptStat("0000");

  TFile fOut(outFileName.c_str(),"RECREATE");

  TSystemDirectory dir("./","./");
  TIter nextFile(dir.GetListOfFiles());
  TSystemFile* file;
  while(file = (TSystemFile*)nextFile.Next()){
    std::string fName=file->GetName();
    std::cout<<fName<<std::endl;
    if(fName.find("mva_")!=std::string::npos){
      fOut.cd();
      std::string caName=fName.substr(0,fName.size()-5);
      TCanvas* ca=new TCanvas(caName.c_str(),caName.c_str());
      TLegend* le=new TLegend(0.2,0.1,0.4,0.3);
      TFile f(fName.c_str());
      int iMethod=0;
      for(std::vector<std::string>::iterator mIter=methods.begin();mIter!=methods.end();++mIter){
        std::string methodFolder;
        if(mIter->find("ANN")!=std::string::npos){
          methodFolder="TMlpANN";
        }
        else{
          methodFolder=*mIter;
        }
        std::cout<<(std::string("Method_")+methodFolder+"/"+*mIter+"/MVA_"+*mIter+"_rejBvsS").c_str()<<std::endl;
        TH1D* h=(TH1D*)f.Get((std::string("Method_")+methodFolder+"/"+*mIter+"/MVA_"+*mIter+"_rejBvsS").c_str());
        le->AddEntry(h);
        if(iMethod==0){
          h->SetTitle(caName.c_str());
          h->SetLineColor(kRed);
          h->Draw();
          h->GetYaxis().SetRangeUser(0,1.2);
          firstMethod=false;
        }
        else{
          h->SetLineColor(iMethod==1?kBlue:kGreen);
          h->Draw("same");
        }
        ++iMethod;
      }
      le->Draw();
      fOut.cd();
      std::cout<<caName.c_str()<<std::endl;
      ca->Write(caName.c_str());
      ca->SaveAs((std::string("./plots/")+caName+".png").c_str());
    }
  }
  fOut.Close();
}
*/