PixelTrackFitting/plugins/PixelTrackProducerFromSoA.cc

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