MultiTrajectory.hpp

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

// This file is part of the Acts project.
//
// Copyright (C) 2019 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/Measurement.hpp"
#include "Acts/EventData/TrackStatePropMask.hpp"
#include "Acts/Geometry/GeometryContext.hpp"
#include "Acts/Utilities/ParameterDefinitions.hpp"
#include "Acts/Utilities/TypeTraits.hpp"

#include <bitset>
#include <cstdint>
#include <type_traits>
#include <vector>

#include <Eigen/Core>

namespace Acts {

enum TrackStateFlag {
  MeasurementFlag = 0,
  ParameterFlag = 1,
  OutlierFlag = 2,
  HoleFlag = 3,
  MaterialFlag = 4,
  NumTrackStateFlags = 5
};

using TrackStateType = std::bitset<TrackStateFlag::NumTrackStateFlags>;

// forward declarations
template <typename source_link_t>
class MultiTrajectory;
class Surface;

namespace detail_lt {
template <typename T, bool select>
using ConstIf = std::conditional_t<select, const T, T>;
template <typename Storage, size_t kSizeIncrement>
struct GrowableColumns {
  auto addCol(size_t n = 1) {
    size_t index = m_size + (n - 1);
    while (capacity() <= index) {
      data.conservativeResize(Eigen::NoChange, data.cols() + kSizeIncrement);
    }
    m_size = index + 1;

    // @TODO: do this or not? If we assume this happens only when something is
    // written, the expectation is that everything is zero
    data.col(index).setZero();

    return data.col(index);
  }

  auto col(size_t index) { return data.col(index); }

  auto col(size_t index) const { return data.col(index); }

  size_t capacity() const { return static_cast<size_t>(data.cols()); }

  size_t size() const { return m_size; }

 private:
  Storage data;
  size_t m_size{0};
};

template <size_t Size, bool ReadOnlyMaps = true>
struct Types {
  enum {
    Flags = Eigen::ColMajor | Eigen::AutoAlign,
    SizeIncrement = 8,
  };
  using Scalar = double;
  // single items
  using Coefficients = Eigen::Matrix<Scalar, Size, 1, Flags>;
  using Covariance = Eigen::Matrix<Scalar, Size, Size, Flags>;
  using CoefficientsMap = Eigen::Map<ConstIf<Coefficients, ReadOnlyMaps>>;
  using CovarianceMap = Eigen::Map<ConstIf<Covariance, ReadOnlyMaps>>;
  // storage of multiple items in flat arrays
  using StorageCoefficients =
      GrowableColumns<Eigen::Array<Scalar, Size, Eigen::Dynamic, Flags>,
                      SizeIncrement>;
  using StorageCovariance =
      GrowableColumns<Eigen::Array<Scalar, Size * Size, Eigen::Dynamic, Flags>,
                      SizeIncrement>;
};

struct IndexData {
  using IndexType = uint16_t;

  static constexpr IndexType kInvalid = UINT16_MAX;

  IndexType irefsurface = kInvalid;
  IndexType iprevious = kInvalid;
  IndexType ipredicted = kInvalid;
  IndexType ifiltered = kInvalid;
  IndexType ismoothed = kInvalid;
  IndexType ijacobian = kInvalid;
  IndexType iprojector = kInvalid;

  double chi2;
  double pathLength;
  TrackStateType typeFlags;

  IndexType iuncalibrated = kInvalid;
  IndexType icalibrated = kInvalid;
  IndexType icalibratedsourcelink = kInvalid;
  IndexType measdim = 0;
};

template <typename source_link_t, size_t M, bool ReadOnly = true>
class TrackStateProxy {
 public:
  using SourceLink = source_link_t;
  using Parameters = typename Types<eBoundSize, ReadOnly>::CoefficientsMap;
  using Covariance = typename Types<eBoundSize, ReadOnly>::CovarianceMap;
  using Measurement = typename Types<M, ReadOnly>::CoefficientsMap;
  using MeasurementCovariance = typename Types<M, ReadOnly>::CovarianceMap;

  // as opposed to the types above, this is an actual Matrix (rather than a
  // map)
  // @TODO: Does not copy flags, because this fails: can't have col major row
  // vector, but that's required for 1xN projection matrices below.
  constexpr static auto ProjectorFlags = Eigen::RowMajor | Eigen::AutoAlign;
  using Projector =
      Eigen::Matrix<typename Covariance::Scalar, M, eBoundSize, ProjectorFlags>;
  using EffectiveProjector =
      Eigen::Matrix<typename Projector::Scalar, Eigen::Dynamic, Eigen::Dynamic,
                    ProjectorFlags, M, eBoundSize>;

  size_t index() const { return m_istate; }

  size_t previous() const { return data().iprevious; }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  IndexData& data() {
    return m_traj->m_index[m_istate];
  }

  TrackStatePropMask getMask() const;

  template <bool RO = ReadOnly, bool ReadOnlyOther,
            typename = std::enable_if<!RO>>
  void copyFrom(const TrackStateProxy<source_link_t, M, ReadOnlyOther>& other,
                TrackStatePropMask mask = TrackStatePropMask::All) {
    using PM = TrackStatePropMask;
    auto dest = getMask();
    auto src = other.getMask() &
               mask;  // combine what we have with what we want to copy
    if (static_cast<std::underlying_type_t<TrackStatePropMask>>((src ^ dest) &
                                                                src) != 0) {
      throw std::runtime_error(
          "Attempt track state copy with incompatible allocations");
    }

    // we're sure now this has correct allocations, so just copy
    if (ACTS_CHECK_BIT(src, PM::Predicted)) {
      predicted() = other.predicted();
      predictedCovariance() = other.predictedCovariance();
    }

    if (ACTS_CHECK_BIT(src, PM::Filtered)) {
      filtered() = other.filtered();
      filteredCovariance() = other.filteredCovariance();
    }

    if (ACTS_CHECK_BIT(src, PM::Smoothed)) {
      smoothed() = other.smoothed();
      smoothedCovariance() = other.smoothedCovariance();
    }

    if (ACTS_CHECK_BIT(src, PM::Uncalibrated)) {
      uncalibrated() = other.uncalibrated();
    }

    if (ACTS_CHECK_BIT(src, PM::Jacobian)) {
      jacobian() = other.jacobian();
    }

    if (ACTS_CHECK_BIT(src, PM::Calibrated)) {
      calibratedSourceLink() = other.calibratedSourceLink();
      calibrated() = other.calibrated();
      calibratedCovariance() = other.calibratedCovariance();
      data().measdim = other.data().measdim;
      setProjectorBitset(other.projectorBitset());
    }

    chi2() = other.chi2();
    pathLength() = other.pathLength();
    typeFlags() = other.typeFlags();

    // can be nullptr, but we just take that
    setReferenceSurface(other.referenceSurfacePointer());
  }

  const IndexData& data() const { return m_traj->m_index[m_istate]; }

  const Surface& referenceSurface() const {
    assert(data().irefsurface != IndexData::kInvalid);
    return *m_traj->m_referenceSurfaces[data().irefsurface];
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void setReferenceSurface(std::shared_ptr<const Surface> srf) {
    m_traj->m_referenceSurfaces[data().irefsurface] = std::move(srf);
  }

  Parameters parameters() const;

  Covariance covariance() const;

  Parameters predicted() const;

  Covariance predictedCovariance() const;

  bool hasPredicted() const { return data().ipredicted != IndexData::kInvalid; }

  Parameters filtered() const;

  Covariance filteredCovariance() const;

  bool hasFiltered() const { return data().ifiltered != IndexData::kInvalid; }

  Parameters smoothed() const;

  Covariance smoothedCovariance() const;

  bool hasSmoothed() const { return data().ismoothed != IndexData::kInvalid; }

  Covariance jacobian() const;

  bool hasJacobian() const { return data().ijacobian != IndexData::kInvalid; }

  Projector projector() const;

  bool hasProjector() const { return data().iprojector != IndexData::kInvalid; }

  EffectiveProjector effectiveProjector() const {
    return projector().topLeftCorner(data().measdim, M);
  }

  template <typename Derived, bool RO = ReadOnly,
            typename = std::enable_if_t<!RO>>
  void setProjector(const Eigen::MatrixBase<Derived>& projector) {
    constexpr int rows = Eigen::MatrixBase<Derived>::RowsAtCompileTime;
    constexpr int cols = Eigen::MatrixBase<Derived>::ColsAtCompileTime;

    static_assert(rows != -1 && cols != -1,
                  "Assignment of dynamic matrices is currently not supported.");

    IndexData& dataref = data();
    assert(dataref.iprojector != IndexData::kInvalid);

    static_assert(rows <= M, "Given projector has too many rows");
    static_assert(cols <= eBoundSize, "Given projector has too many columns");

    // set up full size projector with only zeros
    typename TrackStateProxy::Projector fullProjector =
        decltype(fullProjector)::Zero();

    // assign (potentially) smaller actual projector to matrix, preserving
    // zeroes outside of smaller matrix block.
    fullProjector.template topLeftCorner<rows, cols>() = projector;

    // convert to bitset before storing
    m_traj->m_projectors[dataref.iprojector] = matrixToBitset(fullProjector);
  }

  bool hasUncalibrated() const {
    return data().iuncalibrated != IndexData::kInvalid;
  }

  const SourceLink& uncalibrated() const;

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  SourceLink& uncalibrated() {
    assert(data().iuncalibrated != IndexData::kInvalid);
    return m_traj->m_sourceLinks[data().iuncalibrated];
  }

  bool hasCalibrated() const {
    return data().icalibrated != IndexData::kInvalid;
  }

  const SourceLink& calibratedSourceLink() const;

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  SourceLink& calibratedSourceLink() {
    assert(data().icalibratedsourcelink != IndexData::kInvalid);
    return m_traj->m_sourceLinks[data().icalibratedsourcelink];
  }

  Measurement calibrated() const;

  MeasurementCovariance calibratedCovariance() const;

  auto effectiveCalibrated() const { return calibrated().head(data().measdim); }

  auto effectiveCalibratedCovariance() const {
    return calibratedCovariance().topLeftCorner(data().measdim, data().measdim);
  }

  size_t calibratedSize() const { return data().measdim; }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>,
            BoundIndices... params>
  void setCalibrated(
      const Acts::Measurement<SourceLink, BoundIndices, params...>& meas) {
    IndexData& dataref = data();
    constexpr size_t measdim =
        Acts::Measurement<SourceLink, BoundIndices, params...>::size();

    dataref.measdim = measdim;

    assert(hasCalibrated());
    calibrated().setZero();
    calibrated().template head<measdim>() = meas.parameters();

    calibratedCovariance().setZero();
    calibratedCovariance().template topLeftCorner<measdim, measdim>() =
        meas.covariance();

    setProjector(meas.projector());

    // this shouldn't change
    assert(data().irefsurface != IndexData::kInvalid);
    std::shared_ptr<const Surface>& refSrf =
        m_traj->m_referenceSurfaces[dataref.irefsurface];
    // either unset, or the same, otherwise this is inconsistent assignment
    assert(!refSrf || refSrf.get() == &meas.referenceObject());
    if (!refSrf) {
      // ref surface is not set, set it now
      refSrf = meas.referenceObject().getSharedPtr();
    }

    assert(dataref.icalibratedsourcelink != IndexData::kInvalid);
    calibratedSourceLink() = meas.sourceLink();
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>,
            BoundIndices... params>
  void resetCalibrated(
      const Acts::Measurement<SourceLink, BoundIndices, params...>& meas) {
    IndexData& dataref = data();
    auto& traj = *m_traj;
    // force reallocate, whether currently invalid or shared index
    traj.m_meas.addCol();
    traj.m_measCov.addCol();
    // shared index between meas par
    // and cov
    dataref.icalibrated = traj.m_meas.size() - 1;

    traj.m_sourceLinks.emplace_back();
    dataref.icalibratedsourcelink = traj.m_sourceLinks.size() - 1;

    traj.m_projectors.emplace_back();
    dataref.iprojector = traj.m_projectors.size() - 1;

    // now actually assign to the allocated entries
    setCalibrated(meas);
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  double& chi2() {
    return data().chi2;
  }

  double chi2() const { return data().chi2; }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  double& pathLength() {
    return data().pathLength;
  }

  double pathLength() const { return data().pathLength; }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  TrackStateType& typeFlags() {
    return data().typeFlags;
  }

  TrackStateType typeFlags() const { return data().typeFlags; }

 private:
  // Private since it can only be created by the trajectory.
  TrackStateProxy(ConstIf<MultiTrajectory<SourceLink>, ReadOnly>& trajectory,
                  size_t istate);

  const std::shared_ptr<const Surface>& referenceSurfacePointer() const {
    assert(data().irefsurface != IndexData::kInvalid);
    return m_traj->m_referenceSurfaces[data().irefsurface];
  }

  typename MultiTrajectory<SourceLink>::ProjectorBitset projectorBitset()
      const {
    assert(data().iprojector != IndexData::kInvalid);
    return m_traj->m_projectors[data().iprojector];
  }

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
  void setProjectorBitset(
      typename MultiTrajectory<SourceLink>::ProjectorBitset proj) {
    assert(data().iprojector != IndexData::kInvalid);
    m_traj->m_projectors[data().iprojector] = proj;
  }

  ConstIf<MultiTrajectory<SourceLink>, ReadOnly>* m_traj;
  size_t m_istate;

  friend class Acts::MultiTrajectory<SourceLink>;
};

// implement track state visitor concept
template <typename T, typename TS>
using call_operator_t = decltype(std::declval<T>()(std::declval<TS>()));

template <typename T, typename TS>
constexpr bool VisitorConcept = Concepts ::require<
    Concepts ::either<Concepts ::identical_to<bool, call_operator_t, T, TS>,
                      Concepts ::identical_to<void, call_operator_t, T, TS>>>;

}  // namespace detail_lt

template <typename source_link_t>
class MultiTrajectory {
 public:
  enum {
    MeasurementSizeMax = eBoundSize,
  };
  using SourceLink = source_link_t;
  using ConstTrackStateProxy =
      detail_lt::TrackStateProxy<SourceLink, MeasurementSizeMax, true>;
  using TrackStateProxy =
      detail_lt::TrackStateProxy<SourceLink, MeasurementSizeMax, false>;

  using ProjectorBitset = std::bitset<eBoundSize * MeasurementSizeMax>;

  MultiTrajectory() = default;

  size_t addTrackState(TrackStatePropMask mask = TrackStatePropMask::All,
                       size_t iprevious = SIZE_MAX);

  ConstTrackStateProxy getTrackState(size_t istate) const {
    return {*this, istate};
  }

  TrackStateProxy getTrackState(size_t istate) { return {*this, istate}; }

  template <typename F>
  void visitBackwards(size_t iendpoint, F&& callable) const;

  template <typename F>
  void applyBackwards(size_t iendpoint, F&& callable);

 private:
  std::vector<detail_lt::IndexData> m_index;
  typename detail_lt::Types<eBoundSize>::StorageCoefficients m_params;
  typename detail_lt::Types<eBoundSize>::StorageCovariance m_cov;
  typename detail_lt::Types<MeasurementSizeMax>::StorageCoefficients m_meas;
  typename detail_lt::Types<MeasurementSizeMax>::StorageCovariance m_measCov;
  typename detail_lt::Types<eBoundSize>::StorageCovariance m_jac;
  std::vector<SourceLink> m_sourceLinks;
  std::vector<ProjectorBitset> m_projectors;

  // owning vector of shared pointers to surfaces
  // @TODO: This might be problematic when appending a large number of surfaces
  // trackstates, because vector has to reallocated and thus copy. This might
  // be handled in a smart way by moving but not sure.
  std::vector<std::shared_ptr<const Surface>> m_referenceSurfaces;

  friend class detail_lt::TrackStateProxy<SourceLink, MeasurementSizeMax, true>;
  friend class detail_lt::TrackStateProxy<SourceLink, MeasurementSizeMax,
                                          false>;
};

}  // namespace Acts

#include "Acts/EventData/MultiTrajectory.ipp"