TkSeedGenerator/plugins/SeedProducerFromSoA.cc

0001 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousHost.h"
0002 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0003 #include "DataFormats/GeometrySurface/interface/Plane.h"
0004 #include "DataFormats/TrackerCommon/interface/TrackerTopology.h"
0005 #include "DataFormats/TrackingRecHit/interface/InvalidTrackingRecHit.h"
0006 #include "DataFormats/TrajectorySeed/interface/TrajectorySeedCollection.h"
0007 #include "DataFormats/TrajectoryState/interface/LocalTrajectoryParameters.h"
0008 #include "FWCore/Framework/interface/ConsumesCollector.h"
0009 #include "FWCore/Framework/interface/ESHandle.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/ParameterSet/interface/ConfigurationDescriptions.h"
0015 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0016 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0017 #include "FWCore/PluginManager/interface/ModuleDef.h"
0018 #include "FWCore/Utilities/interface/EDGetToken.h"
0019 #include "FWCore/Utilities/interface/InputTag.h"
0020 #include "Geometry/CommonDetUnit/interface/GeomDet.h"
0021 #include "Geometry/Records/interface/TrackerDigiGeometryRecord.h"
0022 #include "Geometry/Records/interface/TrackerTopologyRcd.h"
0023 #include "Geometry/TrackerGeometryBuilder/interface/TrackerGeometry.h"
0024 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0025 #include "RecoTracker/PixelTrackFitting/interface/FitUtils.h"
0026 #include "TrackingTools/AnalyticalJacobians/interface/JacobianLocalToCurvilinear.h"
0027 #include "TrackingTools/MaterialEffects/interface/PropagatorWithMaterial.h"
0028 #include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
0029 #include "TrackingTools/TrajectoryParametrization/interface/CurvilinearTrajectoryError.h"
0030 #include "TrackingTools/TrajectoryParametrization/interface/GlobalTrajectoryParameters.h"
0031 #include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"
0032
0033 /*
0034   produces seeds directly from cuda produced tuples
0035 */
0036 template <typename TrackerTraits>
0037 class SeedProducerFromSoAT : public edm::global::EDProducer<> {
0038 public:
0039   explicit SeedProducerFromSoAT(const edm::ParameterSet& iConfig);
0040   ~SeedProducerFromSoAT() override = default;
0041
0042   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0043
0044 private:
0045   void produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
0046
0047   // Event data tokens
0048   const edm::EDGetTokenT<reco::BeamSpot> tBeamSpot_;
0049   const edm::EDGetTokenT<TrackSoAHeterogeneousHost<TrackerTraits>> tokenTrack_;
0050   // Event setup tokens
0051   const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> idealMagneticFieldToken_;
0052   const edm::ESGetToken<TrackerGeometry, TrackerDigiGeometryRecord> trackerDigiGeometryToken_;
0053   const edm::ESGetToken<Propagator, TrackingComponentsRecord> trackerPropagatorToken_;
0054   int32_t minNumberOfHits_;
0055 };
0056
0057 template <typename TrackerTraits>
0058 SeedProducerFromSoAT<TrackerTraits>::SeedProducerFromSoAT(const edm::ParameterSet& iConfig)
0059     : tBeamSpot_(consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("beamSpot"))),
0060       tokenTrack_(consumes(iConfig.getParameter<edm::InputTag>("src"))),
0061       idealMagneticFieldToken_(esConsumes()),
0062       trackerDigiGeometryToken_(esConsumes()),
0063       trackerPropagatorToken_(esConsumes(edm::ESInputTag("PropagatorWithMaterial"))),
0064       minNumberOfHits_(iConfig.getParameter<int>("minNumberOfHits"))
0065
0066 {
0067   produces<TrajectorySeedCollection>();
0068 }
0069
0070 template <typename TrackerTraits>
0071 void SeedProducerFromSoAT<TrackerTraits>::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0072   edm::ParameterSetDescription desc;
0073   desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot"));
0074   desc.add<edm::InputTag>("src", edm::InputTag("pixelTrackSoA"));
0075   desc.add<int>("minNumberOfHits", 0);
0076
0077   descriptions.addWithDefaultLabel(desc);
0078 }
0079
0080 template <typename TrackerTraits>
0081 void SeedProducerFromSoAT<TrackerTraits>::produce(edm::StreamID streamID,
0082                                                   edm::Event& iEvent,
0083                                                   const edm::EventSetup& iSetup) const {
0084   // std::cout << "Converting gpu helix to trajectory seed" << std::endl;
0085   auto result = std::make_unique<TrajectorySeedCollection>();
0086
0087   using trackHelper = TracksUtilities<TrackerTraits>;
0088
0089   auto const& fieldESH = iSetup.getHandle(idealMagneticFieldToken_);
0090   auto const& tracker = iSetup.getHandle(trackerDigiGeometryToken_);
0091   auto const& dus = tracker->detUnits();
0092
0093   auto const& propagatorHandle = iSetup.getHandle(trackerPropagatorToken_);
0094   const Propagator* propagator = &(*propagatorHandle);
0095
0096   const auto& bsh = iEvent.get(tBeamSpot_);
0097   // std::cout << "beamspot " << bsh.x0() << ' ' << bsh.y0() << ' ' << bsh.z0() << std::endl;
0098   GlobalPoint bs(bsh.x0(), bsh.y0(), bsh.z0());
0099
0100   auto const& tsoa = iEvent.get(tokenTrack_);
0101
0102   auto const* quality = tsoa.view().quality();
0103   auto const& detIndices = tsoa.view().detIndices();
0104   auto maxTracks = tsoa.view().metadata().size();
0105
0106   for (int32_t it = 0; it < maxTracks; ++it) {
0107     auto nHits = trackHelper::nHits(tsoa.view(), it);
0108     if (nHits == 0)
0109       break;  // this is a guard: maybe we need to move to nTracks...
0110
0111     auto q = quality[it];
0112     if (q != pixelTrack::Quality::loose)
0113       continue;  // FIXME
0114     if (nHits < minNumberOfHits_)
0115       continue;
0116
0117     // fill hits with invalid just to hold the detId
0118     auto b = detIndices.begin(it);
0119     edm::OwnVector<TrackingRecHit> hits;
0120     for (int iHit = 0; iHit < nHits; ++iHit) {
0121       auto const* det = dus[*(b + iHit)];
0122       // FIXME at some point get a proper type ...
0123       hits.push_back(new InvalidTrackingRecHit(*det, TrackingRecHit::bad));
0124     }
0125
0126     // mind: this values are respect the beamspot!
0127
0128     float phi = trackHelper::nHits(tsoa.view(), it);
0129
0130     riemannFit::Vector5d ipar, opar;
0131     riemannFit::Matrix5d icov, ocov;
0132     trackHelper::copyToDense(tsoa.view(), ipar, icov, it);
0133     riemannFit::transformToPerigeePlane(ipar, icov, opar, ocov);
0134
0135     LocalTrajectoryParameters lpar(opar(0), opar(1), opar(2), opar(3), opar(4), 1.);
0136     AlgebraicSymMatrix55 m;
0137     for (int i = 0; i < 5; ++i)
0138       for (int j = i; j < 5; ++j)
0139         m(i, j) = ocov(i, j);
0140
0141     float sp = std::sin(phi);
0142     float cp = std::cos(phi);
0143     Surface::RotationType rot(sp, -cp, 0, 0, 0, -1.f, cp, sp, 0);
0144
0145     Plane impPointPlane(bs, rot);
0146     GlobalTrajectoryParameters gp(impPointPlane.toGlobal(lpar.position()),
0147                                   impPointPlane.toGlobal(lpar.momentum()),
0148                                   lpar.charge(),
0149                                   fieldESH.product());
0150
0151     JacobianLocalToCurvilinear jl2c(impPointPlane, lpar, *fieldESH.product());
0152
0153     AlgebraicSymMatrix55 mo = ROOT::Math::Similarity(jl2c.jacobian(), m);
0154
0155     FreeTrajectoryState fts(gp, CurvilinearTrajectoryError(mo));
0156
0157     auto const& lastHit = hits.back();
0158
0159     TrajectoryStateOnSurface outerState = propagator->propagate(fts, *lastHit.surface());
0160
0161     if (!outerState.isValid()) {
0162       edm::LogError("SeedFromGPU") << " was trying to create a seed from:\n"
0163                                    << fts << "\n propagating to: " << lastHit.geographicalId().rawId();
0164       continue;
0165     }
0166
0167     auto const& pTraj = trajectoryStateTransform::persistentState(outerState, lastHit.geographicalId().rawId());
0168
0169     result->emplace_back(pTraj, hits, alongMomentum);
0170   }
0171
0172   iEvent.put(std::move(result));
0173 }
0174
0175 using SeedProducerFromSoAPhase1 = SeedProducerFromSoAT<pixelTopology::Phase1>;
0176 DEFINE_FWK_MODULE(SeedProducerFromSoAPhase1);
0177
0178 using SeedProducerFromSoAPhase2 = SeedProducerFromSoAT<pixelTopology::Phase2>;
0179 DEFINE_FWK_MODULE(SeedProducerFromSoAPhase2);