File indexing completed on 2024-04-06 12:28:37
0001 #ifndef RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h
0002 #define RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h
0003
0004 #include <alpaka/alpaka.hpp>
0005
0006 #include <Eigen/Core>
0007
0008 #include "DataFormats/Math/interface/choleskyInversion.h"
0009 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
0010 #include "RecoTracker/PixelTrackFitting/interface/alpaka/FitResult.h"
0011
0012 namespace riemannFit {
0013
0014 constexpr double epsilon = 1.e-4;
0015
0016 using VectorXd = Eigen::VectorXd;
0017 using MatrixXd = Eigen::MatrixXd;
0018 template <int N>
0019 using MatrixNd = Eigen::Matrix<double, N, N>;
0020 template <int N>
0021 using MatrixNplusONEd = Eigen::Matrix<double, N + 1, N + 1>;
0022 template <int N>
0023 using ArrayNd = Eigen::Array<double, N, N>;
0024 template <int N>
0025 using Matrix2Nd = Eigen::Matrix<double, 2 * N, 2 * N>;
0026 template <int N>
0027 using Matrix3Nd = Eigen::Matrix<double, 3 * N, 3 * N>;
0028 template <int N>
0029 using Matrix2xNd = Eigen::Matrix<double, 2, N>;
0030 template <int N>
0031 using Array2xNd = Eigen::Array<double, 2, N>;
0032 template <int N>
0033 using MatrixNx3d = Eigen::Matrix<double, N, 3>;
0034 template <int N>
0035 using MatrixNx5d = Eigen::Matrix<double, N, 5>;
0036 template <int N>
0037 using VectorNd = Eigen::Matrix<double, N, 1>;
0038 template <int N>
0039 using VectorNplusONEd = Eigen::Matrix<double, N + 1, 1>;
0040 template <int N>
0041 using Vector2Nd = Eigen::Matrix<double, 2 * N, 1>;
0042 template <int N>
0043 using Vector3Nd = Eigen::Matrix<double, 3 * N, 1>;
0044 template <int N>
0045 using RowVectorNd = Eigen::Matrix<double, 1, 1, N>;
0046 template <int N>
0047 using RowVector2Nd = Eigen::Matrix<double, 1, 2 * N>;
0048
0049 using Matrix2x3d = Eigen::Matrix<double, 2, 3>;
0050
0051 using Matrix3f = Eigen::Matrix3f;
0052 using Vector3f = Eigen::Vector3f;
0053 using Vector4f = Eigen::Vector4f;
0054 using Vector6f = Eigen::Matrix<double, 6, 1>;
0055
0056
0057
0058
0059 template <typename VI5, typename MI5, typename VO5, typename MO5>
0060 inline void transformToPerigeePlane(VI5 const& ip, MI5 const& icov, VO5& op, MO5& ocov) {
0061 auto sinTheta2 = 1. / (1. + ip(3) * ip(3));
0062 auto sinTheta = std::sqrt(sinTheta2);
0063 auto cosTheta = ip(3) * sinTheta;
0064
0065 op(0) = sinTheta * ip(2);
0066 op(1) = 0.;
0067 op(2) = -ip(3);
0068 op(3) = ip(1);
0069 op(4) = -ip(4);
0070
0071 Matrix5d jMat = Matrix5d::Zero();
0072
0073 jMat(0, 2) = sinTheta;
0074 jMat(0, 3) = -sinTheta2 * cosTheta * ip(2);
0075 jMat(1, 0) = 1.;
0076 jMat(2, 3) = -1.;
0077 jMat(3, 1) = 1.;
0078 jMat(4, 4) = -1;
0079
0080 ocov = jMat * icov * jMat.transpose();
0081 }
0082
0083 }
0084
0085 namespace ALPAKA_ACCELERATOR_NAMESPACE {
0086 namespace riemannFit {
0087 using namespace ::riemannFit;
0088
0089 template <typename TAcc, class C>
0090 ALPAKA_FN_ACC void printIt(const TAcc& acc, C* m, const char* prefix = "") {
0091 #ifdef RFIT_DEBUG
0092 for (uint r = 0; r < m->rows(); ++r) {
0093 for (uint c = 0; c < m->cols(); ++c) {
0094 printf("%s Matrix(%d,%d) = %g\n", prefix, r, c, (*m)(r, c));
0095 }
0096 }
0097 #endif
0098 }
0099
0100
0101
0102
0103 template <typename T>
0104 constexpr T sqr(const T a) {
0105 return a * a;
0106 }
0107
0108
0109
0110
0111
0112
0113
0114
0115
0116 template <typename TAcc>
0117 ALPAKA_FN_ACC ALPAKA_FN_INLINE double cross2D(const TAcc& acc, const Vector2d& a, const Vector2d& b) {
0118 return a.x() * b.y() - a.y() * b.x();
0119 }
0120
0121
0122
0123
0124
0125 template <typename TAcc, typename M6xNf, typename M2Nd>
0126 ALPAKA_FN_ACC void loadCovariance2D(const TAcc& acc, M6xNf const& ge, M2Nd& hits_cov) {
0127
0128
0129
0130
0131
0132
0133
0134
0135
0136
0137
0138
0139 constexpr uint32_t hits_in_fit = M6xNf::ColsAtCompileTime;
0140 for (uint32_t i = 0; i < hits_in_fit; ++i) {
0141 {
0142 constexpr uint32_t ge_idx = 0, j = 0, l = 0;
0143 hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
0144 }
0145 {
0146 constexpr uint32_t ge_idx = 2, j = 1, l = 1;
0147 hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
0148 }
0149 {
0150 constexpr uint32_t ge_idx = 1, j = 1, l = 0;
0151 hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
0152 ge.col(i)[ge_idx];
0153 }
0154 }
0155 }
0156
0157 template <typename TAcc, typename M6xNf, typename M3xNd>
0158 ALPAKA_FN_ACC void loadCovariance(const TAcc& acc, M6xNf const& ge, M3xNd& hits_cov) {
0159
0160
0161
0162
0163
0164
0165
0166
0167
0168
0169
0170
0171 constexpr uint32_t hits_in_fit = M6xNf::ColsAtCompileTime;
0172 for (uint32_t i = 0; i < hits_in_fit; ++i) {
0173 {
0174 constexpr uint32_t ge_idx = 0, j = 0, l = 0;
0175 hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
0176 }
0177 {
0178 constexpr uint32_t ge_idx = 2, j = 1, l = 1;
0179 hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
0180 }
0181 {
0182 constexpr uint32_t ge_idx = 5, j = 2, l = 2;
0183 hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) = ge.col(i)[ge_idx];
0184 }
0185 {
0186 constexpr uint32_t ge_idx = 1, j = 1, l = 0;
0187 hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
0188 ge.col(i)[ge_idx];
0189 }
0190 {
0191 constexpr uint32_t ge_idx = 3, j = 2, l = 0;
0192 hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
0193 ge.col(i)[ge_idx];
0194 }
0195 {
0196 constexpr uint32_t ge_idx = 4, j = 2, l = 1;
0197 hits_cov(i + l * hits_in_fit, i + j * hits_in_fit) = hits_cov(i + j * hits_in_fit, i + l * hits_in_fit) =
0198 ge.col(i)[ge_idx];
0199 }
0200 }
0201 }
0202
0203
0204
0205
0206
0207
0208
0209
0210
0211 template <typename TAcc>
0212 ALPAKA_FN_ACC ALPAKA_FN_INLINE void par_uvrtopak(const TAcc& acc,
0213 CircleFit& circle,
0214 const double B,
0215 const bool error) {
0216 Vector3d par_pak;
0217 const double temp0 = circle.par.head(2).squaredNorm();
0218 const double temp1 = alpaka::math::sqrt(acc, temp0);
0219 par_pak << alpaka::math::atan2(acc, circle.qCharge * circle.par(0), -circle.qCharge * circle.par(1)),
0220 circle.qCharge * (temp1 - circle.par(2)), circle.par(2) * B;
0221 if (error) {
0222 const double temp2 = sqr(circle.par(0)) * 1. / temp0;
0223 const double temp3 = 1. / temp1 * circle.qCharge;
0224 Matrix3d j4Mat;
0225 j4Mat << -circle.par(1) * temp2 * 1. / sqr(circle.par(0)), temp2 * 1. / circle.par(0), 0.,
0226 circle.par(0) * temp3, circle.par(1) * temp3, -circle.qCharge, 0., 0., B;
0227 circle.cov = j4Mat * circle.cov * j4Mat.transpose();
0228 }
0229 circle.par = par_pak;
0230 }
0231
0232
0233
0234
0235
0236
0237
0238 template <typename TAcc>
0239 ALPAKA_FN_ACC ALPAKA_FN_INLINE void fromCircleToPerigee(const TAcc& acc, CircleFit& circle) {
0240 Vector3d par_pak;
0241 const double temp0 = circle.par.head(2).squaredNorm();
0242 const double temp1 = alpaka::math::sqrt(acc, temp0);
0243 par_pak << alpaka::math::atan2(acc, circle.qCharge * circle.par(0), -circle.qCharge * circle.par(1)),
0244 circle.qCharge * (temp1 - circle.par(2)), circle.qCharge / circle.par(2);
0245
0246 const double temp2 = sqr(circle.par(0)) * 1. / temp0;
0247 const double temp3 = 1. / temp1 * circle.qCharge;
0248 Matrix3d j4Mat;
0249 j4Mat << -circle.par(1) * temp2 * 1. / sqr(circle.par(0)), temp2 * 1. / circle.par(0), 0., circle.par(0) * temp3,
0250 circle.par(1) * temp3, -circle.qCharge, 0., 0., -circle.qCharge / (circle.par(2) * circle.par(2));
0251 circle.cov = j4Mat * circle.cov * j4Mat.transpose();
0252
0253 circle.par = par_pak;
0254 }
0255
0256 }
0257
0258 }
0259
0260 #endif