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>
 
//BOOST
#include <boost/algorithm/string/predicate.hpp>
 
 
namespace Belle2 {
  class ChargedPidMVAWeights : public TObject {
 
    typedef std::unordered_map<int, std::vector<std::string> > WeightfilesByParticle; 
  public:
 
    ChargedPidMVAWeights() :
      m_energy_unit("energyUnit", Unit::GeV),
      m_ang_unit("angularUnit", Unit::rad)
    {};
 
 
    ~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(TH3F* h)
    {
      m_categories = h;
    }
 
    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.");
        }
 
        Belle2::MVA::Weightfile weightfile;
        if (boost::ends_with(path, ".root")) {
          weightfile = Belle2::MVA::Weightfile::loadFromROOTFile(path);
        } else  if (boost::ends_with(path, ".xml")) {
          weightfile = Belle2::MVA::Weightfile::loadFromXMLFile(path);
        } else {
          B2WARNING("Unknown file extension for file: " << path << ", fallback to xml...");
          weightfile = Belle2::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;
        Belle2::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());
 
        // Conditional expression separator must use square brackets in basf2.
        std::replace(cut.begin(), cut.end(), '(', '[');
        std::replace(cut.begin(), 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);
    }
 
 
    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);
    }
 
 
    unsigned int getMVAWeightIdx(const double& clusterTheta, 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 clusterTheta (visible) bins, along X.
      int nbins_p = m_categories->GetYaxis()->GetNbins(); // nr. of p (visible) bins, along Y.
 
      int glob_bin_idx = findBin(m_categories, clusterTheta / 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 (clusterTheta, p, charge) 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& clusterTheta, const double& p, const double& charge, const int pdg, bool dump_all = false) const
    {
 
      B2INFO("Dumping payload content for:");
      B2INFO("clusterTheta = " << clusterTheta << " [rad], p = " << p << " [GeV/c], charge = " << charge);
 
      if (m_categories) {
        std::string filename = "db_payload_chargedpidmva__clustertheta_p_charge_categories.root";
        B2INFO("\tWriting ROOT file w/ (clusterTheta, p, charge) 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(clusterTheta, 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;
    }
 
 
  private:
 
 
    int findBin(const TH3F* h, const double& x, const double& y, const double& z) const
    {
 
      int nbinsx_vis = h->GetXaxis()->GetNbins();
      int nbinsy_vis = h->GetYaxis()->GetNbins();
      int nbinsz_vis = h->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 < h->GetXaxis()->GetBinLowEdge(1)) { xx = h->GetXaxis()->GetBinCenter(1); }
      if (x >= h->GetXaxis()->GetBinLowEdge(nbinsx_vis + 1)) { xx = h->GetXaxis()->GetBinCenter(nbinsx_vis); }
      if (y < h->GetYaxis()->GetBinLowEdge(1)) { yy = h->GetYaxis()->GetBinCenter(1); }
      if (y >= h->GetYaxis()->GetBinLowEdge(nbinsy_vis + 1)) { yy = h->GetYaxis()->GetBinCenter(nbinsy_vis); }
      if (z < h->GetZaxis()->GetBinLowEdge(1)) { zz = h->GetZaxis()->GetBinCenter(1); }
      if (z >= h->GetZaxis()->GetBinLowEdge(nbinsz_vis + 1)) { zz = h->GetZaxis()->GetBinCenter(nbinsz_vis); }
 
      int nbinsx = h->GetXaxis()->GetNbins() + 2;
      int nbinsy = h->GetYaxis()->GetNbins() + 2;
 
      int j = h->GetXaxis()->FindBin(xx);
      int i = h->GetYaxis()->FindBin(yy);
      int k = h->GetZaxis()->FindBin(zz);
 
      return j + nbinsx * (i + nbinsy * k);
    }
 
 
  private:
 
 
    TParameter<double> m_energy_unit; 
    TParameter<double> m_ang_unit;    
    TH3F* m_categories = nullptr;
 
 
    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>() }
    };
 
 
    ClassDef(ChargedPidMVAWeights, 7);
  };
 
}