PlanarVXDMomentumMeasurement.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 <tracking/trackFitting/measurementCreator/measurements/PlanarMomentumMeasurement.h>
#include <tracking/vxdMomentumEstimation/VXDMomentumEstimation.h>
#include <tracking/vxdMomentumEstimation/VXDMomentumEstimationTools.h>
#include <mdst/dataobjects/MCParticle.h>
#include <tracking/dataobjects/RecoTrack.h>
 
namespace Belle2 {
  template <class HitType>
  class PlanarVXDMomentumMeasurement : public PlanarMomentumMeasurement {
 
  public:
    PlanarVXDMomentumMeasurement(const genfit::PlanarMeasurement& parentElement, HitType* hit, const Belle2::RecoTrack* recoTrack) :
      PlanarMomentumMeasurement(parentElement), m_hit(hit), m_recoTrack(recoTrack)
    {
      rawHitCoords_.ResizeTo(1);
      rawHitCov_.ResizeTo(1, 1);
 
      setHitId(hit->getArrayIndex());
      setDetId(getDetId());
    }
 
    virtual genfit::AbsMeasurement* clone() const override {return new PlanarVXDMomentumMeasurement(*this);}
 
    virtual std::vector<genfit::MeasurementOnPlane*> constructMeasurementsOnPlane(const genfit::StateOnPlane& state) const override;
 
    void setCorrectionFitParameters(const typename VXDMomentumEstimation<HitType>::CorrectionFitParameters& correctionFitParameters)
    {
      m_correctionFitParameters = correctionFitParameters;
    }
 
    void setSigma(double sigma)
    {
      m_sigma = sigma;
    }
 
    void setFitParameters(const typename VXDMomentumEstimation<HitType>::FitParameters& fitParameters)
    {
      m_fitParameters = fitParameters;
    }
 
    void setUseMCInformation(bool useMCInformation)
    {
      m_useMCInformation = useMCInformation;
    }
 
    void setUseThickness(bool useThickness)
    {
      m_useThickness = useThickness;
    }
 
    void setUseTrackingSeeds(bool useTrackingSeeds)
    {
      m_useTrackingSeeds = useTrackingSeeds;
    }
 
    HitType* getHit() const
    {
      return m_hit;
    }
 
  private:
    typename VXDMomentumEstimation<HitType>::FitParameters m_fitParameters;
    typename VXDMomentumEstimation<HitType>::CorrectionFitParameters m_correctionFitParameters;
    bool m_useThickness = false;
    bool m_useMCInformation = true;
    bool m_useTrackingSeeds = false;
    double m_sigma = 0.03;
 
    HitType* m_hit = nullptr;
    const RecoTrack* m_recoTrack = nullptr;
 
    int getDetId() const
    {
      return -1;
    }
  };
 
  template <class HitType>
  std::vector<genfit::MeasurementOnPlane*> PlanarVXDMomentumMeasurement<HitType>::constructMeasurementsOnPlane(
    const genfit::StateOnPlane& state) const
  {
    const VXDMomentumEstimation<HitType>& momentumEstimator = VXDMomentumEstimation<HitType>::getInstance();
    const VXDMomentumEstimationTools<HitType>& momentumEstimationTools = VXDMomentumEstimationTools<HitType>::getInstance();
 
    TVectorD rawHitCoordinates(1);
    TMatrixDSym rawHitCovariance(1);
 
    // Copy the information from the current state
    const B2Vector3D& statePosition(state.getPos());
    const B2Vector3D& stateMomentum(state.getMom());
    short stateCharge = state.getCharge();
 
 
    // Copy the information from the reco track.
    const B2Vector3D& trackPosition(m_recoTrack->getPositionSeed());
    const B2Vector3D& trackMomentum(m_recoTrack->getMomentumSeed());
    short trackCharge = m_recoTrack->getChargeSeed();
 
    // Copy the information from the mc particle (if there is one)
    MCParticle* relatedMCParticle = m_hit->template getRelated<MCParticle>("MCParticles");
 
    ROOT::Math::XYZVector mcMomentum;
    ROOT::Math::XYZVector mcPosition;
    short mcCharge;
 
    if (relatedMCParticle == nullptr) {
      mcMomentum = ROOT::Math::XYZVector(trackMomentum);
      mcPosition = ROOT::Math::XYZVector(trackPosition);
      mcCharge = trackCharge;
    } else {
      mcMomentum = relatedMCParticle->getMomentum();
      mcPosition = relatedMCParticle->getProductionVertex();
      mcCharge = relatedMCParticle->getCharge();
    }
 
    // Copy information from the mc hit
    const ROOT::Math::XYZVector& hitMCMomentum = momentumEstimationTools.getEntryMomentumOfMCParticle(*m_hit);
    const ROOT::Math::XYZVector& hitMCPosition = momentumEstimationTools.getEntryPositionOfMCParticle(*m_hit);
 
 
    if (m_useThickness) {
      rawHitCoordinates(0) = momentumEstimator.estimateQOverPWithThickness(*m_hit, stateCharge, m_fitParameters,
                             m_correctionFitParameters);
    } else {
      if (m_useTrackingSeeds) {
        if (m_useMCInformation) {
          rawHitCoordinates(0) = momentumEstimator.estimateQOverP(*m_hit, mcMomentum, mcPosition, mcCharge,
                                                                  m_fitParameters,
                                                                  m_correctionFitParameters);
 
        } else {
          rawHitCoordinates(0) = momentumEstimator.estimateQOverP(*m_hit, ROOT::Math::XYZVector(trackMomentum),
                                                                  ROOT::Math::XYZVector(trackPosition), trackCharge,
                                                                  m_fitParameters,
                                                                  m_correctionFitParameters);
        }
      } else {
        if (m_useMCInformation) {
          rawHitCoordinates(0) = momentumEstimator.estimateQOverP(*m_hit, hitMCMomentum, hitMCPosition, stateCharge,
                                                                  m_fitParameters,
                                                                  m_correctionFitParameters);
        } else {
          rawHitCoordinates(0) = momentumEstimator.estimateQOverP(*m_hit, ROOT::Math::XYZVector(stateMomentum),
                                                                  ROOT::Math::XYZVector(statePosition), stateCharge,
                                                                  m_fitParameters,
                                                                  m_correctionFitParameters);
        }
      }
    }
 
    rawHitCovariance(0, 0) = m_sigma;
 
    genfit::MeasurementOnPlane* mop = new genfit::MeasurementOnPlane(rawHitCoordinates, rawHitCovariance,
        state.getPlane(), state.getRep(), constructHMatrix(state.getRep()));
    return {mop};
  }
 
  template<>
  int PlanarVXDMomentumMeasurement<PXDCluster>::getDetId() const;
 
  template<>
  int PlanarVXDMomentumMeasurement<SVDCluster>::getDetId() const;
 
}