/cvmfs/atlas.cern.ch/repo/sw/ASG/AnalysisBase/2.4.31/egammaMVACalib/egammaMVACalib/egammaMVACalib.h

#ifndef egammaMVACalibNew_H
#define egammaMVACalibNew_H

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

#include <TString.h>
#include <TObject.h>
#include <TFile.h>
#include "AsgTools/AsgMessaging.h"

class TTreeFormula;
class TList;
class TTree;
class TPRegexp;
class TObjArray;
class TXMLEngine;
class TH2Poly;
class TAxis;
class TF1;

namespace TMVA { class Reader; } // forward decl
namespace MVAUtils { class BDT; } // forward decl

// WARNING: all info to build the key should be defined as a variable or spectator in at least one of the readers

class egammaMVACalib : public asg::AsgMessaging
{
  public:
    enum egammaType {egPHOTON, egELECTRON, NEGAMMATYPES};
    enum ParticleType {UNCONVERTED=0, CONVERTED=1, SiSi=2, ELECTRON=3, NPARTICLETYPES=4, INVALIDPARTICLE=-1};
    enum CalibrationType {correctEcluster, correctEaccordion, fullCalibration, NCalibrationTypes };
    enum ShiftType {NOSHIFT=0, PEAKTOTRUE, MEANTOTRUE, MEDIANTOTRUE, MEAN10TOTRUE, MEAN20TOTRUE, MEDIAN10TOTRUE, MEDIAN20TOTRUE, NSHIFTCORRECTIONS};

    struct ReaderID
    {
    ParticleType particleType;
    int bin; //etaBin; int energyBin;
    bool operator< (const ReaderID &other) const
    {
        if (this->particleType != other.particleType)
        return (this->particleType < other.particleType);
        return (this->bin < other.bin);
    }
    };

    typedef std::map<TString, TString> AdditionalInfoMap;
    typedef std::map< std::pair< egammaMVACalib::ReaderID, egammaMVACalib::ShiftType>, TF1* > ShiftMap;

    struct VarFormula
    {
    float variable;
    TString expression;
    TString infoType;
    TString valueType;
    TTreeFormula* formula;
    bool external;
    };

    egammaMVACalib(int particle,
           bool useNewBDTs = true,
           TString folder = "",
           const TString & method = "BDTG",
           int calibrationType = correctEaccordion,
           bool debug=false,
           const TString & etaBinVar = "",
           const TString & energyBinVar = "",
           const TString & particleTypeVar = "",
           TString filePattern = "",
           bool ignoreSpectators = true);

    virtual ~egammaMVACalib();

    void setPeakCorrection(ShiftType shift_type) { m_shiftType = shift_type; };
    void setPeakCorrection(TString shift_type);

    void useClusterIf0(bool useCluster = true) { m_clusterEif0 = useCluster; }

    void setDefinition(const TString & variable, const TString & expression);

    void InitTree(TTree* tree, bool doNotify = true);

    virtual void LoadVariables(int index); // TODO: move to private

    float getMVAOutput(int index = 0);

    float getMVAEnergy(int index = 0, ShiftType shift_type = NSHIFTCORRECTIONS);

    float getMVAEnergyPhoton(float ph_rawcl_Es0,
                 float ph_rawcl_Es1,
                 float ph_rawcl_Es2,
                 float ph_rawcl_Es3,
                 float ph_cl_eta,
                 float ph_cl_E,
                 float ph_cl_etaCalo,
                 float ph_cl_phiCalo,
                 float ph_ptconv,
                 float ph_pt1conv,
                 float ph_pt2conv,
                 int ph_convtrk1nPixHits,
                 int ph_convtrk1nSCTHits,
                 int ph_convtrk2nPixHits,
                 int ph_convtrk2nSCTHits,
                 float ph_Rconv);

    float getMVAEnergyElectron(float el_rawcl_Es0,
                   float el_rawcl_Es1,
                   float el_rawcl_Es2,
                   float el_rawcl_Es3,
                   float el_cl_eta,
                   float el_cl_E,
                   float el_cl_etaCalo,
                   float el_cl_phiCalo);

    float* getAddress(const TString & input_name);

    void setExternal(const TString & input_name);

    void printValueInput();

    TTree* getMVAResponseTree(TTree *tree, int Nentries = -1,
                  TString branchName = "", TString copyBranches = "input",
                  int first_event = 0, bool flatten = false, int update=10000);

    float getShift(float Et, ReaderID key, ShiftType shift_type) const;

    const TTreeFormula* getFormula(const TString & variable) {
       return (m_formulae.count(variable) ? m_formulae[variable].formula : 0);
    }

    TObjArray* getListOfBranches();

    void activateBranches();

    int getNdata();

    TH2Poly* getTH2Poly() const;

    void writeROOTfile(const TString& directory, int particle = (int) INVALIDPARTICLE);

    TMVA::Reader* getReader(egammaMVACalib::ReaderID key)
    {
    std::map< egammaMVACalib::ReaderID, TMVA::Reader* >::const_iterator it = m_readers.find(key);
    if (it == m_readers.end()) return 0;
    return it->second;
    }

    egammaMVACalib::ReaderID getReaderID();

    MVAUtils::BDT* getBDT(egammaMVACalib::ReaderID key)
    {
    std::map< egammaMVACalib::ReaderID, MVAUtils::BDT* >::const_iterator it = m_BDTs.find(key);
    if (it == m_BDTs.end()) return 0;
    return it->second;
    }

    //virtual bool Notify();


    // Static methods

    static void addBranchesFromFormula(TObjArray* branches, TTreeFormula *formula);

    static TString getString(TObject*);

    static TObjArray* getVariables(const TString &xmlFileName);

    static TMVA::Reader* getDummyReader(const TString &xmlFileName);

    static TString* getVariables(TMVA::Reader *reader);

    static void addReaderInfoToArrays(TMVA::Reader *reader,
                      TObjArray *trees,
                      TObjArray *variables,
                      int index);


 protected:
    TString getCalibTypeString();

    void getReaders(const TString & folder);

    void getBDTs(const std::string& folder);

    struct XmlVariableInfo {
    TString expression;
    TString label;
    TString varType;
    TString nodeName;
    };

    static std::vector<XmlVariableInfo> parseXml(const TString & xmlFilename);

    static std::vector<XmlVariableInfo> parseVariables(TXMLEngine *xml, void* node, const TString & nodeName);

    void setupReader(TMVA::Reader* reader, const TString & xml_filename);

    void setupBDT(const TString& fileName);

    void setupFormulasForReaderID();

    void predefineFormula(const TString & name, const TString & expression,
              const TString & varInfo, const TString & valueType = "F");

    bool parseFileName(const TString & fileName, egammaMVACalib::ReaderID &key);

    bool parseFileName(const TString & fileName, egammaMVACalib::ParticleType &pType);

    TString getBinName( const TString & binField );

    int getBin(float etaMin, float etaMax, float energyMin, float energyMax);

    void printReadersInfo() const;

    TTreeFormula* defineFormula(const TString & varName,
                const TString & expression,
                TTree *tree);

    static egammaMVACalib::ParticleType getParticleType( const TString & s );

    // Functions to define key
    int getParticleType() const;
    int getBin() const;

    void printFormulae() const;

    void checkFormula(TTreeFormula *formula);

    void defineInitialEnergyFormula();

    static AdditionalInfoMap getUserInfo(const TString & xmlfilename);

    TTree* getOutputTree(TString copyBranches, bool deactivateFirst=true);

    bool checkBin(int bin, float etaMin, float etaMax, float etMin, float etMax);

    void defineShiftFormula(ReaderID key);

    const TString& getShiftName(ShiftType shift) const;

    void defineClusterEnergyFormula();

    void checkShowerDepth(TTree *tree);

    int getNreaders() const { return (m_useNewBDTs ? m_BDTs.size() : m_readers.size()); }

    TTree* createInternalTree(egammaType, TTree *tree);


    // Parameters defined in the constructor
    egammaMVACalib::egammaType m_egammaType; 
    bool m_useNewBDTs; 
    TString fMethodName; 
    egammaMVACalib::CalibrationType m_calibrationType; 
    // Variables that define the key and the initial energy (for getMVAEnergy)
    TString m_etaVar, m_energyVar, m_particleTypeVar;
    bool m_ignoreSpectators;
    bool m_hasEnergyBins;
    int m_binMultiplicity;

    bool m_clusterEif0; 

    // Variables that point to the ones that define the key
    // and the initial energy (for getMVAEnergy)
    float *m_particleType, *m_eta, *m_energy, *m_initialEnergy;

    // Bin definitions
    //std::map< ParticleType, TAxis* > m_etaAxis, m_eAxis;
    //TAxis *fAxisEta, *fAxisEnergy;
    TH2Poly* m_hPoly;

    // the tree or chain
    TTree *m_tree, *m_input_tree;

    // use internal TTree
    bool m_useInternalTree;

    // Pattern of xml file names
    TPRegexp *m_fileNamePattern;

    // MVA response
    float m_mvaOutput;

    // Dummy float to be used for spectators if needed
    float m_dummyFloat;

    ShiftType m_shiftType;
    // additional information (mean, median, ...) for each xml file
    std::map< egammaMVACalib::ReaderID, egammaMVACalib::AdditionalInfoMap> m_additional_infos;

    // Map containing TTreeFormulae to get the shifts for each ReaderID, ShiftType
    ShiftMap m_shiftMap;

    // TMVA readers
    std::map< egammaMVACalib::ReaderID, TMVA::Reader* > m_readers;

    // New BDT readers
    std::map< egammaMVACalib::ReaderID, MVAUtils::BDT* > m_BDTs;

    // Variables and formulae for the reader
    std::map< TString, egammaMVACalib::VarFormula > m_formulae;

    // Formula to get cluster energy (std calibration)
    TTreeFormula* m_clusterFormula;

    // fields for internal tree
    float m_ph_rawcl_Es0;
    float m_ph_rawcl_Es1;
    float m_ph_rawcl_Es2;
    float m_ph_rawcl_Es3;
    float m_ph_cl_eta;
    float m_ph_cl_E;
    float m_ph_ptconv;
    float m_ph_pt1conv;
    float m_ph_pt2conv;
    float m_ph_cl_etaCalo;
    float m_ph_cl_phiCalo;
    int m_ph_convtrk1nPixHits;
    int m_ph_convtrk1nSCTHits;
    int m_ph_convtrk2nPixHits;
    int m_ph_convtrk2nSCTHits;
    float m_ph_rawcl_calibHitsShowerDepth;
    float m_ph_Rconv;

    float m_el_rawcl_Es0;
    float m_el_rawcl_Es1;
    float m_el_rawcl_Es2;
    float m_el_rawcl_Es3;
    float m_el_cl_E_TileGap3;
    float m_el_cl_eta;
    float m_el_cl_E;
    float m_el_cl_etaCalo;
    float m_el_cl_phiCalo;
    float m_el_cl_phi;
    float m_el_rawcl_calibHitsShowerDepth;


  public:
    static std::vector<double> get_radius(double eta);
    static double get_shower_depth(double eta,
                   double raw_cl_0,
                   double raw_cl_1,
                   double raw_cl_2,
                   double raw_cl_3);
  private:
    class NotifyDispatcher : public TObject
    {
    public:
      NotifyDispatcher() = default;
      NotifyDispatcher(const std::vector<TObject*>& objs) : m_objs(objs) { }
      void add_object(TObject* obj) { m_objs.push_back(obj); }
      virtual Bool_t Notify() {
         for (auto obj : m_objs) { obj->Notify(); } return true;
      }
    private:
      std::vector<TObject*> m_objs;
    };

    NotifyDispatcher m_notify_dispatcher; 
};

#endif

---