ECLChargedPidPDFs.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
 
#include <framework/logging/Logger.h>
#include <framework/gearbox/Unit.h>
 
#include <unordered_map>
 
#include <TObject.h>
#include <TParameter.h>
#include <TH2F.h>
#include <TF1.h>
 
namespace Belle2 {
  class ECLChargedPidPDFs: public TObject {
 
  public:
 
    ECLChargedPidPDFs() :
      m_energy_unit("energyUnit", Unit::rad),
      m_ang_unit("angularUnit", Unit::GeV),
      m_variablesmap(),
      m_pdfsmap(),
      m_vtsmap()
    {};
 
    ~ECLChargedPidPDFs() {};
 
    enum class InputVar : unsigned int {
      c_NONE = 0,
      c_E1E9 = 1,
      c_E9E21 = 2,
      c_S2 = 3,
      c_E = 4,
      c_EoP = 5,
      c_Z40 = 6,
      c_Z51 = 7,
      c_ZMVA = 8,
      c_PSDMVA = 9,
      c_DeltaL = 10,
      c_LAT = 11
    };
 
    class VarTransfoSettings {
    public:
 
      VarTransfoSettings() : nVars(0), nDivisionsMax(0), ip(0), jth(0), gbin(0) {}
      ~VarTransfoSettings() {}
 
      int nVars; 
      std::string classPath; 
      std::vector<int> nDivisions; 
      int nDivisionsMax; 
      std::vector<double> cumulDist; 
      std::vector<double> x; 
      std::vector<double> covMatrix; 
      unsigned int ip; 
      unsigned int jth; 
      unsigned int gbin; 
    };
 
    typedef std::unordered_map<InputVar, TF1*> PdfsByVariable; 
    typedef std::unordered_map<int, PdfsByVariable > PdfsMapByCategory; 
    typedef std::unordered_map<int, PdfsMapByCategory > PdfsMapByParticle; 
    typedef std::unordered_map<int, std::vector<InputVar> > VariablesMapByCategory; 
    typedef std::unordered_map<int, VariablesMapByCategory > VariablesMapByParticle; 
    typedef std::unordered_map<int, VarTransfoSettings> VTSMapByCategory; 
    typedef std::unordered_map<int, VTSMapByCategory > VTSMapByParticle; 
    void setVars(const unsigned int pdg, const int true_charge, const unsigned int i, const unsigned int j,
                 const std::vector<InputVar>& vars)
    {
      auto ji = m_categories->GetBin(j, i);
 
      m_variablesmap_bycategory[ji] = vars;
 
      const int signed_pdg = pdg * true_charge / std::abs(true_charge);
 
      m_variablesmap[signed_pdg] = m_variablesmap_bycategory;
 
    }
 
    const std::vector<InputVar>* getVars(const unsigned int pdg,  const int charge, const double& p, const double& theta) const
    {
 
      auto gbin = findBin(m_categories, theta, p);
 
      const int signed_pdg = pdg * charge / std::abs(charge);
 
      return &(m_variablesmap.at(signed_pdg).at(gbin));
 
    }
 
    void printPdfMap(const unsigned int pdg = 0,
                     const double& p = -1.0, const double& theta = -999.0,
                     const int true_charge = 1,
                     const InputVar varid = InputVar::c_NONE) const
    {
 
      const int signed_pdg = pdg * true_charge / std::abs(true_charge);
 
      std::cout << "Printing PDF info: " << std::endl;
      if (pdg) std::cout << "-) |pdgId| * true_charge = " << signed_pdg << std::endl;
      if (p != -1.0) std::cout << "-) p = " << p << " [GeV/c]" << std::endl;
      if (theta != -999.0) std::cout << "-) clusterTheta = " << theta << " [rad]" << std::endl;
      if (varid != InputVar::c_NONE) std::cout << "-) varid = " << static_cast<unsigned int>(varid) << std::endl;
 
      int x, y, z;
      for (const auto& pair0 : m_pdfsmap) {
 
        if (signed_pdg && signed_pdg != pair0.first) continue;
 
        for (const auto& pair1 : pair0.second) {
 
          auto ji(-1);
          if (p != -1.0 && theta != -999.0) {
            ji = findBin(m_categories, theta, p);
            m_categories->GetBinXYZ(ji, x, y, z);
          }
          if (ji > 0 && ji != pair1.first) continue;
 
          std::cout << "\tglobal_bin_idx: " << pair1.first << " (" << x << "," << y << ")" << std::endl;
 
          for (const auto& pair2 : pair1.second) {
            if (varid != InputVar::c_NONE && varid != pair2.first) continue;
            std::cout << "\t\tvarid: " << static_cast<unsigned int>(varid) << ", TF1: " << pair2.second->GetName() << std::endl;
          }
 
        }
      }
    }
 
    void setEnergyUnit(const double& unit) { m_energy_unit.SetVal(unit); }
 
    void setAngularUnit(const double& unit) { m_ang_unit.SetVal(unit); }
 
    void setPdfCategories(TH2F* h)
    {
      m_categories = h;
    }
 
    const TH2F* getPdfCategories() const
    {
      return m_categories;
    }
 
    void add(const unsigned int pdg, const int true_charge, const unsigned int i, const unsigned int j, const InputVar varid, TF1* pdf)
    {
 
      auto ji = m_categories->GetBin(j, i);
 
      const int signed_pdg = pdg * true_charge / std::abs(true_charge);
 
      m_pdfsmap[signed_pdg][ji][varid] = pdf;
 
    }
 
    const TF1* getPdf(const unsigned int pdg, const int charge, const double& p, const double& theta, const InputVar varid) const
    {
 
      const int signed_pdg = pdg * charge / std::abs(charge);
 
      double pp = p / m_energy_unit.GetVal();
      double th = TMath::Abs(theta) / m_ang_unit.GetVal();
 
      int gbin = findBin(m_categories, th, pp);
 
      int x, y, z;
      m_categories->GetBinXYZ(gbin, x, y, z);
 
      B2DEBUG(30, "\t\tAngular unit: " <<  m_ang_unit.GetVal());
      B2DEBUG(30, "\t\tEnergy unit: " << m_energy_unit.GetVal());
      B2DEBUG(30, "\t\t|pdgId| * reco_charge = " << signed_pdg << ", clusterTheta = " << th << ", p = " << pp);
      B2DEBUG(30, "\t\tgbin = " << gbin << ", x,y = (" << x << "," << y << ")");
      B2DEBUG(30, "\t\tvariable id = " << static_cast<unsigned int>(varid));
 
      return m_pdfsmap.at(signed_pdg).at(gbin).at(varid);
    }
 
    inline bool doVarsTransfo() const { return m_do_varstransform; }
 
    void storeVarsTransfoSettings(const unsigned int pdg,
                                  const int true_charge,
                                  const unsigned int i, const unsigned int j,
                                  const int nVars,
                                  const std::string& classPath = "",
                                  const std::vector<int>& nDivisions = std::vector<int>(),
                                  const std::vector<double>& cumulDist = std::vector<double>(),
                                  const std::vector<double>& x = std::vector<double>(),
                                  const std::vector<double>& covMatrix = std::vector<double>())
    {
 
      const int signed_pdg = pdg * true_charge / std::abs(true_charge);
 
      m_do_varstransform = true;
 
      VarTransfoSettings vts;
 
      vts.nVars = nVars;
      vts.classPath = classPath;
      vts.nDivisionsMax = (!nDivisions.empty()) ? *(std::max_element(std::begin(nDivisions), std::end(nDivisions))) : 0;
      vts.nDivisions = nDivisions;
      vts.cumulDist = cumulDist;
      vts.x = x;
      vts.covMatrix = covMatrix;
 
      auto ji = m_categories->GetBin(j, i);
      vts.gbin = ji;
      vts.ip   = i;
      vts.jth  = j;
 
      m_vtsmap_bycategory[ji] = vts;
 
      m_vtsmap[signed_pdg] = m_vtsmap_bycategory;
 
    }
 
    const VarTransfoSettings* getVTS(const unsigned int pdg, const int charge, const double& p, const double& theta) const
    {
 
      const int signed_pdg = pdg * charge / std::abs(charge);
 
      auto gbin = findBin(m_categories, theta, p);
 
      return &(m_vtsmap.at(signed_pdg).at(gbin));
 
    }
 
  private:
 
    int findBin(const TH2F* h, const double& x, const double& y) const
    {
 
      int nbinsx_vis = h->GetXaxis()->GetNbins();
      int nbinsy_vis = h->GetYaxis()->GetNbins();
 
      double xx = x;
      double yy = y;
 
      // If x, y are outside of the 2D hogram 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); }
 
      int nbinsx = h->GetXaxis()->GetNbins() + 2;
      int j = h->GetXaxis()->FindBin(xx);
      int i = h->GetYaxis()->FindBin(yy);
 
      return j + nbinsx * i;
    }
 
  private:
 
    TParameter<double> m_energy_unit; 
    TParameter<double> m_ang_unit;    
    TH2F* m_categories = nullptr;
 
    VariablesMapByParticle m_variablesmap;
 
    VariablesMapByCategory m_variablesmap_bycategory;
 
    bool m_do_varstransform = false;
 
    PdfsMapByParticle m_pdfsmap;
 
    PdfsMapByCategory m_pdfsmap_bycategory;
 
    PdfsByVariable m_pdfs_byvariable;
 
    VTSMapByParticle m_vtsmap;
 
    VTSMapByCategory m_vtsmap_bycategory;
 
    ClassDef(ECLChargedPidPDFs, 2); 
  };
} // end namespace Belle2