FEIVariables.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.                  *
 **************************************************************************/
 
// Own header.
#include <analysis/variables/FEIVariables.h>
 
// include VariableManager
#include <analysis/VariableManager/Manager.h>
 
#include <analysis/dataobjects/Particle.h>
#include <analysis/dataobjects/ParticleList.h>
#include <analysis/utility/PCmsLabTransform.h>
#include <framework/gearbox/Const.h>
#include <mdst/dataobjects/MCParticle.h>
#include <framework/datastore/StoreArray.h>
#include <framework/datastore/StoreObjPtr.h>
#include <map>
#include <algorithm>
#include <iterator>
 
namespace Belle2 {
  namespace Variable {
 
    int mostcommonBTagIndex(const Particle* part)
    {
      std::map <int, int> tag_candidates;
      const std::vector<const Particle*>& fsp_tag = part->getFinalStateDaughters();
      for (const Particle* fsp : fsp_tag) {
        const MCParticle* mc_fsp = fsp->getMCParticle();
        if (mc_fsp) {
          int tag_index = finddescendant(mc_fsp);
          if (tag_candidates.find(tag_index) == tag_candidates.end()) {
            tag_candidates[tag_index] = 1;
          } else {
            tag_candidates[tag_index]++;
          }
        }
      }
      int tag_index = -1;
      int tag_N = 0;
      for (const auto& [key, value] : tag_candidates) {
        if (value > tag_N) {
          tag_index = key;
          tag_N = value;
        }
      }
      return tag_index;
    }
    int finddescendant(const MCParticle* mcpart)
    {
      const MCParticle* i_mcpart = mcpart;
      while (i_mcpart) {
        const auto* i_mcpart_mother = i_mcpart->getMother();
        if (i_mcpart_mother) {
          constexpr std::array<int, 2> B_PDG = {511, 521};
          auto result = std::find(B_PDG.begin(), B_PDG.end(), abs(i_mcpart_mother->getPDG()));
          if (result != B_PDG.end()) {
            return i_mcpart_mother->getArrayIndex();
          }
          i_mcpart = i_mcpart_mother;
        } else {
          return -1;
        }
      }
      return -2;
    }
    std::vector<int> truthFSPTag(int BTag_index)
    {
      StoreArray<MCParticle> MC_Particle_list;
      std::vector<int> fsp_truth_index;
      for (const MCParticle& iMCParticle : MC_Particle_list) {
        if ((BTag_index == finddescendant(&iMCParticle)) && (iMCParticle.hasStatus(MCParticle::c_StableInGenerator) == true)
            && (iMCParticle.hasStatus(MCParticle::c_IsISRPhoton) == false) && iMCParticle.hasStatus(MCParticle::c_IsFSRPhoton) == false
            && (iMCParticle.hasStatus(MCParticle::c_IsPHOTOSPhoton) == false)) {
          fsp_truth_index.push_back(iMCParticle.getArrayIndex());
        }
      }
      return fsp_truth_index;
    }
    double percentageMissingParticlesBTag(const Particle* part)
    {
      int index = mostcommonBTagIndex(part);
      if (index == -1) {
        return -1;
      } else {
        std::vector<int> fsp_FEI_index;
        const std::vector<const Particle*>& fsp_tag = part->getFinalStateDaughters();
        for (const Particle* fsp : fsp_tag) {
          const MCParticle* mc_fsp = fsp->getMCParticle();
          if (mc_fsp) {
            int tag_index = finddescendant(mc_fsp);
            if (tag_index == index) {
              fsp_FEI_index.push_back(mc_fsp->getArrayIndex());
            }
          }
        }
        std::vector<int> diff;
        std::vector<int> fsp_truth_index = truthFSPTag(index);
        std::sort(fsp_truth_index.begin(), fsp_truth_index.end());
        std::sort(fsp_FEI_index.begin(), fsp_FEI_index.end());
        std::set_difference(fsp_truth_index.begin(), fsp_truth_index.end(), fsp_FEI_index.begin(), fsp_FEI_index.end(), std::inserter(diff,
                            diff.begin()));
        return double(diff.size()) / double(fsp_truth_index.size());
      }
    }
    double percentageWrongParticlesBTag(const Particle* part)
    {
      int index = mostcommonBTagIndex(part);
      if (index == -1) {
        return -1;
      } else {
        const std::vector<const Particle*>& fsp_tag = part->getFinalStateDaughters();
        int wrong_FEI = 0;
        for (const Particle* fsp : fsp_tag) {
          const MCParticle* mc_fsp = fsp->getMCParticle();
          if (mc_fsp) {
            int tag_index = finddescendant(mc_fsp);
            if (tag_index != index) {
              wrong_FEI ++;
            }
          }
        }
        return double(wrong_FEI) / double(truthFSPTag(index).size());
      }
    }
    double mostcommonBTagDeltaP(const Particle* part)
    {
      int idx = mostcommonBTagIndex(part);
      if (idx < 0) return Const::doubleNaN;
 
      StoreArray<MCParticle> mcParticles;
      PCmsLabTransform T;
      ROOT::Math::XYZVector recoP3 = (T.rotateLabToCms() * part->get4Vector()).Vect();
      ROOT::Math::XYZVector genP3 = (T.rotateLabToCms() * mcParticles[idx]->get4Vector()).Vect();
      return (recoP3 - genP3).R();
    }
 
    double mostcommonBTagPDG(const Particle* part)
    {
      int idx = mostcommonBTagIndex(part);
      if (idx < 0) return Const::doubleNaN;
 
      StoreArray<MCParticle> mcParticles;
      return mcParticles[idx]->getPDG();
    }
 
    Manager::FunctionPtr sigProbRank(const std::vector<std::string>& arguments)
    {
      if (arguments.size() != 2) {
        B2FATAL("sigProbRank requires exactly 2 arguments: bp_list and b0_list. "
                "Received " << arguments.size() << " arguments.");
      }
 
      std::string bpListName = arguments[0];
      std::string b0ListName = arguments[1];
 
      auto func = [bpListName, b0ListName](const Particle * particle) -> int {
        StoreObjPtr<ParticleList> bpList(bpListName);
        StoreObjPtr<ParticleList> b0List(b0ListName);
 
        if (!bpList.isValid() && !b0List.isValid())
          B2FATAL("sigProbRank could not find either input list: bp_list='" << bpListName
                  << "', b0_list='" << b0ListName << "'");
        if (!bpList.isValid())
          B2FATAL("sigProbRank could not find the required bp_list: '" << bpListName << "'");
        if (!b0List.isValid())
          B2FATAL("sigProbRank could not find the required b0_list: '" << b0ListName << "'");
 
        struct RankedCandidate {
          double signalProbability;
          const Particle* particle;
        };
 
        auto collectCandidates = [](const StoreObjPtr<ParticleList>& list,
                                    const std::string & listLabel,
                                    std::vector<RankedCandidate>& rankedCandidates)
        {
          for (unsigned int i = 0; i < list->getListSize(); ++i) {
            const Particle* p = list->getParticle(i);
            if (!p->hasExtraInfo("SignalProbability")) {
              B2FATAL("sigProbRank requires every candidate in " << listLabel
                      << " to provide extraInfo('SignalProbability')");
            }
            rankedCandidates.push_back({p->getExtraInfo("SignalProbability"), p});
          }
        };
 
        std::vector<RankedCandidate> rankedCandidates;
        rankedCandidates.reserve(bpList->getListSize() + b0List->getListSize());
        collectCandidates(bpList, bpListName, rankedCandidates);
        collectCandidates(b0List, b0ListName, rankedCandidates);
 
        std::stable_sort(rankedCandidates.begin(), rankedCandidates.end(),
                         [](const RankedCandidate & lhs, const RankedCandidate & rhs)
        {
          return lhs.signalProbability > rhs.signalProbability;
        });
 
        // Check if particle is a copy made by ParticleCopy::copyParticle
        // (e.g. from setBeamConstrainedMomentum); copies store the original's
        // array index in extraInfo("original_index")
        int originalIndex = -1;
        if (particle->hasExtraInfo("original_index"))
          originalIndex = static_cast<int>(particle->getExtraInfo("original_index"));
 
        for (unsigned int i = 0; i < rankedCandidates.size(); ++i)
        {
          const Particle* cand = rankedCandidates[i].particle;
          if (cand == particle) {
            return static_cast<int>(i) + 1;
          }
          if (originalIndex >= 0 && cand->getArrayIndex() == originalIndex) {
            return static_cast<int>(i) + 1;
          }
        }
 
        B2FATAL("sigProbRank: particle not found in the input lists: "
                "bp_list='" << bpListName << "', b0_list='" << b0ListName << "'");
        return -1;
      };
      return func;
    }
 
    VARIABLE_GROUP("FEIVariables");
    REGISTER_VARIABLE("mostcommonBTagIndex", mostcommonBTagIndex, R"DOC(
Returns the generator-level index of the :math:`B` meson candidate which is assigned the most reconstructed particles. Returns -1
if no :math:`B` meson on the generator level is found. 
 
)DOC");
    REGISTER_VARIABLE("percentageMissingParticlesBTag", percentageMissingParticlesBTag, R"DOC(
Returns the percentage of missing particles calculated by the number of particles not reconstructed by the FEI divided by the number
of generated particles. The latter is determined by using the generator-level index corresponding to the :math:`B` meson returned by
``mostcommonBTagIndex``. Returns -1 if no :math:`B` meson on the generator level is found. 
 
)DOC");
    REGISTER_VARIABLE("percentageWrongParticlesBTag", percentageWrongParticlesBTag, R"DOC(
Returns the percentage of wrong particles calculated by the number of reconstructed particles originating from the other :math:`B` meson divided by the total
number of generated final-state particles. The latter is determined by using the generator-level index corresponding to the :math:`B` meson returned by
``mostcommonBTagIndex``. Returns -1 if no :math:`B` meson on the generator level is found. 
 
)DOC");
    REGISTER_VARIABLE("mostcommonBTagDeltaP", mostcommonBTagDeltaP, R"DOC(
Returns the magnitude of the 3-momentum difference in the CM frame between the reconstructed particles and the generated :math:`B` meson identified using
``mostcommonBTagIndex``. Returns ``NaN`` if no :math:`B` meson on the generator level is found. 
 
)DOC", "GeV/c");
    REGISTER_VARIABLE("mostcommonBTagPDG", mostcommonBTagPDG, R"DOC(
Returns the PDG code of the generated :math:`B` meson identified using ``mostcommonBTagIndex``. Returns ``NaN`` if no :math:`B` meson on the generator level is found.
 
.. note:: This is equivalent to ``genParticle(mostcommonBTagIndex, PDG)``. 
 
)DOC");
    REGISTER_METAVARIABLE("sigProbRank(bp_list, b0_list)", sigProbRank, R"DOC(
Returns the rank (starting at 1) of the given candidate when all candidates from the :math:`B^+` (``bp_list``) and :math:`B^0` (``b0_list``) particle lists
are ordered together by descening signal probability. 
 
.. note:: The signal probability can be accessed using ``extraInfo(SignalProbability)``.
 
)DOC", Manager::VariableDataType::c_int);
  }
}