TrackProxy.hpp

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

// This file is part of the Acts project.
//
// Copyright (C) 2023 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/Definitions/TrackParametrization.hpp"
#include "Acts/EventData/MultiTrajectory.hpp"
#include "Acts/EventData/ParticleHypothesis.hpp"
#include "Acts/EventData/TrackStatePropMask.hpp"
#include "Acts/Utilities/Concepts.hpp"
#include "Acts/Utilities/HashedString.hpp"
#include "Acts/Utilities/UnitVectors.hpp"

#include <iterator>
#include <type_traits>

namespace Acts {

template <ACTS_CONCEPT(Acts::TrackContainerBackend) track_container_t,
          typename traj_t, template <typename> class holder_t>
class TrackContainer;

namespace detail_tc {
template <typename T, bool select>
using ConstIf = std::conditional_t<select, const T, T>;

template <typename container_t, typename proxy_t, bool ReadOnly>
class TrackProxyIterator {
  using ProxyType = proxy_t;
  using IndexType = typename ProxyType::IndexType;
  using ContainerType = container_t;

 public:
  using iterator_category = std::random_access_iterator_tag;
  using value_type = ProxyType;
  using difference_type = std::ptrdiff_t;
  using pointer = void;
  using reference = void;

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  TrackProxyIterator(container_t& container, IndexType itrack)
      : m_container(&container), m_itrack(itrack) {}

  template <bool RO = ReadOnly, typename = std::enable_if_t<RO>>
  TrackProxyIterator(const container_t& container, IndexType itrack)
      : m_container(&container), m_itrack(itrack) {}

  TrackProxyIterator& operator++() {
    m_itrack++;
    return *this;
  }
  TrackProxyIterator& operator--() {
    m_itrack--;
    return *this;
  }

  bool operator==(const TrackProxyIterator& other) const {
    return m_container == other.m_container && m_itrack == other.m_itrack;
  }

  bool operator!=(const TrackProxyIterator& other) const {
    return !(*this == other);
  }

  bool operator<(const TrackProxyIterator& other) const {
    return m_itrack < other.m_itrack;
  }

  bool operator>(const TrackProxyIterator& other) const {
    return m_itrack > other.m_itrack;
  }

  bool operator<=(const TrackProxyIterator& other) const {
    return m_itrack <= other.m_itrack;
  }

  bool operator>=(const TrackProxyIterator& other) const {
    return m_itrack >= other.m_itrack;
  }

  ProxyType operator*() const { return m_container->getTrack(m_itrack); }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  ProxyType operator*() {
    return m_container->getTrack(m_itrack);
  }

  TrackProxyIterator operator[](difference_type n) const {
    TrackProxyIterator copy = *this;
    copy += n;
    return copy;
  };

  TrackProxyIterator& operator+=(difference_type n) {
    m_itrack += n;
    return *this;
  }

  TrackProxyIterator operator-=(difference_type n) {
    m_itrack -= n;
    return *this;
  }

  friend difference_type operator-(const TrackProxyIterator& lhs,
                                   const TrackProxyIterator& rhs) {
    return lhs.m_itrack - rhs.m_itrack;
  }

  friend TrackProxyIterator operator+(const TrackProxyIterator& lhs,
                                      difference_type rhs) {
    TrackProxyIterator copy = lhs;
    copy += rhs;
    return copy;
  }

  friend TrackProxyIterator operator+(difference_type lhs,
                                      const TrackProxyIterator& rhs) {
    return rhs + lhs;
  }

  friend TrackProxyIterator operator-(const TrackProxyIterator& lhs,
                                      difference_type rhs) {
    return lhs + (-rhs);
  }

  friend TrackProxyIterator operator-(difference_type lhs,
                                      const TrackProxyIterator& rhs) {
    return rhs + (-lhs);
  }

 private:
  detail_lt::TransitiveConstPointer<ConstIf<ContainerType, ReadOnly>>
      m_container;
  IndexType m_itrack;
};

}  // namespace detail_tc

template <typename track_container_t, typename trajectory_t,
          template <typename> class holder_t, bool read_only = true>
class TrackProxy {
 public:
  static constexpr bool ReadOnly = read_only;
  using Container = track_container_t;
  using Trajectory = trajectory_t;

  using MutableTrackProxy =
      TrackProxy<track_container_t, trajectory_t, holder_t, false>;
  using ConstTrackProxy =
      TrackProxy<track_container_t, trajectory_t, holder_t, true>;

  using TrackStateProxy = typename Trajectory::TrackStateProxy;
  using ConstTrackStateProxy = typename Trajectory::ConstTrackStateProxy;

  using Parameters =
      typename detail_lt::Types<eBoundSize, false>::CoefficientsMap;
  using ConstParameters =
      typename detail_lt::Types<eBoundSize, true>::CoefficientsMap;

  using Covariance =
      typename detail_lt::Types<eBoundSize, false>::CovarianceMap;
  using ConstCovariance =
      typename detail_lt::Types<eBoundSize, true>::CovarianceMap;

  using IndexType = typename Container::IndexType;
  static constexpr IndexType kInvalid = Container::kInvalid;

#ifndef DOXYGEN
  friend TrackContainer<Container, Trajectory, holder_t>;
  friend MutableTrackProxy;
  friend ConstTrackProxy;
#endif

  TrackProxy(const MutableTrackProxy& other)
      : m_container{other.m_container}, m_index{other.m_index} {}

  TrackProxy& operator=(const MutableTrackProxy& other) {
    m_container = other.m_container;
    m_index = other.m_index;
    return *this;
  }

  template <typename T, HashedString key, bool RO = ReadOnly,
            typename = std::enable_if_t<!RO>>
  constexpr T& component() {
    return m_container->template component<T, key>(m_index);
  }

  template <typename T, bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  constexpr T& component(HashedString key) {
    return m_container->template component<T>(key, m_index);
  }

  template <typename T, bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  constexpr T& component(std::string_view key) {
    return m_container->template component<T>(hashString(key), m_index);
  }

  template <typename T, HashedString key>
  constexpr const T& component() const {
    return m_container->template component<T, key>(m_index);
  }

  template <typename T>
  constexpr const T& component(HashedString key) const {
    return m_container->template component<T>(key, m_index);
  }

  template <typename T>
  constexpr const T& component(std::string_view key) const {
    return m_container->template component<T>(hashString(key), m_index);
  }

  IndexType tipIndex() const {
    return component<IndexType>(hashString("tipIndex"));
  }

  IndexType stemIndex() const {
    return component<IndexType>(hashString("stemIndex"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  IndexType& tipIndex() {
    return component<IndexType>(hashString("tipIndex"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  IndexType& stemIndex() {
    return component<IndexType>(hashString("stemIndex"));
  }

  auto innermostTrackState() const {
    using proxy_t = decltype(m_container->trackStateContainer().getTrackState(
        std::declval<IndexType>()));

    IndexType stem = component<IndexType>(hashString("stemIndex"));
    if (stem == kInvalid) {
      return std::optional<proxy_t>{};
    } else {
      return std::optional<proxy_t>{
          m_container->trackStateContainer().getTrackState(stem)};
    }
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  auto innermostTrackState() {
    using proxy_t = decltype(m_container->trackStateContainer().getTrackState(
        std::declval<IndexType>()));

    IndexType stem = component<IndexType>(hashString("stemIndex"));
    if (stem == kInvalid) {
      return std::optional<proxy_t>{};
    } else {
      return std::optional<proxy_t>{
          m_container->trackStateContainer().getTrackState(stem)};
    }
  }

  const Surface& referenceSurface() const {
    return *m_container->container().referenceSurface_impl(m_index);
  }

  // NOLINTBEGIN(performance-unnecessary-value-param)
  // looks like a false-positive. clang-tidy believes `srf` is not movable.
  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void setReferenceSurface(std::shared_ptr<const Surface> srf) {
    m_container->container().setReferenceSurface_impl(m_index, std::move(srf));
  }
  // NOLINTEND(performance-unnecessary-value-param)

  bool hasReferenceSurface() const {
    // @TODO: This could be more efficient
    return m_container->container().referenceSurface_impl(m_index) != nullptr;
  }

  ConstParameters parameters() const {
    return m_container->parameters(m_index);
  }

  ConstCovariance covariance() const {
    return m_container->covariance(m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  Parameters parameters() {
    return m_container->parameters(m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  Covariance covariance() {
    return m_container->covariance(m_index);
  }

  ActsScalar theta() const {
    return parameters()[eBoundTheta];
  }

  ActsScalar phi() const {
    return parameters()[eBoundPhi];
  }

  ActsScalar loc0() const {
    return parameters()[eBoundLoc0];
  }

  ActsScalar loc1() const {
    return parameters()[eBoundLoc1];
  }

  ActsScalar time() const {
    return parameters()[eBoundTime];
  }

  ActsScalar qOverP() const {
    return parameters()[eBoundQOverP];
  }

  ParticleHypothesis particleHypothesis() const {
    return component<ParticleHypothesis, hashString("particleHypothesis")>();
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void setParticleHypothesis(const ParticleHypothesis& particleHypothesis) {
    m_container->container().setParticleHypothesis_impl(m_index,
                                                        particleHypothesis);
  }

  ActsScalar charge() const {
    return particleHypothesis().qFromQOP(qOverP());
  }

  ActsScalar absoluteMomentum() const {
    return particleHypothesis().extractMomentum(qOverP());
  }

  ActsScalar transverseMomentum() const {
    return std::sin(theta()) * absoluteMomentum();
  }

  Vector3 direction() const {
    return makeDirectionFromPhiTheta(phi(), theta());
  }

  Vector3 momentum() const {
    return absoluteMomentum() * direction();
  }

  auto trackStatesReversed() const {
    return m_container->reverseTrackStateRange(m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  auto trackStatesReversed() {
    return m_container->reverseTrackStateRange(m_index);
  }

  auto trackStates() const {
    return m_container->forwardTrackStateRange(m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  auto trackStates() {
    return m_container->forwardTrackStateRange(m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void linkForward() {
    IndexType last = kInvalid;
    for (auto ts : trackStatesReversed()) {
      ts.template component<IndexType>(hashString("next")) = last;
      last = ts.index();
    }
    stemIndex() = last;
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  auto appendTrackState(TrackStatePropMask mask = TrackStatePropMask::All) {
    auto& tsc = m_container->trackStateContainer();
    auto ts = tsc.getTrackState(tsc.addTrackState(mask, tipIndex()));
    tipIndex() = ts.index();
    return ts;
  }

  unsigned int nTrackStates() const {
    // @TODO: This should probably be cached, distance is expensive
    //        without random access
    if (tipIndex() == kInvalid) {
      // no tip index -> no track states
      return 0;
    }
    auto tsRange = trackStatesReversed();
    return std::distance(tsRange.begin(), tsRange.end());
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  unsigned int& nMeasurements() {
    return component<unsigned int>(hashString("nMeasurements"));
  }

  unsigned int nMeasurements() const {
    return component<unsigned int>(hashString("nMeasurements"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  unsigned int& nHoles() {
    return component<unsigned int>(hashString("nHoles"));
  }

  unsigned int nHoles() const {
    return component<unsigned int>(hashString("nHoles"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  unsigned int& nOutliers() {
    return component<unsigned int>(hashString("nOutliers"));
  }

  unsigned int nOutliers() const {
    return component<unsigned int>(hashString("nOutliers"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  unsigned int& nSharedHits() {
    return component<unsigned int>(hashString("nSharedHits"));
  }

  unsigned int nSharedHits() const {
    return component<unsigned int>(hashString("nSharedHits"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  float& chi2() {
    return component<float>(hashString("chi2"));
  }

  float chi2() const {
    return component<float>(hashString("chi2"));
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  unsigned int& nDoF() {
    return component<unsigned int>(hashString("ndf"));
  }

  unsigned int nDoF() const {
    return component<unsigned int>(hashString("ndf"));
  }

  IndexType index() const {
    return m_index;
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  auto& container() {
    return *m_container;
  }

  template <typename track_proxy_t, bool RO = ReadOnly,
            typename = std::enable_if_t<!RO>>
  void copyFrom(const track_proxy_t& other, bool copyTrackStates = true) {
    // @TODO: Add constraint on which track proxies are allowed,
    // this is only implicit right now

    if (copyTrackStates) {
      // append track states (cheap), but they're in the wrong order
      for (const auto& srcTrackState : other.trackStatesReversed()) {
        auto destTrackState = appendTrackState(srcTrackState.getMask());
        if (srcTrackState.hasCalibrated()) {
          destTrackState.allocateCalibrated(srcTrackState.calibratedSize());
        }
        destTrackState.copyFrom(srcTrackState, Acts::TrackStatePropMask::All,
                                true);
      }

      // reverse using standard linked list reversal algorithm
      reverseTrackStates();
    }

    parameters() = other.parameters();
    covariance() = other.covariance();
    setParticleHypothesis(other.particleHypothesis());
    if (other.hasReferenceSurface()) {
      setReferenceSurface(other.referenceSurface().getSharedPtr());
    }
    nMeasurements() = other.nMeasurements();
    nHoles() = other.nHoles();
    nOutliers() = other.nOutliers();
    nSharedHits() = other.nSharedHits();
    chi2() = other.chi2();
    nDoF() = other.nDoF();

    // This will only be valid if the backends match and support this operation
    m_container->copyDynamicFrom(m_index, other.m_container->container(),
                                 other.m_index);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void reverseTrackStates(bool invertJacobians = false) {
    IndexType current = tipIndex();
    IndexType next = kInvalid;
    IndexType prev = kInvalid;

    stemIndex() = tipIndex();

    // @TODO: Maybe refactor to not need this variable if invertJacobians == false
    BoundMatrix nextJacobian;

    while (current != kInvalid) {
      auto ts = m_container->trackStateContainer().getTrackState(current);
      prev = ts.previous();
      ts.template component<IndexType>(hashString("next")) = prev;
      ts.previous() = next;
      if (invertJacobians) {
        if (next != kInvalid) {
          BoundMatrix curJacobian = ts.jacobian();
          ts.jacobian() = nextJacobian.inverse();
          nextJacobian = curJacobian;
        } else {
          nextJacobian = ts.jacobian();
          ts.jacobian().setZero();
        }
      }
      next = current;
      tipIndex() = current;
      current = prev;
    }
  }

  const auto& container() const {
    return *m_container;
  }

  bool operator==(const TrackProxy& other) const {
    return &(*m_container) == &(*other.m_container) && m_index == other.m_index;
  }

  // Track proxies are friends, not food!
  template <typename A, typename B, template <typename> class H, bool R>
  friend class TrackProxy;

 private:
  TrackProxy(detail_tc::ConstIf<TrackContainer<Container, Trajectory, holder_t>,
                                ReadOnly>& container,
             IndexType itrack)
      : m_container{&container}, m_index{itrack} {}

  detail_lt::TransitiveConstPointer<detail_tc::ConstIf<
      TrackContainer<Container, Trajectory, holder_t>, ReadOnly>>
      m_container;
  IndexType m_index;
};
}  // namespace Acts