PixelVertexFinding/plugins/PixelVertexProducerFromSoAAlpaka.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/VertexReco/interface/Vertex.h"
0007 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0008 #include "DataFormats/VertexSoA/interface/ZVertexHost.h"
0009 #include "FWCore/Framework/interface/Event.h"
0010 #include "FWCore/Framework/interface/EventSetup.h"
0011 #include "FWCore/Framework/interface/MakerMacros.h"
0012 #include "FWCore/Framework/interface/global/EDProducer.h"
0013 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0015 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0016 #include "FWCore/PluginManager/interface/ModuleDef.h"
0017 #include "FWCore/Utilities/interface/EDGetToken.h"
0018 #include "FWCore/Utilities/interface/InputTag.h"
0019 #include "Geometry/Records/interface/TrackerTopologyRcd.h"
0020 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0021
0022 class PixelVertexProducerFromSoAAlpaka : public edm::global::EDProducer<> {
0023 public:
0024   using IndToEdm = std::vector<uint32_t>;
0025
0026   explicit PixelVertexProducerFromSoAAlpaka(const edm::ParameterSet &iConfig);
0027   ~PixelVertexProducerFromSoAAlpaka() override = default;
0028
0029   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0030
0031 private:
0032   void produce(edm::StreamID streamID, edm::Event &iEvent, const edm::EventSetup &iSetup) const override;
0033
0034   edm::EDGetTokenT<ZVertexHost> tokenVertex_;
0035   edm::EDGetTokenT<reco::BeamSpot> tokenBeamSpot_;
0036   edm::EDGetTokenT<reco::TrackCollection> tokenTracks_;
0037   edm::EDGetTokenT<IndToEdm> tokenIndToEdm_;
0038 };
0039
0040 PixelVertexProducerFromSoAAlpaka::PixelVertexProducerFromSoAAlpaka(const edm::ParameterSet &conf)
0041     : tokenVertex_(consumes(conf.getParameter<edm::InputTag>("src"))),
0042       tokenBeamSpot_(consumes(conf.getParameter<edm::InputTag>("beamSpot"))),
0043       tokenTracks_(consumes(conf.getParameter<edm::InputTag>("TrackCollection"))),
0044       tokenIndToEdm_(consumes(conf.getParameter<edm::InputTag>("TrackCollection"))) {
0045   produces<reco::VertexCollection>();
0046 }
0047
0048 void PixelVertexProducerFromSoAAlpaka::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0049   edm::ParameterSetDescription desc;
0050
0051   desc.add<edm::InputTag>("TrackCollection", edm::InputTag("pixelTracks"));
0052   desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot"));
0053   desc.add<edm::InputTag>("src", edm::InputTag("pixelVerticesAlpaka"));
0054
0055   descriptions.add("pixelVertexFromSoAAlpaka", desc);
0056 }
0057
0058 void PixelVertexProducerFromSoAAlpaka::produce(edm::StreamID streamID,
0059                                                edm::Event &iEvent,
0060                                                const edm::EventSetup &) const {
0061   auto vertexes = std::make_unique<reco::VertexCollection>();
0062
0063   auto tracksHandle = iEvent.getHandle(tokenTracks_);
0064   auto tracksSize = tracksHandle->size();
0065   auto const &indToEdm = iEvent.get(tokenIndToEdm_);
0066   auto bsHandle = iEvent.getHandle(tokenBeamSpot_);
0067
0068   float x0 = 0, y0 = 0, z0 = 0, dxdz = 0, dydz = 0;
0069   std::vector<int32_t> itrk;
0070   itrk.reserve(64);  // avoid first relocations
0071   if (!bsHandle.isValid()) {
0072     edm::LogWarning("PixelVertexProducer") << "No beamspot found. returning vertexes with (0,0,Z) ";
0073   } else {
0074     const reco::BeamSpot &bs = *bsHandle;
0075     x0 = bs.x0();
0076     y0 = bs.y0();
0077     z0 = bs.z0();
0078     dxdz = bs.dxdz();
0079     dydz = bs.dydz();
0080   }
0081
0082   auto const &soa = iEvent.get(tokenVertex_);
0083
0084   int nv = soa.view().nvFinal();
0085
0086 #ifdef PIXVERTEX_DEBUG_PRODUCE
0087   std::cout << "converting " << nv << " vertices "
0088             << " from " << indToEdm.size() << " tracks" << std::endl;
0089 #endif  // PIXVERTEX_DEBUG_PRODUCE
0090
0091   std::set<uint32_t> uind;  // for verifing index consistency
0092   for (int j = nv - 1; j >= 0; --j) {
0093     auto i = soa.view()[j].sortInd();  // on gpu sorted in ascending order....
0094     assert(i < nv);
0095     uind.insert(i);
0096     assert(itrk.empty());
0097     auto z = soa.view()[i].zv();
0098     auto x = x0 + dxdz * z;
0099     auto y = y0 + dydz * z;
0100     z += z0;
0101     reco::Vertex::Error err;
0102     err(2, 2) = 1.f / soa.view()[i].wv();
0103     err(2, 2) *= 2.;  // artifically inflate error
0104     //Copy also the tracks (no intention to be efficient....)
0105     for (auto k = 0U; k < indToEdm.size(); ++k) {
0106       if (soa.view()[k].idv() == int16_t(i))
0107         itrk.push_back(k);
0108     }
0109     auto nt = itrk.size();
0110     if (nt == 0) {
0111 #ifdef PIXVERTEX_DEBUG_PRODUCE
0112       std::cout << "vertex " << i << " with no tracks..." << std::endl;
0113 #endif  // PIXVERTEX_DEBUG_PRODUCE
0114       continue;
0115     }
0116     if (nt < 2) {
0117       itrk.clear();
0118       continue;
0119     }  // remove outliers
0120     (*vertexes).emplace_back(reco::Vertex::Point(x, y, z), err, soa.view()[i].chi2(), soa.view()[i].ndof(), nt);
0121     auto &v = (*vertexes).back();
0122     v.reserve(itrk.size());
0123     for (auto it : itrk) {
0124       assert(it < int(indToEdm.size()));
0125       auto k = indToEdm[it];
0126       if (k > tracksSize) {
0127         edm::LogWarning("PixelVertexProducer") << "oops track " << it << " does not exists on CPU " << k;
0128         continue;
0129       }
0130       auto tk = reco::TrackRef(tracksHandle, k);
0131       v.add(tk);
0132     }
0133     itrk.clear();
0134   }
0135
0136   LogDebug("PixelVertexProducer") << ": Found " << vertexes->size() << " vertexes\n";
0137   for (unsigned int i = 0; i < vertexes->size(); ++i) {
0138     LogDebug("PixelVertexProducer") << "Vertex number " << i << " has " << (*vertexes)[i].tracksSize()
0139                                     << " tracks with a position of " << (*vertexes)[i].z() << " +- "
0140                                     << std::sqrt((*vertexes)[i].covariance(2, 2));
0141   }
0142
0143   // legacy logic....
0144   if (vertexes->empty() && bsHandle.isValid()) {
0145     const reco::BeamSpot &bs = *bsHandle;
0146
0147     GlobalError bse(bs.rotatedCovariance3D());
0148     if ((bse.cxx() <= 0.) || (bse.cyy() <= 0.) || (bse.czz() <= 0.)) {
0149       AlgebraicSymMatrix33 we;
0150       we(0, 0) = 10000;
0151       we(1, 1) = 10000;
0152       we(2, 2) = 10000;
0153       vertexes->push_back(reco::Vertex(bs.position(), we, 0., 0., 0));
0154
0155       edm::LogInfo("PixelVertexProducer") << "No vertices found. Beamspot with invalid errors " << bse.matrix()
0156                                           << "\nWill put Vertex derived from dummy-fake BeamSpot into Event.\n"
0157                                           << (*vertexes)[0].x() << "\n"
0158                                           << (*vertexes)[0].y() << "\n"
0159                                           << (*vertexes)[0].z() << "\n";
0160     } else {
0161       vertexes->push_back(reco::Vertex(bs.position(), bs.rotatedCovariance3D(), 0., 0., 0));
0162
0163       edm::LogInfo("PixelVertexProducer") << "No vertices found. Will put Vertex derived from BeamSpot into Event:\n"
0164                                           << (*vertexes)[0].x() << "\n"
0165                                           << (*vertexes)[0].y() << "\n"
0166                                           << (*vertexes)[0].z() << "\n";
0167     }
0168   } else if (vertexes->empty() && !bsHandle.isValid()) {
0169     edm::LogWarning("PixelVertexProducer") << "No beamspot and no vertex found. No vertex returned.";
0170   }
0171
0172   iEvent.put(std::move(vertexes));
0173 }
0174
0175 DEFINE_FWK_MODULE(PixelVertexProducerFromSoAAlpaka);