ChargedPidMVAWeights.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/gearbox/Const.h>
#include <framework/gearbox/Unit.h>
#include <framework/logging/Logger.h>
 
// MVA
#include <mva/interface/Weightfile.h>
 
// ROOT
#include <TObject.h>
#include <TH3F.h>
#include <TParameter.h>
#include <TFile.h>
 
namespace Belle2 {

  class ChargedPidMVAWeights : public TObject {
 
    typedef std::unordered_map<int, std::vector<std::string> > WeightfilesByParticle; 
    typedef std::map<std::string, std::string> VariablesByAlias; 
 
  public:

    ChargedPidMVAWeights() :
      m_energy_unit("energyUnit", Unit::GeV),
      m_ang_unit("angularUnit", Unit::rad),
      m_thetaVarName("clusterTheta"),
      m_implicitNaNmasking(false)
    {};
 

    ChargedPidMVAWeights(const double& energyUnit, const double& angUnit,
                         const std::string& thetaVarName = "clusterTheta",
                         bool implictNaNmasking = false)
    {
      setEnergyUnit(energyUnit);
      setAngularUnit(angUnit);
      m_thetaVarName = thetaVarName;
      m_implicitNaNmasking = implictNaNmasking;
    }

    ~ChargedPidMVAWeights() {};

    enum class ChargedPidMVATrainingMode : unsigned int {
      c_Classification = 0,
      c_Multiclass = 1,
      c_ECL_Classification = 2,
      c_ECL_Multiclass = 3,
      c_PSD_Classification = 4,
      c_PSD_Multiclass = 5,
      c_ECL_PSD_Classification = 6,
      c_ECL_PSD_Multiclass = 7
    };
 

    void setEnergyUnit(const double& unit) { m_energy_unit.SetVal(unit); }
 

    void setAngularUnit(const double& unit) { m_ang_unit.SetVal(unit); }

    void setWeightCategories(const double* clusterThetaBins, const int nClusterThetaBins,
                             const double* pBins, const int nPBins,
                             const double* chargeBins, const int nChargeBins)
    {
 
      m_categories = std::make_unique<TH3F>("clustertheta_p_charge_binsgrid",
                                            ";ECL cluster #theta;p_{lab};Q",
                                            nClusterThetaBins, clusterThetaBins,
                                            nPBins, pBins,
                                            nChargeBins, chargeBins);
    }

    void storeMVAWeights(const int pdg, const std::vector<std::string>& filepaths,
                         const std::vector<std::tuple<double, double, double>>& categoryBinCentres)
    {
 
      if (!isValidPdg(pdg)) {
        B2FATAL("PDG: " << pdg << " is not that of a valid charged particle! Aborting...");
      }
 
      unsigned int idx(0);
      for (const auto& path : filepaths) {
 
        // Index consistency check.
        auto bin_centres_tuple = categoryBinCentres.at(idx);
 
        auto theta_bin_centre = std::get<0>(bin_centres_tuple);
        auto p_bin_centre = std::get<1>(bin_centres_tuple);
        auto charge_bin_centre = std::get<2>(bin_centres_tuple);
 
        auto h_idx = getMVAWeightIdx(theta_bin_centre, p_bin_centre, charge_bin_centre);
        if (idx != h_idx) {
          B2FATAL("xml file:\n" << path << "\nindex in input vector:\n" << idx << "\ndoes not correspond to:\n" << h_idx <<
                  "\n, i.e. the linearised index of the 3D bin centered in (clusterTheta, p, charge) = (" << theta_bin_centre << ", " << p_bin_centre
                  << ", " <<
                  charge_bin_centre <<
                  ")\nPlease check how the input xml file list is being filled.");
        }
 
        MVA::Weightfile weightfile;
        if (path.ends_with(".root")) {
          weightfile = MVA::Weightfile::loadFromROOTFile(path);
        } else  if (path.ends_with(".xml")) {
          weightfile = MVA::Weightfile::loadFromXMLFile(path);
        } else {
          B2WARNING("Unknown file extension for file: " << path << ", fallback to xml...");
          weightfile = MVA::Weightfile::loadFromXMLFile(path);
        }
 
        // Serialize the MVA::Weightfile object into a string for storage in the database,
        // otherwise there are issues w/ dictionary generation for the payload class...
        std::stringstream ss;
        MVA::Weightfile::saveToStream(weightfile, ss);
        m_weightfiles[pdg].push_back(ss.str());
 
        ++idx;
      }
 
    }
 

    void storeMVAWeightsMultiClass(const std::vector<std::string>& filepaths,
                                   const std::vector<std::tuple<double, double, double>>& categoryBinCentres)
    {
      storeMVAWeights(0, filepaths, categoryBinCentres);
    }
 

    void storeCuts(const int pdg, const std::vector<std::string>& cutfiles,
                   const std::vector<std::tuple<double, double, double>>& categoryBinCentres)
    {
 
      if (!isValidPdg(pdg)) {
        B2FATAL("PDG: " << pdg << " is not that of a valid charged particle! Aborting...");
      }
 
      unsigned int idx(0);
      for (const auto& cutfile : cutfiles) {
 
        auto bin_centres_tuple = categoryBinCentres.at(idx);
 
        auto theta_bin_centre = std::get<0>(bin_centres_tuple);
        auto p_bin_centre = std::get<1>(bin_centres_tuple);
        auto charge_bin_centre = std::get<2>(bin_centres_tuple);
 
        auto h_idx = getMVAWeightIdx(theta_bin_centre, p_bin_centre, charge_bin_centre);
        if (idx != h_idx) {
          B2FATAL("Cut file:\n" << cutfile << "\nindex in input vector:\n" << idx << "\ndoes not correspond to:\n" << h_idx <<
                  "\n, i.e. the linearised index of the 3D bin centered in (clusterTheta, p, charge) = (" << theta_bin_centre << ", " << p_bin_centre
                  << ", " <<
                  charge_bin_centre <<
                  ")\nPlease check how the input cut file list is being filled.");
        }
 
        std::ifstream ifs(cutfile);
        std::string cut((std::istreambuf_iterator<char>(ifs)), (std::istreambuf_iterator<char>()));
 
        // Strip trailing newline.
        cut.erase(std::remove(cut.begin(), cut.end(), '\n'), cut.end());
 
        m_cuts[pdg].push_back(cut);
 
        ++idx;
      }
 
    }

    void storeCutsMultiClass(const std::vector<std::string>& cutfiles,
                             const std::vector<std::tuple<double, double, double>>& categoryBinCentres)
    {
      storeCuts(0, cutfiles, categoryBinCentres);
    }
 

    void storeAliases(const VariablesByAlias& aliases)
    {
      m_aliases = VariablesByAlias(aliases);
    }
 

    const TH3F* getWeightCategories() const
    {
      return m_categories.get();
    }
 

    const std::vector<std::string>* getMVAWeights(const int pdg) const
    {
      return &(m_weightfiles.at(pdg));
    }
 

    const std::vector<std::string>* getMVAWeightsMulticlass() const
    {
      return getMVAWeights(0);
    }
 

    const std::vector<std::string>* getCuts(const int pdg) const
    {
      return &(m_cuts.at(pdg));
    }
 

    const std::vector<std::string>* getCutsMulticlass() const
    {
      return getCuts(0);
    }
 

    const VariablesByAlias* getAliases() const
    {
      return &m_aliases;
    }
 

    unsigned int getMVAWeightIdx(const double& theta, const double& p, const double& charge, int& idx_theta, int& idx_p,
                                 int& idx_charge) const
    {
 
      if (!m_categories) {
        B2FATAL("No (clusterTheta, p, charge) TH3 grid was found in the DB payload. Most likely, you are using a GT w/ an old payload which is no longer compatible with the DB object class implementation. This should not happen! Abort...");
      }
 
      int nbins_th = m_categories->GetXaxis()->GetNbins(); // nr. of theta (visible) bins, along X.
      int nbins_p = m_categories->GetYaxis()->GetNbins(); // nr. of p (visible) bins, along Y.
 
      int glob_bin_idx = findBin(theta / m_ang_unit.GetVal(), p / m_energy_unit.GetVal(), charge);
      m_categories->GetBinXYZ(glob_bin_idx, idx_theta, idx_p, idx_charge);
 
      // The index of the linearised 3D m_categories.
      // The unit offset is b/c ROOT sets global bin idx also for overflows and underflows.
      return (idx_theta - 1) + nbins_th * ((idx_p - 1) + nbins_p * (idx_charge - 1));
    }

    unsigned int getMVAWeightIdx(const double& theta, const double& p, const double& charge) const
    {
      int idx_theta, idx_p, idx_charge;
      return getMVAWeightIdx(theta, p, charge, idx_theta, idx_p, idx_charge);
    }
 

    void dumpPayload(const double& theta, const double& p, const double& charge, const int pdg, bool dump_all = false) const
    {
 
      B2INFO("Dumping payload content for:");
      B2INFO("clusterTheta(theta) = " << theta << " [rad], p = " << p << " [GeV/c], charge = " << charge);
 
      if (m_categories) {
        std::string filename = "db_payload_chargedpidmva__theta_p_charge_categories.root";
        B2INFO("\tWriting ROOT file w/ TH3F grid that defines categories:" << filename);
        auto f = std::make_unique<TFile>(filename.c_str(), "RECREATE");
        m_categories->Write();
        f->Close();
      } else {
        B2WARNING("\tThe TH3F object that defines categories is a nullptr!");
      }
 
      for (const auto& [pdgId, weights] : m_weightfiles) {
 
        if (!dump_all && pdg != pdgId) continue;
 
        auto idx = getMVAWeightIdx(theta, p, charge);
 
        auto serialized_weightfile = weights.at(idx);
 
        std::string filename = "db_payload_chargedpidmva__weightfile_pdg_" + std::to_string(pdgId) +
                               "_glob_bin_" + std::to_string(idx + 1) + ".xml";
 
        auto cutstr = getCuts(pdgId)->at(idx);
 
        B2INFO("\tpdgId = " << pdgId);
        B2INFO("\tCut: " << cutstr);
        B2INFO("\tWriting weight file: " << filename);
 
        std::ofstream weightfile;
        weightfile.open(filename.c_str(), std::ios::out);
        weightfile << serialized_weightfile << std::endl;
        weightfile.close();
 
      }
 
    };
 

    void dumpPayloadMulticlass(const double& theta, const double& p, const double& charge) const
    {
      dumpPayload(theta, p, charge, 0);
    }
 

    bool isValidPdg(const int pdg) const
    {
      bool isValid = (Const::chargedStableSet.find(pdg) != Const::invalidParticle) || (pdg == 0);
      return isValid;
    }

    std::string getThetaVarName() const
    {
      return m_thetaVarName;
    }
 

    bool hasImplicitNaNmasking() const
    {
      return m_implicitNaNmasking;
    }
 
 
  private:
 

    int findBin(const double& x, const double& y, const double& z) const
    {
 
      int nbinsx_vis = m_categories->GetXaxis()->GetNbins();
      int nbinsy_vis = m_categories->GetYaxis()->GetNbins();
      int nbinsz_vis = m_categories->GetZaxis()->GetNbins();
 
      double xx = x;
      double yy = y;
      double zz = z;
 
      // If x, y, z are outside of the 3D grid (visible) range, set their value to
      // fall in the last (first) bin before (after) overflow (underflow).
      if (x < m_categories->GetXaxis()->GetBinLowEdge(1)) { xx = m_categories->GetXaxis()->GetBinCenter(1); }
      if (x >= m_categories->GetXaxis()->GetBinLowEdge(nbinsx_vis + 1)) { xx = m_categories->GetXaxis()->GetBinCenter(nbinsx_vis); }
      if (y < m_categories->GetYaxis()->GetBinLowEdge(1)) { yy = m_categories->GetYaxis()->GetBinCenter(1); }
      if (y >= m_categories->GetYaxis()->GetBinLowEdge(nbinsy_vis + 1)) { yy = m_categories->GetYaxis()->GetBinCenter(nbinsy_vis); }
      if (z < m_categories->GetZaxis()->GetBinLowEdge(1)) { zz = m_categories->GetZaxis()->GetBinCenter(1); }
      if (z >= m_categories->GetZaxis()->GetBinLowEdge(nbinsz_vis + 1)) { zz = m_categories->GetZaxis()->GetBinCenter(nbinsz_vis); }
 
      int nbinsx = m_categories->GetXaxis()->GetNbins() + 2;
      int nbinsy = m_categories->GetYaxis()->GetNbins() + 2;
 
      int j = m_categories->GetXaxis()->FindBin(xx);
      int i = m_categories->GetYaxis()->FindBin(yy);
      int k = m_categories->GetZaxis()->FindBin(zz);
 
      return j + nbinsx * (i + nbinsy * k);
    }
 
 
  private:
 
 
    TParameter<double> m_energy_unit; 
    TParameter<double> m_ang_unit;    
    std::string
    m_thetaVarName; 
    bool m_implicitNaNmasking; 
 

    std::unique_ptr<TH3F> m_categories;
 

    WeightfilesByParticle m_weightfiles = {
      { 0, std::vector<std::string>() },
      { Const::electron.getPDGCode(), std::vector<std::string>() },
      { Const::muon.getPDGCode(), std::vector<std::string>() },
      { Const::pion.getPDGCode(), std::vector<std::string>() },
      { Const::kaon.getPDGCode(), std::vector<std::string>() },
      { Const::proton.getPDGCode(), std::vector<std::string>() },
      { Const::deuteron.getPDGCode(), std::vector<std::string>() }
    };
 

    WeightfilesByParticle m_cuts = {
      { 0, std::vector<std::string>() },
      { Const::electron.getPDGCode(), std::vector<std::string>() },
      { Const::muon.getPDGCode(), std::vector<std::string>() },
      { Const::pion.getPDGCode(), std::vector<std::string>() },
      { Const::kaon.getPDGCode(), std::vector<std::string>() },
      { Const::proton.getPDGCode(), std::vector<std::string>() },
      { Const::deuteron.getPDGCode(), std::vector<std::string>() }
    };
 

    VariablesByAlias m_aliases;
 
 
    ClassDef(ChargedPidMVAWeights, 10);
  };

}