chargedParticleIdentificator.cc

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2019 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Marco Milesi                                             *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#include <framework/logging/Logger.h>
#include <framework/gearbox/Const.h>
#include <framework/gearbox/Unit.h>
 
#include <analysis/dbobjects/ChargedPidMVAWeights.h>
 
#include <gtest/gtest.h>
 
#include <TH2F.h>
 
namespace Belle2 {
  class ChargedParticleIdentificatorTest : public ::testing::Test {
 
  public:
 
    ChargedPidMVAWeights m_dbrep;
 
    Const::ChargedStable m_testHypo = Const::electron;
 
    std::unique_ptr<TH2F> m_grid;
 
    std::vector<float> m_thetabins = {0.2164208, 0.5480334, 0.561996, 2.2462387, 2.2811453, 2.7070057};
 
    std::vector<float> m_pbins = {0.0, 0.5, 0.75, 1.0, 3.0, 100.0};
 
    std::vector<int> m_eclregbins = {1, 11, 2, 13, 3};
 
    std::string m_basename = "dummy_weightfile";
 
    std::vector<std::string> m_dummyfiles;
 
  protected:
 
    void SetUp() override
    {
 
      m_grid = std::make_unique<TH2F>("theta_p_binsgrid",
                                      ";ECL cluster #theta;p_{lab}",
                                      m_thetabins.size() - 1,
                                      m_thetabins.data(),
                                      m_pbins.size() - 1,
                                      m_pbins.data());
 
      m_dbrep.setWeightCategories(m_grid.get());
 
      m_dbrep.setAngularUnit(Unit::rad);
      m_dbrep.setEnergyUnit(Unit::GeV);
 
      std::vector<std::pair<float, float>> gridBinCentres;
 
      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
                       + "__p__" + std::to_string(m_pbins.at(ip)) + "_" + std::to_string(m_pbins.at(ip + 1))
                       + "__clusterTheta__" + std::to_string(m_thetabins.at(jth)) + "_" + std::to_string(m_thetabins.at(jth + 1));
          std::replace(fname.begin(), fname.end(), '.', '_');
          fname += ".xml";
          std::ofstream dummyfile(fname);
          dummyfile.close();
          m_dummyfiles.push_back(fname);
 
          auto centre = std::make_pair(th_bin_centre, p_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 value for (clusterTheta, p) in the grid.
    int binx = 3;
    int biny = 4;
    auto theta = m_grid->GetXaxis()->GetBinCenter(binx);
    auto p = m_grid->GetYaxis()->GetBinCenter(biny);
 
    int jth, ip;
    auto ji = m_dbrep.getMVAWeightIdx(theta, p, jth, ip);
 
    EXPECT_EQ(jth, binx);
    EXPECT_EQ(ip, biny);
 
    auto thisfname = m_basename
                     + "__p__" + std::to_string(m_pbins.at(ip - 1)) + "_" + std::to_string(m_pbins.at(ip))
                     + "__clusterTheta__" + std::to_string(m_thetabins.at(jth - 1)) + "_" + std::to_string(m_thetabins.at(jth));
    std::replace(thisfname.begin(), thisfname.end(), '.', '_');
    thisfname += ".xml";
 
    EXPECT_EQ(thisfname, m_dummyfiles.at(ji));
 
    auto matchitr = std::find(m_dummyfiles.begin(), m_dummyfiles.end(), thisfname);
    auto thisidx = std::distance(m_dummyfiles.begin(), matchitr);
 
    EXPECT_EQ(thisidx, ji);
 
  }
 
}  // namespace