V0Fitter.cc

#include <tracking/v0Finding/fitter/V0Fitter.h>
 
#include <framework/logging/Logger.h>
#include <framework/datastore/StoreArray.h>
#include <tracking/v0Finding/dataobjects/VertexVector.h>
#include <tracking/dataobjects/RecoTrack.h>
#include <tracking/trackFitting/fitter/base/TrackFitter.h>
#include <tracking/trackFitting/trackBuilder/factories/TrackBuilder.h>
 
#include <mdst/dataobjects/HitPatternVXD.h>
#include <mdst/dataobjects/HitPatternCDC.h>
#include <mdst/dataobjects/Track.h>
#include <mdst/dataobjects/TrackFitResult.h>
#include <cdc/dataobjects/CDCRecoHit.h>
#include <pxd/reconstruction/PXDRecoHit.h>
#include <svd/reconstruction/SVDRecoHit.h>
#include <svd/reconstruction/SVDRecoHit2D.h>
 
#include <genfit/Track.h>
#include <genfit/TrackPoint.h>
#include <genfit/MeasuredStateOnPlane.h>
#include <genfit/GFRaveVertexFactory.h>
#include <genfit/GFRaveVertex.h>
#include <genfit/FieldManager.h>
#include <genfit/MaterialEffects.h>
#include <genfit/FitStatus.h>
#include <genfit/KalmanFitStatus.h>
#include <genfit/KalmanFitterInfo.h>
 
#include <framework/utilities/IOIntercept.h>
 
using namespace Belle2;
 
V0Fitter::V0Fitter(const std::string& trackFitResultsName, const std::string& v0sName,
                   const std::string& v0ValidationVerticesName, const std::string& recoTracksName,
                   const std::string& copiedRecoTracksName, bool enableValidation)
  : m_validation(enableValidation), m_recoTracksName(recoTracksName), m_v0FitterMode(1), m_forcestore(false),
    m_useOnlyOneSVDHitPair(true)
{
  m_trackFitResults.isRequired(trackFitResultsName);
  m_v0s.registerInDataStore(v0sName, DataStore::c_WriteOut | DataStore::c_ErrorIfAlreadyRegistered);
  //Relation to RecoTracks from Tracks is already tested at the module level.
 
  if (m_validation) {
    B2DEBUG(24, "Register DataStore for validation.");
    m_validationV0s.registerInDataStore(v0ValidationVerticesName, DataStore::c_ErrorIfAlreadyRegistered);
    m_v0s.registerRelationTo(m_validationV0s);
  }
 
  m_recoTracks.isRequired(m_recoTracksName);
 
  m_copiedRecoTracks.registerInDataStore(copiedRecoTracksName,
                                         DataStore::c_ErrorIfAlreadyRegistered);
  RecoTrack::registerRequiredRelations(m_copiedRecoTracks);
 
  m_copiedRecoTracks.registerRelationTo(m_recoTracks);
 
 
  B2ASSERT(genfit::MaterialEffects::getInstance()->isInitialized(),
           "Material effects not set up.  Please use SetupGenfitExtrapolationModule.");
  B2ASSERT(genfit::FieldManager::getInstance()->isInitialized(),
           "Magnetic field not set up.  Please use SetupGenfitExtrapolationModule.");
}
 
void V0Fitter::setFitterMode(int fitterMode)
{
  if (not(0 <= fitterMode && fitterMode <= 2)) {
    B2FATAL("Invalid fitter mode!");
  } else {
    m_v0FitterMode = fitterMode;
    if (fitterMode == 0)
      m_forcestore = true;
    else
      m_forcestore = false;
    if (fitterMode == 2)
      m_useOnlyOneSVDHitPair = false;
    else
      m_useOnlyOneSVDHitPair = true;
  }
}
 
void V0Fitter::initializeCuts(double beamPipeRadius,
                              double vertexChi2CutOutside)
{
  m_beamPipeRadius = beamPipeRadius;
  m_vertexChi2CutOutside = vertexChi2CutOutside;
}
 
 
bool V0Fitter::rejectCandidate(genfit::MeasuredStateOnPlane& stPlus, genfit::MeasuredStateOnPlane& stMinus)
{
  const TVector3& posPlus = stPlus.getPos();
  const TVector3& posMinus = stMinus.getPos();
 
  const double Rstart = std::min(posPlus.Perp(), posMinus.Perp());
  try {
    stPlus.extrapolateToCylinder(Rstart);
    stMinus.extrapolateToCylinder(Rstart);
  } catch (...) {
    B2DEBUG(22, "Extrapolation to cylinder failed.");
    return true;
  }
 
  return false;
}
 
 
bool V0Fitter::fitGFRaveVertex(genfit::Track& trackPlus, genfit::Track& trackMinus, genfit::GFRaveVertex& vertex)
{
  VertexVector vertexVector;
  std::vector<genfit::Track*> trackPair {&trackPlus, &trackMinus};
 
  try {
    IOIntercept::OutputToLogMessages
    logCapture("V0Fitter GFRaveVertexFactory", LogConfig::c_Debug, LogConfig::c_Debug);
    logCapture.start();
 
    genfit::GFRaveVertexFactory vertexFactory;
    vertexFactory.findVertices(&vertexVector.v, trackPair);
  } catch (...) {
    B2ERROR("Exception during vertex fit.");
    return false;
  }
 
  if (vertexVector.size() != 1) {
    B2DEBUG(21, "Vertex fit failed. Size of vertexVector not 1, but: " << vertexVector.size());
    return false;
  }
 
  if ((*vertexVector[0]).getNTracks() != 2) {
    B2DEBUG(20, "Wrong number of tracks in vertex.");
    return false;
  }
 
  vertex = *vertexVector[0];
  return true;
}
 
bool V0Fitter::extrapolateToVertex(genfit::MeasuredStateOnPlane& stPlus, genfit::MeasuredStateOnPlane& stMinus,
                                   const TVector3& vertexPosition, unsigned int& hasInnerHitStatus)
{
  hasInnerHitStatus = 0;
 
  try {
    double extralengthPlus  =  stPlus.extrapolateToPoint(vertexPosition);
    double extralengthMinus = stMinus.extrapolateToPoint(vertexPosition);
    if (extralengthPlus  > 0) hasInnerHitStatus |= 0x1; 
    if (extralengthMinus > 0) hasInnerHitStatus |= 0x2; 
  } catch (...) {
    B2DEBUG(22, "Could not extrapolate track to vertex.");
    return false;
  }
  return true;
}
 
double V0Fitter::getBzAtVertex(const TVector3& vertexPosition)
{
  double Bx, By, Bz;
  genfit::FieldManager::getInstance()->getFieldVal(vertexPosition.X(), vertexPosition.Y(), vertexPosition.Z(),
                                                   Bx, By, Bz);
  Bz = Bz / 10;
  return Bz;
}
 
 
TrackFitResult* V0Fitter::buildTrackFitResult(const genfit::Track& track, const RecoTrack* recoTrack,
                                              const genfit::MeasuredStateOnPlane& msop, const double Bz,
                                              const Const::ParticleType& trackHypothesis)
{
  const uint64_t hitPatternCDCInitializer = TrackBuilder::getHitPatternCDCInitializer(*recoTrack);
  const uint32_t hitPatternVXDInitializer = TrackBuilder::getHitPatternVXDInitializer(*recoTrack);
 
  TrackFitResult* v0TrackFitResult
    = m_trackFitResults.appendNew(msop.getPos(), msop.getMom(),
                                  msop.get6DCov(), msop.getCharge(),
                                  trackHypothesis,
                                  track.getFitStatus()->getPVal(),
                                  Bz, hitPatternCDCInitializer, hitPatternVXDInitializer, track.getFitStatus()->getNdf());
  return v0TrackFitResult;
}
 
std::pair<Const::ParticleType, Const::ParticleType> V0Fitter::getTrackHypotheses(const Const::ParticleType& v0Hypothesis) const
{
  if (v0Hypothesis == Const::Kshort) {
    return std::make_pair(Const::pion, Const::pion);
  } else if (v0Hypothesis == Const::photon) {
    return std::make_pair(Const::electron, Const::electron);
  } else if (v0Hypothesis == Const::Lambda) {
    return std::make_pair(Const::proton, Const::pion);
  } else if (v0Hypothesis == Const::antiLambda) {
    return std::make_pair(Const::pion, Const::proton);
  }
  B2FATAL("Given V0Hypothesis not available.");
  return std::make_pair(Const::invalidParticle, Const::invalidParticle); // return something to avoid triggering cppcheck
}
 
bool V0Fitter::fitAndStore(const Track* trackPlus, const Track* trackMinus,
                           const Const::ParticleType& v0Hypothesis)
{
  RecoTrack* recoTrackPlus  =  trackPlus->getRelated<RecoTrack>(m_recoTracksName);
  RecoTrack* recoTrackMinus = trackMinus->getRelated<RecoTrack>(m_recoTracksName);
 
  unsigned int hasInnerHitStatus = 0;
 
  TVector3 vertexPos(0, 0, 0);
 
  if (not vertexFitWithRecoTracks(trackPlus, trackMinus, recoTrackPlus, recoTrackMinus, v0Hypothesis,
                                  hasInnerHitStatus, vertexPos, m_forcestore))
    return false;
 
  if ((hasInnerHitStatus == 0) or m_forcestore)
    return true;
 
  bool failflag = false;
  unsigned int nRemoveHitsPlus  = 0;
  unsigned int nRemoveHitsMinus = 0;
 
  const auto trackHypotheses = getTrackHypotheses(v0Hypothesis);
  const int pdg_trackPlus   = trackPlus->getTrackFitResultWithClosestMass(
                                trackHypotheses.first)->getParticleType().getPDGCode();
  const int pdg_trackMinus  = trackMinus->getTrackFitResultWithClosestMass(
                                trackHypotheses.second)->getParticleType().getPDGCode();
 
  RecoTrack* recoTrackPlus_forRefit  = copyRecoTrack(recoTrackPlus, pdg_trackPlus);
  RecoTrack* recoTrackMinus_forRefit = copyRecoTrack(recoTrackMinus, pdg_trackMinus);
 
  if ((recoTrackPlus_forRefit == nullptr) or (recoTrackMinus_forRefit == nullptr))
    return false;
 
  while (hasInnerHitStatus != 0) {
 
    if (hasInnerHitStatus & 0x1) {
      ++nRemoveHitsPlus;
      if (not removeInnerHits(recoTrackPlus, recoTrackPlus_forRefit, pdg_trackPlus, nRemoveHitsPlus)) {
        failflag = true;
        break;
      }
    }
 
    if (hasInnerHitStatus & 0x2) {
      ++nRemoveHitsMinus;
      if (not removeInnerHits(recoTrackMinus, recoTrackMinus_forRefit, -1 * pdg_trackMinus, nRemoveHitsMinus)) {
        failflag = true;
        break;
      }
    }
 
    hasInnerHitStatus = 0;
    if (not vertexFitWithRecoTracks(trackPlus, trackMinus, recoTrackPlus_forRefit, recoTrackMinus_forRefit,
                                    v0Hypothesis, hasInnerHitStatus, vertexPos, false)) {
      //B2WARNING("vertex refit failed.");
      failflag = true;
      break;
    }
  }
 
  if (failflag) {
    bool forcestore = true;
    if (not vertexFitWithRecoTracks(trackPlus, trackMinus, recoTrackPlus, recoTrackMinus, v0Hypothesis,
                                    hasInnerHitStatus, vertexPos, forcestore)) {
      B2WARNING("original vertex fit fails. this should not happen.");
      return false;
    }
  }
 
  return true;
}
 
bool V0Fitter::vertexFitWithRecoTracks(const Track* trackPlus, const Track* trackMinus,
                                       RecoTrack* recoTrackPlus, RecoTrack* recoTrackMinus,
                                       const Const::ParticleType& v0Hypothesis,
                                       unsigned int& hasInnerHitStatus, TVector3& vertexPos,
                                       const bool forceStore)
{
  const auto trackHypotheses = getTrackHypotheses(v0Hypothesis);
 
  genfit::Track gfTrackPlus = RecoTrackGenfitAccess::getGenfitTrack(
                                *recoTrackPlus);
  const int pdgTrackPlus = trackPlus->getTrackFitResultWithClosestMass(trackHypotheses.first)->getParticleType().getPDGCode();
  genfit::AbsTrackRep* plusRepresentation = recoTrackPlus->getTrackRepresentationForPDG(pdgTrackPlus);
  if ((plusRepresentation == nullptr) or (not recoTrackPlus->wasFitSuccessful(plusRepresentation))) {
    B2ERROR("Track hypothesis with closest mass not available. Should never happen, but I can continue savely anyway.");
    return false;
  }
 
  genfit::Track gfTrackMinus = RecoTrackGenfitAccess::getGenfitTrack(
                                 *recoTrackMinus);
  const int pdgTrackMinus = trackMinus->getTrackFitResultWithClosestMass(trackHypotheses.second)->getParticleType().getPDGCode();
  genfit::AbsTrackRep* minusRepresentation = recoTrackMinus->getTrackRepresentationForPDG(pdgTrackMinus);
  if ((minusRepresentation == nullptr) or (not recoTrackMinus->wasFitSuccessful(minusRepresentation))) {
    B2ERROR("Track hypothesis with closest mass not available. Should never happen, but I can continue savely anyway.");
    return false;
  }
 
  std::vector<genfit::AbsTrackRep*> repsPlus = gfTrackPlus.getTrackReps();
  std::vector<genfit::AbsTrackRep*> repsMinus = gfTrackMinus.getTrackReps();
  if (repsPlus.size() == repsMinus.size()) {
    for (unsigned int id = 0; id < repsPlus.size(); id++) {
      if (abs(repsPlus[id]->getPDG()) == pdgTrackPlus)
        gfTrackPlus.setCardinalRep(id);
      if (abs(repsMinus[id]->getPDG()) == pdgTrackMinus)
        gfTrackMinus.setCardinalRep(id);
    }
  }
  else {
    for (unsigned int id = 0; id < repsPlus.size(); id++) {
      if (abs(repsPlus[id]->getPDG()) == pdgTrackPlus)
        gfTrackPlus.setCardinalRep(id);
    }
    for (unsigned int id = 0; id < repsMinus.size(); id++) {
      if (abs(repsMinus[id]->getPDG()) == pdgTrackMinus)
        gfTrackMinus.setCardinalRep(id);
    }
  }
 
  genfit::MeasuredStateOnPlane stPlus  = recoTrackPlus->getMeasuredStateOnPlaneFromFirstHit(plusRepresentation);
  genfit::MeasuredStateOnPlane stMinus = recoTrackMinus->getMeasuredStateOnPlaneFromFirstHit(minusRepresentation);
 
  if (rejectCandidate(stPlus, stMinus)) {
    return false;
  }
 
  genfit::GFRaveVertex vert;
  if (not fitGFRaveVertex(gfTrackPlus, gfTrackMinus, vert)) {
    return false;
  }
 
  stPlus  = gfTrackPlus.getFittedState();
  stMinus = gfTrackMinus.getFittedState();
 
  const TVector3& posVert(vert.getPos());
  vertexPos = posVert;
 
  if (posVert.Perp() < m_beamPipeRadius) {
    return false;
  } else {
    if (vert.getChi2() > m_vertexChi2CutOutside) {
      B2DEBUG(22, "Vertex outside beam pipe, chi^2 too large.");
      return false;
    }
  }
 
  B2DEBUG(22, "Vertex accepted.");
 
  if (not extrapolateToVertex(stPlus, stMinus, posVert, hasInnerHitStatus)) {
    return false;
  }
 
  if (forceStore || hasInnerHitStatus == 0) {
    const TVector3 origin(0, 0, 0);
    const double Bz = getBzAtVertex(origin);
 
    TrackFitResult* tfrPlusVtx =  buildTrackFitResult(gfTrackPlus, recoTrackPlus, stPlus,  Bz, trackHypotheses.first);
    TrackFitResult* tfrMinusVtx = buildTrackFitResult(gfTrackMinus, recoTrackMinus, stMinus, Bz, trackHypotheses.second);
 
    B2DEBUG(20, "Creating new V0.");
    auto v0 = m_v0s.appendNew(std::make_pair(trackPlus, tfrPlusVtx),
                              std::make_pair(trackMinus, tfrMinusVtx));
 
    if (m_validation) {
      B2DEBUG(24, "Create StoreArray and Output for validation.");
      const genfit::GFRaveTrackParameters* tr0 = vert.getParameters(0);
      const genfit::GFRaveTrackParameters* tr1 = vert.getParameters(1);
      TLorentzVector lv0, lv1;
      lv0.SetVectM(tr0->getMom(), trackHypotheses.first.getMass());
      lv1.SetVectM(tr1->getMom(), trackHypotheses.second.getMass());
      auto validationV0 = m_validationV0s.appendNew(
                            std::make_pair(trackPlus, tfrPlusVtx),
                            std::make_pair(trackMinus, tfrMinusVtx),
                            vert.getPos(),
                            vert.getCov(),
                            (lv0 + lv1).P(),
                            (lv0 + lv1).M(),
                            vert.getChi2()
                          );
      v0->addRelationTo(validationV0);
    }
  }
  return true;
}
 
RecoTrack* V0Fitter::copyRecoTrack(RecoTrack* origRecoTrack, const int trackPDG)
{
  Const::ChargedStable particleUsedForFitting(std::abs(trackPDG));
  const genfit::AbsTrackRep* origTrackRep = origRecoTrack->getTrackRepresentationForPDG(std::abs(
                                              trackPDG));
 
  RecoTrack* newRecoTrack = origRecoTrack->copyToStoreArray(m_copiedRecoTracks);
  newRecoTrack->addHitsFromRecoTrack(origRecoTrack);
  newRecoTrack->addRelationTo(origRecoTrack);
 
  TrackFitter fitter;
  if (not fitter.fit(*newRecoTrack, particleUsedForFitting)) {
    B2WARNING("track fit failed for copied RecoTrack.");
    if (not origRecoTrack->wasFitSuccessful(origTrackRep))
      B2WARNING("\t original track fit was also failed.");
    return nullptr;
  }
 
  return newRecoTrack;
}
 
bool V0Fitter::removeInnerHits(RecoTrack* origRecoTrack, RecoTrack* recoTrack,
                               const int trackPDG, unsigned int& nRemoveHits)
{
  Const::ChargedStable particleUsedForFitting(std::abs(trackPDG));
  const genfit::AbsTrackRep* origTrackRep = origRecoTrack->getTrackRepresentationForPDG(std::abs(
                                              trackPDG));
 
  const std::vector<RecoHitInformation*>& recoHitInformations = recoTrack->getRecoHitInformations(true);
  if (recoHitInformations.size() < nRemoveHits) {
    B2WARNING("N removed hits exceeds N hits in the track.");
    return false;
  }
  for (unsigned int i = 0 ; i < nRemoveHits ; ++i) {
    recoHitInformations[i]->setUseInFit(false);
  }
  if (recoHitInformations[nRemoveHits - 1]->getTrackingDetector() == RecoHitInformation::RecoHitDetector::c_SVD) {
    if (recoHitInformations.size() < nRemoveHits + 1) { 
      B2WARNING("N removed hits exceeds N hits in the track.");
      return false;
    }
    if (recoHitInformations[nRemoveHits]->getTrackingDetector() == RecoHitInformation::RecoHitDetector::c_SVD) {
      const SVDCluster* lastRemovedSVDHit = recoHitInformations[nRemoveHits - 1]->getRelatedTo<SVDCluster>();
      const SVDCluster* nextSVDHit        = recoHitInformations[nRemoveHits]  ->getRelatedTo<SVDCluster>();
      if (lastRemovedSVDHit->getSensorID() == nextSVDHit->getSensorID() &&
          lastRemovedSVDHit->isUCluster() && !(nextSVDHit->isUCluster())) {
        recoHitInformations[nRemoveHits]->setUseInFit(false);
        ++nRemoveHits;
      }
    }
  }
 
  if (!m_useOnlyOneSVDHitPair &&
      recoHitInformations[nRemoveHits - 1]->getTrackingDetector() == RecoHitInformation::RecoHitDetector::c_SVD &&
      recoHitInformations.size() > nRemoveHits + 2) {
    if (recoHitInformations[nRemoveHits  ]->getTrackingDetector() == RecoHitInformation::RecoHitDetector::c_SVD && 
        recoHitInformations[nRemoveHits + 1]->getTrackingDetector() == RecoHitInformation::RecoHitDetector::c_SVD && 
        recoHitInformations[nRemoveHits + 2]->getTrackingDetector() != RecoHitInformation::RecoHitDetector::c_SVD) { 
      recoHitInformations[nRemoveHits  ]->setUseInFit(false);
      recoHitInformations[nRemoveHits + 1]->setUseInFit(false);
      nRemoveHits += 2;
    }
  }
 
  TrackFitter fitter;
  if (not fitter.fit(*recoTrack, particleUsedForFitting)) {
    B2WARNING("track fit failed.");
    if (not origRecoTrack->wasFitSuccessful(origTrackRep))
      B2WARNING("\t original track fit was also failed.");
    return false;
  }
 
  return true;
}