CbmKfTrackFitter.cxx

Go to the documentation of this file.

/* Copyright (C) 2023-2025 GSI Helmholtzzentrum fuer Schwerionenforschung, Darmstadt
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Sergey Gorbunov [committer] */
 
#include "CbmKfTrackFitter.h"
 
#include "CbmGlobalTrack.h"
#include "CbmKfTrackingGeoSetupContainer.h"
#include "CbmMuchPixelHit.h"
#include "CbmMuchTrack.h"
#include "CbmMvdHit.h"
#include "CbmRecoSetupManager.h"
#include "CbmStsAddress.h"
#include "CbmStsHit.h"
#include "CbmStsSetup.h"
#include "CbmStsTrack.h"
#include "CbmTofHit.h"
#include "CbmTofTrack.h"
#include "CbmTrdHit.h"
#include "CbmTrdTrack.h"
#include "FairRootManager.h"
#include "KFParticleDatabase.h"
#include "KfDefs.h"
#include "KfToolsField.h"
#include "KfTrackKalmanFilter.h"
#include "KfTrackParam.h"
#include "TClonesArray.h"
#include "TDatabasePDG.h"
 
using std::vector;
using namespace std;
using namespace cbm::algo;
 
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
CbmKfTrackFitter<FlagFitTime>::CbmKfTrackFitter()
{
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
CbmKfTrackFitter<FlagFitTime>::~CbmKfTrackFitter()
{
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::Init()
{
  FairRootManager* ioman = FairRootManager::Instance();
 
  if (!ioman) {
    LOG(fatal) << fDebugPrefix << "no FairRootManager";
  }
 
  // Get hits
 
  fInputMvdHits  = dynamic_cast<const TClonesArray*>(ioman->GetObject("MvdHit"));
  fInputStsHits  = dynamic_cast<const TClonesArray*>(ioman->GetObject("StsHit"));
  fInputTrdHits  = dynamic_cast<const TClonesArray*>(ioman->GetObject("TrdHit"));
  fInputMuchHits = dynamic_cast<const TClonesArray*>(ioman->GetObject("MuchHit"));
  fInputTofHits  = dynamic_cast<const TClonesArray*>(ioman->GetObject("TofHit"));
 
  // Get global tracks
  fInputGlobalTracks = dynamic_cast<const TClonesArray*>(ioman->GetObject("GlobalTrack"));
 
  // Get detector tracks
  fInputStsTracks  = dynamic_cast<const TClonesArray*>(ioman->GetObject("StsTrack"));
  fInputMuchTracks = dynamic_cast<const TClonesArray*>(ioman->GetObject("MuchTrack"));
  fInputTrdTracks  = dynamic_cast<const TClonesArray*>(ioman->GetObject("TrdTrack"));
  fInputTofTracks  = dynamic_cast<const TClonesArray*>(ioman->GetObject("TofTrack"));
 
  fKfSetup = cbm::kf::TrackingGeoSetupContainer::Instance().GetSetup();
 
  // Check Detector interfaces
  if (!cbm::RecoSetupManager::Instance()->IsInitialized()) {
    LOG(fatal) << "cbm::RecoSetupManager not initialized";
  }
  const auto& recoSetup = cbm::RecoSetupManager::Instance()->GetSetup();
 
  if (fInputMvdHits && !recoSetup.GetMvd()) {
    LOG(fatal) << "RecoSetupUnit for MVD was not provided";
  }
  if (fInputStsHits && !recoSetup.GetSts()) {
    LOG(fatal) << "RecoSetupUnit for STS was not provided";
  }
  if (fInputMuchHits && !recoSetup.GetMuch()) {
    LOG(fatal) << "RecoSetupUnit for MUCH was not provided";
  }
  if (fInputTrdHits && !recoSetup.GetTrd()) {
    LOG(fatal) << "RecoSetupUnit for TRD was not provided";
  }
  if (fInputTofHits && !recoSetup.GetTof()) {
    LOG(fatal) << "RecoSetupUnit for TOF was not provided";
  }
 
  //auto CheckoutDetectorUnit = [] (const auto* node, auto moduleId) {
  //  if (node && !cbm::RecoSetupManager::Instance()->GetSetup().Get<moduleId>()) {
  //    LOG(fatal) << "RecoSetupUnit for " << moduleId << " was not provided";
  //  }
  //};
 
  //CheckoutDetectorUnit(fInputMvdHits, ECbmModuleId::kMvd);
  //CheckoutDetectorUnit(fInputStsHits, ECbmModuleId::kSts);
  //CheckoutDetectorUnit(fInputMuchHits, ECbmModuleId::kMuch);
  //CheckoutDetectorUnit(fInputTrdHits, ECbmModuleId::kTrd);
  //CheckoutDetectorUnit(fInputTofHits, ECbmModuleId::kTof);
 
  cbm::kf::tools::SetOriginalCbmField();
 
  fIsInitialized = true;
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::SetParticleHypothesis(int pdg)
{
  TParticlePDG* particlePDG = TDatabasePDG::Instance()->GetParticle(pdg);
  if (!particlePDG) {
    LOG(fatal) << fDebugPrefix << "particle PDG " << pdg << " is not in the data base, please set the mass manually";
    return;
  }
  fMass       = particlePDG->Mass();
  fIsElectron = (abs(pdg) == 11);
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::SetMassHypothesis(double mass)
{
  assert(mass >= 0.);
  fMass = mass;
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::SetElectronFlag(bool isElectron)
{
  fIsElectron = isElectron;
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::CreateFromMvdStsTrack(CbmKfTrackFitter<FlagFitTime>::Trajectory& trajectory,
                                                          int stsTrackIndex, bool addMaterialNodesOutsideHitRange)
{
  trajectory.Clear();
 
  if (!fIsInitialized) {
    LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
    return false;
  }
 
  if (!fInputStsTracks) {
    LOG(error) << fDebugPrefix << "Sts track array not found!";
    return false;
  }
  if (stsTrackIndex >= fInputStsTracks->GetEntriesFast()) {
    LOG(error) << fDebugPrefix << "Sts track index " << stsTrackIndex << " is out of range!";
    return false;
  }
  auto* stsTrack = dynamic_cast<const CbmStsTrack*>(fInputStsTracks->At(stsTrackIndex));
  if (!stsTrack) {
    LOG(error) << fDebugPrefix << "Sts track is null!";
    return false;
  }
 
  CbmGlobalTrack globalTrack;
  globalTrack.SetStsTrackIndex(stsTrackIndex);
  globalTrack.SetParamFirst(*dynamic_cast<const FairTrackParam*>(stsTrack->GetParamFirst()));
 
  return CreateFromGlobalTrack(trajectory, globalTrack, addMaterialNodesOutsideHitRange);
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::CreateFromGlobalTrack(CbmKfTrackFitter<FlagFitTime>::Trajectory& trajectory,
                                                          int globalTrackIndex, bool addMaterialNodesOutsideHitRange)
{
  trajectory.Clear();
 
  if (!fIsInitialized) {
    LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
    return false;
  }
 
  if (!fInputGlobalTracks) {
    LOG(error) << fDebugPrefix << "Global track array not found!";
    return false;
  }
 
  if (globalTrackIndex >= fInputGlobalTracks->GetEntriesFast()) {
    LOG(error) << fDebugPrefix << "Global track index " << globalTrackIndex << " is out of range!";
    return false;
  }
 
  auto* globalTrack = dynamic_cast<const CbmGlobalTrack*>(fInputGlobalTracks->At(globalTrackIndex));
  if (!globalTrack) {
    LOG(error) << fDebugPrefix << "Global track is null!";
    return false;
  }
 
  return CreateFromGlobalTrack(trajectory, *globalTrack, addMaterialNodesOutsideHitRange);
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::CreateFromGlobalTrack(Trajectory& trajectory, const CbmGlobalTrack& globalTrack,
                                                          bool addMaterialNodesOutsideHitRange)
{
  trajectory.Clear();
  if (!fIsInitialized) {
    LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
    return false;
  }
 
  Trajectory t(fKfSetup);
  {
    t.fNodes.resize(fKfSetup->GetNofLayers());
    for (int i = 0; i < fKfSetup->GetNofLayers(); i++) {
      auto& n         = t.fNodes[i];
      n               = {};
      n.materialLayer = i;
      n.z             = fKfSetup->GetMaterial(i).GetZref();
      n.radThick      = 0.;
      n.isFitted      = false;
      n.isXySet       = false;
      n.isTimeSet     = false;
      Trajectory::SetHitSystemId(n, ECbmModuleId::kNotExist);
      Trajectory::SetHitAddress(n, -1);
      Trajectory::SetHitIndex(n, -1);
    }
  }
 
  // lambda to set the node from the pixel hit
  auto setNode = [&](const CbmPixelHit& h, int stIdx, int hitIdx, ca::EDetectorID detId) -> typename Trajectory::Node* {
    int iLayer = fKfSetup->GetIndexMap().LocalToGlobal(detId, stIdx);
    if (iLayer < 0) {
      return nullptr;
    }
    assert(iLayer >= 0 && iLayer < fKfSetup->GetNofLayers());
    auto& n = t.fNodes[iLayer];
    n.z     = h.GetZ();
    n.measurementXY.SetX(h.GetX());
    n.measurementXY.SetY(h.GetY());
    n.measurementXY.SetDx2(h.GetDx() * h.GetDx());
    n.measurementXY.SetDy2(h.GetDy() * h.GetDy());
    n.measurementXY.SetDxy(h.GetDxy());
    n.measurementXY.SetNdfX(1);
    n.measurementXY.SetNdfY(1);
    n.isXySet = true;
    n.measurementTime.SetT(h.GetTime());
    n.measurementTime.SetDt2(h.GetTimeError() * h.GetTimeError());
    n.measurementTime.SetNdfT(1);
    n.isTimeSet = (detId != ca::EDetectorID::kMvd);
    n.radThick  = 0.;
    Trajectory::SetHitSystemId(n, ca::ToCbmModuleId(detId));
    Trajectory::SetHitAddress(n, h.GetAddress());
    Trajectory::SetHitIndex(n, hitIdx);
    return &n;
  };
 
  // Read MVD & STS hits
 
  if (globalTrack.GetStsTrackIndex() >= 0) {
 
    int stsTrackIndex = globalTrack.GetStsTrackIndex();
 
    if (!fInputStsTracks) {
      LOG(error) << fDebugPrefix << "Sts track array not found!";
      return false;
    }
    if (stsTrackIndex >= fInputStsTracks->GetEntriesFast()) {
      LOG(error) << fDebugPrefix << "Sts track index " << stsTrackIndex << " is out of range!";
      return false;
    }
    auto* stsTrack = dynamic_cast<const CbmStsTrack*>(fInputStsTracks->At(stsTrackIndex));
    if (!stsTrack) {
      LOG(error) << fDebugPrefix << "Sts track is null!";
      return false;
    }
 
    // Read MVD hits
 
    int nMvdHits = stsTrack->GetNofMvdHits();
    if (nMvdHits > 0) {
      if (!fInputMvdHits) {
        LOG(error) << fDebugPrefix << " Mvd hit array not found!";
        return false;
      }
      const auto* pRecoUnit = cbm::RecoSetupManager::Instance()->GetSetup().GetMvd();
      for (int ih = 0; ih < nMvdHits; ih++) {
        int hitIndex = stsTrack->GetMvdHitIndex(ih);
        if (hitIndex >= fInputMvdHits->GetEntriesFast()) {
          LOG(error) << fDebugPrefix << " Mvd hit index " << hitIndex << " is out of range!";
          return false;
        }
        const auto& hit = *dynamic_cast<const CbmMvdHit*>(fInputMvdHits->At(hitIndex));
        setNode(hit, pRecoUnit->GetTrackingStationId(hit.GetAddress()), hitIndex, ca::EDetectorID::kMvd);
      }
    }
 
    // Read STS hits
 
    int nStsHits = stsTrack->GetNofStsHits();
    if (nStsHits > 0) {
      if (!fInputStsHits) {
        LOG(error) << fDebugPrefix << " Sts hit array not found!";
        return false;
      }
      const auto* pRecoUnit = cbm::RecoSetupManager::Instance()->GetSetup().GetSts();
      for (int ih = 0; ih < nStsHits; ih++) {
        int hitIndex = stsTrack->GetStsHitIndex(ih);
        if (hitIndex >= fInputStsHits->GetEntriesFast()) {
          LOG(error) << fDebugPrefix << " Sts hit index " << hitIndex << " is out of range!";
          return false;
        }
        const auto& hit = *dynamic_cast<const CbmStsHit*>(fInputStsHits->At(hitIndex));
        setNode(hit, pRecoUnit->GetTrackingStationId(hit.GetAddress()), hitIndex, ca::EDetectorID::kSts);
      }
    }
  }  // MVD & STS hits
 
 
  // Read Much hits
 
  if (globalTrack.GetMuchTrackIndex() >= 0) {
    int muchTrackIndex = globalTrack.GetMuchTrackIndex();
    if (!fInputMuchTracks) {
      LOG(error) << fDebugPrefix << "Much track array not found!";
      return false;
    }
    if (muchTrackIndex >= fInputMuchTracks->GetEntriesFast()) {
      LOG(error) << fDebugPrefix << "Much track index " << muchTrackIndex << " is out of range!";
      return false;
    }
    auto* track = dynamic_cast<const CbmMuchTrack*>(fInputMuchTracks->At(muchTrackIndex));
    if (!track) {
      LOG(error) << fDebugPrefix << "Much track is null!";
      return false;
    }
    int nHits = track->GetNofHits();
    if (nHits > 0) {
      if (!fInputMuchHits) {
        LOG(error) << fDebugPrefix << "Much hit array not found!";
        return false;
      }
      const auto* pRecoUnit = cbm::RecoSetupManager::Instance()->GetSetup().GetMuch();
      for (int ih = 0; ih < nHits; ih++) {
        int hitIndex = track->GetHitIndex(ih);
        if (hitIndex >= fInputMuchHits->GetEntriesFast()) {
          LOG(error) << fDebugPrefix << "Much hit index " << hitIndex << " is out of range!";
          return false;
        }
        const auto& hit = *dynamic_cast<const CbmMuchPixelHit*>(fInputMuchHits->At(hitIndex));
        setNode(hit, pRecoUnit->GetTrackingStationId(hit.GetAddress()), hitIndex, ca::EDetectorID::kMuch);
      }
    }
  }
 
  // Read TRD hits
 
  if (globalTrack.GetTrdTrackIndex() >= 0) {
    int trdTrackIndex = globalTrack.GetTrdTrackIndex();
    if (!fInputTrdTracks) {
      LOG(error) << fDebugPrefix << "Trd track array not found!";
      return false;
    }
    if (trdTrackIndex >= fInputTrdTracks->GetEntriesFast()) {
      LOG(error) << fDebugPrefix << "Trd track index " << trdTrackIndex << " is out of range!";
      return false;
    }
    auto* track = dynamic_cast<const CbmTrdTrack*>(fInputTrdTracks->At(trdTrackIndex));
    if (!track) {
      LOG(error) << fDebugPrefix << "Trd track is null!";
      return false;
    }
    int nHits = track->GetNofHits();
    if (nHits > 0) {
      if (!fInputTrdHits) {
        LOG(error) << fDebugPrefix << "Trd hit array not found!";
        return false;
      }
      const auto* pRecoUnit = cbm::RecoSetupManager::Instance()->GetSetup().GetTrd();
      for (int ih = 0; ih < nHits; ih++) {
        int hitIndex = track->GetHitIndex(ih);
        if (hitIndex >= fInputTrdHits->GetEntriesFast()) {
          LOG(error) << fDebugPrefix << "Trd hit index " << hitIndex << " is out of range!";
          return false;
        }
        const auto& hit = *dynamic_cast<const CbmTrdHit*>(fInputTrdHits->At(hitIndex));
        auto* node = setNode(hit, pRecoUnit->GetTrackingStationId(hit.GetAddress()), hitIndex, ca::EDetectorID::kTrd);
 
        if (node != nullptr && hit.GetClassType() == 1) {
          Trajectory::SetHitSystemId(*node, ECbmModuleId::kTrd2d);
        }
      }
    }
  }
 
 
  // Read TOF hits
 
  if (globalTrack.GetTofTrackIndex() >= 0) {
    int tofTrackIndex = globalTrack.GetTofTrackIndex();
    if (!fInputTofTracks) {
      LOG(error) << fDebugPrefix << "Tof track array not found!";
      return false;
    }
    if (tofTrackIndex >= fInputTofTracks->GetEntriesFast()) {
      LOG(error) << fDebugPrefix << "Tof track index " << tofTrackIndex << " is out of range!";
      return false;
    }
    auto* track = dynamic_cast<const CbmTofTrack*>(fInputTofTracks->At(tofTrackIndex));
    if (!track) {
      LOG(error) << fDebugPrefix << "Tof track is null!";
      return false;
    }
 
    int nHits = track->GetNofHits();
    if (nHits > 0) {
      if (!fInputTofHits) {
        LOG(error) << fDebugPrefix << "Tof hit array not found!";
        return false;
      }
      const auto* pRecoUnit = cbm::RecoSetupManager::Instance()->GetSetup().GetTof();
      for (int ih = 0; ih < nHits; ih++) {
        int hitIndex = track->GetHitIndex(ih);
        if (hitIndex >= fInputTofHits->GetEntriesFast()) {
          LOG(error) << fDebugPrefix << "Tof hit index " << hitIndex << " is out of range!";
          return false;
        }
        const auto& hit = *dynamic_cast<const CbmTofHit*>(fInputTofHits->At(hitIndex));
        setNode(hit, pRecoUnit->GetTrackingStationId(hit.GetAddress()), hitIndex, ca::EDetectorID::kTof);
      }
    }
  }
 
  t.MakeConsistent();
 
  if (!addMaterialNodesOutsideHitRange) {  // TODO: dont create the nodes if they are not needed
    int firstHitNode = t.fFirstMeasurementNodeId;
    int lastHitNode  = t.fLastMeasurementNodeId;
    assert(firstHitNode >= 0 && lastHitNode >= 0);
    t.fNodes.erase(t.fNodes.begin() + lastHitNode + 1, t.fNodes.end());
    t.fNodes.erase(t.fNodes.begin(), t.fNodes.begin() + firstHitNode);
    t.MakeConsistent();
  }
 
  trajectory = t;
  return true;
}
 
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::FilterFirstMeasurement(const cbm::algo::kf::Trajectory<>::Node& n)
{
  // a special routine to filter the first measurement.
  // the measurement errors are simply copied to the track covariance matrix
 
  assert(n.isXySet);
 
  const auto& mxy = n.measurementXY;
  const auto& mt  = n.measurementTime;
 
  auto& tr = fFit.Tr();
 
  if (fabs(tr.GetZ() - n.z) > 1.e-15) {
    LOG(fatal) << fDebugPrefix << "Z mismatch: fitted track " << tr.GetZ() << " != node " << n.z;
  }
 
  tr.ResetErrors(mxy.Dx2(), mxy.Dy2(), 100., 100., 10., 1.e4, 1.e2);
  tr.SetC10(mxy.Dxy());
  tr.SetX(mxy.X());
  tr.SetY(mxy.Y());
  tr.SetZ(n.z);
 
  if (fIgnorePoorCoordinates) {  // TODO: take C10 correlation into account
    if (mxy.NdfX() == 0) {
      tr.SetX(n.paramDn.GetX());
      tr.SetC00(1.e4);
    }
    if (mxy.NdfY() == 0) {
      tr.SetY(n.paramDn.GetY());
      tr.SetC11(1.e4);
      tr.SetC10(0.);
    }
  }
 
  tr.SetChiSq(0.);
  tr.SetChiSqTime(0.);
  tr.SetNdf(-5. + mxy.NdfX() + mxy.NdfY());
  if (n.isTimeSet) {
    tr.SetTime(mt.T());
    tr.SetC55(mt.Dt2());
    tr.SetNdfTime(-2 + 1);
  }
  else {
    tr.SetNdfTime(-2);
  }
  tr.Vi() = cbm::algo::kf::defs::SpeedOfLightInv<>;
  tr.InitVelocityRange(0.5);
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
void CbmKfTrackFitter<FlagFitTime>::AddMaterialEffects(const LinearizationAtNode& l, kf::FitDirection direction)
{
  // add material effects
  if (l.radThick <= 0.) {
    return;
  }
 
  double tx   = 0.5 * (l.linearizationDn.tx + l.linearizationUp.tx);
  double ty   = 0.5 * (l.linearizationDn.ty + l.linearizationUp.ty);
  double msQp = 0.5 * (l.linearizationDn.qp + l.linearizationUp.qp);
 
  if (fEnforceDefaultMomentumForMs) {
    msQp = fDefaultQpForMs;
  }
 
  if (!fIgnoreMultipleScattering) {
    fFit.MultipleScattering(l.radThick, tx, ty, msQp);
  }
 
  if (!fIgnoreEnergyLosses && !fEnforceDefaultMomentumForMs) {
    if (direction == kf::FitDirection::kDownstream) {
      fFit.Linearization().qp = l.linearizationUp.qp;
    }
    else {
      fFit.Linearization().qp = l.linearizationDn.qp;
    }
    fFit.EnergyLossCorrection(l.radThick, direction);
  }
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::FitTrajectory(cbm::algo::kf::Trajectory<>& t)
{
  // (re)fit the track
 
 
  if (fVerbosityLevel > 0) {
    LOG(info) << fDebugPrefix << "-------- FitTrajectory ... ";
  }
 
 
  {  // check input trajectory consistency
 
    if (!fIsInitialized) {
      LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
      return false;
    }
 
 
    if (t.IsFitted()) {
      for (int in = t.GetFirstMeasurementNodeId(); in <= t.GetLastMeasurementNodeId(); in++) {
        const auto& n = t.fNodes[in];
        if (!n.isFitted) {
          LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fitted, but its node " << in
                     << " in the measured region is not fitted!";
        }
        if (n.materialLayer >= 0 && n.radThick < 0.) {
          LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fitted, but its node " << in
                     << " with material layer " << n.materialLayer
                     << " in the measured region does not have material budget set!";
        }
      }
    }
 
    if (t.IsFullyExtrapolated()) {
      if (!t.IsFitted()) {
        LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but it is not fitted!";
      }
      for (unsigned int in = 0; in < t.fNodes.size(); in++) {
        const auto& n = t.fNodes[in];
        if (!n.isFitted) {
          LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                     << " is not fitted!";
        }
        if (n.materialLayer >= 0 && n.radThick < 0.) {
          LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                     << " with material layer " << n.materialLayer << " does not have material budget set!";
        }
      }
    }
 
    if (fIsMaterialFixed) {  // check that the material budget is already set
      if (!t.IsFitted()) {
        LOG(fatal) << fDebugPrefix << "Material fixed flag is set, but the input trajectory is not fitted.";
      }
      if (!t.fIsFullyExtrapolated) {
        LOG(fatal) << fDebugPrefix << "Material fixed flag is set, but the input trajectory is not extrapolated.";
      }
    }
 
  }  // check input trajectory consistency
 
 
  bool ok = true;
 
  bool wasFitted = t.IsFitted();
 
  t.fIsFitted            = false;
  t.fIsFullyExtrapolated = false;
 
  int nNodes = t.fNodes.size();
 
  if (nNodes <= 0) {
    LOG(warning) << fDebugPrefix << " no nodes found!";
    return false;
  }
 
  if (t.GetNofMeasurements() <= 0) {
    LOG(warning) << fDebugPrefix << "no nodes with measurements found!";
    return false;
  }
 
  fFit.SetParticleMass(fMass);
 
  kf::FieldRegion<double> field(kf::GlobalField::fgOriginalFieldType, kf::GlobalField::fgOriginalField);
 
  std::vector<LinearizationAtNode> linearization(nNodes);
 
  int firstHitNode = t.fFirstMeasurementNodeId;
  int lastHitNode  = t.fLastMeasurementNodeId;
 
  // linearization initialization
 
  if (wasFitted) {  // take the linearization from the previously fitted trajectory
    for (int i = firstHitNode; i <= lastHitNode; i++) {
      linearization[i].linearizationDn = t.fNodes[i].paramDn;
      linearization[i].linearizationUp = t.fNodes[i].paramUp;
      linearization[i].radThick        = t.fNodes[i].radThick;
    }
  }
  else {  // first approximation with straight line segments connecting the nodes
    for (int i1 = firstHitNode, i2 = firstHitNode + 1; i2 <= lastHitNode; i2++) {
      auto& n1 = t.fNodes[i1];
      auto& n2 = t.fNodes[i2];
      if (!n2.isXySet) {
        continue;
      }
      double dz  = n2.z - n1.z;
      double dzi = (fabs(dz) > 1.e-4) ? 1. / dz : 0.;
      kf::LinearizationFull<double> l;
      double lz = n1.z;
      l.x       = n1.measurementXY.X();
      l.y       = n1.measurementXY.Y();
      l.tx      = (n2.measurementXY.X() - n1.measurementXY.X()) * dzi;
      l.ty      = (n2.measurementXY.Y() - n1.measurementXY.Y()) * dzi;
      l.qp      = 0.;
      l.time    = 0.;
      l.vi      = cbm::algo::kf::defs::SpeedOfLightInv<double>;
 
      for (int i = i1; i <= i2; i++) {  // fill the nodes inbetween
        auto& n  = t.fNodes[i];
        auto& nl = linearization[i];
        l.x += l.tx * (n.z - lz);
        l.y += l.ty * (n.z - lz);
        lz = n.z;
        if (i < i2 || i == lastHitNode) {  // downstream parameters for i2 will be set later, except i2 == last hit node
          nl.linearizationDn = l;
        }
        if (i > i1 || i == firstHitNode) {  // upstream parameters for i1 are already set, except i1 == first hit node
          nl.linearizationUp = l;
        }
      }
      i1 = i2;
    }
    // compute material budget for each node
    for (int in = firstHitNode; in <= lastHitNode; in++) {
      auto& n = t.fNodes[in];
      auto& l = linearization[in];
      if (n.materialLayer >= 0) {
        const kf::MaterialMap& map = fKfSetup->GetMaterial(n.materialLayer);
        l.radThick                 = map.GetThicknessX0(l.linearizationDn.x, l.linearizationDn.y);
      }
      else {
        l.radThick = 0.;
      }
    }
  }  // linearization initialization
 
  t.fQaChi2Downstream = -1.;
 
  auto printNode = [&](std::string str, int node) {
    if (fVerbosityLevel > 0) {
      LOG(info) << fDebugPrefix << str << ": node " << node << " chi2 " << fFit.Tr().GetChiSq() << " x "
                << fFit.Tr().GetX() << " y " << fFit.Tr().GetY() << " z " << fFit.Tr().GetZ() << " tx "
                << fFit.Tr().GetTx() << " ty " << fFit.Tr().GetTy() << " qp " << fFit.Tr().GetQp() << " time "
                << fFit.Tr().GetTime();
    }
  };
 
  // fit trajectory in the measured area in several iterations
 
  for (int iter = 0; iter < fNofIterations && ok; iter++) {
 
    if (fVerbosityLevel > 0) {
      LOG(info) << fDebugPrefix << "---- Fit iteration " << iter << " ----";
    }
 
    {  // fit downstream from the first to the last measurement
 
      {  // initialisation at the first measurement
        auto& n = t.fNodes[firstHitNode];
        assert(n.isXySet);
        fFit.SetTrack(n.z, linearization[firstHitNode].linearizationDn);
        FilterFirstMeasurement(n);
        printNode("fit downstream", firstHitNode);
      }
 
      for (int iNode = firstHitNode + 1; iNode <= lastHitNode && ok; iNode++) {
 
        auto& n = t.fNodes[iNode];
 
        // transport the partially fitted track to the current node
        fFit.SetLinearization(linearization[iNode - 1].linearizationDn);
        fFit.Extrapolate(n.z, field);
 
        // measure the track at the current node
        if (n.isXySet) {
          fFit.FilterXY(n.measurementXY, fIgnorePoorCoordinates);
        }
        if (n.isTimeSet) {
          fFit.FilterTime(n.measurementTime);
        }
        printNode("fit downstream", iNode);
 
        // store the partially fitted track before the scattering
        n.paramUp = fFit.Tr();
 
        // add material effects
        if (n.materialLayer >= 0) {
          AddMaterialEffects(linearization[iNode], kf::FitDirection::kDownstream);
        }
 
        // store the partially fitted track after the scattering
        n.paramDn = fFit.Tr();
      }
    }  // fit downstream
 
    if (!ok) {
      break;
    }
 
    // store the chi2 for debugging purposes
    t.fQaChi2Downstream = fFit.Tr().GetChiSq();
 
    {  // fit upstream with the Kalman Filter smoothing
 
      {  // initialization at the last measurement
        auto& n = t.fNodes[lastHitNode];
        assert(n.isXySet);
        n.isFitted = true;
        fFit.SetTrack(n.z, linearization[lastHitNode].linearizationUp);
        FilterFirstMeasurement(n);
        printNode("fit upstream", lastHitNode);
      }
 
      for (int iNode = lastHitNode - 1; ok && (iNode > firstHitNode); iNode--) {
 
        auto& n = t.fNodes[iNode];
 
        // transport the partially fitted track to the current node
        fFit.SetLinearization(linearization[iNode + 1].linearizationUp);
        fFit.Extrapolate(n.z, field);
 
        // combine partially fitted downstream and upstream tracks
        ok = ok && Smooth(n.paramDn, fFit.Tr());
 
        if (n.materialLayer >= 0) {
          AddMaterialEffects(linearization[iNode], kf::FitDirection::kUpstream);
        }
 
        // combine partially fitted downstream and upstream tracks
        ok = ok && Smooth(n.paramUp, fFit.Tr());
        if (ok) {
          n.isFitted = true;
        }
 
        // measure the track at the current node
        if (n.isXySet) {
          fFit.FilterXY(n.measurementXY, fIgnorePoorCoordinates);
        }
        if (n.isTimeSet) {
          fFit.FilterTime(n.measurementTime);
        }
        printNode("fit upstream", iNode);
      }
 
      if (!ok) {
        break;
      }
 
      {  // the first measurement node
        int iNode = firstHitNode;
 
        auto& n = t.fNodes[iNode];
 
        // transport the partially fitted track to the current node
        fFit.SetLinearization(linearization[iNode + 1].linearizationUp);
        fFit.Extrapolate(n.z, field);
 
        // measure the track at the current node
        if (n.isXySet) {
          fFit.FilterXY(n.measurementXY, fIgnorePoorCoordinates);
        }
        if (n.isTimeSet) {
          fFit.FilterTime(n.measurementTime);
        }
        printNode("fit upstream", iNode);
        n.paramDn = fFit.Tr();
 
        if (n.materialLayer >= 0) {
          AddMaterialEffects(linearization[iNode], kf::FitDirection::kUpstream);
        }
        n.paramUp  = fFit.Tr();
        n.isFitted = true;
      }
    }  // fit upstream
 
    if (!ok) {
      t.fIsFitted = false;
      for (auto& n : t.fNodes) {
        n.isFitted = false;
      }
      return false;
    }
 
    t.fIsFitted = true;
 
    if (!fIsMaterialFixed) {
      for (int in = firstHitNode; in <= lastHitNode; in++) {
        auto& n = t.fNodes[in];
        if (n.materialLayer >= 0) {
          const kf::MaterialMap& map = fKfSetup->GetMaterial(n.materialLayer);
          n.radThick                 = map.GetThicknessX0(n.paramDn.GetX(), n.paramDn.GetY());
        }
        else {
          n.radThick = 0.;
        }
      }
    }
  }  // fit iterations
 
 
  {  // extrapolate the fitted trajectory outside the measured area downstream
    fFit.SetTrack(t.fNodes[lastHitNode].paramDn);
    for (int i = lastHitNode + 1; i < nNodes; i++) {
      auto& n = t.fNodes[i];
      fFit.Extrapolate(n.z, field);
      if (n.materialLayer >= 0) {
        n.radThick = fKfSetup->GetMaterial(n.materialLayer).GetThicknessX0(fFit.Tr().GetX(), fFit.Tr().GetY());
      }
      n.paramUp = fFit.Tr();
      if (n.materialLayer >= 0) {
        LinearizationAtNode l;
        l.linearizationUp = fFit.Tr();
        l.linearizationDn = fFit.Tr();
        l.radThick        = n.radThick;
        AddMaterialEffects(l, kf::FitDirection::kDownstream);
      }
      n.paramDn  = fFit.Tr();
      n.isFitted = true;
    }
  }
 
  // TODO: check when the extrapolation is failing
 
  {  // extrapolate the fitted trajectory outside the measured area upstream
    fFit.SetTrack(t.fNodes[firstHitNode].paramUp);
    for (int i = firstHitNode - 1; i >= 0; i--) {
      auto& n = t.fNodes[i];
      fFit.Extrapolate(n.z, field);
      n.paramDn = fFit.Tr();
      if (n.materialLayer >= 0) {
        n.radThick = fKfSetup->GetMaterial(n.materialLayer).GetThicknessX0(fFit.Tr().GetX(), fFit.Tr().GetY());
        LinearizationAtNode l;
        l.linearizationDn = fFit.Tr();
        l.linearizationUp = fFit.Tr();
        l.radThick        = n.radThick;
        AddMaterialEffects(l, kf::FitDirection::kUpstream);
      }
      n.paramUp  = fFit.Tr();
      n.isFitted = true;
    }
  }
 
  // distribute the final chi2, ndf to all nodes
 
  const auto& tt = fFit.Tr();
 
  for (int iNode = 0; iNode < nNodes; iNode++) {
    auto& n = t.fNodes[iNode];
    n.paramDn.SetNdf(tt.GetNdf());
    n.paramDn.SetNdfTime(tt.GetNdfTime());
    n.paramDn.SetChiSq(tt.GetChiSq());
    n.paramDn.SetChiSqTime(tt.GetChiSqTime());
    n.paramUp.SetNdf(tt.GetNdf());
    n.paramUp.SetNdfTime(tt.GetNdfTime());
    n.paramUp.SetChiSq(tt.GetChiSq());
    n.paramUp.SetChiSqTime(tt.GetChiSqTime());
  }
 
  t.fIsFullyExtrapolated = true;
 
  return ok;
}
 
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::FitTrajectoryUpstream(const cbm::algo::kf::Trajectory<>& t,
                                                          kf::TrackParamD& outputParam)
{
 
  // fit the track upstream using linearization stored in the trajectory
 
  outputParam = kf::TrackParamD();
 
  if (fVerbosityLevel > 0) {
    LOG(info) << fDebugPrefix << "-------- FitTrajectoryUpstream ... ";
  }
 
  {  // check input trajectory consistency
 
    if (!fIsInitialized) {
      LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
      return false;
    }
 
    if (!t.IsFitted()) {
      LOG(fatal) << fDebugPrefix << "the input trajectory is not fitted!";
    }
 
    if (!t.IsFullyExtrapolated()) {
      LOG(fatal) << fDebugPrefix << "the input trajectory is not fully extrapolated!";
    }
 
    for (unsigned int in = 0; in < t.fNodes.size(); in++) {
      const auto& n = t.fNodes[in];
      if (!n.isFitted) {
        LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                   << " is not fitted!";
      }
      if (n.materialLayer >= 0 && n.radThick < 0.) {
        LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                   << " with material layer " << n.materialLayer << " does not have material budget set!";
      }
    }
  }  // check input trajectory consistency
 
  bool ok = true;
 
  int nNodes = t.fNodes.size();
 
  if (nNodes <= 0) {
    LOG(warning) << fDebugPrefix << " no nodes found!";
    return false;
  }
 
  if (t.GetNofMeasurements() <= 0) {
    LOG(warning) << fDebugPrefix << "no nodes with measurements found!";
    return false;
  }
 
  fFit.SetParticleMass(fMass);
 
  kf::FieldRegion<double> field(kf::GlobalField::fgOriginalFieldType, kf::GlobalField::fgOriginalField);
 
  std::vector<LinearizationAtNode> linearization(nNodes);
 
  int lastHitNode = t.fLastMeasurementNodeId;
 
  // linearization initialization
  for (int i = 0; i <= lastHitNode; i++) {
    linearization[i].linearizationDn = t.fNodes[i].paramDn;
    linearization[i].linearizationUp = t.fNodes[i].paramUp;
    linearization[i].radThick        = t.fNodes[i].radThick;
  }
 
  auto printNode = [&](std::string str, int node) {
    if (fVerbosityLevel > 0) {
      LOG(info) << fDebugPrefix << str << ": node " << node << " chi2 " << fFit.Tr().GetChiSq() << " x "
                << fFit.Tr().GetX() << " y " << fFit.Tr().GetY() << " z " << fFit.Tr().GetZ() << " tx "
                << fFit.Tr().GetTx() << " ty " << fFit.Tr().GetTy() << " qp " << fFit.Tr().GetQp() << " time "
                << fFit.Tr().GetTime();
    }
  };
 
  // fit upstream
 
  {  // initialization at the last measurement
    const auto& n = t.fNodes[lastHitNode];
    assert(n.isXySet);
    fFit.SetTrack(n.paramUp);
    FilterFirstMeasurement(n);
    printNode("fit upstream", lastHitNode);
  }
 
  for (int iNode = lastHitNode - 1; ok && (iNode >= 0); iNode--) {
 
    const auto& n = t.fNodes[iNode];
 
    // transport the partially fitted track to the current node
    fFit.SetLinearization(linearization[iNode + 1].linearizationUp);
    fFit.Extrapolate(n.z, field);
 
    if (n.materialLayer >= 0) {
      AddMaterialEffects(linearization[iNode], kf::FitDirection::kUpstream);
    }
 
    // measure the track at the current node
    if (n.isXySet) {
      fFit.FilterXY(n.measurementXY, fIgnorePoorCoordinates);
    }
    if (n.isTimeSet) {
      fFit.FilterTime(n.measurementTime);
    }
    printNode("fit upstream", iNode);
  }
 
  if (ok) {
    outputParam = fFit.Tr();
  }
 
  return ok;
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::FitTrajectoryDownstream(const cbm::algo::kf::Trajectory<>& t,
                                                            kf::TrackParamD& outputParam)
{
 
  // fit the track downstream using linearization stored in the trajectory
 
  outputParam = kf::TrackParamD();
 
  if (fVerbosityLevel > 0) {
    LOG(info) << fDebugPrefix << "-------- FitTrajectoryDownstream ... ";
  }
 
  {  // check input trajectory consistency
 
    if (!fIsInitialized) {
      LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
      return false;
    }
 
    if (!t.IsFitted()) {
      LOG(fatal) << fDebugPrefix << "the input trajectory is not fitted!";
    }
 
    if (!t.IsFullyExtrapolated()) {
      LOG(fatal) << fDebugPrefix << "the input trajectory is not fully extrapolated!";
    }
 
    for (unsigned int in = 0; in < t.fNodes.size(); in++) {
      const auto& n = t.fNodes[in];
      if (!n.isFitted) {
        LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                   << " is not fitted!";
      }
      if (n.materialLayer >= 0 && n.radThick < 0.) {
        LOG(fatal) << fDebugPrefix << "the input trajectory is marked as fully extrapolated, but its node " << in
                   << " with material layer " << n.materialLayer << " does not have material budget set!";
      }
    }
  }  // check input trajectory consistency
 
  bool ok = true;
 
  int nNodes = t.fNodes.size();
 
  if (nNodes <= 0) {
    LOG(warning) << fVerbosityLevel << " no nodes found!";
    return false;
  }
 
  if (t.GetNofMeasurements() <= 0) {
    LOG(warning) << fVerbosityLevel << "no nodes with measurements found!";
    return false;
  }
 
  fFit.SetParticleMass(fMass);
 
  kf::FieldRegion<double> field(kf::GlobalField::fgOriginalFieldType, kf::GlobalField::fgOriginalField);
 
  std::vector<LinearizationAtNode> linearization(nNodes);
 
  int firstHitNode = t.fFirstMeasurementNodeId;
 
  // linearization initialization
  for (int i = firstHitNode; i < nNodes; i++) {
    linearization[i].linearizationDn = t.fNodes[i].paramDn;
    linearization[i].linearizationUp = t.fNodes[i].paramUp;
    linearization[i].radThick        = t.fNodes[i].radThick;
  }
 
  auto printNode = [&](std::string str, int node) {
    if (fVerbosityLevel > 0) {
      LOG(info) << fDebugPrefix << str << ": node " << node << " chi2 " << fFit.Tr().GetChiSq() << " x "
                << fFit.Tr().GetX() << " y " << fFit.Tr().GetY() << " z " << fFit.Tr().GetZ() << " tx "
                << fFit.Tr().GetTx() << " ty " << fFit.Tr().GetTy() << " qp " << fFit.Tr().GetQp() << " time "
                << fFit.Tr().GetTime();
    }
  };
 
  // fit downstream
 
  {  // initialization at the first measurement
    const auto& n = t.fNodes[firstHitNode];
    assert(n.isXySet);
    fFit.SetTrack(n.paramDn);
    FilterFirstMeasurement(n);
    printNode("fit downstream", firstHitNode);
  }
 
  for (int iNode = firstHitNode + 1; ok && (iNode < nNodes); iNode++) {
 
    const auto& n = t.fNodes[iNode];
 
    // transport the partially fitted track to the current node
    fFit.SetLinearization(linearization[iNode - 1].linearizationDn);
    fFit.Extrapolate(n.z, field);
 
    if (n.materialLayer >= 0) {
      AddMaterialEffects(linearization[iNode], kf::FitDirection::kDownstream);
    }
 
    // measure the track at the current node
    if (n.isXySet) {
      fFit.FilterXY(n.measurementXY, fIgnorePoorCoordinates);
    }
    if (n.isTimeSet) {
      fFit.FilterTime(n.measurementTime);
    }
    printNode("fit downstream", iNode);
  }
 
  if (ok) {
    outputParam = fFit.Tr();
  }
  return ok;
}
 
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::Smooth(kf::TrackParamD& t1, const kf::TrackParamD& t2)
{
  // TODO: move to the CaTrackFit class
 
  constexpr int nPar = kf::TrackParamD::kNtrackParam;
  constexpr int nCov = kf::TrackParamD::kNcovParam;
 
  double r[nPar] = {t1.X(), t1.Y(), t1.Tx(), t1.Ty(), t1.Qp(), t1.Time(), t1.Vi()};
  double m[nPar] = {t2.X(), t2.Y(), t2.Tx(), t2.Ty(), t2.Qp(), t2.Time(), t2.Vi()};
 
  double S[nCov], S1[nCov], Tmp[nCov];
  for (int i = 0; i < nCov; i++) {
    S[i] = (t1.GetCovMatrix()[i] + t2.GetCovMatrix()[i]);
  }
 
  int nullty = 0;
  int ifault = 0;
  cbm::algo::kf::utils::math::SymInv(S, nPar, S1, Tmp, &nullty, &ifault);
 
  if (fVerbosityLevel > 0 && ((0 != ifault) || (0 != nullty))) {
    if (fVerbosityLevel > 0) {
      LOG(error) << fDebugPrefix << "matrix inversion failed! ifault " << ifault << " nullty " << nullty;
      for (int i = 0; i < nCov; i++) {
        LOG(error) << fDebugPrefix << "Cov1 [" << i << "] = " << t1.GetCovMatrix()[i] << " Cov2[" << i
                   << "] = " << t2.GetCovMatrix()[i];
      }
      for (int i = 0; i < nCov; i++) {
        LOG(error) << fDebugPrefix << "S[" << i << "] = " << S[i] << " S1[" << i << "] = " << S1[i];
      }
    }
    return false;
  }
 
  double dzeta[nPar];
 
  for (int i = 0; i < nPar; i++) {
    dzeta[i] = m[i] - r[i];
  }
 
  double C[nPar][nPar];
  double Si[nPar][nPar];
  for (int i = 0, k = 0; i < nPar; i++) {
    for (int j = 0; j <= i; j++, k++) {
      C[i][j]  = t1.GetCovMatrix()[k];
      Si[i][j] = S1[k];
      C[j][i]  = C[i][j];
      Si[j][i] = Si[i][j];
    }
  }
 
  // K = C * S^{-1}
  double K[nPar][nPar];
  for (int i = 0; i < nPar; i++) {
    for (int j = 0; j < nPar; j++) {
      K[i][j] = 0.;
      for (int k = 0; k < nPar; k++) {
        K[i][j] += C[i][k] * Si[k][j];
      }
    }
  }
 
  for (int i = 0, k = 0; i < nPar; i++) {
    for (int j = 0; j <= i; j++, k++) {
      double kc = 0.;  // K*C[i][j]
      for (int l = 0; l < nPar; l++) {
        kc += K[i][l] * C[l][j];
      }
      t1.CovMatrix()[k] = t1.CovMatrix()[k] - kc;
    }
  }
 
  for (int i = 0; i < nPar; i++) {
    double kd = 0.;
    for (int j = 0; j < nPar; j++) {
      kd += K[i][j] * dzeta[j];
    }
    r[i] += kd;
  }
  t1.X()    = r[0];
  t1.Y()    = r[1];
  t1.Tx()   = r[2];
  t1.Ty()   = r[3];
  t1.Qp()   = r[4];
  t1.Time() = r[5];
  t1.Vi()   = r[6];
 
  double chi2     = 0.;
  double chi2Time = 0.;
  for (int i = 0; i < nPar; i++) {
    double SiDzeta = 0.;
    for (int j = 0; j < nPar; j++) {
      SiDzeta += Si[i][j] * dzeta[j];
    }
    if (i < 5) {
      chi2 += dzeta[i] * SiDzeta;
    }
    else {
      chi2Time += dzeta[i] * SiDzeta;
    }
  }
  t1.SetChiSq(chi2 + t1.GetChiSq() + t2.GetChiSq());
  t1.SetChiSqTime(chi2Time + t1.GetChiSqTime() + t2.GetChiSqTime());
  t1.SetNdf(5 + t1.GetNdf() + t2.GetNdf());
  t1.SetNdfTime(2 + t1.GetNdfTime() + t2.GetNdfTime());
  return true;
}
 
template<cbm::algo::kf::DoFitTime FlagFitTime>
bool CbmKfTrackFitter<FlagFitTime>::PropagateToZ(cbm::algo::kf::TrackParamD& track, double z)
{
  // propagate the track to the given Z position
 
  if (!fIsInitialized) {
    LOG(error) << fDebugPrefix << "CbmKfTrackFitter is not initialized!";
    return false;
  }
 
  kf::FieldRegion<double> field(kf::GlobalField::fgOriginalFieldType, kf::GlobalField::fgOriginalField);
  fFit.SetTrack(track);
  auto& tr = fFit.Tr();
 
  kf::FitDirection direction = (tr.GetZ() <= z) ? kf::FitDirection::kDownstream : kf::FitDirection::kUpstream;
 
  auto processMaterialLayer = [&](int i) {
    double zMaterial = fKfSetup->GetMaterial(i).GetZref();
 
    fFit.Extrapolate(zMaterial, field);
 
    double radThick = fKfSetup->GetMaterial(i).GetThicknessX0(tr.GetX(), tr.GetY());
    double msQp     = tr.GetQp();
    if (fEnforceDefaultMomentumForMs) {
      msQp = fDefaultQpForMs;
    }
    if (!fIgnoreMultipleScattering) {
      fFit.MultipleScattering(radThick, tr.GetTx(), tr.GetTy(), msQp);
    }
    if (!fEnforceDefaultMomentumForMs) {
      fFit.EnergyLossCorrection(radThick, direction);
    }
  };
 
  if (direction == kf::FitDirection::kDownstream) {
    for (int i = 0; i < fKfSetup->GetNofLayers(); i++) {
      double zMaterial = fKfSetup->GetMaterial(i).GetZref();
      if (zMaterial <= tr.GetZ()) {
        continue;  // skip upstream material layers
      }
      processMaterialLayer(i);
      if (zMaterial > z) {
        break;
      }
    }
  }
  else {  // upstream
    for (int i = fKfSetup->GetNofLayers() - 1; i >= 0; i--) {
      double zMaterial = fKfSetup->GetMaterial(i).GetZref();
      if (zMaterial >= tr.GetZ()) {
        continue;  // skip downstream material layers
      }
      if (zMaterial < z) {
        break;
      }
      processMaterialLayer(i);
    }
  }
 
  fFit.Extrapolate(z, field);
  track = fFit.Tr();
  return true;
}
 
 
template class CbmKfTrackFitter<cbm::algo::kf::DoFitTime::Y>;
template class CbmKfTrackFitter<cbm::algo::kf::DoFitTime::N>;