FitUtils.h

Go to the documentation of this file.

#ifndef RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h
#define RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h

#include <alpaka/alpaka.hpp>

#include <Eigen/Core>

#include "DataFormats/Math/interface/choleskyInversion.h"
#include "HeterogeneousCore/AlpakaInterface/interface/config.h"
#include "RecoTracker/PixelTrackFitting/interface/alpaka/FitResult.h"

namespace riemannFit {

  constexpr double epsilon = 1.e-4;  

  using VectorXd = Eigen::VectorXd;
  using MatrixXd = Eigen::MatrixXd;
  template <int N>
  using MatrixNd = Eigen::Matrix<double, N, N>;
  template <int N>
  using MatrixNplusONEd = Eigen::Matrix<double, N + 1, N + 1>;
  template <int N>
  using ArrayNd = Eigen::Array<double, N, N>;
  template <int N>
  using Matrix2Nd = Eigen::Matrix<double, 2 * N, 2 * N>;
  template <int N>
  using Matrix3Nd = Eigen::Matrix<double, 3 * N, 3 * N>;
  template <int N>
  using Matrix2xNd = Eigen::Matrix<double, 2, N>;
  template <int N>
  using Array2xNd = Eigen::Array<double, 2, N>;
  template <int N>
  using MatrixNx3d = Eigen::Matrix<double, N, 3>;
  template <int N>
  using MatrixNx5d = Eigen::Matrix<double, N, 5>;
  template <int N>
  using VectorNd = Eigen::Matrix<double, N, 1>;
  template <int N>
  using VectorNplusONEd = Eigen::Matrix<double, N + 1, 1>;
  template <int N>
  using Vector2Nd = Eigen::Matrix<double, 2 * N, 1>;
  template <int N>
  using Vector3Nd = Eigen::Matrix<double, 3 * N, 1>;
  template <int N>
  using RowVectorNd = Eigen::Matrix<double, 1, 1, N>;
  template <int N>
  using RowVector2Nd = Eigen::Matrix<double, 1, 2 * N>;

  using Matrix2x3d = Eigen::Matrix<double, 2, 3>;

  using Matrix3f = Eigen::Matrix3f;
  using Vector3f = Eigen::Vector3f;
  using Vector4f = Eigen::Vector4f;
  using Vector6f = Eigen::Matrix<double, 6, 1>;
  // transformation between the "perigee" to cmssw localcoord frame
  // the plane of the latter is the perigee plane...
  // from   //!<(phi,Tip,q/pt,cotan(theta)),Zip)
  // to q/p,dx/dz,dy/dz,x,z
  template <typename VI5, typename MI5, typename VO5, typename MO5>
  inline void transformToPerigeePlane(VI5 const& ip, MI5 const& icov, VO5& op, MO5& ocov) {
    auto sinTheta2 = 1. / (1. + ip(3) * ip(3));
    auto sinTheta = std::sqrt(sinTheta2);
    auto cosTheta = ip(3) * sinTheta;

    op(0) = sinTheta * ip(2);
    op(1) = 0.;
    op(2) = -ip(3);
    op(3) = ip(1);
    op(4) = -ip(4);

    Matrix5d jMat = Matrix5d::Zero();

    jMat(0, 2) = sinTheta;
    jMat(0, 3) = -sinTheta2 * cosTheta * ip(2);
    jMat(1, 0) = 1.;
    jMat(2, 3) = -1.;
    jMat(3, 1) = 1.;
    jMat(4, 4) = -1;

    ocov = jMat * icov * jMat.transpose();
  }

}  // namespace riemannFit

namespace ALPAKA_ACCELERATOR_NAMESPACE {
  namespace riemannFit {
    using namespace ::riemannFit;

    template <typename TAcc, class C>
    ALPAKA_FN_ACC void printIt(const TAcc& acc, C* m, const char* prefix = "") {
#ifdef RFIT_DEBUG
      for (uint r = 0; r < m->rows(); ++r) {
        for (uint c = 0; c < m->cols(); ++c) {
          printf("%s Matrix(%d,%d) = %g\n", prefix, r, c, (*m)(r, c));
        }
      }
#endif
    }

    template <typename T>
    constexpr T sqr(const T a) {
      return a * a;
    }

    template <typename TAcc>
    ALPAKA_FN_ACC ALPAKA_FN_INLINE double cross2D(const TAcc& acc, const Vector2d& a, const Vector2d& b) {
      return a.x() * b.y() - a.y() * b.x();
    }

    template <typename TAcc, typename M6xNf, typename M2Nd>
    ALPAKA_FN_ACC void loadCovariance2D(const TAcc& acc, M6xNf const& ge, M2Nd& hits_cov) {
      // Index numerology:
      // i: index of the hits/point (0,..,3)
      // j: index of space component (x,y,z)
      // l: index of space components (x,y,z)
      // ge is always in sync with the index i and is formatted as:
      // ge[] ==> [xx, xy, yy, xz, yz, zz]
      // in (j,l) notation, we have:
      // ge[] ==> [(0,0), (0,1), (1,1), (0,2), (1,2), (2,2)]
      // so the index ge_idx corresponds to the matrix elements:
      // | 0  1  3 |
      // | 1  2  4 |
      // | 3  4  5 |
      constexpr uint32_t hits_in_fit = M6xNf::ColsAtCompileTime;
      for (uint32_t i = 0; i < hits_in_fit; ++i) {
        {
          constexpr uint32_t ge_idx = 0, j = 0, l = 0;
          hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 2, j = 1, l = 1;
          hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 1, j = 1, l = 0;
          hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
              ge.col(i)[ge_idx];
        }
      }
    }

    template <typename TAcc, typename M6xNf, typename M3xNd>
    ALPAKA_FN_ACC void loadCovariance(const TAcc& acc, M6xNf const& ge, M3xNd& hits_cov) {
      // Index numerology:
      // i: index of the hits/point (0,..,3)
      // j: index of space component (x,y,z)
      // l: index of space components (x,y,z)
      // ge is always in sync with the index i and is formatted as:
      // ge[] ==> [xx, xy, yy, xz, yz, zz]
      // in (j,l) notation, we have:
      // ge[] ==> [(0,0), (0,1), (1,1), (0,2), (1,2), (2,2)]
      // so the index ge_idx corresponds to the matrix elements:
      // | 0  1  3 |
      // | 1  2  4 |
      // | 3  4  5 |
      constexpr uint32_t hits_in_fit = M6xNf::ColsAtCompileTime;
      for (uint32_t i = 0; i < hits_in_fit; ++i) {
        {
          constexpr uint32_t ge_idx = 0, j = 0, l = 0;
          hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 2, j = 1, l = 1;
          hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 5, j = 2, l = 2;
          hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 1, j = 1, l = 0;
          hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
              ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 3, j = 2, l = 0;
          hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
              ge.col(i)[ge_idx];
        }
        {
          constexpr uint32_t ge_idx = 4, j = 2, l = 1;
          hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
              ge.col(i)[ge_idx];
        }
      }
    }

    template <typename TAcc>
    ALPAKA_FN_ACC ALPAKA_FN_INLINE void par_uvrtopak(const TAcc& acc,
                                                     CircleFit& circle,
                                                     const double B,
                                                     const bool error) {
      Vector3d par_pak;
      const double temp0 = circle.par.head(2).squaredNorm();
      const double temp1 = alpaka::math::sqrt(acc, temp0);
      par_pak << alpaka::math::atan2(acc, circle.qCharge * circle.par(0), -circle.qCharge * circle.par(1)),
          circle.qCharge * (temp1 - circle.par(2)), circle.par(2) * B;
      if (error) {
        const double temp2 = sqr(circle.par(0)) * 1. / temp0;
        const double temp3 = 1. / temp1 * circle.qCharge;
        Matrix3d j4Mat;
        j4Mat << -circle.par(1) * temp2 * 1. / sqr(circle.par(0)), temp2 * 1. / circle.par(0), 0.,
            circle.par(0) * temp3, circle.par(1) * temp3, -circle.qCharge, 0., 0., B;
        circle.cov = j4Mat * circle.cov * j4Mat.transpose();
      }
      circle.par = par_pak;
    }

    template <typename TAcc>
    ALPAKA_FN_ACC ALPAKA_FN_INLINE void fromCircleToPerigee(const TAcc& acc, CircleFit& circle) {
      Vector3d par_pak;
      const double temp0 = circle.par.head(2).squaredNorm();
      const double temp1 = alpaka::math::sqrt(acc, temp0);
      par_pak << alpaka::math::atan2(acc, circle.qCharge * circle.par(0), -circle.qCharge * circle.par(1)),
          circle.qCharge * (temp1 - circle.par(2)), circle.qCharge / circle.par(2);

      const double temp2 = sqr(circle.par(0)) * 1. / temp0;
      const double temp3 = 1. / temp1 * circle.qCharge;
      Matrix3d j4Mat;
      j4Mat << -circle.par(1) * temp2 * 1. / sqr(circle.par(0)), temp2 * 1. / circle.par(0), 0., circle.par(0) * temp3,
          circle.par(1) * temp3, -circle.qCharge, 0., 0., -circle.qCharge / (circle.par(2) * circle.par(2));
      circle.cov = j4Mat * circle.cov * j4Mat.transpose();

      circle.par = par_pak;
    }

  }  // namespace riemannFit

}  // namespace ALPAKA_ACCELERATOR_NAMESPACE

#endif  // RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h