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,
                      "By giving e.g. a FEI B meson candidate the B meson index on generator level is determined, where most reconstructed particles can be assigned to. If no B meson found on generator level -1 is returned.");
    REGISTER_VARIABLE("percentageMissingParticlesBTag", percentageMissingParticlesBTag,
                      "Get the percentage of missing particles by using the mostcommonBTagIndex. So the number of particles not reconstructed by e.g. the FEI are determined and divided by the number of generated particles using the given index of the B meson. If no B meson found on generator level -1 is returned.");
    REGISTER_VARIABLE("percentageWrongParticlesBTag", percentageWrongParticlesBTag,
                      "Get the percentage of wrong particles by using the mostcommonBTagIndex. In this context wrong means that the reconstructed particles originated from the other B meson. The absolute number is divided by the total number of generated FSP from the given B meson index. If no B meson found on generator level -1 is returned.");
    REGISTER_VARIABLE("mostcommonBTagDeltaP", mostcommonBTagDeltaP,
                      "Returns the magnitude of the 3-momentum difference in CMS frame between the "
                      "reconstructed particle and the generated B meson identified by mostcommonBTagIndex. "
                      "Returns NaN if no B meson is found on generator level.", "GeV/c");
    REGISTER_VARIABLE("mostcommonBTagPDG", mostcommonBTagPDG,
                      "Returns the PDG code of the generated B meson identified by mostcommonBTagIndex. "
                      "Returns NaN if no B meson is found on generator level. "
                      "Note: this is equivalent to ``genParticle(mostcommonBTagIndex, PDG)``. "
                      "Other variables can be accessed the same way by replacing ``PDG`` with any variable.");
    REGISTER_METAVARIABLE("sigProbRank(bp_list, b0_list)", sigProbRank,
                          "Returns the rank (starting at 1) of the candidate when all candidates from the B+ list (bp_list) "
                          "and B0 list (b0_list) are ordered together by descending ``extraInfo(SignalProbability)``.", Manager::VariableDataType::c_int);
  }
}