TrackSoA/interface/TracksSoA.h

0001 #ifndef DataFormats_TrackSoA_interface_TracksSoA_h
0002 #define DataFormats_TrackSoA_interface_TracksSoA_h
0003
0004 #include <alpaka/alpaka.hpp>
0005
0006 #include <Eigen/Core>
0007
0008 #include "HeterogeneousCore/AlpakaInterface/interface/OneToManyAssoc.h"
0009 #include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
0010 #include "DataFormats/SoATemplate/interface/SoALayout.h"
0011 #include "DataFormats/TrackSoA/interface/TrackDefinitions.h"
0012
0013 namespace reco {
0014
0015   // Aliases in order to not confuse the GENERATE_SOA_LAYOUT
0016   // macro with weird colons and angled brackets.
0017   using Vector5f = Eigen::Matrix<float, 5, 1>;
0018   using Vector15f = Eigen::Matrix<float, 15, 1>;
0019   using Quality = pixelTrack::Quality;
0020
0021   GENERATE_SOA_LAYOUT(TrackLayout,
0022                       SOA_COLUMN(Quality, quality),
0023                       SOA_COLUMN(float, chi2),
0024                       SOA_COLUMN(int8_t, nLayers),
0025                       SOA_COLUMN(float, eta),
0026                       SOA_COLUMN(float, pt),
0027                       // state at the beam spot: {phi, tip, 1/pt, cotan(theta), zip}
0028                       SOA_EIGEN_COLUMN(Vector5f, state),
0029                       SOA_EIGEN_COLUMN(Vector15f, covariance),
0030                       SOA_SCALAR(int, nTracks),
0031                       SOA_COLUMN(uint32_t, hitOffsets))
0032
0033   GENERATE_SOA_LAYOUT(TrackHitsLayout, SOA_COLUMN(uint32_t, id), SOA_COLUMN(uint32_t, detId))
0034
0035   using TrackSoA = TrackLayout<>;
0036   using TrackSoAView = TrackSoA::View;
0037   using TrackSoAConstView = TrackSoA::ConstView;
0038
0039   using TrackHitSoA = TrackHitsLayout<>;
0040   using TrackHitSoAView = TrackHitSoA::View;
0041   using TrackHitSoAConstView = TrackHitSoA::ConstView;
0042
0043   // All these below were constexpr. Now I get this:
0044   // note: non-literal type 'reco::TrackLayout<128, false>::ConstViewTemplateFreeParams<128, false, true, true>::const_element'
0045   // cannot be used in a constant expression
0046
0047   // TODO: move to use the layer gaps defined in CAGeometry
0048   ALPAKA_FN_HOST_ACC inline int nLayers(const TrackSoAConstView &tracks,
0049                                         const TrackHitSoAConstView &hits,
0050                                         uint16_t maxLayers,
0051                                         uint32_t const *__restrict__ layerStarts,
0052                                         int32_t i) {
0053     auto start = (i == 0) ? 0 : tracks[i - 1].hitOffsets();
0054     auto end = tracks[i].hitOffsets();
0055     auto hitId = hits[start].id();
0056     int nl = 1;
0057     int ol = 0;
0058     while (hitId >= layerStarts[ol + 1] and ol < maxLayers)
0059       ++ol;
0060     ++start;
0061     for (; start < end; ++start) {
0062       hitId = hits[start].id();
0063       int il = 0;
0064       while (hitId >= layerStarts[il + 1] and il < maxLayers)
0065         ++il;
0066       if (il != ol)
0067         ++nl;
0068       ol = il;
0069     }
0070     return nl;
0071   }
0072
0073   ALPAKA_FN_HOST_ACC inline float charge(const TrackSoAConstView &tracks, int32_t i) {
0074     //was: std::copysign(1.f, tracks[i].state()(2)). Will be constexpr with C++23
0075     float v = tracks[i].state()(2);
0076     return float((0.0f < v) - (v < 0.0f));
0077   }
0078
0079   ALPAKA_FN_HOST_ACC inline float phi(const TrackSoAConstView &tracks, int32_t i) { return tracks[i].state()(0); }
0080
0081   ALPAKA_FN_HOST_ACC inline float tip(const TrackSoAConstView &tracks, int32_t i) { return tracks[i].state()(1); }
0082
0083   ALPAKA_FN_HOST_ACC inline float zip(const TrackSoAConstView &tracks, int32_t i) { return tracks[i].state()(4); }
0084
0085   ALPAKA_FN_HOST_ACC inline bool isTriplet(const TrackSoAConstView &tracks, int32_t i) {
0086     return tracks[i].nLayers() == 3;
0087   }
0088
0089   ALPAKA_FN_HOST_ACC inline int nHits(const TrackSoAConstView &tracks, int i) {
0090     auto start = (i == 0) ? 0 : tracks[i - 1].hitOffsets();
0091     return tracks[i].hitOffsets() - start;
0092   }
0093
0094   // state at the beam spot: { phi, tip, 1/pt, cotan(theta), zip }
0095
0096   template <typename V3, typename M3, typename V2, typename M2>
0097   ALPAKA_FN_HOST_ACC inline void copyFromCircle(
0098       TrackSoAView &tracks, V3 const &cp, M3 const &ccov, V2 const &lp, M2 const &lcov, float b, int32_t i) {
0099     tracks[i].state() << cp.template cast<float>(), lp.template cast<float>();
0100
0101     tracks[i].state()(2) = tracks[i].state()(2) * b;
0102     auto cov = tracks[i].covariance();
0103     cov(0) = ccov(0, 0);
0104     cov(1) = ccov(0, 1);
0105     cov(2) = b * float(ccov(0, 2));
0106     cov(4) = cov(3) = 0;
0107     cov(5) = ccov(1, 1);
0108     cov(6) = b * float(ccov(1, 2));
0109     cov(8) = cov(7) = 0;
0110     cov(9) = b * b * float(ccov(2, 2));
0111     cov(11) = cov(10) = 0;
0112     cov(12) = lcov(0, 0);
0113     cov(13) = lcov(0, 1);
0114     cov(14) = lcov(1, 1);
0115   }
0116
0117   template <typename V5, typename M5>
0118   ALPAKA_FN_HOST_ACC inline void copyFromDense(TrackSoAView &tracks, V5 const &v, M5 const &cov, int32_t i) {
0119     tracks[i].state() = v.template cast<float>();
0120     for (int j = 0, ind = 0; j < 5; ++j)
0121       for (auto k = j; k < 5; ++k)
0122         tracks[i].covariance()(ind++) = cov(j, k);
0123   }
0124
0125   template <typename V5, typename M5>
0126   ALPAKA_FN_HOST_ACC inline void copyToDense(const TrackSoAConstView &tracks, V5 &v, M5 &cov, int32_t i) {
0127     v = tracks[i].state().template cast<typename V5::Scalar>();
0128     for (int j = 0, ind = 0; j < 5; ++j) {
0129       cov(j, j) = tracks[i].covariance()(ind++);
0130       for (auto k = j + 1; k < 5; ++k)
0131         cov(k, j) = cov(j, k) = tracks[i].covariance()(ind++);
0132     }
0133   }
0134
0135 }  // namespace reco
0136
0137 #endif  // DataFormats_TrackSoA_interface_TracksSoA_h