PixelTrackFitting/plugins/PixelTrackProducerFromSoAAlpaka.cc

0001 #include <algorithm>
0002 #include <cmath>
0003 #include <iostream>
0004 #include <memory>
0005 #include <numeric>
0006 #include <string>
0007 #include <utility>
0008 #include <vector>
0009
0010 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0011 #include "DataFormats/GeometrySurface/interface/Plane.h"
0012 #include "DataFormats/SiPixelClusterSoA/interface/ClusteringConstants.h"
0013 #include "DataFormats/TrackReco/interface/Track.h"
0014 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0015 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0016 #include "DataFormats/TrackSoA/interface/TracksHost.h"
0017 #include "DataFormats/TrackSoA/interface/alpaka/TrackUtilities.h"
0018 #include "DataFormats/TrackerCommon/interface/TrackerTopology.h"
0019 #include "DataFormats/TrackerRecHit2D/interface/SiPixelRecHitCollection.h"
0020 #include "DataFormats/TrajectoryState/interface/LocalTrajectoryParameters.h"
0021 #include "FWCore/Framework/interface/ConsumesCollector.h"
0022 #include "FWCore/Framework/interface/Event.h"
0023 #include "FWCore/Framework/interface/EventSetup.h"
0024 #include "FWCore/Framework/interface/global/EDProducer.h"
0025 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0026 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0027 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0028 #include "FWCore/Utilities/interface/EDGetToken.h"
0029 #include "FWCore/Utilities/interface/InputTag.h"
0030 #include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
0031 #include "Geometry/Records/interface/TrackerTopologyRcd.h"
0032 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0033 #include "RecoTracker/PixelTrackFitting/interface/alpaka/FitUtils.h"
0034 #include "TrackingTools/AnalyticalJacobians/interface/JacobianLocalToCurvilinear.h"
0035 #include "TrackingTools/TrajectoryParametrization/interface/CurvilinearTrajectoryError.h"
0036 #include "TrackingTools/TrajectoryParametrization/interface/GlobalTrajectoryParameters.h"
0037
0038 #include "storeTracks.h"
0039
0040 /**
0041  * This class creates "legacy" reco::Track
0042  * objects from the output of SoA CA.
0043  */
0044
0045 //#define GPU_DEBUG
0046
0047 template <typename TrackerTraits>
0048 class PixelTrackProducerFromSoAAlpaka : public edm::global::EDProducer<> {
0049   using TrackSoAHost = TracksHost<TrackerTraits>;
0050   using TracksHelpers = TracksUtilities<TrackerTraits>;
0051   using HMSstorage = std::vector<uint32_t>;
0052   using IndToEdm = std::vector<uint32_t>;
0053
0054 public:
0055   explicit PixelTrackProducerFromSoAAlpaka(const edm::ParameterSet &iConfig);
0056   ~PixelTrackProducerFromSoAAlpaka() override = default;
0057
0058   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0059
0060 private:
0061   void produce(edm::StreamID streamID, edm::Event &iEvent, const edm::EventSetup &iSetup) const override;
0062
0063   // Event Data tokens
0064   const edm::EDGetTokenT<reco::BeamSpot> tBeamSpot_;
0065   const edm::EDGetTokenT<TrackSoAHost> tokenTrack_;
0066   const edm::EDGetTokenT<SiPixelRecHitCollectionNew> cpuHits_;
0067   const edm::EDGetTokenT<HMSstorage> hmsToken_;
0068   // Event Setup tokens
0069   const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> idealMagneticFieldToken_;
0070   const edm::ESGetToken<TrackerTopology, TrackerTopologyRcd> ttTopoToken_;
0071
0072   int32_t const minNumberOfHits_;
0073   pixelTrack::Quality const minQuality_;
0074 };
0075
0076 template <typename TrackerTraits>
0077 PixelTrackProducerFromSoAAlpaka<TrackerTraits>::PixelTrackProducerFromSoAAlpaka(const edm::ParameterSet &iConfig)
0078     : tBeamSpot_(consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("beamSpot"))),
0079       tokenTrack_(consumes(iConfig.getParameter<edm::InputTag>("trackSrc"))),
0080       cpuHits_(consumes<SiPixelRecHitCollectionNew>(iConfig.getParameter<edm::InputTag>("pixelRecHitLegacySrc"))),
0081       hmsToken_(consumes<HMSstorage>(iConfig.getParameter<edm::InputTag>("pixelRecHitLegacySrc"))),
0082       idealMagneticFieldToken_(esConsumes()),
0083       ttTopoToken_(esConsumes()),
0084       minNumberOfHits_(iConfig.getParameter<int>("minNumberOfHits")),
0085       minQuality_(pixelTrack::qualityByName(iConfig.getParameter<std::string>("minQuality"))) {
0086   if (minQuality_ == pixelTrack::Quality::notQuality) {
0087     throw cms::Exception("PixelTrackConfiguration")
0088         << iConfig.getParameter<std::string>("minQuality") + " is not a pixelTrack::Quality";
0089   }
0090   if (minQuality_ < pixelTrack::Quality::dup) {
0091     throw cms::Exception("PixelTrackConfiguration")
0092         << iConfig.getParameter<std::string>("minQuality") + " not supported";
0093   }
0094   produces<TrackingRecHitCollection>();
0095   produces<reco::TrackExtraCollection>();
0096   // TrackCollection refers to TrackingRechit and TrackExtra
0097   // collections, need to declare its production after them to work
0098   // around a rare race condition in framework scheduling
0099   produces<reco::TrackCollection>();
0100   produces<IndToEdm>();
0101 }
0102
0103 template <typename TrackerTraits>
0104 void PixelTrackProducerFromSoAAlpaka<TrackerTraits>::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0105   edm::ParameterSetDescription desc;
0106   desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot"));
0107   desc.add<edm::InputTag>("trackSrc", edm::InputTag("pixelTracksAlpaka"));
0108   desc.add<edm::InputTag>("pixelRecHitLegacySrc", edm::InputTag("siPixelRecHitsPreSplittingLegacy"));
0109   desc.add<int>("minNumberOfHits", 0);
0110   desc.add<std::string>("minQuality", "loose");
0111   descriptions.addWithDefaultLabel(desc);
0112 }
0113
0114 template <typename TrackerTraits>
0115 void PixelTrackProducerFromSoAAlpaka<TrackerTraits>::produce(edm::StreamID streamID,
0116                                                              edm::Event &iEvent,
0117                                                              const edm::EventSetup &iSetup) const {
0118   // enum class Quality : uint8_t { bad = 0, edup, dup, loose, strict, tight, highPurity };
0119   reco::TrackBase::TrackQuality recoQuality[] = {reco::TrackBase::undefQuality,
0120                                                  reco::TrackBase::undefQuality,
0121                                                  reco::TrackBase::discarded,
0122                                                  reco::TrackBase::loose,
0123                                                  reco::TrackBase::tight,
0124                                                  reco::TrackBase::tight,
0125                                                  reco::TrackBase::highPurity};
0126   assert(reco::TrackBase::highPurity == recoQuality[int(pixelTrack::Quality::highPurity)]);
0127
0128 #ifdef GPU_DEBUG
0129   std::cout << "Converting soa helix in reco tracks" << std::endl;
0130 #endif
0131
0132   auto indToEdmP = std::make_unique<IndToEdm>();
0133   auto &indToEdm = *indToEdmP;
0134
0135   auto const &idealField = iSetup.getData(idealMagneticFieldToken_);
0136
0137   pixeltrackfitting::TracksWithRecHits tracks;
0138
0139   auto const &httopo = iSetup.getData(ttTopoToken_);
0140
0141   const auto &bsh = iEvent.get(tBeamSpot_);
0142   GlobalPoint bs(bsh.x0(), bsh.y0(), bsh.z0());
0143
0144   auto const &rechits = iEvent.get(cpuHits_);
0145   std::vector<TrackingRecHit const *> hitmap;
0146   auto const &rcs = rechits.data();
0147   auto const nhits = rcs.size();
0148
0149   hitmap.resize(nhits, nullptr);
0150
0151   auto const &hitsModuleStart = iEvent.get(hmsToken_);
0152
0153   for (auto const &hit : rcs) {
0154     auto const &thit = static_cast<BaseTrackerRecHit const &>(hit);
0155     auto const detI = thit.det()->index();
0156     auto const &clus = thit.firstClusterRef();
0157     assert(clus.isPixel());
0158     auto const idx = hitsModuleStart[detI] + clus.pixelCluster().originalId();
0159     if (idx >= hitmap.size())
0160       hitmap.resize(idx + 256, nullptr);  // only in case of hit overflow in one module
0161
0162     assert(nullptr == hitmap[idx]);
0163     hitmap[idx] = &hit;
0164   }
0165
0166   std::vector<const TrackingRecHit *> hits;
0167   hits.reserve(5);
0168
0169   auto const &tsoa = iEvent.get(tokenTrack_);
0170   auto const *quality = tsoa.view().quality();
0171   auto const &hitIndices = tsoa.view().hitIndices();
0172   auto nTracks = tsoa.view().nTracks();
0173
0174   tracks.reserve(nTracks);
0175
0176   int32_t nt = 0;
0177
0178   //sort index by pt
0179   std::vector<int32_t> sortIdxs(nTracks);
0180   std::iota(sortIdxs.begin(), sortIdxs.end(), 0);
0181   // sort good-quality tracks by pt, keep bad-quality tracks at the bottom
0182   std::sort(sortIdxs.begin(), sortIdxs.end(), [&](int32_t const i1, int32_t const i2) {
0183     if (quality[i1] >= minQuality_ && quality[i2] >= minQuality_)
0184       return tsoa.view()[i1].pt() > tsoa.view()[i2].pt();
0185     else
0186       return quality[i1] > quality[i2];
0187   });
0188
0189   //store the index of the SoA: indToEdm[index_SoAtrack] -> index_edmTrack (if it exists)
0190   indToEdm.resize(sortIdxs.size(), -1);
0191   for (const auto &it : sortIdxs) {
0192     auto nHits = TracksHelpers::nHits(tsoa.view(), it);
0193     assert(nHits >= 3);
0194     auto q = quality[it];
0195
0196     if (q < minQuality_)
0197       continue;
0198     if (nHits < minNumberOfHits_)  //move to nLayers?
0199       continue;
0200     indToEdm[it] = nt;
0201     ++nt;
0202
0203     hits.resize(nHits);
0204     auto b = hitIndices.begin(it);
0205     for (int iHit = 0; iHit < nHits; ++iHit)
0206       hits[iHit] = hitmap[*(b + iHit)];
0207
0208     // mind: this values are respect the beamspot!
0209     float chi2 = tsoa.view()[it].chi2();
0210     float phi = reco::phi(tsoa.view(), it);
0211
0212     riemannFit::Vector5d ipar, opar;
0213     riemannFit::Matrix5d icov, ocov;
0214     TracksHelpers::template copyToDense<riemannFit::Vector5d, riemannFit::Matrix5d>(tsoa.view(), ipar, icov, it);
0215     riemannFit::transformToPerigeePlane(ipar, icov, opar, ocov);
0216
0217     LocalTrajectoryParameters lpar(opar(0), opar(1), opar(2), opar(3), opar(4), 1.);
0218     AlgebraicSymMatrix55 m;
0219     for (int i = 0; i < 5; ++i)
0220       for (int j = i; j < 5; ++j)
0221         m(i, j) = ocov(i, j);
0222
0223     float sp = std::sin(phi);
0224     float cp = std::cos(phi);
0225     Surface::RotationType rot(sp, -cp, 0, 0, 0, -1.f, cp, sp, 0);
0226
0227     Plane impPointPlane(bs, rot);
0228     GlobalTrajectoryParameters gp(
0229         impPointPlane.toGlobal(lpar.position()), impPointPlane.toGlobal(lpar.momentum()), lpar.charge(), &idealField);
0230     JacobianLocalToCurvilinear jl2c(impPointPlane, lpar, idealField);
0231
0232     AlgebraicSymMatrix55 mo = ROOT::Math::Similarity(jl2c.jacobian(), m);
0233
0234     int ndof = 2 * hits.size() - 5;
0235     chi2 = chi2 * ndof;
0236     GlobalPoint vv = gp.position();
0237     math::XYZPoint pos(vv.x(), vv.y(), vv.z());
0238     GlobalVector pp = gp.momentum();
0239     math::XYZVector mom(pp.x(), pp.y(), pp.z());
0240
0241     auto track = std::make_unique<reco::Track>(chi2, ndof, pos, mom, gp.charge(), CurvilinearTrajectoryError(mo));
0242
0243     // bad and edup not supported as fit not present or not reliable
0244     auto tkq = recoQuality[int(q)];
0245     track->setQuality(tkq);
0246     // loose,tight and HP are inclusive
0247     if (reco::TrackBase::highPurity == tkq) {
0248       track->setQuality(reco::TrackBase::tight);
0249       track->setQuality(reco::TrackBase::loose);
0250     } else if (reco::TrackBase::tight == tkq) {
0251       track->setQuality(reco::TrackBase::loose);
0252     }
0253     track->setQuality(tkq);
0254     // filter???
0255     tracks.emplace_back(track.release(), hits);
0256   }
0257 #ifdef GPU_DEBUG
0258   std::cout << "processed " << nt << " good tuples " << tracks.size() << " out of " << indToEdm.size() << std::endl;
0259 #endif
0260
0261   // store tracks
0262   storeTracks(iEvent, tracks, httopo);
0263   iEvent.put(std::move(indToEdmP));
0264 }
0265
0266 using PixelTrackProducerFromSoAAlpakaPhase1 = PixelTrackProducerFromSoAAlpaka<pixelTopology::Phase1>;
0267 using PixelTrackProducerFromSoAAlpakaPhase2 = PixelTrackProducerFromSoAAlpaka<pixelTopology::Phase2>;
0268 using PixelTrackProducerFromSoAAlpakaHIonPhase1 = PixelTrackProducerFromSoAAlpaka<pixelTopology::HIonPhase1>;
0269
0270 #include "FWCore/Framework/interface/MakerMacros.h"
0271 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAAlpakaPhase1);
0272 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAAlpakaPhase2);
0273 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAAlpakaHIonPhase1);