ECLChargedPIDMVAWeights.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
// FRAMEWORK
#include <framework/logging/Logger.h>
 
// MVA
#include <mva/interface/Weightfile.h>
 
// ROOT
#include <TObject.h>
#include <TF1.h>
#include <TH1.h>
#include <TParameter.h>
 
// C++
#include <unordered_map>
#include <algorithm>
 
namespace Belle2 {

  class ECLChargedPIDPhasespaceBinning : public TObject {
 
  public:

    ECLChargedPIDPhasespaceBinning() {};

    explicit ECLChargedPIDPhasespaceBinning(const std::vector<std::vector<float>>& binEdges) :
      m_binEdges(binEdges)
    {
      for (const auto& dimensionBinEdges : binEdges) {
        m_nBins.push_back(dimensionBinEdges.size() - 1);
      }
    }

    ~ECLChargedPIDPhasespaceBinning() {};

    int getLinearisedBinIndex(const std::vector<float>& values)
    {
      int globalBin(-1);
      std::vector<int> binIndices = getBinIndices(values);
      for (unsigned int i = 0; i < binIndices.size(); i++) {
        if (binIndices[i] < 0) return -1;
        if (i == 0) {
          globalBin = binIndices[i];
        } else {
          globalBin = globalBin * m_nBins[i] + binIndices[i];
        }
      }
      return globalBin;
    }
 
  private:

    std::vector<int> getBinIndices(const std::vector<float>& values)
    {
      std::vector<int> binIndices(m_binEdges.size());
 
      for (unsigned int i = 0; i < m_binEdges.size(); i++) {
        std::vector<float> dimBinEdges = m_binEdges[i];
        auto it = std::upper_bound(dimBinEdges.begin(), dimBinEdges.end(), values[i]);
        if (it == dimBinEdges.end()) {
          binIndices[i] = -1;
        } else {
          int index = std::distance(dimBinEdges.begin(), it) - 1;
          binIndices[i] = index;
        }
      }
      return binIndices;
    }

    std::vector<std::vector<float>> m_binEdges;

    std::vector<int> m_nBins;
 
    // 1: First class implementation.
    ClassDef(ECLChargedPIDPhasespaceBinning, 1); 
  };

  class ECLChargedPIDPhasespaceCategory : public TObject {
 
  public:
    enum class MVAResponseTransformMode : unsigned int {
      c_LogTransform = 0,
      c_LogTransformSingle = 1,
      c_GaussianTransform = 2,
      c_DecorrelationTransform = 3,
      c_DirectMVAResponse = 4,
      c_LogMVAResponse = 5
    };

    ECLChargedPIDPhasespaceCategory() :
      m_log_transform_offset("logTransformOffset", 1e-15),
      m_max_possible_response_value("maxPossibleResponseValue", 1.0),
      m_temperature("temperature", 1.0)
    {};

    ECLChargedPIDPhasespaceCategory(const std::string& weightfilePath,
                                    const MVAResponseTransformMode& mvaResponeTransformMode,
                                    const std::unordered_map<unsigned int, unsigned int>& mvaIndexForHypothesis) :
      m_log_transform_offset("logTransformOffset", 1e-15),
      m_max_possible_response_value("maxPossibleResponseValue", 1.0),
      m_temperature("temperature", 1.0)
    {
      // Load and serialize the MVA::Weightfile object into a string for storage in the database,
      // otherwise there are issues w/ dictionary generation for the payload class...
      Belle2::MVA::Weightfile weightfile;
      if (weightfilePath.ends_with(".root")) {
        weightfile = Belle2::MVA::Weightfile::loadFromROOTFile(weightfilePath);
      } else if (weightfilePath.ends_with(".xml")) {
        weightfile = Belle2::MVA::Weightfile::loadFromXMLFile(weightfilePath);
      } else {
        B2WARNING("Unknown file extension for file: " << weightfilePath << ", fallback to xml...");
        weightfile = Belle2::MVA::Weightfile::loadFromXMLFile(weightfilePath);
      }
      std::stringstream ss;
      Belle2::MVA::Weightfile::saveToStream(weightfile, ss);
 
      // store
      m_weight = ss.str();
      m_mvaResponseTransformMode = mvaResponeTransformMode;
      m_mvaIndexForHypothesis = mvaIndexForHypothesis;
    }

    ~ECLChargedPIDPhasespaceCategory() {};

    const std::string& getSerialisedWeight() const {return m_weight;}

    void setSerialisedWeight(const std::string& weight) { m_weight = weight; }

    MVAResponseTransformMode getTransformMode() const {return m_mvaResponseTransformMode;}

    const TF1* getPDF(const unsigned int iMVAResponse, const unsigned int hypoPDG) const
    {
      return &m_pdfs.at(iMVAResponse).at(hypoPDG);
    }

    const TH1F* getCDF(const unsigned int iMVAResponse, const int hypoPDG) const
    {
      return &m_cdfs.at(iMVAResponse).at(hypoPDG);
    }

    const std::vector<float>* getDecorrelationMatrix(const int hypoPDG) const
    {
      return &m_decorrelationMatrices.at(hypoPDG);
    }

    void setCDFs(const std::vector<std::unordered_map<unsigned int, TH1F>>& cdfs) {m_cdfs = cdfs;}
 

    void setPDFs(const std::vector<std::unordered_map<unsigned int, TF1>>& pdfs) {m_pdfs = pdfs;}
 

    void setDecorrelationMatrixMap(const std::unordered_map<unsigned int, std::vector<float>>& decorrelationMatrices)
    {
      m_decorrelationMatrices = decorrelationMatrices;
    }

    void setlogTransformOffset(const float& offset)
    {
      m_log_transform_offset.SetVal(offset);
    }

    float getLogTransformOffset() const
    {
      return m_log_transform_offset.GetVal();
    }

    void setTemperature(const float& temperature)
    {
      m_temperature.SetVal(temperature);
    }

    float getTemperature() const
    {
      return m_temperature.GetVal();
    }

    void setMaxPossibleResponseValue(const float& offset)
    {
      m_max_possible_response_value.SetVal(offset);
    }

    float getMaxPossibleResponseValue() const
    {
      return m_max_possible_response_value.GetVal();
    }

    unsigned int getMVAIndexForHypothesis(const unsigned int hypoPDG) const
    {
      return m_mvaIndexForHypothesis.at(hypoPDG);
    }
 
  private:
 
    TParameter<float> m_log_transform_offset; 
    TParameter<float> m_max_possible_response_value; 
    TParameter<float> m_temperature; 

    std::string m_weight;

    MVAResponseTransformMode m_mvaResponseTransformMode = MVAResponseTransformMode::c_LogTransform;

    std::vector<std::unordered_map<unsigned int, TF1>> m_pdfs;

    std::unordered_map<unsigned int, unsigned int> m_mvaIndexForHypothesis;

    std::vector<std::unordered_map<unsigned int, TH1F>> m_cdfs;

    std::unordered_map<unsigned int, std::vector<float>> m_decorrelationMatrices;
 
    // 1: First class implementation.
    ClassDef(ECLChargedPIDPhasespaceCategory, 1); 
  };

  class ECLChargedPIDMVAWeights : public TObject {
  public:
    ECLChargedPIDMVAWeights() {};

    ~ECLChargedPIDMVAWeights() {};
 

    void setWeightCategories(ECLChargedPIDPhasespaceBinning* h) {m_categories = h;}

    void storeMVAWeights(const std::unordered_map<unsigned int, ECLChargedPIDPhasespaceCategory>& phasespaceCategories)
    {
      m_phasespaceCategories = phasespaceCategories;
    }

    const ECLChargedPIDPhasespaceCategory* getPhasespaceCategory(const unsigned int idx)  const {return &m_phasespaceCategories.at(idx);}

    const std::unordered_map<unsigned int, ECLChargedPIDPhasespaceCategory>* getPhasespaceCategories() const {return &m_phasespaceCategories;}

    bool isPhasespaceCovered(const int linearBinIndex) const
    {
      // if the vector of values passed falls outside the defined phasespace.
      if (linearBinIndex < 0) return false;
      // if the vector is within the defined phasespace but we do not provide an ECLChargedPIDPhasespaceCategory object for this bin.
      if (m_phasespaceCategories.count(linearBinIndex) == 0) return false;
      return true;
    }

 
    unsigned int getLinearisedCategoryIndex(const std::vector<float>& values) const
    {
      if (!m_categories) {
        B2FATAL("No N dimensional grid was found in the ECLChargedPIDMVA DB payload. This should not happen! Abort...");
      }
      return m_categories->getLinearisedBinIndex(values);
    }

    const std::vector<std::string>& getBinningVariables() const {return m_binningVariables;}

    void setBinningVariables(const std::vector<std::string>& binningVariables) {m_binningVariables = binningVariables;}
  private:
    ECLChargedPIDPhasespaceBinning* m_categories = nullptr;

    std::unordered_map<unsigned int, ECLChargedPIDPhasespaceCategory> m_phasespaceCategories;

    std::vector<std::string> m_binningVariables;
 
    // 1: First class implementation.
    ClassDef(ECLChargedPIDMVAWeights, 1); 
 
  }; // class ECLChargedPIDMVAWeights

} // Belle 2 Namespace