GaussianSumFitter.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file GaussianSumFitter.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2021 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once

#include "Acts/EventData/TrackHelpers.hpp"
#include "Acts/EventData/VectorMultiTrajectory.hpp"
#include "Acts/Propagator/EigenStepper.hpp"
#include "Acts/Propagator/MultiStepperAborters.hpp"
#include "Acts/Propagator/StandardAborters.hpp"
#include "Acts/TrackFitting/GsfOptions.hpp"
#include "Acts/TrackFitting/KalmanFitter.hpp"
#include "Acts/TrackFitting/detail/GsfActor.hpp"
#include "Acts/Utilities/Logger.hpp"

#include <fstream>

namespace Acts {

namespace detail {

template <typename T>
struct IsMultiComponentBoundParameters : public std::false_type {};

template <>
struct IsMultiComponentBoundParameters<MultiComponentBoundTrackParameters>
    : public std::true_type {};

}  // namespace detail

template <typename propagator_t, typename bethe_heitler_approx_t,
          typename traj_t>
struct GaussianSumFitter {
  GaussianSumFitter(propagator_t&& propagator, bethe_heitler_approx_t&& bha,
                    std::unique_ptr<const Logger> _logger =
                        getDefaultLogger("GSF", Logging::INFO))
      : m_propagator(std::move(propagator)),
        m_betheHeitlerApproximation(std::move(bha)),
        m_logger{std::move(_logger)},
        m_actorLogger(m_logger->cloneWithSuffix("Actor")) {}

  propagator_t m_propagator;

  bethe_heitler_approx_t m_betheHeitlerApproximation;

  std::unique_ptr<const Logger> m_logger;
  std::unique_ptr<const Logger> m_actorLogger;

  const Logger& logger() const { return *m_logger; }

  using GsfNavigator = typename propagator_t::Navigator;

  using GsfActor = detail::GsfActor<bethe_heitler_approx_t, traj_t>;

  template <typename source_link_it_t, typename start_parameters_t,
            typename track_container_t, template <typename> class holder_t>
  auto fit(source_link_it_t begin, source_link_it_t end,
           const start_parameters_t& sParameters,
           const GsfOptions<traj_t>& options,
           const std::vector<const Surface*>& sSequence,
           TrackContainer<track_container_t, traj_t, holder_t>& trackContainer)
      const {
    // Check if we have the correct navigator
    static_assert(
        std::is_same_v<DirectNavigator, typename propagator_t::Navigator>);

    // Initialize the forward propagation with the DirectNavigator
    auto fwdPropInitializer = [&sSequence, this](const auto& opts) {
      using Actors = ActionList<GsfActor, DirectNavigator::Initializer>;
      using Aborters = AbortList<>;

      PropagatorOptions<Actors, Aborters> propOptions(opts.geoContext,
                                                      opts.magFieldContext);

      propOptions.setPlainOptions(opts.propagatorPlainOptions);

      propOptions.actionList.template get<DirectNavigator::Initializer>()
          .navSurfaces = sSequence;
      propOptions.actionList.template get<GsfActor>()
          .m_cfg.bethe_heitler_approx = &m_betheHeitlerApproximation;

      return propOptions;
    };

    // Initialize the backward propagation with the DirectNavigator
    auto bwdPropInitializer = [&sSequence, this](const auto& opts) {
      using Actors = ActionList<GsfActor, DirectNavigator::Initializer>;
      using Aborters = AbortList<>;

      std::vector<const Surface*> backwardSequence(
          std::next(sSequence.rbegin()), sSequence.rend());
      backwardSequence.push_back(opts.referenceSurface);

      PropagatorOptions<Actors, Aborters> propOptions(opts.geoContext,
                                                      opts.magFieldContext);

      propOptions.setPlainOptions(opts.propagatorPlainOptions);

      propOptions.actionList.template get<DirectNavigator::Initializer>()
          .navSurfaces = std::move(backwardSequence);
      propOptions.actionList.template get<GsfActor>()
          .m_cfg.bethe_heitler_approx = &m_betheHeitlerApproximation;

      return propOptions;
    };

    return fit_impl(begin, end, sParameters, options, fwdPropInitializer,
                    bwdPropInitializer, trackContainer);
  }

  template <typename source_link_it_t, typename start_parameters_t,
            typename track_container_t, template <typename> class holder_t>
  auto fit(source_link_it_t begin, source_link_it_t end,
           const start_parameters_t& sParameters,
           const GsfOptions<traj_t>& options,
           TrackContainer<track_container_t, traj_t, holder_t>& trackContainer)
      const {
    // Check if we have the correct navigator
    static_assert(std::is_same_v<Navigator, typename propagator_t::Navigator>);

    // Initialize the forward propagation with the DirectNavigator
    auto fwdPropInitializer = [this](const auto& opts) {
      using Actors = ActionList<GsfActor>;
      using Aborters = AbortList<EndOfWorldReached>;

      PropagatorOptions<Actors, Aborters> propOptions(opts.geoContext,
                                                      opts.magFieldContext);
      propOptions.setPlainOptions(opts.propagatorPlainOptions);
      propOptions.actionList.template get<GsfActor>()
          .m_cfg.bethe_heitler_approx = &m_betheHeitlerApproximation;

      return propOptions;
    };

    // Initialize the backward propagation with the DirectNavigator
    auto bwdPropInitializer = [this](const auto& opts) {
      using Actors = ActionList<GsfActor>;
      using Aborters = AbortList<EndOfWorldReached>;

      PropagatorOptions<Actors, Aborters> propOptions(opts.geoContext,
                                                      opts.magFieldContext);

      propOptions.setPlainOptions(opts.propagatorPlainOptions);

      propOptions.actionList.template get<GsfActor>()
          .m_cfg.bethe_heitler_approx = &m_betheHeitlerApproximation;

      return propOptions;
    };

    return fit_impl(begin, end, sParameters, options, fwdPropInitializer,
                    bwdPropInitializer, trackContainer);
  }

  template <typename source_link_it_t, typename start_parameters_t,
            typename fwd_prop_initializer_t, typename bwd_prop_initializer_t,
            typename track_container_t, template <typename> class holder_t>
  Acts::Result<
      typename TrackContainer<track_container_t, traj_t, holder_t>::TrackProxy>
  fit_impl(source_link_it_t begin, source_link_it_t end,
           const start_parameters_t& sParameters,
           const GsfOptions<traj_t>& options,
           const fwd_prop_initializer_t& fwdPropInitializer,
           const bwd_prop_initializer_t& bwdPropInitializer,
           TrackContainer<track_container_t, traj_t, holder_t>& trackContainer)
      const {
    // return or abort utility
    auto return_error_or_abort = [&](auto error) {
      if (options.abortOnError) {
        std::abort();
      }
      return error;
    };

    // Define directions based on input propagation direction. This way we can
    // refer to 'forward' and 'backward' regardless of the actual direction.
    const auto gsfForward = options.propagatorPlainOptions.direction;
    const auto gsfBackward = gsfForward.invert();

    // Check if the start parameters are on the start surface
    auto intersectionStatusStartSurface =
        sParameters.referenceSurface()
            .intersect(GeometryContext{},
                       sParameters.position(GeometryContext{}),
                       sParameters.direction(), true)
            .closest()
            .status();

    if (intersectionStatusStartSurface != Intersection3D::Status::onSurface) {
      ACTS_ERROR(
          "Surface intersection of start parameters with bound-check failed");
      return GsfError::StartParametersNotOnStartSurface;
    }

    // To be able to find measurements later, we put them into a map
    // We need to copy input SourceLinks anyway, so the map can own them.
    ACTS_VERBOSE("Preparing " << std::distance(begin, end)
                              << " input measurements");
    std::map<GeometryIdentifier, SourceLink> inputMeasurements;
    for (auto it = begin; it != end; ++it) {
      SourceLink sl = *it;
      inputMeasurements.emplace(
          options.extensions.surfaceAccessor(sl)->geometryId(), std::move(sl));
    }

    ACTS_VERBOSE(
        "Gsf: Final measurement map size: " << inputMeasurements.size());

    if (sParameters.covariance() == std::nullopt) {
      return GsfError::StartParametersHaveNoCovariance;
    }

    // Forward pass
    ACTS_VERBOSE("+-----------------------------+");
    ACTS_VERBOSE("| Gsf: Do forward propagation |");
    ACTS_VERBOSE("+-----------------------------+");

    auto fwdResult = [&]() {
      auto fwdPropOptions = fwdPropInitializer(options);

      // Catch the actor and set the measurements
      auto& actor = fwdPropOptions.actionList.template get<GsfActor>();
      actor.setOptions(options);
      actor.m_cfg.inputMeasurements = &inputMeasurements;
      actor.m_cfg.numberMeasurements = inputMeasurements.size();
      actor.m_cfg.inReversePass = false;
      actor.m_cfg.logger = m_actorLogger.get();

      fwdPropOptions.direction = gsfForward;

      // If necessary convert to MultiComponentBoundTrackParameters
      using IsMultiParameters =
          detail::IsMultiComponentBoundParameters<start_parameters_t>;

      typename propagator_t::template action_list_t_result_t<
          MultiComponentCurvilinearTrackParameters,
          decltype(fwdPropOptions.actionList)>
          inputResult;

      auto& r = inputResult.template get<typename GsfActor::result_type>();

      r.fittedStates = &trackContainer.trackStateContainer();

      // This allows the initialization with single- and multicomponent start
      // parameters
      if constexpr (not IsMultiParameters::value) {
        MultiComponentBoundTrackParameters params(
            sParameters.referenceSurface().getSharedPtr(),
            sParameters.parameters(), *sParameters.covariance(),
            sParameters.particleHypothesis());

        return m_propagator.propagate(params, fwdPropOptions, false,
                                      std::move(inputResult));
      } else {
        return m_propagator.propagate(sParameters, fwdPropOptions, false,
                                      std::move(inputResult));
      }
    }();

    if (!fwdResult.ok()) {
      return return_error_or_abort(fwdResult.error());
    }

    auto& fwdGsfResult =
        fwdResult->template get<typename GsfActor::result_type>();

    if (!fwdGsfResult.result.ok()) {
      return return_error_or_abort(fwdGsfResult.result.error());
    }

    if (fwdGsfResult.measurementStates == 0) {
      return return_error_or_abort(GsfError::NoMeasurementStatesCreatedForward);
    }

    ACTS_VERBOSE("Finished forward propagation");
    ACTS_VERBOSE("- visited surfaces: " << fwdGsfResult.visitedSurfaces.size());
    ACTS_VERBOSE("- processed states: " << fwdGsfResult.processedStates);
    ACTS_VERBOSE("- measurement states: " << fwdGsfResult.measurementStates);

    std::size_t nInvalidBetheHeitler = fwdGsfResult.nInvalidBetheHeitler;

    // Backward pass
    ACTS_VERBOSE("+------------------------------+");
    ACTS_VERBOSE("| Gsf: Do backward propagation |");
    ACTS_VERBOSE("+------------------------------+");

    auto bwdResult = [&]() {
      auto bwdPropOptions = bwdPropInitializer(options);

      auto& actor = bwdPropOptions.actionList.template get<GsfActor>();
      actor.setOptions(options);
      actor.m_cfg.inputMeasurements = &inputMeasurements;
      actor.m_cfg.inReversePass = true;
      actor.m_cfg.logger = m_actorLogger.get();
      actor.setOptions(options);

      bwdPropOptions.direction = gsfBackward;

      const Surface& target = options.referenceSurface
                                  ? *options.referenceSurface
                                  : sParameters.referenceSurface();

      using PM = TrackStatePropMask;

      typename propagator_t::template action_list_t_result_t<
          MultiComponentBoundTrackParameters,
          decltype(bwdPropOptions.actionList)>
          inputResult;

      // Unfortunately we must construct the result type here to be able to
      // return an error code
      using ResultType =
          decltype(m_propagator.template propagate<
                   MultiComponentBoundTrackParameters, decltype(bwdPropOptions),
                   MultiStepperSurfaceReached>(
              std::declval<MultiComponentBoundTrackParameters>(),
              std::declval<Acts::Surface&>(),
              std::declval<decltype(bwdPropOptions)>(),
              std::declval<decltype(inputResult)>()));

      auto& r = inputResult.template get<typename GsfActor::result_type>();

      r.fittedStates = &trackContainer.trackStateContainer();

      assert(
          (fwdGsfResult.lastMeasurementTip != MultiTrajectoryTraits::kInvalid &&
           "tip is invalid"));

      auto proxy =
          r.fittedStates->getTrackState(fwdGsfResult.lastMeasurementTip);
      proxy.shareFrom(TrackStatePropMask::Filtered,
                      TrackStatePropMask::Smoothed);

      r.currentTip = fwdGsfResult.lastMeasurementTip;
      r.visitedSurfaces.push_back(&proxy.referenceSurface());
      r.surfacesVisitedBwdAgain.push_back(&proxy.referenceSurface());
      r.measurementStates++;
      r.processedStates++;

      const auto& params = *fwdGsfResult.lastMeasurementState;

      return m_propagator.template propagate<std::decay_t<decltype(params)>,
                                             decltype(bwdPropOptions),
                                             MultiStepperSurfaceReached>(
          params, target, bwdPropOptions, std::move(inputResult));
    }();

    if (!bwdResult.ok()) {
      return return_error_or_abort(bwdResult.error());
    }

    auto& bwdGsfResult =
        bwdResult->template get<typename GsfActor::result_type>();

    if (!bwdGsfResult.result.ok()) {
      return return_error_or_abort(bwdGsfResult.result.error());
    }

    if (bwdGsfResult.measurementStates == 0) {
      return return_error_or_abort(
          GsfError::NoMeasurementStatesCreatedBackward);
    }

    nInvalidBetheHeitler += bwdGsfResult.nInvalidBetheHeitler;

    if (nInvalidBetheHeitler > 0) {
      ACTS_WARNING("Encountered "
                   << nInvalidBetheHeitler
                   << " cases where the material thickness exceeds the range "
                      "of the Bethe-Heitler-Approximation. Enable DEBUG output "
                      "for more information.");
    }

    // Create Kalman Result
    ACTS_VERBOSE("Gsf - States summary:");
    ACTS_VERBOSE("- Fwd measurement states: " << fwdGsfResult.measurementStates
                                              << ", holes: "
                                              << fwdGsfResult.measurementHoles);
    ACTS_VERBOSE("- Bwd measurement states: " << bwdGsfResult.measurementStates
                                              << ", holes: "
                                              << bwdGsfResult.measurementHoles);

    // TODO should this be warning level? it happens quite often... Investigate!
    if (bwdGsfResult.measurementStates != fwdGsfResult.measurementStates) {
      ACTS_DEBUG("Fwd and bwd measurement states do not match");
    }

    // Go through the states and assign outliers / unset smoothed if surface not
    // passed in backward pass
    const auto& foundBwd = bwdGsfResult.surfacesVisitedBwdAgain;
    std::size_t measurementStatesFinal = 0;

    for (auto state : fwdGsfResult.fittedStates->reverseTrackStateRange(
             fwdGsfResult.currentTip)) {
      const bool found = std::find(foundBwd.begin(), foundBwd.end(),
                                   &state.referenceSurface()) != foundBwd.end();
      if (not found && state.typeFlags().test(MeasurementFlag)) {
        state.typeFlags().set(OutlierFlag);
        state.typeFlags().reset(MeasurementFlag);
        state.unset(TrackStatePropMask::Smoothed);
      }

      measurementStatesFinal +=
          static_cast<std::size_t>(state.typeFlags().test(MeasurementFlag));
    }

    if (measurementStatesFinal == 0) {
      return return_error_or_abort(GsfError::NoMeasurementStatesCreatedFinal);
    }

    auto track = trackContainer.getTrack(trackContainer.addTrack());
    track.tipIndex() = fwdGsfResult.lastMeasurementTip;

    if (options.referenceSurface) {
      const auto& params = *bwdResult->endParameters;

      const auto [finalPars, finalCov] = Acts::reduceGaussianMixture(
          params.components(), params.referenceSurface(),
          options.stateReductionMethod, [](auto& t) {
            return std::tie(std::get<0>(t), std::get<1>(t), *std::get<2>(t));
          });

      track.parameters() = finalPars;
      track.covariance() = finalCov;

      track.setReferenceSurface(params.referenceSurface().getSharedPtr());

      if (trackContainer.hasColumn(
              hashString(GsfConstants::kFinalMultiComponentStateColumn))) {
        ACTS_DEBUG("Add final multi-component state to track")
        track.template component<GsfConstants::FinalMultiComponentState>(
            GsfConstants::kFinalMultiComponentStateColumn) = std::move(params);
      }
    }

    calculateTrackQuantities(track);

    track.nMeasurements() = measurementStatesFinal;
    track.nHoles() = fwdGsfResult.measurementHoles;

    return track;
  }
};

}  // namespace Acts