PixelTrackFitting/plugins/PixelTrackProducerFromSoA.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 "CUDADataFormats/Common/interface/HostProduct.h"
0011 #include "CUDADataFormats/SiPixelCluster/interface/gpuClusteringConstants.h"
0012 #include "CUDADataFormats/Track/interface/PixelTrackUtilities.h"
0013 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousDevice.h"
0014 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousHost.h"
0015 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0016 #include "DataFormats/Common/interface/OrphanHandle.h"
0017 #include "DataFormats/GeometrySurface/interface/Plane.h"
0018 #include "DataFormats/TrackReco/interface/Track.h"
0019 #include "DataFormats/TrackReco/interface/TrackExtra.h"
0020 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0021 #include "DataFormats/TrackerCommon/interface/TrackerTopology.h"
0022 #include "DataFormats/TrackerRecHit2D/interface/SiPixelRecHitCollection.h"
0023 #include "DataFormats/TrajectoryState/interface/LocalTrajectoryParameters.h"
0024 #include "FWCore/Framework/interface/ConsumesCollector.h"
0025 #include "FWCore/Framework/interface/Event.h"
0026 #include "FWCore/Framework/interface/EventSetup.h"
0027 #include "FWCore/Framework/interface/MakerMacros.h"
0028 #include "FWCore/Framework/interface/global/EDProducer.h"
0029 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0030 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0031 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0032 #include "FWCore/PluginManager/interface/ModuleDef.h"
0033 #include "FWCore/Utilities/interface/EDGetToken.h"
0034 #include "FWCore/Utilities/interface/InputTag.h"
0035 #include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
0036 #include "Geometry/Records/interface/TrackerTopologyRcd.h"
0037 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0038 #include "RecoTracker/PixelTrackFitting/interface/FitUtils.h"
0039 #include "TrackingTools/AnalyticalJacobians/interface/JacobianLocalToCurvilinear.h"
0040 #include "TrackingTools/TrajectoryParametrization/interface/CurvilinearTrajectoryError.h"
0041 #include "TrackingTools/TrajectoryParametrization/interface/GlobalTrajectoryParameters.h"
0042
0043 #include "storeTracks.h"
0044
0045 /**
0046  * This class creates "legacy" reco::Track
0047  * objects from the output of SoA CA.
0048  */
0049 template <typename TrackerTraits>
0050 class PixelTrackProducerFromSoAT : public edm::global::EDProducer<> {
0051   using TrackSoAHost = TrackSoAHeterogeneousHost<TrackerTraits>;
0052   using TracksHelpers = TracksUtilities<TrackerTraits>;
0053   using HMSstorage = HostProduct<uint32_t[]>;
0054   using IndToEdm = std::vector<uint32_t>;
0055
0056 public:
0057   explicit PixelTrackProducerFromSoAT(const edm::ParameterSet &iConfig);
0058   ~PixelTrackProducerFromSoAT() override = default;
0059
0060   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0061
0062 private:
0063   void produce(edm::StreamID streamID, edm::Event &iEvent, const edm::EventSetup &iSetup) const override;
0064
0065   // Event Data tokens
0066   const edm::EDGetTokenT<reco::BeamSpot> tBeamSpot_;
0067   const edm::EDGetTokenT<TrackSoAHost> tokenTrack_;
0068   const edm::EDGetTokenT<SiPixelRecHitCollectionNew> cpuHits_;
0069   const edm::EDGetTokenT<HMSstorage> hmsToken_;
0070   // Event Setup tokens
0071   const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> idealMagneticFieldToken_;
0072   const edm::ESGetToken<TrackerTopology, TrackerTopologyRcd> ttTopoToken_;
0073
0074   int32_t const minNumberOfHits_;
0075   pixelTrack::Quality const minQuality_;
0076 };
0077
0078 template <typename TrackerTraits>
0079 PixelTrackProducerFromSoAT<TrackerTraits>::PixelTrackProducerFromSoAT(const edm::ParameterSet &iConfig)
0080     : tBeamSpot_(consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("beamSpot"))),
0081       tokenTrack_(consumes(iConfig.getParameter<edm::InputTag>("trackSrc"))),
0082       cpuHits_(consumes<SiPixelRecHitCollectionNew>(iConfig.getParameter<edm::InputTag>("pixelRecHitLegacySrc"))),
0083       hmsToken_(consumes<HMSstorage>(iConfig.getParameter<edm::InputTag>("pixelRecHitLegacySrc"))),
0084       idealMagneticFieldToken_(esConsumes()),
0085       ttTopoToken_(esConsumes()),
0086       minNumberOfHits_(iConfig.getParameter<int>("minNumberOfHits")),
0087       minQuality_(pixelTrack::qualityByName(iConfig.getParameter<std::string>("minQuality"))) {
0088   if (minQuality_ == pixelTrack::Quality::notQuality) {
0089     throw cms::Exception("PixelTrackConfiguration")
0090         << iConfig.getParameter<std::string>("minQuality") + " is not a pixelTrack::Quality";
0091   }
0092   if (minQuality_ < pixelTrack::Quality::dup) {
0093     throw cms::Exception("PixelTrackConfiguration")
0094         << iConfig.getParameter<std::string>("minQuality") + " not supported";
0095   }
0096   produces<TrackingRecHitCollection>();
0097   produces<reco::TrackExtraCollection>();
0098   // TrackCollection refers to TrackingRechit and TrackExtra
0099   // collections, need to declare its production after them to work
0100   // around a rare race condition in framework scheduling
0101   produces<reco::TrackCollection>();
0102   produces<IndToEdm>();
0103 }
0104
0105 template <typename TrackerTraits>
0106 void PixelTrackProducerFromSoAT<TrackerTraits>::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0107   edm::ParameterSetDescription desc;
0108   desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot"));
0109   desc.add<edm::InputTag>("trackSrc", edm::InputTag("pixelTracksSoA"));
0110   desc.add<edm::InputTag>("pixelRecHitLegacySrc", edm::InputTag("siPixelRecHitsPreSplittingLegacy"));
0111   desc.add<int>("minNumberOfHits", 0);
0112   desc.add<std::string>("minQuality", "loose");
0113   descriptions.addWithDefaultLabel(desc);
0114 }
0115
0116 template <typename TrackerTraits>
0117 void PixelTrackProducerFromSoAT<TrackerTraits>::produce(edm::StreamID streamID,
0118                                                         edm::Event &iEvent,
0119                                                         const edm::EventSetup &iSetup) const {
0120   // enum class Quality : uint8_t { bad = 0, edup, dup, loose, strict, tight, highPurity };
0121   reco::TrackBase::TrackQuality recoQuality[] = {reco::TrackBase::undefQuality,
0122                                                  reco::TrackBase::undefQuality,
0123                                                  reco::TrackBase::discarded,
0124                                                  reco::TrackBase::loose,
0125                                                  reco::TrackBase::tight,
0126                                                  reco::TrackBase::tight,
0127                                                  reco::TrackBase::highPurity};
0128   assert(reco::TrackBase::highPurity == recoQuality[int(pixelTrack::Quality::highPurity)]);
0129
0130   // std::cout << "Converting gpu helix in reco tracks" << std::endl;
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_).get();
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 in 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 = TracksHelpers::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   // std::cout << "processed " << nt << " good tuples " << tracks.size() << "out of " << indToEdm.size() << std::endl;
0258
0259   // store tracks
0260   storeTracks(iEvent, tracks, httopo);
0261   iEvent.put(std::move(indToEdmP));
0262 }
0263
0264 using PixelTrackProducerFromSoAPhase1 = PixelTrackProducerFromSoAT<pixelTopology::Phase1>;
0265 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAPhase1);
0266
0267 using PixelTrackProducerFromSoAPhase2 = PixelTrackProducerFromSoAT<pixelTopology::Phase2>;
0268 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAPhase2);
0269
0270 using PixelTrackProducerFromSoAHIonPhase1 = PixelTrackProducerFromSoAT<pixelTopology::HIonPhase1>;
0271 DEFINE_FWK_MODULE(PixelTrackProducerFromSoAHIonPhase1);