MultiEigenStepperLoop.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file MultiEigenStepperLoop.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

// Workaround for building on clang+libstdc++
#include "Acts/Utilities/detail/ReferenceWrapperAnyCompat.hpp"

#include "Acts/Definitions/Algebra.hpp"
#include "Acts/Definitions/Direction.hpp"
#include "Acts/Definitions/TrackParametrization.hpp"
#include "Acts/Definitions/Units.hpp"
#include "Acts/EventData/MultiComponentTrackParameters.hpp"
#include "Acts/EventData/TrackParameters.hpp"
#include "Acts/EventData/detail/CorrectedTransformationFreeToBound.hpp"
#include "Acts/MagneticField/MagneticFieldProvider.hpp"
#include "Acts/Propagator/ConstrainedStep.hpp"
#include "Acts/Propagator/EigenStepper.hpp"
#include "Acts/Propagator/EigenStepperError.hpp"
#include "Acts/Propagator/Propagator.hpp"
#include "Acts/Surfaces/Surface.hpp"
#include "Acts/Utilities/GaussianMixtureReduction.hpp"
#include "Acts/Utilities/Intersection.hpp"
#include "Acts/Utilities/Result.hpp"

#include <algorithm>
#include <cmath>
#include <cstddef>
#include <functional>
#include <limits>
#include <numeric>
#include <sstream>
#include <vector>

#include <boost/container/small_vector.hpp>

#include "MultiStepperError.hpp"

namespace Acts {

using namespace Acts::UnitLiterals;

struct WeightedComponentReducerLoop {
  template <typename component_range_t>
  static Vector3 toVector3(const component_range_t& comps,
                           const FreeIndices i) {
    return std::accumulate(
        comps.begin(), comps.end(), Vector3{Vector3::Zero()},
        [i](const auto& sum, const auto& cmp) -> Vector3 {
          return sum + cmp.weight * cmp.state.pars.template segment<3>(i);
        });
  }

  template <typename stepper_state_t>
  static Vector3 position(const stepper_state_t& s) {
    return toVector3(s.components, eFreePos0);
  }

  template <typename stepper_state_t>
  static Vector3 direction(const stepper_state_t& s) {
    return toVector3(s.components, eFreeDir0).normalized();
  }

  // TODO: Maybe we can cache this value and only update it when the parameters
  // change
  template <typename stepper_state_t>
  static ActsScalar qOverP(const stepper_state_t& s) {
    return std::accumulate(
        s.components.begin(), s.components.end(), ActsScalar{0.},
        [](const auto& sum, const auto& cmp) -> ActsScalar {
          return sum + cmp.weight * cmp.state.pars[eFreeQOverP];
        });
  }

  template <typename stepper_state_t>
  static ActsScalar absoluteMomentum(const stepper_state_t& s) {
    return std::accumulate(
        s.components.begin(), s.components.end(), ActsScalar{0.},
        [&s](const auto& sum, const auto& cmp) -> ActsScalar {
          return sum + cmp.weight * s.particleHypothesis.extractMomentum(
                                        cmp.state.pars[eFreeQOverP]);
        });
  }

  template <typename stepper_state_t>
  static Vector3 momentum(const stepper_state_t& s) {
    return std::accumulate(
        s.components.begin(), s.components.end(), Vector3::Zero().eval(),
        [&s](const auto& sum, const auto& cmp) -> Vector3 {
          return sum + cmp.weight *
                           s.particleHypothesis.extractMomentum(
                               cmp.state.pars[eFreeQOverP]) *
                           cmp.state.pars.template segment<3>(eFreeDir0);
        });
  }

  template <typename stepper_state_t>
  static ActsScalar charge(const stepper_state_t& s) {
    return std::accumulate(
        s.components.begin(), s.components.end(), ActsScalar{0.},
        [&s](const auto& sum, const auto& cmp) -> ActsScalar {
          return sum + cmp.weight * s.particleHypothesis.extractCharge(
                                        cmp.state.pars[eFreeQOverP]);
        });
  }

  template <typename stepper_state_t>
  static ActsScalar time(const stepper_state_t& s) {
    return std::accumulate(
        s.components.begin(), s.components.end(), ActsScalar{0.},
        [](const auto& sum, const auto& cmp) -> ActsScalar {
          return sum + cmp.weight * cmp.state.pars[eFreeTime];
        });
  }

  template <typename stepper_state_t>
  static FreeVector pars(const stepper_state_t& s) {
    return std::accumulate(s.components.begin(), s.components.end(),
                           FreeVector{FreeVector::Zero()},
                           [](const auto& sum, const auto& cmp) -> FreeVector {
                             return sum + cmp.weight * cmp.state.pars;
                           });
  }

  template <typename stepper_state_t>
  static FreeVector cov(const stepper_state_t& s) {
    return std::accumulate(s.components.begin(), s.components.end(),
                           FreeMatrix{FreeMatrix::Zero()},
                           [](const auto& sum, const auto& cmp) -> FreeMatrix {
                             return sum + cmp.weight * cmp.state.cov;
                           });
  }
};

struct MaxMomentumReducerLoop {
  template <typename component_range_t>
  static const auto& maxAbsoluteMomentumIt(const component_range_t& cmps) {
    return *std::max_element(cmps.begin(), cmps.end(),
                             [&](const auto& a, const auto& b) {
                               return std::abs(a.state.pars[eFreeQOverP]) >
                                      std::abs(b.state.pars[eFreeQOverP]);
                             });
  }

  template <typename stepper_state_t>
  static Vector3 position(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components)
        .state.pars.template segment<3>(eFreePos0);
  }

  template <typename stepper_state_t>
  static Vector3 direction(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components)
        .state.pars.template segment<3>(eFreeDir0);
  }

  template <typename stepper_state_t>
  static ActsScalar qOverP(const stepper_state_t& s) {
    const auto& cmp = maxAbsoluteMomentumIt(s.components);
    return cmp.state.pars[eFreeQOverP];
  }

  template <typename stepper_state_t>
  static ActsScalar absoluteMomentum(const stepper_state_t& s) {
    const auto& cmp = maxAbsoluteMomentumIt(s.components);
    return std::abs(cmp.state.absCharge / cmp.state.pars[eFreeQOverP]);
  }

  template <typename stepper_state_t>
  static Vector3 momentum(const stepper_state_t& s) {
    const auto& cmp = maxAbsoluteMomentumIt(s.components);
    return std::abs(cmp.state.absCharge / cmp.state.pars[eFreeQOverP]) *
           cmp.state.pars.template segment<3>(eFreeDir0);
  }

  template <typename stepper_state_t>
  static ActsScalar charge(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components).state.absCharge;
  }

  template <typename stepper_state_t>
  static ActsScalar time(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components).state.pars[eFreeTime];
  }

  template <typename stepper_state_t>
  static FreeVector pars(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components).state.pars;
  }

  template <typename stepper_state_t>
  static FreeVector cov(const stepper_state_t& s) {
    return maxAbsoluteMomentumIt(s.components).state.cov;
  }
};

template <typename extensionlist_t = StepperExtensionList<DefaultExtension>,
          typename component_reducer_t = WeightedComponentReducerLoop,
          typename auctioneer_t = detail::VoidAuctioneer>
class MultiEigenStepperLoop
    : public EigenStepper<extensionlist_t, auctioneer_t> {
  std::size_t m_stepLimitAfterFirstComponentOnSurface = 50;

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

  template <typename T>
  using SmallVector = boost::container::small_vector<T, 16>;

 public:
  using SingleStepper = EigenStepper<extensionlist_t, auctioneer_t>;

  using SingleState = typename SingleStepper::State;

  using typename SingleStepper::Covariance;
  using typename SingleStepper::Jacobian;

  using BoundState =
      std::tuple<MultiComponentBoundTrackParameters, Jacobian, ActsScalar>;

  using CurvilinearState = std::tuple<MultiComponentCurvilinearTrackParameters,
                                      Jacobian, ActsScalar>;

  using Reducer = component_reducer_t;

  static constexpr int maxComponents = std::numeric_limits<int>::max();

  MultiEigenStepperLoop(std::shared_ptr<const MagneticFieldProvider> bField,
                        std::unique_ptr<const Logger> logger =
                            getDefaultLogger("GSF", Logging::INFO))
      : EigenStepper<extensionlist_t, auctioneer_t>(std::move(bField)),
        m_logger(std::move(logger)) {}

  struct State {
    struct Component {
      SingleState state;
      ActsScalar weight;
      Intersection3D::Status status;
    };

    ParticleHypothesis particleHypothesis = ParticleHypothesis::pion();

    SmallVector<Component> components;

    bool covTransport = false;
    double pathAccumulated = 0.;
    std::size_t steps = 0;

    std::reference_wrapper<const GeometryContext> geoContext;

    std::reference_wrapper<const MagneticFieldContext> magContext;

    std::optional<std::size_t> stepCounterAfterFirstComponentOnSurface;

    State() = delete;

    explicit State(const GeometryContext& gctx,
                   const MagneticFieldContext& mctx,
                   const std::shared_ptr<const MagneticFieldProvider>& bfield,
                   const MultiComponentBoundTrackParameters& multipars,
                   double ssize = std::numeric_limits<double>::max())
        : particleHypothesis(multipars.particleHypothesis()),
          geoContext(gctx),
          magContext(mctx) {
      if (multipars.components().empty()) {
        throw std::invalid_argument(
            "Cannot construct MultiEigenStepperLoop::State with empty "
            "multi-component parameters");
      }

      const auto surface = multipars.referenceSurface().getSharedPtr();

      for (auto i = 0ul; i < multipars.components().size(); ++i) {
        const auto& [weight, singlePars] = multipars[i];
        components.push_back(
            {SingleState(gctx, bfield->makeCache(mctx), singlePars, ssize),
             weight, Intersection3D::Status::onSurface});
      }

      if (std::get<2>(multipars.components().front())) {
        covTransport = true;
      }
    }
  };

  State makeState(std::reference_wrapper<const GeometryContext> gctx,
                  std::reference_wrapper<const MagneticFieldContext> mctx,
                  const MultiComponentBoundTrackParameters& par,
                  double ssize = std::numeric_limits<double>::max()) const {
    return State(gctx, mctx, SingleStepper::m_bField, par, ssize);
  }

  void resetState(
      State& state, const BoundVector& boundParams,
      const BoundSquareMatrix& cov, const Surface& surface,
      const double stepSize = std::numeric_limits<double>::max()) const {
    for (auto& component : state.components) {
      SingleStepper::resetState(component.state, boundParams, cov, surface,
                                stepSize);
    }
  }

  template <typename stepping_t, typename navigation_t, typename options_t,
            typename geoctx_t>
  struct SinglePropState {
    stepping_t& stepping;
    navigation_t& navigation;
    options_t& options;
    geoctx_t& geoContext;

    SinglePropState(stepping_t& s, navigation_t& n, options_t& o, geoctx_t& g)
        : stepping(s), navigation(n), options(o), geoContext(g) {}
  };

  template <typename component_t>
  struct ComponentProxyBase {
    static_assert(std::is_same_v<std::remove_const_t<component_t>,
                                 typename State::Component>);

    component_t& cmp;

    ComponentProxyBase(component_t& c) : cmp(c) {}

    // These are the const accessors, which are shared between the mutable
    // ComponentProxy and the ConstComponentProxy
    auto status() const { return cmp.status; }
    auto weight() const { return cmp.weight; }
    auto pathAccumulated() const { return cmp.state.pathAccumulated; }
    const auto& pars() const { return cmp.state.pars; }
    const auto& derivative() const { return cmp.state.derivative; }
    const auto& jacTransport() const { return cmp.state.jacTransport; }
    const auto& cov() const { return cmp.state.cov; }
    const auto& jacobian() const { return cmp.state.jacobian; }
    const auto& jacToGlobal() const { return cmp.state.jacToGlobal; }

    template <typename propagator_state_t>
    auto singleState(const propagator_state_t& state) const {
      using DeducedStepping = decltype(state.stepping.components.front().state);
      static_assert(std::is_same_v<SingleState, DeducedStepping>);

      return SinglePropState<
          const SingleState, const decltype(state.navigation),
          const decltype(state.options), const decltype(state.geoContext)>(
          cmp.state, state.navigation, state.options, state.geoContext);
    }

    const auto& singleStepper(const MultiEigenStepperLoop& stepper) const {
      return static_cast<const SingleStepper&>(stepper);
    }
  };

  using ConstComponentProxy =
      ComponentProxyBase<const typename State::Component>;

  struct ComponentProxy : ComponentProxyBase<typename State::Component> {
    using Base = ComponentProxyBase<typename State::Component>;

    // Import the const accessors from ComponentProxyBase
    using Base::cmp;
    using Base::cov;
    using Base::derivative;
    using Base::jacobian;
    using Base::jacToGlobal;
    using Base::jacTransport;
    using Base::pars;
    using Base::pathAccumulated;
    using Base::singleState;
    using Base::singleStepper;
    using Base::status;
    using Base::weight;

    // The multi-component state of the stepper
    const State& all_state;

    ComponentProxy(typename State::Component& c, const State& s)
        : Base(c), all_state(s) {}

    // These are the mutable accessors, the const ones are inherited from the
    // ComponentProxyBase
    auto& status() { return cmp.status; }
    auto& weight() { return cmp.weight; }
    auto& pathAccumulated() { return cmp.state.pathAccumulated; }
    auto& pars() { return cmp.state.pars; }
    auto& derivative() { return cmp.state.derivative; }
    auto& jacTransport() { return cmp.state.jacTransport; }
    auto& cov() { return cmp.state.cov; }
    auto& jacobian() { return cmp.state.jacobian; }
    auto& jacToGlobal() { return cmp.state.jacToGlobal; }

    template <typename propagator_state_t>
    auto singleState(propagator_state_t& state) {
      using DeducedStepping = decltype(state.stepping.components.front().state);
      static_assert(std::is_same_v<SingleState, DeducedStepping>);

      return SinglePropState<SingleState, decltype(state.navigation),
                             decltype(state.options),
                             decltype(state.geoContext)>(
          cmp.state, state.navigation, state.options, state.geoContext);
    }

    Result<typename SingleStepper::BoundState> boundState(
        const Surface& surface, bool transportCov,
        const FreeToBoundCorrection& freeToBoundCorrection) {
      return detail::boundState(
          all_state.geoContext, cov(), jacobian(), jacTransport(), derivative(),
          jacToGlobal(), pars(), all_state.particleHypothesis,
          all_state.covTransport && transportCov, cmp.state.pathAccumulated,
          surface, freeToBoundCorrection);
    }

    void update(const FreeVector& freeParams, const BoundVector& boundParams,
                const Covariance& covariance, const Surface& surface) {
      cmp.state.pars = freeParams;
      cmp.state.cov = covariance;
      cmp.state.jacToGlobal =
          surface.boundToFreeJacobian(all_state.geoContext, boundParams);
    }
  };

  auto componentIterable(State& state) const {
    struct Iterator {
      using difference_type = std::ptrdiff_t;
      using value_type = ComponentProxy;
      using reference = ComponentProxy;
      using pointer = void;
      using iterator_category = std::forward_iterator_tag;

      typename decltype(state.components)::iterator it;
      const State& s;

      // clang-format off
      auto& operator++() { ++it; return *this; }
      auto operator!=(const Iterator& other) const { return it != other.it; }
      auto operator==(const Iterator& other) const { return it == other.it; }
      auto operator*() const { return ComponentProxy(*it, s); }
      // clang-format on
    };

    struct Iterable {
      using iterator = Iterator;

      State& s;

      // clang-format off
      auto begin() { return Iterator{s.components.begin(), s}; }
      auto end() { return Iterator{s.components.end(), s}; }
      // clang-format on
    };

    return Iterable{state};
  }

  auto constComponentIterable(const State& state) const {
    struct ConstIterator {
      using difference_type = std::ptrdiff_t;
      using value_type = ConstComponentProxy;
      using reference = ConstComponentProxy;
      using pointer = void;
      using iterator_category = std::forward_iterator_tag;

      typename decltype(state.components)::const_iterator it;
      const State& s;

      // clang-format off
      auto& operator++() { ++it; return *this; }
      auto operator!=(const ConstIterator& other) const { return it != other.it; }
      auto operator==(const ConstIterator& other) const { return it == other.it; }
      auto operator*() const { return ConstComponentProxy{*it}; }
      // clang-format on
    };

    struct Iterable {
      using iterator = ConstIterator;
      const State& s;

      // clang-format off
      auto begin() const { return ConstIterator{s.components.cbegin(), s}; }
      auto end() const { return ConstIterator{s.components.cend(), s}; }
      // clang-format on
    };

    return Iterable{state};
  }

  std::size_t numberComponents(const State& state) const {
    return state.components.size();
  }

  void removeMissedComponents(State& state) const {
    auto new_end = std::remove_if(
        state.components.begin(), state.components.end(), [](const auto& cmp) {
          return cmp.status == Intersection3D::Status::missed;
        });

    state.components.erase(new_end, state.components.end());
  }

  void reweightComponents(State& state) const {
    ActsScalar sumOfWeights = 0.0;
    for (const auto& cmp : state.components) {
      sumOfWeights += cmp.weight;
    }
    for (auto& cmp : state.components) {
      cmp.weight /= sumOfWeights;
    }
  }

  void clearComponents(State& state) const { state.components.clear(); }

  Result<ComponentProxy> addComponent(State& state,
                                      const BoundTrackParameters& pars,
                                      double weight) const {
    state.components.push_back(
        {SingleState(state.geoContext,
                     SingleStepper::m_bField->makeCache(state.magContext),
                     pars),
         weight, Intersection3D::Status::onSurface});

    return ComponentProxy{state.components.back(), state};
  }

  Result<Vector3> getField(State& state, const Vector3& pos) const {
    return SingleStepper::getField(state.components.front().state, pos);
  }

  Vector3 position(const State& state) const {
    return Reducer::position(state);
  }

  Vector3 direction(const State& state) const {
    return Reducer::direction(state);
  }

  double qOverP(const State& state) const { return Reducer::qOverP(state); }

  double absoluteMomentum(const State& state) const {
    return Reducer::absoluteMomentum(state);
  }

  Vector3 momentum(const State& state) const {
    return Reducer::momentum(state);
  }

  double charge(const State& state) const { return Reducer::charge(state); }

  ParticleHypothesis particleHypothesis(const State& state) const {
    return state.particleHypothesis;
  }

  double time(const State& state) const { return Reducer::time(state); }

  Intersection3D::Status updateSurfaceStatus(
      State& state, const Surface& surface, Direction navDir,
      const BoundaryCheck& bcheck,
      ActsScalar surfaceTolerance = s_onSurfaceTolerance,
      const Logger& logger = getDummyLogger()) const {
    using Status = Intersection3D::Status;

    std::array<int, 4> counts = {0, 0, 0, 0};

    for (auto& component : state.components) {
      component.status = detail::updateSingleSurfaceStatus<SingleStepper>(
          *this, component.state, surface, navDir, bcheck, surfaceTolerance,
          logger);
      ++counts[static_cast<std::size_t>(component.status)];
    }

    // If at least one component is on a surface, we can remove all missed
    // components before the step. If not, we must keep them for the case that
    // all components miss and we need to retarget
    if (counts[static_cast<std::size_t>(Status::onSurface)] > 0) {
      removeMissedComponents(state);
      reweightComponents(state);
    }

    ACTS_VERBOSE("Component status wrt "
                 << surface.geometryId() << " at {"
                 << surface.center(state.geoContext).transpose() << "}:\t"
                 << [&]() {
                      std::stringstream ss;
                      for (auto& component : state.components) {
                        ss << component.status << "\t";
                      }
                      return ss.str();
                    }());

    // Switch on stepCounter if one or more components reached a surface, but
    // some are still in progress of reaching the surface
    if (!state.stepCounterAfterFirstComponentOnSurface &&
        counts[static_cast<std::size_t>(Status::onSurface)] > 0 &&
        counts[static_cast<std::size_t>(Status::reachable)] > 0) {
      state.stepCounterAfterFirstComponentOnSurface = 0;
      ACTS_VERBOSE("started stepCounterAfterFirstComponentOnSurface");
    }

    // If there are no components onSurface, but the counter is switched on
    // (e.g., if the navigator changes the target surface), we need to switch it
    // off again
    if (state.stepCounterAfterFirstComponentOnSurface &&
        counts[static_cast<std::size_t>(Status::onSurface)] == 0) {
      state.stepCounterAfterFirstComponentOnSurface.reset();
      ACTS_VERBOSE("switch off stepCounterAfterFirstComponentOnSurface");
    }

    // This is a 'any_of' criterium. As long as any of the components has a
    // certain state, this determines the total state (in the order of a
    // somewhat importance)
    if (counts[static_cast<std::size_t>(Status::reachable)] > 0) {
      return Status::reachable;
    } else if (counts[static_cast<std::size_t>(Status::onSurface)] > 0) {
      state.stepCounterAfterFirstComponentOnSurface.reset();
      return Status::onSurface;
    } else if (counts[static_cast<std::size_t>(Status::unreachable)] > 0) {
      return Status::unreachable;
    } else {
      return Status::missed;
    }
  }

  template <typename object_intersection_t>
  void updateStepSize(State& state, const object_intersection_t& oIntersection,
                      Direction direction, bool release = true) const {
    const Surface& surface = *oIntersection.representation();

    for (auto& component : state.components) {
      auto intersection = surface.intersect(
          component.state.geoContext, SingleStepper::position(component.state),
          direction * SingleStepper::direction(component.state),
          true)[oIntersection.index()];

      SingleStepper::updateStepSize(component.state, intersection, direction,
                                    release);
    }
  }

  void setStepSize(State& state, double stepSize,
                   ConstrainedStep::Type stype = ConstrainedStep::actor,
                   bool release = true) const {
    for (auto& component : state.components) {
      SingleStepper::setStepSize(component.state, stepSize, stype, release);
    }
  }

  double getStepSize(const State& state, ConstrainedStep::Type stype) const {
    return std::min_element(state.components.begin(), state.components.end(),
                            [=](const auto& a, const auto& b) {
                              return std::abs(a.state.stepSize.value(stype)) <
                                     std::abs(b.state.stepSize.value(stype));
                            })
        ->state.stepSize.value(stype);
  }

  void releaseStepSize(State& state) const {
    for (auto& component : state.components) {
      SingleStepper::releaseStepSize(component.state);
    }
  }

  std::string outputStepSize(const State& state) const {
    std::stringstream ss;
    for (const auto& component : state.components) {
      ss << component.state.stepSize.toString() << " || ";
    }

    return ss.str();
  }

  double overstepLimit(const State& state) const {
    // A dynamic overstep limit could sit here
    return SingleStepper::overstepLimit(state.components.front().state);
  }

  Result<BoundState> boundState(
      State& state, const Surface& surface, bool transportCov = true,
      const FreeToBoundCorrection& freeToBoundCorrection =
          FreeToBoundCorrection(false)) const;

  CurvilinearState curvilinearState(State& state,
                                    bool transportCov = true) const;

  void transportCovarianceToCurvilinear(State& state) const {
    for (auto& component : state.components) {
      SingleStepper::transportCovarianceToCurvilinear(component.state);
    }
  }

  void transportCovarianceToBound(
      State& state, const Surface& surface,
      const FreeToBoundCorrection& freeToBoundCorrection =
          FreeToBoundCorrection(false)) const {
    for (auto& component : state.components) {
      SingleStepper::transportCovarianceToBound(component.state, surface,
                                                freeToBoundCorrection);
    }
  }

  template <typename propagator_state_t, typename navigator_t>
  Result<double> step(propagator_state_t& state,
                      const navigator_t& navigator) const;
};

}  // namespace Acts

#include "MultiEigenStepperLoop.ipp"