TrackMatchLookUp.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 <mdst/dataobjects/MCParticle.h>
#include <mdst/dataobjects/TrackFitResult.h>
 
#include <tracking/dataobjects/RecoTrack.h>
 
namespace Belle2 {
  class TrackMatchLookUp {
 
  public:
    enum class MCToPRMatchInfo { c_undefined, c_matched, c_matchedWrongCharge, c_merged, c_mergedWrongCharge, c_missing };
 
    enum class PRToMCMatchInfo { c_undefined, c_matched, c_matchedWrongCharge, c_clone, c_cloneWrongCharge, c_background, c_ghost };
 
  public:
    explicit TrackMatchLookUp(const std::string& mcRecoTrackStoreArrayName,
                              const std::string& prRecoTrackStoreArrayName = "");
 
  private:
    MCToPRMatchInfo extractMCToPRMatchInfo(const RecoTrack& mcRecoTrack,
                                           const RecoTrack* prRecoTrack,
                                           const float& efficiency) const;
 
    PRToMCMatchInfo extractPRToMCMatchInfo(const RecoTrack& prRecoTrack,
                                           const RecoTrack* mcRecoTrack,
                                           const float& purity) const;
 
    const RecoTrack* getMCRecoTrackWithStatus(const RecoTrack& prRecoTrack,
                                              const PRToMCMatchInfo matchingStatus = PRToMCMatchInfo::c_matched) const;
 
    const RecoTrack* getPRRecoTrackWithStatus(const RecoTrack& mcRecoTrack,
                                              const MCToPRMatchInfo matchingStatus = MCToPRMatchInfo::c_matched) const;
 
  public:
    const RecoTrack* getCorrectChargeMatchedMCRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getMCRecoTrackWithStatus(prRecoTrack, PRToMCMatchInfo::c_matched);
    }
 
    const RecoTrack* getWrongChargeMatchedMCRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getMCRecoTrackWithStatus(prRecoTrack, PRToMCMatchInfo::c_matchedWrongCharge);
    }
 
    const RecoTrack* getAnyChargeMatchedMCRecoTrack(const RecoTrack& prRecoTrack) const
    {
      const RecoTrack* anyMatchedRecoTrack = getMCRecoTrackWithStatus(prRecoTrack, PRToMCMatchInfo::c_matched);
      if (anyMatchedRecoTrack == nullptr) {
        anyMatchedRecoTrack = getMCRecoTrackWithStatus(prRecoTrack, PRToMCMatchInfo::c_matchedWrongCharge);
      }
      return anyMatchedRecoTrack;
    }
 
    const RecoTrack* getCorrectChargeMatchedPRRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getPRRecoTrackWithStatus(mcRecoTrack, MCToPRMatchInfo::c_matched);
    }
 
    const RecoTrack* getWrongChargeMatchedPRRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getPRRecoTrackWithStatus(mcRecoTrack, MCToPRMatchInfo::c_matchedWrongCharge);
    }
 
    const RecoTrack* getAnyChargeMatchedPRRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      const RecoTrack* anyMatchedRecoTrack = getPRRecoTrackWithStatus(mcRecoTrack, MCToPRMatchInfo::c_matched);
      if (anyMatchedRecoTrack == nullptr) {
        anyMatchedRecoTrack = getPRRecoTrackWithStatus(mcRecoTrack, MCToPRMatchInfo::c_matchedWrongCharge);
      }
      return anyMatchedRecoTrack;
    }
 
  public:
    float getMatchedPurity(const RecoTrack& recoTrack) const;
 
    float getMatchedEfficiency(const RecoTrack& recoTrack) const;
 
 
    bool isChargeMatched(const RecoTrack& recoTrack) const;
 
  public:
    MCToPRMatchInfo getMCToPRMatchInfo(const RecoTrack& mcRecoTrack) const
    {
      float efficiency = NAN;
      const RecoTrack* prRecoTrack = getRelatedPRRecoTrack(mcRecoTrack, efficiency);
      return extractMCToPRMatchInfo(mcRecoTrack, prRecoTrack, efficiency);
    }
 
    bool isCorrectChargeMatchedPRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_matched;
    }
 
    bool isWrongChargeMatchedPRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_matchedWrongCharge;
    }
 
    bool isAnyChargeMatchedPRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_matched or
             getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_matchedWrongCharge;
    }
 
    bool isCorrectChargeClonePRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_clone;
    }
 
    bool isWrongChargeClonePRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_cloneWrongCharge;
    }
 
    bool isAnyChargeClonePRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_clone or
             getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_cloneWrongCharge;
    }
 
    bool isBackgroundPRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_background;
    }
 
    bool isGhostPRRecoTrack(const RecoTrack& prRecoTrack) const
    {
      return getPRToMCMatchInfo(prRecoTrack) == PRToMCMatchInfo::c_ghost;
    }
 
  public:
    PRToMCMatchInfo getPRToMCMatchInfo(const RecoTrack& prRecoTrack) const
    {
      float purity = NAN;
      const RecoTrack* mcRecoTrack = getRelatedMCRecoTrack(prRecoTrack, purity);
      return extractPRToMCMatchInfo(prRecoTrack, mcRecoTrack, purity);
    }
 
    bool isCorrectChargeMatchedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_matched;
    }
 
    bool isWrongChargeMatchedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_matchedWrongCharge;
    }
 
    bool isAnyChargeMatchedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_matched or
             getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_matchedWrongCharge;
    }
 
    bool isCorrectChargeMergedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_merged;
    }
 
    bool isWrongChargeMergedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_mergedWrongCharge;
    }
 
    bool isAnyChargeMergedMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_merged or
             getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_mergedWrongCharge;
    }
 
    bool isMissingMCRecoTrack(const RecoTrack& mcRecoTrack) const
    {
      return getMCToPRMatchInfo(mcRecoTrack) == MCToPRMatchInfo::c_missing;
    }
 
  public:
    const MCParticle* getRelatedMCParticle(const RecoTrack& recoTrack) const;
 
    const TrackFitResult*
    getRelatedTrackFitResult(const RecoTrack& prRecoTrack,
                             Const::ChargedStable chargedStable = Const::pion) const;
 
    const RecoTrack* getRelatedMCRecoTrack(const RecoTrack& prRecoTrack) const;
 
    const RecoTrack* getRelatedPRRecoTrack(const RecoTrack& mcRecoTrack) const;
 
  public:
    float getRelatedPurity(const RecoTrack& prRecoTrack) const;
 
    float getRelatedEfficiency(const RecoTrack& mcRecoTrack) const;
 
    const RecoTrack* getRelatedMCRecoTrack(const RecoTrack& prRecoTrack, float& purity) const;
 
    const RecoTrack* getRelatedPRRecoTrack(const RecoTrack& mcRecoTrack, float& efficiency) const;
 
  public:
    bool isMCRecoTrack(const RecoTrack& recoTrack) const;
 
    bool isPRRecoTrack(const RecoTrack& recoTrack) const;
 
  public:
    const std::string& getMCTracksStoreArrayName() const
    {
      return m_mcTracksStoreArrayName;
    }
 
    const std::string& getPRTracksStoreArrayName() const
    {
      return m_prTracksStoreArrayName;
    }
 
  private:
    std::string m_mcTracksStoreArrayName;
 
    std::string m_prTracksStoreArrayName;
  };
}