Back to home page

Project CMSSW displayed by LXR

 
 

    


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;  //!< used in numerical derivative (J2 in Circle_fit())
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   // transformation between the "perigee" to cmssw localcoord frame
0056   // the plane of the latter is the perigee plane...
0057   // from   //!<(phi,Tip,q/pt,cotan(theta)),Zip)
0058   // to q/p,dx/dz,dy/dz,x,z
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 }  // namespace riemannFit
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     \brief raise to square.
0102   */
0103     template <typename T>
0104     constexpr T sqr(const T a) {
0105       return a * a;
0106     }
0107 
0108     /*!
0109     \brief Compute cross product of two 2D vector (assuming z component 0),
0110     returning z component of the result.
0111     \param a first 2D vector in the product.
0112     \param b second 2D vector in the product.
0113     \return z component of the cross product.
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    *  load error in CMSSW format to our formalism
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       // Index numerology:
0128       // i: index of the hits/point (0,..,3)
0129       // j: index of space component (x,y,z)
0130       // l: index of space components (x,y,z)
0131       // ge is always in sync with the index i and is formatted as:
0132       // ge[] ==> [xx, xy, yy, xz, yz, zz]
0133       // in (j,l) notation, we have:
0134       // ge[] ==> [(0,0), (0,1), (1,1), (0,2), (1,2), (2,2)]
0135       // so the index ge_idx corresponds to the matrix elements:
0136       // | 0  1  3 |
0137       // | 1  2  4 |
0138       // | 3  4  5 |
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       // Index numerology:
0160       // i: index of the hits/point (0,..,3)
0161       // j: index of space component (x,y,z)
0162       // l: index of space components (x,y,z)
0163       // ge is always in sync with the index i and is formatted as:
0164       // ge[] ==> [xx, xy, yy, xz, yz, zz]
0165       // in (j,l) notation, we have:
0166       // ge[] ==> [(0,0), (0,1), (1,1), (0,2), (1,2), (2,2)]
0167       // so the index ge_idx corresponds to the matrix elements:
0168       // | 0  1  3 |
0169       // | 1  2  4 |
0170       // | 3  4  5 |
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     \brief Transform circle parameter from (X0,Y0,R) to (phi,Tip,p_t) and
0205     consequently covariance matrix.
0206     \param circle_uvr parameter (X0,Y0,R), covariance matrix to
0207     be transformed and particle charge.
0208     \param B magnetic field in Gev/cm/c unit.
0209     \param error flag for errors computation.
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     \brief Transform circle parameter from (X0,Y0,R) to (phi,Tip,q/R) and
0234     consequently covariance matrix.
0235     \param circle_uvr parameter (X0,Y0,R), covariance matrix to
0236     be transformed and particle charge.
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   }  // namespace riemannFit
0257 
0258 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0259 
0260 #endif  // RecoTracker_PixelTrackFitting_interface_alpaka_FitUtils_h