LSTCore/interface/LSTPrepareInput.h

0001 #ifndef RecoTracker_LSTCore_interface_LSTPrepareInput_h
0002 #define RecoTracker_LSTCore_interface_LSTPrepareInput_h
0003
0004 #include <memory>
0005 #include <Math/Vector3D.h>
0006 #include <Math/VectorUtil.h>
0007
0008 #include "RecoTracker/LSTCore/interface/Common.h"
0009 #include "RecoTracker/LSTCore/interface/LSTInputHostCollection.h"
0010
0011 namespace lst {
0012
0013   inline ROOT::Math::XYZVector calculateR3FromPCA(const ROOT::Math::XYZVector& p3, float dxy, float dz) {
0014     const float pt = p3.rho();
0015     const float p = p3.r();
0016     const float vz = dz * pt * pt / p / p;
0017
0018     const float vx = -dxy * p3.y() / pt - p3.x() / p * p3.z() / p * dz;
0019     const float vy = dxy * p3.x() / pt - p3.y() / p * p3.z() / p * dz;
0020     return {vx, vy, vz};
0021   }
0022
0023   template <typename TQueue>
0024   inline LSTInputHostCollection prepareInput(std::vector<float> const& see_px,
0025                                              std::vector<float> const& see_py,
0026                                              std::vector<float> const& see_pz,
0027                                              std::vector<float> const& see_dxy,
0028                                              std::vector<float> const& see_dz,
0029                                              std::vector<float> const& see_ptErr,
0030                                              std::vector<float> const& see_etaErr,
0031                                              std::vector<float> const& see_stateTrajGlbX,
0032                                              std::vector<float> const& see_stateTrajGlbY,
0033                                              std::vector<float> const& see_stateTrajGlbZ,
0034                                              std::vector<float> const& see_stateTrajGlbPx,
0035                                              std::vector<float> const& see_stateTrajGlbPy,
0036                                              std::vector<float> const& see_stateTrajGlbPz,
0037                                              std::vector<int> const& see_q,
0038                                              std::vector<std::vector<int>> const& see_hitIdx,
0039                                              std::vector<unsigned int> const& see_algo,
0040                                              std::vector<unsigned int> const& ph2_detId,
0041                                              std::vector<float> const& ph2_x,
0042                                              std::vector<float> const& ph2_y,
0043                                              std::vector<float> const& ph2_z,
0044 #ifndef LST_STANDALONE
0045                                              std::vector<TrackingRecHit const*> const& ph2_hits,
0046 #endif
0047                                              float const ptCut,
0048                                              TQueue const& queue) {
0049     std::vector<float> trkX;
0050     std::vector<float> trkY;
0051     std::vector<float> trkZ;
0052     std::vector<unsigned int> hitId;
0053     std::vector<unsigned int> hitIdxs;
0054     std::vector<Params_pLS::ArrayUxHits> hitIndices_vec;
0055     std::vector<float> deltaPhi_vec;
0056     std::vector<float> ptIn_vec;
0057     std::vector<float> ptErr_vec;
0058     std::vector<float> px_vec;
0059     std::vector<float> py_vec;
0060     std::vector<float> pz_vec;
0061     std::vector<float> eta_vec;
0062     std::vector<float> etaErr_vec;
0063     std::vector<float> phi_vec;
0064     std::vector<int> charge_vec;
0065     std::vector<unsigned int> seedIdx_vec;
0066     std::vector<int> superbin_vec;
0067     std::vector<PixelType> pixelType_vec;
0068     std::vector<char> isQuad_vec;
0069
0070     const int hit_size = ph2_x.size();
0071
0072     unsigned int count = 0;
0073     auto n_see = see_stateTrajGlbPx.size();
0074     px_vec.reserve(n_see);
0075     py_vec.reserve(n_see);
0076     pz_vec.reserve(n_see);
0077     hitIndices_vec.reserve(n_see);
0078     ptIn_vec.reserve(n_see);
0079     ptErr_vec.reserve(n_see);
0080     etaErr_vec.reserve(n_see);
0081     eta_vec.reserve(n_see);
0082     phi_vec.reserve(n_see);
0083     charge_vec.reserve(n_see);
0084     seedIdx_vec.reserve(n_see);
0085     deltaPhi_vec.reserve(n_see);
0086     trkX.reserve(4 * n_see);
0087     trkY.reserve(4 * n_see);
0088     trkZ.reserve(4 * n_see);
0089     hitId.reserve(4 * n_see);
0090     hitIdxs.reserve(hit_size + 4 * n_see);
0091     hitIdxs.resize(hit_size);
0092
0093     std::iota(hitIdxs.begin(), hitIdxs.end(), 0);
0094
0095     for (size_t iSeed = 0; iSeed < n_see; iSeed++) {
0096       // Only needed for standalone
0097       bool good_seed_type = see_algo.empty() || see_algo[iSeed] == 4 || see_algo[iSeed] == 22;
0098       if (!good_seed_type)
0099         continue;
0100
0101       ROOT::Math::XYZVector p3LH(see_stateTrajGlbPx[iSeed], see_stateTrajGlbPy[iSeed], see_stateTrajGlbPz[iSeed]);
0102       float ptIn = p3LH.rho();
0103       float eta = p3LH.eta();
0104       float ptErr = see_ptErr[iSeed];
0105
0106       if ((ptIn > ptCut - 2 * ptErr)) {
0107         ROOT::Math::XYZVector r3LH(see_stateTrajGlbX[iSeed], see_stateTrajGlbY[iSeed], see_stateTrajGlbZ[iSeed]);
0108         ROOT::Math::XYZVector p3PCA(see_px[iSeed], see_py[iSeed], see_pz[iSeed]);
0109         ROOT::Math::XYZVector r3PCA(calculateR3FromPCA(p3PCA, see_dxy[iSeed], see_dz[iSeed]));
0110
0111         // The charge could be used directly in the line below
0112         float pixelSegmentDeltaPhiChange = ROOT::Math::VectorUtil::DeltaPhi(p3LH, r3LH);
0113         float etaErr = see_etaErr[iSeed];
0114         float px = p3LH.x();
0115         float py = p3LH.y();
0116         float pz = p3LH.z();
0117
0118         int charge = see_q[iSeed];
0119         PixelType pixtype = PixelType::kInvalid;
0120
0121         if (ptIn >= 2.0)
0122           pixtype = PixelType::kHighPt;
0123         else if (ptIn >= (ptCut - 2 * ptErr) and ptIn < 2.0) {
0124           if (pixelSegmentDeltaPhiChange >= 0)
0125             pixtype = PixelType::kLowPtPosCurv;
0126           else
0127             pixtype = PixelType::kLowPtNegCurv;
0128         } else
0129           continue;
0130
0131         unsigned int hitIdx0 = hit_size + count;
0132         count++;
0133         unsigned int hitIdx1 = hit_size + count;
0134         count++;
0135         unsigned int hitIdx2 = hit_size + count;
0136         count++;
0137         unsigned int hitIdx3;
0138         if (see_hitIdx[iSeed].size() <= 3)
0139           hitIdx3 = hitIdx2;
0140         else {
0141           hitIdx3 = hit_size + count;
0142           count++;
0143         }
0144
0145         trkX.push_back(r3PCA.x());
0146         trkY.push_back(r3PCA.y());
0147         trkZ.push_back(r3PCA.z());
0148         trkX.push_back(p3PCA.rho());
0149         float p3PCA_Eta = p3PCA.eta();
0150         trkY.push_back(p3PCA_Eta);
0151         float p3PCA_Phi = p3PCA.phi();
0152         trkZ.push_back(p3PCA_Phi);
0153         trkX.push_back(r3LH.x());
0154         trkY.push_back(r3LH.y());
0155         trkZ.push_back(r3LH.z());
0156         hitId.push_back(1);
0157         hitId.push_back(1);
0158         hitId.push_back(1);
0159         if (see_hitIdx[iSeed].size() > 3) {
0160           trkX.push_back(r3LH.x());
0161           trkY.push_back(see_dxy[iSeed]);
0162           trkZ.push_back(see_dz[iSeed]);
0163           hitId.push_back(1);
0164         }
0165         px_vec.push_back(px);
0166         py_vec.push_back(py);
0167         pz_vec.push_back(pz);
0168
0169         hitIndices_vec.push_back({{hitIdx0, hitIdx1, hitIdx2, hitIdx3}});
0170         ptIn_vec.push_back(ptIn);
0171         ptErr_vec.push_back(ptErr);
0172         etaErr_vec.push_back(etaErr);
0173         eta_vec.push_back(eta);
0174         float phi = p3LH.phi();
0175         phi_vec.push_back(phi);
0176         charge_vec.push_back(charge);
0177         seedIdx_vec.push_back(iSeed);
0178         deltaPhi_vec.push_back(pixelSegmentDeltaPhiChange);
0179
0180         hitIdxs.push_back(see_hitIdx[iSeed][0]);
0181         hitIdxs.push_back(see_hitIdx[iSeed][1]);
0182         hitIdxs.push_back(see_hitIdx[iSeed][2]);
0183         char isQuad = false;
0184         if (see_hitIdx[iSeed].size() > 3) {
0185           isQuad = true;
0186           hitIdxs.push_back(see_hitIdx[iSeed][3]);
0187         }
0188         float neta = 25.;
0189         float nphi = 72.;
0190         float nz = 25.;
0191         int etabin = (p3PCA_Eta + 2.6) / ((2 * 2.6) / neta);
0192         int phibin = (p3PCA_Phi + std::numbers::pi_v<float>) / ((2. * std::numbers::pi_v<float>) / nphi);
0193         int dzbin = (std::clamp(see_dz[iSeed], -30.f, 30.f) + 30) / (2 * 30 / nz);
0194         int isuperbin = (nz * nphi) * etabin + (nz)*phibin + dzbin;
0195         superbin_vec.push_back(isuperbin);
0196         pixelType_vec.push_back(pixtype);
0197         isQuad_vec.push_back(isQuad);
0198       }
0199     }
0200
0201     // Build the SoAs
0202     int nHitsOT = ph2_x.size();
0203     int nHitsIT = trkX.size();
0204     int nPixelSeeds = ptIn_vec.size();
0205     if (static_cast<unsigned int>(nPixelSeeds) > n_max_pixel_segments_per_module) {
0206       nPixelSeeds = n_max_pixel_segments_per_module;
0207     }
0208
0209     std::array<int, 2> const soa_sizes{{nHitsIT + nHitsOT, nPixelSeeds}};
0210     LSTInputHostCollection lstInputHC(soa_sizes, queue);
0211
0212     auto hits = lstInputHC.view<HitsBaseSoA>();
0213     std::memcpy(hits.xs(), ph2_x.data(), nHitsOT * sizeof(float));
0214     std::memcpy(hits.ys(), ph2_y.data(), nHitsOT * sizeof(float));
0215     std::memcpy(hits.zs(), ph2_z.data(), nHitsOT * sizeof(float));
0216     std::memcpy(hits.detid(), ph2_detId.data(), nHitsOT * sizeof(unsigned int));
0217 #ifndef LST_STANDALONE
0218     std::memcpy(hits.hits(), ph2_hits.data(), nHitsOT * sizeof(TrackingRecHit const*));
0219 #endif
0220
0221     std::memcpy(hits.xs() + nHitsOT, trkX.data(), nHitsIT * sizeof(float));
0222     std::memcpy(hits.ys() + nHitsOT, trkY.data(), nHitsIT * sizeof(float));
0223     std::memcpy(hits.zs() + nHitsOT, trkZ.data(), nHitsIT * sizeof(float));
0224     std::memcpy(hits.detid() + nHitsOT, hitId.data(), nHitsIT * sizeof(unsigned int));
0225 #ifndef LST_STANDALONE
0226     std::memset(hits.hits() + nHitsOT, 0, nHitsIT * sizeof(TrackingRecHit const*));
0227 #endif
0228
0229     std::memcpy(hits.idxs(), hitIdxs.data(), (nHitsIT + nHitsOT) * sizeof(unsigned int));
0230
0231     auto pixelSeeds = lstInputHC.view<PixelSeedsSoA>();
0232     std::memcpy(pixelSeeds.hitIndices(), hitIndices_vec.data(), nPixelSeeds * sizeof(Params_pLS::ArrayUxHits));
0233     std::memcpy(pixelSeeds.deltaPhi(), deltaPhi_vec.data(), nPixelSeeds * sizeof(float));
0234     std::memcpy(pixelSeeds.ptIn(), ptIn_vec.data(), nPixelSeeds * sizeof(float));
0235     std::memcpy(pixelSeeds.ptErr(), ptErr_vec.data(), nPixelSeeds * sizeof(float));
0236     std::memcpy(pixelSeeds.px(), px_vec.data(), nPixelSeeds * sizeof(float));
0237     std::memcpy(pixelSeeds.py(), py_vec.data(), nPixelSeeds * sizeof(float));
0238     std::memcpy(pixelSeeds.pz(), pz_vec.data(), nPixelSeeds * sizeof(float));
0239     std::memcpy(pixelSeeds.etaErr(), etaErr_vec.data(), nPixelSeeds * sizeof(float));
0240     std::memcpy(pixelSeeds.isQuad(), isQuad_vec.data(), nPixelSeeds * sizeof(char));
0241     std::memcpy(pixelSeeds.eta(), eta_vec.data(), nPixelSeeds * sizeof(float));
0242     std::memcpy(pixelSeeds.phi(), phi_vec.data(), nPixelSeeds * sizeof(float));
0243     std::memcpy(pixelSeeds.charge(), charge_vec.data(), nPixelSeeds * sizeof(int));
0244     std::memcpy(pixelSeeds.seedIdx(), seedIdx_vec.data(), nPixelSeeds * sizeof(unsigned int));
0245     std::memcpy(pixelSeeds.superbin(), superbin_vec.data(), nPixelSeeds * sizeof(int));
0246     std::memcpy(pixelSeeds.pixelType(), pixelType_vec.data(), nPixelSeeds * sizeof(PixelType));
0247
0248     return lstInputHC;
0249   }
0250
0251 }  // namespace lst
0252
0253 #endif