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 <bitset>
#include <cstdint>
#include <type_traits>
#include <vector>

#include <Eigen/Core>

#include "Acts/EventData/TrackState.hpp"
#include "Acts/Geometry/GeometryContext.hpp"
#include "Acts/Utilities/ParameterDefinitions.hpp"
#include "Acts/Utilities/TypeTraits.hpp"

namespace Acts {

// 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 N, size_t M, bool ReadOnly = true>
class TrackStateProxy {
 public:
  using SourceLink = source_link_t;
  using Parameters = typename Types<N, ReadOnly>::CoefficientsMap;
  using Covariance = typename Types<N, 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, N, ProjectorFlags>;
  using EffectiveProjector =
      Eigen::Matrix<typename Projector::Scalar, Eigen::Dynamic, Eigen::Dynamic,
                    ProjectorFlags, M, N>;

  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];
  }

  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;

  BoundParameters predictedParameters(const Acts::GeometryContext& gctx) const;

  Covariance predictedCovariance() const;

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

  Parameters filtered() const;

  Covariance filteredCovariance() const;

  BoundParameters filteredParameters(const Acts::GeometryContext& gctx) const;

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

  Parameters smoothed() const;

  Covariance smoothedCovariance() const;

  BoundParameters smoothedParameters(const Acts::GeometryContext& gctx) 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);

    constexpr int max_measdim = MultiTrajectory<SourceLink>::MeasurementSizeMax;

    static_assert(rows <= max_measdim, "Given projector has too many rows");
    static_assert(cols <= max_measdim, "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, max_measdim>() = 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>,
            ParID_t... params>
  void setCalibrated(const Acts::Measurement<SourceLink, params...>& meas) {
    IndexData& dataref = data();
    constexpr size_t measdim = Acts::Measurement<SourceLink, 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.referenceSurface());
    if (!refSrf) {
      // ref surface is not set, set it now
      refSrf = meas.referenceSurface().getSharedPtr();
    }

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

  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>,
            ParID_t... params>
  void resetCalibrated(const Acts::Measurement<SourceLink, 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);

  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 = concept ::require<
    concept ::either<concept ::identical_to<bool, call_operator_t, T, TS>,
                     concept ::identical_to<void, call_operator_t, T, TS>>>;

}  // namespace detail_lt

namespace TrackStatePropMask {
using Type = std::bitset<8>;

constexpr Type Predicted{1 << 0};
constexpr Type Filtered{1 << 1};
constexpr Type Smoothed{1 << 2};
constexpr Type Jacobian{1 << 3};

constexpr Type Uncalibrated{1 << 4};
constexpr Type Calibrated{1 << 5};

// Initialize to all 1s. This only works as long as number of bits <= 64
constexpr Type All{static_cast<unsigned long long>(-1)};
constexpr Type None{0};
}  // namespace TrackStatePropMask

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

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

  MultiTrajectory() = default;

  template <typename parameters_t>
  size_t addTrackState(const TrackState<SourceLink, parameters_t>& ts,
                       size_t iprevious = SIZE_MAX);

  size_t addTrackState(
      const TrackStatePropMask::Type& 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<ParametersSize>::StorageCoefficients m_params;
  typename detail_lt::Types<ParametersSize>::StorageCovariance m_cov;
  typename detail_lt::Types<MeasurementSizeMax>::StorageCoefficients m_meas;
  typename detail_lt::Types<MeasurementSizeMax>::StorageCovariance m_measCov;
  typename detail_lt::Types<ParametersSize>::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, ParametersSize,
                                          MeasurementSizeMax, true>;
  friend class detail_lt::TrackStateProxy<SourceLink, ParametersSize,
                                          MeasurementSizeMax, false>;
};

}  // namespace Acts

#include "Acts/EventData/MultiTrajectory.ipp"