NewV0Fitter.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 <framework/datastore/StoreArray.h>
#include <framework/gearbox/Const.h>
#include <mdst/dataobjects/Track.h>
#include <mdst/dataobjects/TrackFitResult.h>
#include <mdst/dataobjects/V0.h>
#include <tracking/dataobjects/V0ValidationVertex.h>
#include <tracking/dataobjects/RecoTrack.h>
#include <genfit/Track.h>
#include <genfit/GFRaveVertex.h>
#include <map>
 
namespace Belle2 {
  class NewV0Fitter {
 
  public:
 
    struct FittedTrack {
      const RecoTrack* recoTrack = nullptr; 
      Const::ParticleType ptype = Const::invalidParticle; 
      genfit::MeasuredStateOnPlane state; 
      double pValue = 0; 
      int Ndf = 0; 
      void set(const RecoTrack* recoTrk, const Const::ParticleType& hypo, const genfit::MeasuredStateOnPlane& mSoP,
               const genfit::AbsTrackRep* rep)
      {
        recoTrack = recoTrk;
        ptype = hypo;
        state = mSoP;
        const auto* fitStatus = recoTrk->getGenfitTrack().getFitStatus(rep);
        pValue = fitStatus->getPVal();
        Ndf = fitStatus->getNdf();
      }
    };
 
 
    NewV0Fitter(const std::string& trackFitResultsName = "",
                const std::string& v0sName = "",
                const std::string& v0ValidationVerticesName = "",
                const std::string& recoTracksName = "",
                const std::string& copiedRecoTracksName = "CopiedRecoTracks",
                bool enableValidation = false);
 
    void initializeCuts(double vertexDistanceCut,
                        double vertexChi2Cut,
                        const std::tuple<double, double>& invMassRangeKshort,
                        const std::tuple<double, double>& invMassRangeLambda,
                        const std::tuple<double, double>& invMassRangePhoton);
 
    void setFitterMode(int fitterMode) {m_fitterMode = fitterMode;}
 
    bool fitAndStore(const Track* trackPlus, const Track* trackMinus, const Const::ParticleType& v0Hypothesis,
                     bool& isForceStored, bool& isHitRemoved);
 
    static std::pair<Const::ParticleType, Const::ParticleType> getTrackHypotheses(const Const::ParticleType& v0Hypothesis);
 
 
  private:
 
    enum ErrorStatus {
      c_NoTrackRepresentation = -4, 
      c_VertexFitFailed = -3,       
      c_ExtrapolationFailed = -2,   
      c_NotSelected = -1            
    };
 
 
    enum EInnerHitBits {
      c_BitTrackPlus = 0x1, 
      c_BitTrackMinus = 0x2 
    };
 
 
    int vertexFit(const RecoTrack* recoTrackPlus, const RecoTrack* recoTrackMinus,
                  int pdgTrackPlus, int pdgTrackMinus, const Const::ParticleType& v0Hypothesis);
 
    const genfit::AbsTrackRep* getTrackRepresentation(const RecoTrack* recoTrack, int pdgCode);
 
    bool setCardinalRep(genfit::Track& gfTrack, int pdgCode);
 
    bool fitGFRaveVertex(genfit::Track& trackPlus, genfit::Track& trackMinus, genfit::GFRaveVertex& vertex);
 
    int extrapolateToVertex(genfit::MeasuredStateOnPlane& statePlus, genfit::MeasuredStateOnPlane& stateMinus,
                            const genfit::GFRaveVertex& vertex);
 
    RecoTrack* copyRecoTrack(const RecoTrack* origRecoTrack, const genfit::MeasuredStateOnPlane& state);
 
    const RecoTrack* removeHitsAndRefit(const RecoTrack* origRecoTrack, const RecoTrack* lastRecoTrack,
                                        const Const::ParticleType& ptype);
 
    int isInnermostClusterShared(const RecoTrack* recoTrackPlus, const RecoTrack* recoTrackMinus);
 
    const TrackFitResult* saveTrackFitResult(const FittedTrack& trk, int sharedInnermostCluster);
 
 
    std::string m_recoTracksName;                    
    StoreArray<RecoTrack> m_recoTracks;              
    StoreArray<TrackFitResult> m_trackFitResults;    
    StoreArray<V0> m_v0s;                            
    StoreArray<V0ValidationVertex> m_validationV0s;  
    StoreArray<RecoTrack> m_copiedRecoTracks;        
    double m_vertexDistanceCut = 0;  
    double m_vertexChi2Cut = 0;   
    std::map<int, std::pair<double, double> > m_invMassCuts; 
    int m_fitterMode = 1;  
    bool m_validation = false; 
    // temporary storage
 
    genfit::GFRaveVertex m_fittedVertex; 
    double m_momentum = 0; 
    double m_invMass = 0; 
    FittedTrack m_trkPlus; 
    FittedTrack m_trkMinus; 
  };
 
}