RecoTrack.h

/**************************************************************************
 * BASF2 (Belle Analysis Framework 2)                                     *
 * Copyright(C) 2014 - Belle II Collaboration                             *
 *                                                                        *
 * Author: The Belle II Collaboration                                     *
 * Contributors: Martin Heck, Nils Braun, Thomas Hauth                    *
 *                                                                        *
 * This software is provided "as is" without any warranty.                *
 **************************************************************************/
 
#pragma once
 
#include <framework/datastore/StoreArray.h>
#include <framework/datastore/RelationsObject.h>
#include <framework/core/FrameworkExceptions.h>
 
#include <genfit/Track.h>
 
#include <tracking/dataobjects/RecoHitInformation.h>
 
#include <boost/optional.hpp>
 
#include <string>
#include <vector>
 
namespace genfit {
  class TrackCand;
  class AbsTrackRep;
}
 
namespace Belle2 {
  class RecoTrackGenfitAccess;
 
  BELLE2_DEFINE_EXCEPTION(NoTrackFitResult, "No track fit result available for this hit (e.g. DAF has removed it).")
 
  
  class RecoTrack : public RelationsObject {
    friend class RecoTrackGenfitAccess;
 
  private:
    typedef RecoHitInformation::RightLeftInformation RightLeftInformation;
    typedef RecoHitInformation::RecoHitDetector TrackingDetector;
    typedef RecoHitInformation::OriginTrackFinder OriginTrackFinder;
    typedef RecoHitInformation::UsedCDCHit UsedCDCHit;
    typedef RecoHitInformation::UsedSVDHit UsedSVDHit;
    typedef RecoHitInformation::UsedPXDHit UsedPXDHit;
    typedef RecoHitInformation::UsedBKLMHit UsedBKLMHit;
    typedef RecoHitInformation::UsedEKLMHit UsedEKLMHit;
 
  public:
    enum MatchingStatus {
      c_undefined,
      c_matched,
      c_clone,
      c_background,
      c_ghost
    };
 
    static void registerRequiredRelations(
      StoreArray<RecoTrack>& recoTracks,
      std::string const& pxdHitsStoreArrayName = "",
      std::string const& svdHitsStoreArrayName = "",
      std::string const& cdcHitsStoreArrayName = "",
      std::string const& bklmHitsStoreArrayName = "",
      std::string const& eklmHitsStoreArrayName = "",
      std::string const& recoHitInformationStoreArrayName = "");
 
    RecoTrack() { }
 
    RecoTrack(const TVector3& seedPosition, const TVector3& seedMomentum, const short int seedCharge,
              const std::string& storeArrayNameOfCDCHits = "",
              const std::string& storeArrayNameOfSVDHits = "",
              const std::string& storeArrayNameOfPXDHits = "",
              const std::string& storeArrayNameOfBKLMHits = "",
              const std::string& storeArrayNameOfEKLMHits = "",
              const std::string& storeArrayNameOfRecoHitInformation = "");
 
    RecoTrack(const RecoTrack&) = delete;
    RecoTrack& operator=(RecoTrack const&) = delete;
 
    static RecoTrack* createFromTrackCand(const genfit::TrackCand& trackCand,
                                          const std::string& storeArrayNameOfRecoTracks = "",
                                          const std::string& storeArrayNameOfCDCHits = "",
                                          const std::string& storeArrayNameOfSVDHits = "",
                                          const std::string& storeArrayNameOfPXDHits = "",
                                          const std::string& storeArrayNameOfBKLMHits = "",
                                          const std::string& storeArrayNameOfEKLMHits = "",
                                          const std::string& storeArrayNameOfRecoHitInformation = "",
                                          const bool recreateSortingParameters = false
                                         );
 
    genfit::TrackCand createGenfitTrackCand() const;
 
    RecoTrack* copyToStoreArrayUsing(StoreArray<RecoTrack>& storeArray, const TVector3& position,
                                     const TVector3& momentum, short charge,
                                     const TMatrixDSym& covariance, double timeSeed) const;
 
    RecoTrack* copyToStoreArrayUsingSeeds(StoreArray<RecoTrack>& storeArray) const;
 
    RecoTrack* copyToStoreArray(StoreArray<RecoTrack>& storeArray) const;
 
    size_t addHitsFromRecoTrack(const RecoTrack* recoTrack, unsigned int sortingParameterOffset = 0,
                                bool reversed = false, boost::optional<double> optionalMinimalWeight = boost::none);
 
    bool addCDCHit(const UsedCDCHit* cdcHit, const unsigned int sortingParameter,
                   RightLeftInformation rightLeftInformation = RightLeftInformation::c_undefinedRightLeftInformation,
                   OriginTrackFinder foundByTrackFinder = OriginTrackFinder::c_undefinedTrackFinder)
    {
      return addHit(cdcHit, rightLeftInformation, foundByTrackFinder, sortingParameter);
    }
 
    bool addPXDHit(const UsedPXDHit* pxdHit, const unsigned int sortingParameter,
                   OriginTrackFinder foundByTrackFinder = OriginTrackFinder::c_undefinedTrackFinder)
    {
      return addHit(pxdHit, foundByTrackFinder, sortingParameter);
    }
 
    bool addSVDHit(const UsedSVDHit* svdHit, const unsigned int sortingParameter,
                   OriginTrackFinder foundByTrackFinder = OriginTrackFinder::c_undefinedTrackFinder)
    {
      return addHit(svdHit, foundByTrackFinder, sortingParameter);
    }
 
    bool addBKLMHit(const UsedBKLMHit* bklmHit, const unsigned int sortingParameter,
                    OriginTrackFinder foundByTrackFinder = OriginTrackFinder::c_undefinedTrackFinder)
    {
      return addHit(bklmHit, foundByTrackFinder, sortingParameter);
    }
 
    bool addEKLMHit(const UsedEKLMHit* eklmHit, const unsigned int sortingParameter,
                    OriginTrackFinder foundByTrackFinder = OriginTrackFinder::c_undefinedTrackFinder)
    {
      return addHit(eklmHit, foundByTrackFinder, sortingParameter);
    }
 
    template<class HitType>
    RecoHitInformation* getRecoHitInformation(HitType* hit) const
    {
      RelationVector<RecoHitInformation> relatedHitInformationToHit = hit->template getRelationsFrom<RecoHitInformation>
      (m_storeArrayNameOfRecoHitInformation);
 
      for (RecoHitInformation& recoHitInformation : relatedHitInformationToHit) {
        // cppcheck-suppress useStlAlgorithm
        if (recoHitInformation.getRelatedFrom<RecoTrack>(this->getArrayName()) == this) {
          return &recoHitInformation;
        }
      }
 
      return nullptr;
    }
 
    // Hits Information Questioning
    template <class HitType>
    TrackingDetector getTrackingDetector(const HitType* hit) const
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      return recoHitInformation->getTrackingDetector();
    }
 
    template <class HitType>
    RightLeftInformation getRightLeftInformation(const HitType* hit) const
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      return recoHitInformation->getRightLeftInformation();
    }
 
    template <class HitType>
    OriginTrackFinder getFoundByTrackFinder(const HitType* hit) const
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      return recoHitInformation->getFoundByTrackFinder();
    }
 
    template <class HitType>
    unsigned int getSortingParameter(const HitType* hit) const
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      return recoHitInformation->getSortingParameter();
    }
 
    template <class HitType>
    void setRightLeftInformation(const HitType* hit, RightLeftInformation rightLeftInformation)
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      recoHitInformation->setRightLeftInformation(rightLeftInformation);
      setDirtyFlag();
    }
 
    template <class HitType>
    void setFoundByTrackFinder(const HitType* hit, OriginTrackFinder originTrackFinder)
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      recoHitInformation->setFoundByTrackFinder(originTrackFinder);
    }
 
    template <class HitType>
    void setSortingParameter(const HitType* hit, unsigned int sortingParameter)
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformationSafely(hit);
      recoHitInformation->setSortingParameter(sortingParameter);
      setDirtyFlag();
    }
 
    const genfit::TrackPoint* getCreatedTrackPoint(const RecoHitInformation* recoHitInformation) const;
 
    // Hits Added Questioning
    bool hasCDCHits() const { return getRelatedFrom<UsedCDCHit>(m_storeArrayNameOfCDCHits) != nullptr; }
 
    bool hasSVDHits() const { return getRelatedFrom<UsedSVDHit>(m_storeArrayNameOfSVDHits) != nullptr; }
 
    bool hasPXDHits() const { return getRelatedFrom<UsedPXDHit>(m_storeArrayNameOfPXDHits) != nullptr; }
 
    bool hasBKLMHits() const { return getRelatedFrom<UsedBKLMHit>(m_storeArrayNameOfBKLMHits) != nullptr; }
 
    bool hasEKLMHits() const { return getRelatedFrom<UsedEKLMHit>(m_storeArrayNameOfEKLMHits) != nullptr; }
 
    template <class HitType>
    bool hasHit(const HitType* hit) const
    {
      const RelationVector<RecoTrack>& relatedTracksToHit = hit->template getRelationsTo<RecoTrack>(getArrayName());
      return std::find_if(relatedTracksToHit.begin(), relatedTracksToHit.end(), [this](const RecoTrack & recoTrack) {
        return &recoTrack == this;
      }) != relatedTracksToHit.end();
    }
 
    // Hits Questioning
    unsigned int getNumberOfCDCHits() const { return getNumberOfHitsOfGivenType<UsedCDCHit>(m_storeArrayNameOfCDCHits); }
 
    unsigned int getNumberOfSVDHits() const { return getNumberOfHitsOfGivenType<UsedSVDHit>(m_storeArrayNameOfSVDHits); }
 
    unsigned int getNumberOfPXDHits() const { return getNumberOfHitsOfGivenType<UsedPXDHit>(m_storeArrayNameOfPXDHits); }
 
    unsigned int getNumberOfBKLMHits() const { return getNumberOfHitsOfGivenType<UsedBKLMHit>(m_storeArrayNameOfBKLMHits); }
 
    unsigned int getNumberOfEKLMHits() const { return getNumberOfHitsOfGivenType<UsedEKLMHit>(m_storeArrayNameOfEKLMHits); }
 
    unsigned int getNumberOfTotalHits() const
    {
      return getNumberOfCDCHits() + getNumberOfPXDHits() +
             getNumberOfSVDHits() + getNumberOfBKLMHits() + getNumberOfEKLMHits();
    }
 
    unsigned int getNumberOfTrackingHits() const
    {
      return getNumberOfCDCHits() + getNumberOfPXDHits() +
             getNumberOfSVDHits();
    }
 
    std::vector<Belle2::RecoTrack::UsedCDCHit*> getCDCHitList() const { return getHitList<UsedCDCHit>(m_storeArrayNameOfCDCHits); }
 
    std::vector<Belle2::RecoTrack::UsedSVDHit*> getSVDHitList() const { return getHitList<UsedSVDHit>(m_storeArrayNameOfSVDHits); }
 
    std::vector<Belle2::RecoTrack::UsedPXDHit*> getPXDHitList() const { return getHitList<UsedPXDHit>(m_storeArrayNameOfPXDHits); }
 
    std::vector<Belle2::RecoTrack::UsedBKLMHit*> getBKLMHitList() const { return getHitList<UsedBKLMHit>(m_storeArrayNameOfBKLMHits); }
 
    std::vector<Belle2::RecoTrack::UsedEKLMHit*> getEKLMHitList() const { return getHitList<UsedEKLMHit>(m_storeArrayNameOfEKLMHits); }
 
    std::vector<Belle2::RecoTrack::UsedCDCHit*> getSortedCDCHitList() const { return getSortedHitList<UsedCDCHit>(m_storeArrayNameOfCDCHits); }
 
    std::vector<Belle2::RecoTrack::UsedSVDHit*> getSortedSVDHitList() const { return getSortedHitList<UsedSVDHit>(m_storeArrayNameOfSVDHits); }
 
    std::vector<Belle2::RecoTrack::UsedPXDHit*> getSortedPXDHitList() const { return getSortedHitList<UsedPXDHit>(m_storeArrayNameOfPXDHits); }
 
    std::vector<Belle2::RecoTrack::UsedBKLMHit*> getSortedBKLMHitList() const { return getSortedHitList<UsedBKLMHit>(m_storeArrayNameOfBKLMHits); }
 
    std::vector<Belle2::RecoTrack::UsedEKLMHit*> getSortedEKLMHitList() const { return getSortedHitList<UsedEKLMHit>(m_storeArrayNameOfEKLMHits); }
 
    // Seed Helix Functionality
    TVector3 getPositionSeed() const
    {
      const TVectorD& seed = m_genfitTrack.getStateSeed();
      return TVector3(seed(0), seed(1), seed(2));
    }
 
    TVector3 getMomentumSeed() const
    {
      const TVectorD& seed = m_genfitTrack.getStateSeed();
      return TVector3(seed(3), seed(4), seed(5));
    }
 
    const TVectorD& getStateSeed() const
    {
      return m_genfitTrack.getStateSeed();
    }
 
    short int getChargeSeed() const { return m_charge; }
 
    double getTimeSeed() const { return m_genfitTrack.getTimeSeed(); }
 
    std::tuple<TVector3, TVector3, short> extractTrackState() const;
 
    void setPositionAndMomentum(const TVector3& positionSeed, const TVector3& momentumSeed)
    {
      m_genfitTrack.setStateSeed(positionSeed, momentumSeed);
      deleteFittedInformation();
    }
 
    void setChargeSeed(const short int chargeSeed)
    {
      m_charge = chargeSeed;
      deleteFittedInformation();
    }
 
    void setTimeSeed(const double timeSeed)
    {
      m_genfitTrack.setTimeSeed(timeSeed);
      deleteFittedInformation();
    }
 
    const TMatrixDSym& getSeedCovariance() const { return m_genfitTrack.getCovSeed(); }
 
    void setSeedCovariance(const TMatrixDSym& seedCovariance) { m_genfitTrack.setCovSeed(seedCovariance); }
 
    // Fitting
    bool wasFitSuccessful(const genfit::AbsTrackRep* representation = nullptr) const;
 
    const genfit::FitStatus* getTrackFitStatus(const genfit::AbsTrackRep* representation = nullptr) const
    {
      checkDirtyFlag();
      return m_genfitTrack.getFitStatus(representation);
    }
 
    bool hasTrackFitStatus(const genfit::AbsTrackRep* representation = nullptr) const;
 
    genfit::AbsTrackRep* getCardinalRepresentation() const
    {
      checkDirtyFlag();
      return m_genfitTrack.getCardinalRep();
    }
 
    const std::vector<genfit::AbsTrackRep*>& getRepresentations() const
    {
      checkDirtyFlag();
      return m_genfitTrack.getTrackReps();
    }
 
    genfit::AbsTrackRep* getTrackRepresentationForPDG(int pdgCode);
 
    std::vector<RecoHitInformation*> getRecoHitInformations(bool getSorted = false) const;
 
    const genfit::MeasuredStateOnPlane& getMeasuredStateOnPlaneFromFirstHit(const genfit::AbsTrackRep* representation = nullptr) const;
 
    const genfit::MeasuredStateOnPlane& getMeasuredStateOnPlaneFromLastHit(const genfit::AbsTrackRep* representation = nullptr) const;
 
    const genfit::MeasuredStateOnPlane& getMeasuredStateOnPlaneFromRecoHit(const RecoHitInformation* recoHitInfo,
        const genfit::AbsTrackRep* representation = nullptr) const;
 
    const genfit::MeasuredStateOnPlane& getMeasuredStateOnPlaneClosestTo(const TVector3& closestPoint,
        const genfit::AbsTrackRep* representation = nullptr);
 
    void prune();
 
    const std::vector<genfit::TrackPoint*>& getHitPointsWithMeasurement() const
    {
      checkDirtyFlag();
      return m_genfitTrack.getPointsWithMeasurement();
    }
 
    bool getDirtyFlag() const { return m_dirtyFlag; }
 
 
    void setDirtyFlag(const bool& dirtyFlag = true)
    {
      m_dirtyFlag = dirtyFlag;
      if (dirtyFlag) {
        deleteFittedInformation();
      }
    }
 
    // Store Array Names
    const std::string& getStoreArrayNameOfCDCHits() const { return m_storeArrayNameOfCDCHits; }
 
    const std::string& getStoreArrayNameOfSVDHits() const { return m_storeArrayNameOfSVDHits; }
 
    const std::string& getStoreArrayNameOfPXDHits() const { return m_storeArrayNameOfPXDHits; }
 
    const std::string& getStoreArrayNameOfBKLMHits() const { return m_storeArrayNameOfBKLMHits; }
 
    const std::string& getStoreArrayNameOfEKLMHits() const { return m_storeArrayNameOfEKLMHits; }
 
    const std::string& getStoreArrayNameOfRecoHitInformation() const { return m_storeArrayNameOfRecoHitInformation; }
 
    template<class HitType>
    void mapOnHits(const std::string& storeArrayNameOfHits,
                   std::function<void(RecoHitInformation&, HitType*)> const&   mapFunction,
                   std::function<bool(const RecoHitInformation&, const HitType*)> const& pickFunction)
    {
      RelationVector<RecoHitInformation> relatedHitInformation = getRelationsTo<RecoHitInformation>
                                                                 (m_storeArrayNameOfRecoHitInformation);
 
      for (RecoHitInformation& hitInformation : relatedHitInformation) {
        HitType* const hit = hitInformation.getRelatedTo<HitType>(storeArrayNameOfHits);
        if (hit != nullptr && pickFunction(hitInformation, hit)) {
          mapFunction(hitInformation, hit);
        }
      }
    }
 
    template<class HitType>
    void mapOnHits(const std::string& storeArrayNameOfHits,
                   std::function<void(const RecoHitInformation&, const HitType*)> const& mapFunction,
                   std::function<bool(const RecoHitInformation&, const HitType*)> const& pickFunction) const
    {
      RelationVector<RecoHitInformation> relatedHitInformation = getRelationsTo<RecoHitInformation>
                                                                 (m_storeArrayNameOfRecoHitInformation);
 
      for (const RecoHitInformation& hitInformation : relatedHitInformation) {
        const HitType* const hit = hitInformation.getRelatedTo<HitType>(storeArrayNameOfHits);
        if (hit != nullptr && pickFunction(hitInformation, hit)) {
          mapFunction(hitInformation, hit);
        }
      }
    }
 
    template<class HitType>
    void mapOnHits(const std::string& storeArrayNameOfHits,
                   std::function<void(RecoHitInformation&, HitType*)> const& mapFunction)
    {
      mapOnHits<HitType>(storeArrayNameOfHits, mapFunction, [](const RecoHitInformation&, const HitType*) -> bool { return true; });
    }
 
    template<class HitType>
    void mapOnHits(const std::string& storeArrayNameOfHits,
                   std::function<void(const RecoHitInformation&, const HitType*)> const&   mapFunction) const
    {
      mapOnHits<HitType>(storeArrayNameOfHits, mapFunction, [](const RecoHitInformation&, const HitType*) -> bool { return true; });
    }
 
    // Matching status
    MatchingStatus getMatchingStatus() const
    {
      return m_matchingStatus;
    }
 
    void setMatchingStatus(MatchingStatus matchingStatus)
    {
      m_matchingStatus = matchingStatus;
    }
 
    float getQualityIndicator() const
    {
      return m_qualityIndicator;
    }
 
    void setQualityIndicator(const float qualityIndicator)
    {
      m_qualityIndicator = qualityIndicator;
    }
 
    void deleteFittedInformation();
 
    void deleteFittedInformationForRepresentation(const genfit::AbsTrackRep* rep);
 
    virtual std::string getInfoHTML() const;
 
  private:
    genfit::Track m_genfitTrack;
    short int m_charge = 1;
    std::string m_storeArrayNameOfCDCHits = "";
    std::string m_storeArrayNameOfSVDHits = "";
    std::string m_storeArrayNameOfPXDHits = "";
    std::string m_storeArrayNameOfBKLMHits = "";
    std::string m_storeArrayNameOfEKLMHits = "";
    std::string m_storeArrayNameOfRecoHitInformation = "";
    bool m_dirtyFlag = true;
    MatchingStatus m_matchingStatus = MatchingStatus::c_undefined;
    float m_qualityIndicator = NAN;
 
    template<class HitType, class ...Args>
    bool addHit(const HitType* hit, Args&& ... params)
    {
      if (hasHit(hit)) {
        return false;
      }
 
      StoreArray<RecoHitInformation> recoHitInformations(m_storeArrayNameOfRecoHitInformation);
      RecoHitInformation* newRecoHitInformation = recoHitInformations.appendNew(hit, params...);
 
      addHitWithHitInformation(hit, newRecoHitInformation);
 
      return true;
    }
 
    template <class HitType>
    void addHitWithHitInformation(const HitType* hit, RecoHitInformation* recoHitInformation)
    {
      hit->addRelationTo(this);
      addRelationTo(recoHitInformation);
 
      setDirtyFlag();
    }
 
    template <class HitType>
    RecoHitInformation* getRecoHitInformationSafely(HitType* hit) const
    {
      RecoHitInformation* recoHitInformation = getRecoHitInformation(hit);
      if (recoHitInformation == nullptr) {
        B2FATAL("Queried hit is not in the reco track! Did you prune it?");
      } else {
        return recoHitInformation;
      }
    }
 
    template <class HitType>
    unsigned int getNumberOfHitsOfGivenType(const std::string& storeArrayNameOfHits) const
    {
      return getRelationsFrom<HitType>(storeArrayNameOfHits).size();
    }
 
    template<class HitType>
    std::vector<HitType*> getSortedHitList(const std::string& storeArrayNameOfHits) const
    {
      const RelationVector<RecoHitInformation>& relatedHitInformation = getRelationsTo<RecoHitInformation>
          (m_storeArrayNameOfRecoHitInformation);
 
      std::vector<const RecoHitInformation*> relatedHitInformationAsVector;
      relatedHitInformationAsVector.reserve(relatedHitInformation.size());
      for (const RecoHitInformation& hitInformation : relatedHitInformation) {
        // cppcheck-suppress useStlAlgorithm
        relatedHitInformationAsVector.push_back(&hitInformation);
      }
      std::sort(relatedHitInformationAsVector.begin(), relatedHitInformationAsVector.end(), [](const RecoHitInformation * a,
      const RecoHitInformation * b) -> bool {
        return a->getSortingParameter() < b->getSortingParameter();
      });
 
      std::vector<HitType*> hitList;
      hitList.reserve(relatedHitInformationAsVector.size());
      for (const RecoHitInformation* hitInformation : relatedHitInformationAsVector) {
        HitType* relatedHit = hitInformation->getRelatedTo<HitType>(storeArrayNameOfHits);
        if (relatedHit != nullptr) {
          hitList.push_back(relatedHit);
        }
      }
 
      return hitList;
    }
 
    template<class HitType>
    std::vector<HitType*> getHitList(const std::string& storeArrayNameOfHits) const
    {
      RelationVector<HitType> relatedHits = getRelationsFrom<HitType>(storeArrayNameOfHits);
      std::vector<HitType*> hitList;
      hitList.reserve(relatedHits.size());
      for (HitType& hit : relatedHits) {
        // cppcheck-suppress useStlAlgorithm
        hitList.push_back(&hit);
      }
 
      return hitList;
    }
 
    void checkDirtyFlag() const
    {
      if (m_dirtyFlag) {
        B2DEBUG(100, "Dirty flag is set. The result may not be in sync with the latest changes. Refit the track to be sure.");
      }
    }
 
    ClassDef(RecoTrack, 9);
  };
 
  class RecoTrackGenfitAccess {
  public:
    static genfit::Track& getGenfitTrack(RecoTrack& recoTrack);
 
    static genfit::AbsTrackRep* createOrReturnRKTrackRep(RecoTrack& recoTrack, int PDGcode);
  };
 
}