TMVA.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2016 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Thomas Keck                                              *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#pragma once
#ifndef INCLUDE_GUARD_BELLE2_MVA_TMVA_HEADER
#define INCLUDE_GUARD_BELLE2_MVA_TMVA_HEADER
 
#include <mva/interface/Options.h>
#include <mva/interface/Teacher.h>
#include <mva/interface/Expert.h>
 
#include <TMVA/Factory.h>
#include <TMVA/Tools.h>
#include <TMVA/Reader.h>
#if ROOT_VERSION_CODE >= ROOT_VERSION(6,8,0)
#include <TMVA/DataLoader.h>
#endif
 
#include <memory>
 
namespace Belle2 {
  namespace MVA {
 
    class TMVAOptions : public SpecificOptions {
 
    public:
      virtual void load(const boost::property_tree::ptree& pt) override;
 
      virtual void save(boost::property_tree::ptree& pt) const override;
 
      virtual po::options_description getDescription() override;
 
      virtual std::string getMethod() const override { return "TMVA"; }
 
    public:
      std::string m_method = "BDT"; 
      std::string m_type = "BDT"; 
      std::string m_config =
        "!H:!V:CreateMVAPdfs:NTrees=400:BoostType=Grad:Shrinkage=0.1:UseBaggedBoost:BaggedSampleFraction=0.5:nCuts=1024:MaxDepth=3:IgnoreNegWeightsInTraining";
      //std::string method = "FastBDT";
      //std::string type = "Plugins";
      //std::string config = "!H:!V:CreateMVAPdfs:NTrees=400:Shrinkage=0.10:RandRatio=0.5:NCutLevel=8:NTreeLayers=3";
      std::string m_factoryOption = "!V:!Silent:Color:DrawProgressBar"; 
      std::string m_prepareOption = "SplitMode=random:!V"; 
      std::string m_workingDirectory = ""; 
      std::string m_prefix = "TMVA"; 
    };
 
    class TMVAOptionsClassification : public TMVAOptions {
 
    public:
      TMVAOptionsClassification()
      {
        m_factoryOption += ":AnalysisType=Classification";
      }
 
      virtual void load(const boost::property_tree::ptree& pt) override;
 
      virtual void save(boost::property_tree::ptree& pt) const override;
 
      virtual po::options_description getDescription() override;
 
      virtual std::string getMethod() const override { return "TMVAClassification"; }
 
    public:
      bool transform2probability = true; 
    };
 
 
    class TMVAOptionsMulticlass : public TMVAOptions {
 
    public:
      TMVAOptionsMulticlass()
      {
        m_factoryOption += ":AnalysisType=Multiclass";
      }
 
      virtual void load(const boost::property_tree::ptree& pt) override;
 
      virtual void save(boost::property_tree::ptree& pt) const override;
 
      virtual po::options_description getDescription() override;
 
      virtual std::string getMethod() const override { return "TMVAMulticlass"; }
 
    public:
 
      std::vector<std::string> m_classes; 
    };
 
 
    class TMVAOptionsRegression : public TMVAOptions {
 
    public:
      TMVAOptionsRegression()
      {
        m_factoryOption += ":AnalysisType=Regression";
      }
 
      virtual std::string getMethod() const override { return "TMVARegression"; }
    };
 
 
    class TMVATeacher : public Teacher {
 
    public:
      TMVATeacher(const GeneralOptions& general_options, const TMVAOptions& _specific_options);
 
#if ROOT_VERSION_CODE >= ROOT_VERSION(6,8,0)
 
      Weightfile trainFactory(TMVA::Factory& factory, TMVA::DataLoader& data_loader, std::string& jobName) const;
#else
 
      Weightfile trainFactory(TMVA::Factory& factory, std::string& jobName) const;
#endif
 
    private:
      TMVAOptions specific_options; 
    };
 
    class TMVATeacherClassification : public TMVATeacher {
 
    public:
      TMVATeacherClassification(const GeneralOptions& general_options, const TMVAOptionsClassification& _specific_options);
 
      virtual Weightfile train(Dataset& training_data) const override;
 
    protected:
      TMVAOptionsClassification specific_options; 
    };
 
    class TMVATeacherMulticlass : public TMVATeacher {
 
    public:
      TMVATeacherMulticlass(const GeneralOptions& general_options, const TMVAOptionsMulticlass& _specific_options);
 
      virtual Weightfile train(Dataset& training_data) const override;
 
    protected:
      TMVAOptionsMulticlass specific_options; 
    };
 
    class TMVATeacherRegression : public TMVATeacher {
 
    public:
      TMVATeacherRegression(const GeneralOptions& general_options, const TMVAOptionsRegression& _specific_options);
 
      virtual Weightfile train(Dataset& training_data) const override;
 
    protected:
      TMVAOptionsRegression specific_options; 
    };
 
 
    class TMVAExpert : public MVA::Expert {
 
    public:
      virtual void load(Weightfile& weightfile) override;
 
    protected:
      std::unique_ptr<TMVA::Reader> m_expert; 
      mutable std::vector<float>
      m_input_cache; 
      mutable std::vector<float>
      m_spectators_cache; 
    };
 
    class TMVAExpertClassification : public TMVAExpert {
 
    public:
      virtual void load(Weightfile& weightfile) override;
 
      virtual std::vector<float> apply(Dataset& test_data) const override;
 
    protected:
      TMVAOptionsClassification specific_options; 
      float expert_signalFraction; 
    };
 
    class TMVAExpertMulticlass : public TMVAExpert {
 
    public:
      virtual void load(Weightfile& weightfile) override;
 
      virtual std::vector<float> apply(Dataset& test_data) const override
      {
        (void) test_data;
        return std::vector<float>();
      };
 
      virtual std::vector<float> apply(Dataset& test_data, const unsigned int classID) const override;
 
    protected:
      TMVAOptionsMulticlass specific_options; 
    };
 
    class TMVAExpertRegression : public TMVAExpert {
 
    public:
      virtual void load(Weightfile& weightfile) override;
 
      virtual std::vector<float> apply(Dataset& test_data) const override;
 
    protected:
      TMVAOptionsRegression specific_options; 
    };
 
  }
}
#endif