GaussianSumFitter.cxx

Go to the documentation of this file.

/*§
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/

 
#include "TrkGaussianSumFilter/GaussianSumFitter.h"
//
#include "TrkGaussianSumFilterUtils/GsfConstants.h"
#include "TrkGaussianSumFilterUtils/TwoStateCombiner.h"
//
#include "TrkCaloCluster_OnTrack/CaloCluster_OnTrack.h"
#include "TrkEventPrimitives/FitQuality.h"
#include "TrkEventUtils/MeasurementBaseComparisonFunction.h"
#include "TrkEventUtils/PrepRawDataComparisonFunction.h"
#include "TrkMaterialOnTrack/EstimatedBremOnTrack.h"
#include "TrkParameters/TrackParameters.h"
#include "TrkPrepRawData/PrepRawData.h"
#include "TrkPseudoMeasurementOnTrack/PseudoMeasurementOnTrack.h"
#include "TrkRIO_OnTrack/RIO_OnTrack.h"
#include "TrkSurfaces/PerigeeSurface.h"
#include "TrkTrack/Track.h"
#include "TrkTrack/TrackInfo.h"
//
#include <algorithm>
#include <vector>
 
namespace {

std::unique_ptr<Trk::FitQuality> buildFitQuality(
    const Trk::GaussianSumFitter::GSFTrajectory& smoothedTrajectory) {
  double chiSquared = 0.;
  int numberDoF = -5;
  Trk::GaussianSumFitter::GSFTrajectory::const_iterator stateOnSurface =
    smoothedTrajectory.begin();
  for (; stateOnSurface != smoothedTrajectory.end(); ++stateOnSurface) {
    if (!stateOnSurface->typeFlags.test(Trk::TrackStateOnSurface::Measurement)) {
      continue;
    }
    if (!stateOnSurface->fitQualityOnSurface) {
      continue;
    }
    chiSquared += stateOnSurface->fitQualityOnSurface.chiSquared();
    numberDoF += stateOnSurface->fitQualityOnSurface.numberDoF();
  }
  if (std::isnan(chiSquared) || chiSquared <= 0.) {
    return nullptr;
  }
  return std::make_unique<Trk::FitQuality>(chiSquared, numberDoF);
}

GSFTsos smootherHelper(Trk::MultiComponentState&& updatedState,
                       std::unique_ptr<Trk::MeasurementBase>&& measurement,
                       const Trk::FitQualityOnSurface& fitQuality,
                       bool combineToSingle, bool useMode) {
  if (combineToSingle) {
    auto combinedLastState = Trk::MultiComponentStateCombiner::combineToSingle(
        updatedState, useMode);
    if (combinedLastState) {
      return {fitQuality, std::move(measurement), std::move(combinedLastState),
              std::move(updatedState)};
    }
  }
  return {fitQuality, std::move(measurement), nullptr, std::move(updatedState)};
}

Trk::MeasurementSet stripMeasurements(
    const Trk::Track& inputTrack, const Trk::MeasurementSet* addMeasurementSet,
    bool reintegrateOutliers) {
  Trk::MeasurementSet measurementSet;
  for (const auto* tsos : *(inputTrack.trackStateOnSurfaces())) {
    if (!tsos) {
      continue;
    }
    bool use =
        tsos->type(Trk::TrackStateOnSurface::Measurement) ||
        (reintegrateOutliers && tsos->type(Trk::TrackStateOnSurface::Outlier));
    if (!use) {
      continue;
    }
    const Trk::MeasurementBase* meas = tsos->measurementOnTrack();
    if (meas) {
      measurementSet.push_back(meas);
    }
  }
 
  if (addMeasurementSet) {
    for (const auto* meas : *addMeasurementSet) {
      measurementSet.push_back(meas);
    }
  }
  return measurementSet;
}

Trk::PrepRawDataSet stripPrepRawData(
    const Trk::Track& inputTrack, const Trk::PrepRawDataSet* addPrepRawDataSet,
    bool reintegrateOutliers) {
 
  Trk::PrepRawDataSet prepRawDataSet;
  for (const auto* tsos : *(inputTrack.trackStateOnSurfaces())) {
    if (!tsos) {
      continue;
    }
    bool use =
        tsos->type(Trk::TrackStateOnSurface::Measurement) ||
        (reintegrateOutliers && tsos->type(Trk::TrackStateOnSurface::Outlier));
    if (!use) {
      continue;
    }
    const Trk::MeasurementBase* meas = tsos->measurementOnTrack();
    if (!meas) {
      continue;
    }
    const Trk::RIO_OnTrack* rioOnTrack = nullptr;
    if (meas->type(Trk::MeasurementBaseType::RIO_OnTrack)) {
      rioOnTrack = static_cast<const Trk::RIO_OnTrack*>(meas);
    }
    if (!rioOnTrack) {
      continue;
    }
    const Trk::PrepRawData* prepRawData = rioOnTrack->prepRawData();
    if (prepRawData) {
      prepRawDataSet.push_back(prepRawData);
    }
  }
  if (addPrepRawDataSet) {
    for (const auto* prepRawData : *addPrepRawDataSet) {
      prepRawDataSet.push_back(prepRawData);
    }
  }
  return prepRawDataSet;
}
 
} // end of anonymous namespace
 
Trk::GaussianSumFitter::GaussianSumFitter(const std::string& type,
                                          const std::string& name,
                                          const IInterface* parent)
  : AthAlgTool(type, name, parent)
  , m_trkParametersComparisonFunction{}
{
  declareInterface<ITrackFitter>(this);
}
 
StatusCode
Trk::GaussianSumFitter::initialize()
{
 
  if (m_maximumNumberOfComponents > GSFConstants::maxNumberofStateComponents) {
    ATH_MSG_FATAL("Requested MaximumNumberOfComponents > "
                  << GSFConstants::maxNumberofStateComponents);
    return StatusCode::FAILURE;
  }
  // GSF extrapolator
  ATH_CHECK(m_extrapolator.retrieve());
  // We need a to create RIO_OnTrack (calibrated/corrected)
  // measurements only if we start from PrePRawData.
  if (!m_refitOnMeasurementBase) {
    ATH_CHECK(m_rioOnTrackCreator.retrieve());
  } else {
    m_rioOnTrackCreator.disable();
  }
  // Initialise the closest track parameters search algorithm
  const Amg::Vector3D referencePosition(0, 0, 0);
  m_trkParametersComparisonFunction = Trk::TrkParametersComparisonFunction(referencePosition);
  m_particleHypothesis = m_extrapolator->particleHypothesis();
  return StatusCode::SUCCESS;
}

std::unique_ptr<Trk::Track>
Trk::GaussianSumFitter::fit(
  const EventContext& ctx,
  const Trk::Track& inputTrack,
  const Trk::RunOutlierRemoval outlierRemoval,
  const Trk::ParticleHypothesis particleHypothesis) const
{
  // Check that the input track has well defined parameters
  if (inputTrack.trackParameters()->empty()) {
    ATH_MSG_FATAL("No track parameters on Track!");
    return nullptr;
  }
  // Check that the input track has associated MeasurementBase objects
  if (inputTrack.trackStateOnSurfaces()->empty()) {
    ATH_MSG_FATAL("Empty MeasurementBase ");
    return nullptr;
  }
  // Retrieve the set of track parameters closest to the reference point
  const Trk::TrackParameters* paramNearestRef =
    *(std::min_element(inputTrack.trackParameters()->begin(),
                       inputTrack.trackParameters()->end(),
                       m_trkParametersComparisonFunction));
 
  // Rrefit using measurement base then
  if (m_refitOnMeasurementBase) {
    MeasurementSet measurementSet =
        stripMeasurements(inputTrack, nullptr, m_reintegrateOutliers);
    return fit(ctx, measurementSet, *paramNearestRef, outlierRemoval,
               particleHypothesis);
  }
  // Refit using PrepRawData level then
  PrepRawDataSet prepRawDataSet =
      stripPrepRawData(inputTrack, nullptr, m_reintegrateOutliers);
  return fit(ctx, prepRawDataSet, *paramNearestRef, outlierRemoval,
             particleHypothesis);
}

std::unique_ptr<Trk::Track>
Trk::GaussianSumFitter::fit(const EventContext& ctx,
                            const Track& intrk,
                            const PrepRawDataSet& addPrdColl,
                            const RunOutlierRemoval runOutlier,
                            const ParticleHypothesis matEffects) const
{
  // Check that the input track has well defined parameters
  if (intrk.trackParameters()->empty()) {
    ATH_MSG_FATAL("No track parameters on Track!");
    return nullptr;
  }
  // Check that the input track has associated MeasurementBase objects
  if (intrk.trackStateOnSurfaces()->empty()) {
    ATH_MSG_FATAL("Empty MeasurementBase ");
    return nullptr;
  }
  // determine the Track Parameter closest to the reference
  const TrackParameters* paramNearestRef = *(std::min_element(
      intrk.trackParameters()->begin(), intrk.trackParameters()->end(),
      m_trkParametersComparisonFunction));
  PrepRawDataSet prepRawDataSet =
      stripPrepRawData(intrk, &addPrdColl, m_reintegrateOutliers);
  return fit(ctx, prepRawDataSet, *paramNearestRef, runOutlier, matEffects);
}

std::unique_ptr<Trk::Track>
Trk::GaussianSumFitter::fit(const EventContext& ctx,
                            const Track& inputTrack,
                            const MeasurementSet& addSet,
                            const RunOutlierRemoval runOutlier,
                            const ParticleHypothesis matEffects) const
{
  // Check that the input track has well defined parameters
  if (inputTrack.trackParameters()->empty()) {
    ATH_MSG_FATAL("No estimation of track parameters near origin!");
    return nullptr;
  }
  // Check that the input track has associated MeasurementBase objects
  if (inputTrack.trackStateOnSurfaces()->empty()) {
    ATH_MSG_FATAL(
        "Attempting to fit track to empty MeasurementBase "
        "collection!");
    return nullptr;
  }
  // Retrieve the set of track parameters closest to the reference point
  const Trk::TrackParameters* paramNearestRef = *(std::min_element(
      inputTrack.trackParameters()->begin(),
      inputTrack.trackParameters()->end(), m_trkParametersComparisonFunction));
 
  MeasurementSet measurementSet =
      stripMeasurements(inputTrack, &addSet, m_reintegrateOutliers);
  return fit(ctx, measurementSet, *paramNearestRef, runOutlier, matEffects);
}

std::unique_ptr<Trk::Track> Trk::GaussianSumFitter::fit(
    const EventContext& ctx,
    const Track& intrk1,
    const Track& intrk2,
    const RunOutlierRemoval runOutlier,
    const ParticleHypothesis matEffects) const {
  // protection against not having measurements on the input tracks
  if (!intrk1.trackStateOnSurfaces() || !intrk2.trackStateOnSurfaces() ||
      intrk1.trackStateOnSurfaces()->size() < 2) {
    ATH_MSG_WARNING(
        "called to refit empty track or track with too little "
        "information, reject fit");
    return nullptr;
  }
  if (!intrk1.trackParameters() || intrk1.trackParameters()->empty()) {
    ATH_MSG_WARNING(
        "input #1 fails to provide track parameters for "
        "seeding the GXF, reject fit");
    return nullptr;
  }
  // Here we assume that the reference paramrs are the beginning
  // of the 1st track.
  const TrackParameters* minPar = *(intrk1.trackParameters()->begin());
  MeasurementSet measurementSet1 =
      stripMeasurements(intrk1, nullptr, m_reintegrateOutliers);
  MeasurementSet measurementSet2 =
      stripMeasurements(intrk2, nullptr, m_reintegrateOutliers);
 
  for (const auto* meas : measurementSet2) {
    measurementSet1.push_back(meas);
  }
 
  return fit(ctx, measurementSet1, *minPar, runOutlier, matEffects);
}

std::unique_ptr<Trk::Track>
Trk::GaussianSumFitter::fit(
  const EventContext& ctx,
  const Trk::PrepRawDataSet& prepRawDataSet,
  const Trk::TrackParameters& estimatedParametersNearOrigin,
  const Trk::RunOutlierRemoval /* Not used*/,
  const Trk::ParticleHypothesis /* Not used*/) const
{
  if (prepRawDataSet.size() < 3) {
    ATH_MSG_WARNING("Requesting Fit with less than three prep Raw Data!!!");
    return nullptr;
  }
  // We need a sorted PrepRawDataSet
  Trk::PrepRawDataSet sortedPrepRawDataSet = PrepRawDataSet(prepRawDataSet);
  const Trk::PrepRawDataComparisonFunction prdComparisonFunction =
      Trk::PrepRawDataComparisonFunction(
          estimatedParametersNearOrigin.position(),
          estimatedParametersNearOrigin.momentum());
 
  std::sort(sortedPrepRawDataSet.begin(), sortedPrepRawDataSet.end(),
            prdComparisonFunction);
 
  // Create Extrapolator cache that holds material effects cache;
  Trk::IMultiStateExtrapolator::Cache extrapolatorCache;
  // Perform GSF forward fit
  GSFTrajectory forwardTrajectory =
      forwardFit(ctx, extrapolatorCache, sortedPrepRawDataSet,
                 estimatedParametersNearOrigin);
 
  if (forwardTrajectory.empty()) {
    return nullptr;
  }
  // Perform GSF smoother operation
  GSFTrajectory smoothedTrajectory = smootherFit(
      ctx, extrapolatorCache, forwardTrajectory);
  if (smoothedTrajectory.empty()) {
    return nullptr;
  }
 
  // Fit quality
  std::unique_ptr<FitQuality> fitQuality = buildFitQuality(smoothedTrajectory);
  if (!fitQuality) {
    return nullptr;
  }
 
  // Create parameters at perigee if needed
  auto perigeeMultiStateOnSurface = makePerigee(
      ctx, extrapolatorCache, smoothedTrajectory);
  if (!perigeeMultiStateOnSurface.multiComponentState.empty()) {
    smoothedTrajectory.push_back(std::move(perigeeMultiStateOnSurface));
  } else {
    return nullptr;
  }
 
  // Reverse the order of the TSOS's after the smoother
  // to make be order flow from inside to out
  std::reverse(smoothedTrajectory.begin(), smoothedTrajectory.end());
 
  // Create Trk::Track
  Trk::TrackInfo info(Trk::TrackInfo::GaussianSumFilter, m_particleHypothesis);
  info.setTrackProperties(TrackInfo::BremFit);
  info.setTrackProperties(TrackInfo::BremFitSuccessful);
  return std::make_unique<Track>(info, convertTrajToTrack(smoothedTrajectory),
                                 std::move(fitQuality));
}

std::unique_ptr<Trk::Track>
Trk::GaussianSumFitter::fit(
  const EventContext& ctx,
  const Trk::MeasurementSet& measurementSet,
  const Trk::TrackParameters& estimatedParametersNearOrigin,
  const Trk::RunOutlierRemoval /* Not used*/,
  const Trk::ParticleHypothesis /*Not used*/) const
{
  if (measurementSet.size() < 3) {
    ATH_MSG_WARNING("Requesting fit with less than 3 Measurements!!!");
    return nullptr;
  }
  // We need to separate the possible CaloCluster on Track
  // measurement from the rest
  const Trk::CaloCluster_OnTrack* ccot(nullptr);
 
  Trk::MeasurementSet cleanedMeasurementSet;
  cleanedMeasurementSet.reserve(measurementSet.size());
 
  for (const auto* meas : measurementSet) {
    if (!meas) {
      ATH_MSG_WARNING("There is an empty MeasurementBase object ! ");
      continue;
    }
    if (meas->type(Trk::MeasurementBaseType::CaloCluster_OnTrack)) {
      ccot = static_cast<const Trk::CaloCluster_OnTrack*>(meas);
    } else {
      cleanedMeasurementSet.push_back(meas);
    }
  }
 
  // We need a sorted measurement set
  Trk::MeasurementSet sortedMeasurementSet =
      MeasurementSet(cleanedMeasurementSet);
 
  const Trk::MeasurementBaseComparisonFunction
      measurementBaseComparisonFunction(
          estimatedParametersNearOrigin.position(),
          estimatedParametersNearOrigin.momentum());
  sort(sortedMeasurementSet.begin(), sortedMeasurementSet.end(),
       measurementBaseComparisonFunction);
 
  // Create Extrapolator cache that holds material effects cache;
  Trk::IMultiStateExtrapolator::Cache extrapolatorCache;
 
  // Perform GSF forwards fit
  GSFTrajectory forwardTrajectory =
      forwardFit(ctx, extrapolatorCache, sortedMeasurementSet,
                 estimatedParametersNearOrigin);
 
  if (forwardTrajectory.empty()) {
    return nullptr;
  }
 
  // Perform GSF smoother operation
  GSFTrajectory smoothedTrajectory = smootherFit(
      ctx, extrapolatorCache, forwardTrajectory, ccot);
  if (smoothedTrajectory.empty()) {
    return nullptr;
  }
 
  // fit quality
  std::unique_ptr<FitQuality> fitQuality = buildFitQuality(smoothedTrajectory);
  if (!fitQuality) {
    return nullptr;
  }
 
  // Create parameters at perigee if needed
  auto perigeeMultiStateOnSurface = makePerigee(
      ctx, extrapolatorCache, smoothedTrajectory);
  if (!perigeeMultiStateOnSurface.multiComponentState.empty()) {
    smoothedTrajectory.push_back(std::move(perigeeMultiStateOnSurface));
  } else {
    return nullptr;
  }
 
  // Reverse the order of the TSOS's to make be order flow from inside to out
  std::reverse(smoothedTrajectory.begin(), smoothedTrajectory.end());
 
  // Create track
  Trk::TrackInfo info(Trk::TrackInfo::GaussianSumFilter, m_particleHypothesis);
  info.setTrackProperties(TrackInfo::BremFit);
  info.setTrackProperties(TrackInfo::BremFitSuccessful);
  return std::make_unique<Track>(info, convertTrajToTrack(smoothedTrajectory),
                                 std::move(fitQuality));
}

std::unique_ptr<MultiComponentStateOnSurfaceDV> Trk::GaussianSumFitter::convertTrajToTrack(
    GSFTrajectory& trajectory) const{
  const bool slimTransientMTSOS = m_slimTransientMTSOS;
  auto MTSOS = std::make_unique<MultiComponentStateOnSurfaceDV>();
  MTSOS->reserve(trajectory.size());
  for (GSFTsos& state : trajectory) {
    MTSOS->push_back(state.convert(slimTransientMTSOS));
  }
  return MTSOS;
}

GSFTsos
Trk::GaussianSumFitter::makePerigee(
  const EventContext& ctx,
  Trk::IMultiStateExtrapolator::Cache& extrapolatorCache,
  const GSFTrajectory& smoothedTrajectory) const
{
 
  // Start at the end of the smoothed trajectory
  // we should be the closest to perigee/origin.
  const GSFTsos&
    multiComponentStateOnSurfaceNearestOrigin = smoothedTrajectory.back();
  const Trk::MultiComponentState* multiComponentState =
    &(multiComponentStateOnSurfaceNearestOrigin.multiComponentState);
 
  // Extrapolate to perigee
  const Trk::PerigeeSurface perigeeSurface;
  Trk::MultiComponentState stateExtrapolatedToPerigee =
    m_extrapolator->extrapolate(ctx,
                                extrapolatorCache,
                                *multiComponentState,
                                perigeeSurface,
                                Trk::oppositeMomentum,
                                false);
 
  if (stateExtrapolatedToPerigee.empty()) {
    return {};
  }
 
  // Determine the combined state
  std::unique_ptr<Trk::TrackParameters> combinedPerigee =
    MultiComponentStateCombiner::combineToSingle(stateExtrapolatedToPerigee, m_useMode);
 
  // Perigee is an additional MultiComponentStateOnSurface
  std::bitset<Trk::TrackStateOnSurface::NumberOfTrackStateOnSurfaceTypes>
    pattern(0);
  pattern.set(Trk::TrackStateOnSurface::Perigee);
 
  if(std::abs(combinedPerigee->position().z())>5000.) {
    ATH_MSG_WARNING("Pathological perigee well outside of tracking detector!! Returning {}");
    return {};
  }
 
  if (std::abs(combinedPerigee->parameters()[Trk::qOverP]) > 1e8) {
    ATH_MSG_WARNING(
      "makePerigee() about to return with 0 momentum!! Returning {}");
    return {};
  }
 
  return {Trk::FitQualityOnSurface{},
          nullptr,
          std::move(combinedPerigee),
          std::move(stateExtrapolatedToPerigee),
          pattern};
}
 
/*
 * Implementation of the smoothing of the trajectory.
 * 1. We start from the forward trajectory containing
 * the measurement and prediction from the fwd fit.
 * 2. We do an update for the last state of the forward fit,
 * as this is special for the EDM
 * 3. Then we inflate the last state covariance matrix and we
 * run a filter in the backward direction.
 * 4. We optionally can use a two filter "Bayessian" smoother
 * ,see "Track fitting with non-Gaussian noise" R.Fruhwirth,
 * since we kept the "predictedStates" in the forward pass.
 *
 * This is also our chance to do any special treatement
 * needed for the ATLAS EDM.
 */
Trk::GaussianSumFitter::GSFTrajectory
Trk::GaussianSumFitter::smootherFit(
  const EventContext& ctx,
  Trk::IMultiStateExtrapolator::Cache& extrapolatorCache,
  GSFTrajectory& forwardTrajectory,
  const Trk::CaloCluster_OnTrack* ccot) const
{
  if (forwardTrajectory.empty()) {
    ATH_MSG_ERROR(
      "Attempting to smooth an empty forward trajectory... Exiting!");
    return {};
  }
 
  GSFTrajectory smoothedTrajectory;
  smoothedTrajectory.reserve(forwardTrajectory.size());
  // For the smoother we start from the end and we go
  // back. We need to find the the first track
  // state on surface  in this reverse direction
  auto trackStateOnSurfaceItr = forwardTrajectory.rbegin();
  bool foundMeasurement = false;
  for (; trackStateOnSurfaceItr != forwardTrajectory.rend(); ++trackStateOnSurfaceItr) {
    if (!(*trackStateOnSurfaceItr).typeFlags.test(TrackStateOnSurface::Measurement)) {
      smoothedTrajectory.emplace_back(std::move(*trackStateOnSurfaceItr));
    } else {
      foundMeasurement = true;
      break;
    }
  }
  if(!foundMeasurement){
    return {};
  }
  // This is the 1st track state on surface for a measurement
  // in the reverse direction. Our starting point for the smoother
  auto smootherPredictionMultiState = std::move(trackStateOnSurfaceItr->multiComponentState);
  // Perform an update with the measurement.
  // This will be the first entry in the trajectory
  std::unique_ptr<Trk::MeasurementBase> firstSmootherMeasurementOnTrack =
      std::move(trackStateOnSurfaceItr->measurementOnTrack);
  if (!firstSmootherMeasurementOnTrack) {
    ATH_MSG_WARNING(
        "Initial state on surface in smoother does not have an associated "
        "MeasurementBase object");
    return {};
  }
  Trk::FitQualityOnSurface fitQuality;
  Trk::MultiComponentState firstSmoothedState =
    Trk::GsfMeasurementUpdator::update(std::move(smootherPredictionMultiState),
                                       *firstSmootherMeasurementOnTrack,
                                       fitQuality);
  if (firstSmoothedState.empty()) {
    return {};
  }
  if (!MultiComponentStateHelpers::allHaveCovariance(firstSmoothedState)) {
    ATH_MSG_WARNING(
      "Not all components have covariance. Rejecting smoothed state.");
    return {};
  }
  // The first in reverse (last in normal order) TSOS is also special
  // for the EDM. Sso we also  a proper collapse
  // of the multi component to single TrackParameter
  std::unique_ptr<Trk::TrackParameters> combinedFirstSmoothedState =
      MultiComponentStateCombiner::combineToSingle(firstSmoothedState,
                                                   m_useMode);
  smoothedTrajectory.emplace_back(
      fitQuality, std::move(firstSmootherMeasurementOnTrack),
      std::move(combinedFirstSmoothedState),
      MultiComponentStateHelpers::clone(firstSmoothedState));
  const auto& updatedFirstStateOnSurface = smoothedTrajectory.back();
  //
  // Now we are ready to proceed
  // Generate a prediction with an inflated covariance of all components
  // in the first smoothed state.
  // This way there is no dependance on error of prediction
  //
  // NB local Y and theta need care due to TRT.
  //
  Trk::MultiComponentState smoothedStateWithScaledError =
      MultiComponentStateHelpers::WithScaledError(
          std::move(firstSmoothedState), m_smootherCovFactors[0],
          m_smootherCovFactors[1], m_smootherCovFactors[2],
          m_smootherCovFactors[3], m_smootherCovFactors[4]);
  // Do the 1st update using the state with inflated covariance.
  Trk::FitQualityOnSurface fitQualityWithScaledErrors;
  Trk::MultiComponentState updatedState = Trk::GsfMeasurementUpdator::update(
    std::move(smoothedStateWithScaledError),
    *(updatedFirstStateOnSurface.measurementOnTrack),
    fitQualityWithScaledErrors);
  if (updatedState.empty()) {
    ATH_MSG_WARNING("Smoother prediction could not be determined");
    return {};
  }
  // loopUpdatedState points to the most recent predicted MultiComponentState.
  Trk::MultiComponentState* loopUpdatedState = &updatedState;
  // Increase reverse iterator by one.
  ++trackStateOnSurfaceItr;
  // The is the last one we will see as we go back.
  auto lasttrackStateOnSurface = forwardTrajectory.rend() - 1;
  // And this is the pre-last
  auto secondLastTrackStateOnSurface = forwardTrajectory.rend() - 2;
  //Loop
  for (; trackStateOnSurfaceItr != forwardTrajectory.rend();++trackStateOnSurfaceItr) {
    auto& trackStateOnSurface = (*trackStateOnSurfaceItr);
    // Retrieve the MeasurementBase object from the TrackStateOnSurface object
    auto measurement = std::move(trackStateOnSurface.measurementOnTrack);
    if (!measurement) {
      ATH_MSG_WARNING("MeasurementBase object could not be extracted from a "
                      "measurement TSOS...continuing");
      continue;
    }
    // Create prediction for the next measurement surface.
    // For the smoother the direction of propagation
    // is opposite to the direction of momentum
    Trk::MultiComponentState extrapolatedState = m_extrapolator->extrapolate(
        ctx, extrapolatorCache, (*loopUpdatedState),
        measurement->associatedSurface(), Trk::oppositeMomentum, false);
 
    if (extrapolatedState.empty()) {
      return {};
    }
    // Handle the case where Original measurement was flagged as an outlier.
    if (!trackStateOnSurface.typeFlags.test(TrackStateOnSurface::Measurement)) {
      std::bitset<TrackStateOnSurface::NumberOfTrackStateOnSurfaceTypes> type(
          0);
      type.set(TrackStateOnSurface::Outlier);
      smoothedTrajectory.emplace_back(FitQualityOnSurface(1, 1),
                                      std::move(measurement), nullptr,
                                      std::move(extrapolatedState), type);
      loopUpdatedState = &(smoothedTrajectory.back().multiComponentState);
      continue;
    }
    // Update with the measurement
    updatedState = Trk::GsfMeasurementUpdator::update(
        std::move(extrapolatedState), *measurement, fitQuality);
    if (updatedState.empty()) {
      ATH_MSG_WARNING("Could not update the multi-component state");
      return {};
    }
    // last in reverse (first in normal order) is special for the EDM
    const bool islast = (trackStateOnSurfaceItr == lasttrackStateOnSurface);
    if (m_combineWithFitter) {
      // Optional combine smoother state with fitter state
      const Trk::MultiComponentState& forwardsMultiState =
          trackStateOnSurface.multiComponentState;
      Trk::MultiComponentState combinedfitterState =
          Trk::TwoStateCombiner::combine(forwardsMultiState, updatedState,
                                         m_maximumNumberOfComponents);
      if (combinedfitterState.empty()) {
        ATH_MSG_WARNING(
            "Could not combine state from forward fit with "
            "smoother state");
        return {};
      }
      auto combinedFitQuality = Trk::GsfMeasurementUpdator::fitQuality(
          combinedfitterState, *measurement);
      smoothedTrajectory.emplace_back(
          smootherHelper(std::move(combinedfitterState), std::move(measurement),
                         combinedFitQuality, islast, m_useMode));
    } else {
      // If combination with forwards state is not done
      smoothedTrajectory.emplace_back(
          smootherHelper(std::move(updatedState), std::move(measurement),
                         fitQuality, islast, m_useMode));
    }
    // Handle adding measurement from calo if it is present
    if (ccot && trackStateOnSurfaceItr == secondLastTrackStateOnSurface) {
      if (!addCCOT(ctx, ccot, smoothedTrajectory)) {
        ATH_MSG_WARNING("Could not add Calo Cluster On Track Measurement");
        return {};
      }
    }
    // For the next iteration start from last added
    loopUpdatedState = &(smoothedTrajectory.back().multiComponentState);
  }  // End for loop over all components
  return smoothedTrajectory;
}

bool
Trk::GaussianSumFitter::addCCOT(
  const EventContext& ctx,
  const Trk::CaloCluster_OnTrack* ccot,
  GSFTrajectory& smoothedTrajectory) const
{
  const GSFTsos& currentMultiStateOS = smoothedTrajectory.back();
  if (!ccot) {
    return false;
  }
  const auto& currentMultiComponentState = currentMultiStateOS.multiComponentState;
  const Trk::MeasurementBase* measurement = currentMultiStateOS.measurementOnTrack.get();
  if (!measurement) {
     return false;
  }
  const Trk::Surface& currentSurface = measurement->associatedSurface();
 
  //Create a ccot that we will own. So we use it to build our own TSOS
  auto ownCCOT = ccot->uniqueClone();
  // Extrapolate to the Calo to get prediction
  Trk::MultiComponentState extrapolatedToCaloState = m_extrapolator->extrapolateDirectly(
    ctx, currentMultiComponentState, ownCCOT->associatedSurface(),
    Trk::alongMomentum, false);
 
  if (extrapolatedToCaloState.empty()) {
    return false;
  }
  // Update newly extrapolated state with measurement
  Trk::FitQualityOnSurface fitQuality;
  Trk::MultiComponentState updatedStateAtCalo =
      Trk::GsfMeasurementUpdator::update(std::move(extrapolatedToCaloState),
                                         *ownCCOT, fitQuality);
  if (updatedStateAtCalo.empty()) {
    return false;
  }
 
  // Extrapolate back to the surface near the origin
  auto improvedState = m_extrapolator->extrapolateDirectly(
      ctx, updatedStateAtCalo, currentSurface, Trk::oppositeMomentum,
      false);
 
  if (improvedState.empty()) {
    return false;
  }
 // Combine the improved state after extrapolating back from the calo
 // and find the mode of the distribution
  std::unique_ptr<Trk::TrackParameters> combinedSingleState =
      MultiComponentStateCombiner::combineToSingle(improvedState, m_useMode);
 
  // Now build a dummy measurement for the improved estimation
  AmgSymMatrix(5) covMatrix;
  covMatrix.setZero();
  covMatrix(0, 0) = 1e6;
  const Trk::DefinedParameter locX(0, Trk::locX);
  Trk::LocalParameters locpars(locX);
  auto pseudoMeasurement = std::make_unique<Trk::PseudoMeasurementOnTrack>(
    std::move(locpars), std::move(covMatrix), currentSurface);
 
  // Build TSOS at the surface of calo
  smoothedTrajectory.emplace_back(
      fitQuality,
      std::move(ownCCOT),
      nullptr,
      std::move(updatedStateAtCalo));
 
  // Build a TSOS using the dummy measurement and and the final combined state
  smoothedTrajectory.emplace_back(
      FitQualityOnSurface{},
      // We do not add the fitquality again. As this would be the one we
      // add at calo. Here not a real measurement
      std::move(pseudoMeasurement),
      std::move(combinedSingleState),
      std::move(improvedState));
 
  return true;
}