chargedParticleIdentificator.cc

/**************************************************************************
 * 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.                  *
 **************************************************************************/
 
#include <framework/logging/Logger.h>
#include <framework/gearbox/Const.h>
 
#include <analysis/dbobjects/ChargedPidMVAWeights.h>
 
#include <gtest/gtest.h>
#include <random>
 
namespace Belle2 {

  class ChargedParticleIdentificatorTest : public ::testing::Test {
 
  public:

    ChargedPidMVAWeights m_dbrep;

    Const::ChargedStable m_testHypo = Const::electron;

    std::vector<double> m_thetabins = {0.2164208, 0.5480334, 0.561996, 2.2462387, 2.2811453, 2.7070057};

    std::vector<double> m_pbins = {0.2, 0.6, 0.75, 1.0, 3.0, 7.0};

    std::vector<double> m_chbins = { -1.5, -0.5, 0.5, 1.5};

    std::string m_basename = "dummy_weightfile";

    std::vector<std::string> m_dummyfiles;
 
  protected:

    void SetUp() override
    {
 
      m_dbrep.setWeightCategories(m_thetabins.data(), m_thetabins.size() - 1,
                                  m_pbins.data(), m_pbins.size() - 1,
                                  m_chbins.data(), m_chbins.size() - 1);
 
      std::vector<std::tuple<double, double, double>> gridBinCentres;
 
      for (unsigned int kch(0); kch < m_chbins.size() - 1; kch++) {
        auto ch_bin_centre = (m_chbins.at(kch) + m_chbins.at(kch + 1)) / 2.0;
        for (unsigned int ip(0); ip < m_pbins.size() - 1; ip++) {
          auto p_bin_centre = (m_pbins.at(ip) + m_pbins.at(ip + 1)) / 2.0;
          for (unsigned int jth(0); jth < m_thetabins.size() - 1; jth++) {
            auto th_bin_centre = (m_thetabins.at(jth) + m_thetabins.at(jth + 1)) / 2.0;
            auto fname = m_basename
                         + "__clusterTheta__" + std::to_string(m_thetabins.at(jth)) + "_" + std::to_string(m_thetabins.at(jth + 1))
                         + "__p__" + std::to_string(m_pbins.at(ip)) + "_" + std::to_string(m_pbins.at(ip + 1))
                         + "__charge__" + std::to_string(ch_bin_centre);
 
            std::replace(fname.begin(), fname.end(), '.', '_');
            fname += ".xml";
            std::ofstream dummyfile(fname);
            dummyfile.close();
            m_dummyfiles.push_back(fname);
 
            auto centre = std::make_tuple(th_bin_centre, p_bin_centre, ch_bin_centre);
            gridBinCentres.push_back(centre);
          }
        }
      }
      m_dbrep.storeMVAWeights(m_testHypo.getPDGCode(), m_dummyfiles, gridBinCentres);
 
    }

    void TearDown() override
    {
 
      // Delete all dummy files.
      for (const auto& fname : m_dummyfiles) {
        if (remove(fname.c_str())) {
          B2ERROR("Couldn't remove file: " << fname);
        }
      }
 
    }
 
  };

  TEST_F(ChargedParticleIdentificatorTest, TestDBRep)
  {
 
    // Pick a random index in the clusterTheta, p, charge bins arrays.
    // Exclude underflow and overflow.
    std::random_device rd; // non-deterministic uniform int rand generator.
    std::uniform_int_distribution<int> binx_idx_distr(1, m_thetabins.size() - 1);
    std::uniform_int_distribution<int> biny_idx_distr(1, m_pbins.size() - 1);
    std::uniform_int_distribution<int> binz_idx_distr(1, m_chbins.size() - 1);
    int binx = binx_idx_distr(rd);
    int biny = biny_idx_distr(rd);
    int binz = binz_idx_distr(rd);
 
    // Pick each axis' bin centre as a test value for (clusterTheta, p, charge).
    auto theta = m_dbrep.getWeightCategories()->GetXaxis()->GetBinCenter(binx);
    auto p = m_dbrep.getWeightCategories()->GetYaxis()->GetBinCenter(biny);
    auto charge = m_dbrep.getWeightCategories()->GetZaxis()->GetBinCenter(binz);
 
    int jth, ip, kch;
    auto jik = m_dbrep.getMVAWeightIdx(theta, p, charge, jth, ip, kch);
 
    EXPECT_EQ(jth, binx);
    EXPECT_EQ(ip, biny);
    EXPECT_EQ(kch, binz);
 
    auto thisfname = m_basename
                     + "__clusterTheta__" + std::to_string(m_thetabins.at(jth - 1)) + "_" + std::to_string(m_thetabins.at(jth))
                     + "__p__" + std::to_string(m_pbins.at(ip - 1)) + "_" + std::to_string(m_pbins.at(ip))
                     + "__charge__" + std::to_string(charge);
    std::replace(thisfname.begin(), thisfname.end(), '.', '_');
    thisfname += ".xml";
 
    EXPECT_EQ(thisfname, m_dummyfiles.at(jik));
 
    auto matchitr = std::find(m_dummyfiles.begin(), m_dummyfiles.end(), thisfname);
    auto thisidx = std::distance(m_dummyfiles.begin(), matchitr);
 
    EXPECT_EQ(thisidx, jik);
 
  }

}  // namespace