|
| 1 | +#include "DataFormats/BeamSpot/interface/BeamSpot.h" |
| 2 | +#include "DataFormats/Common/interface/OrphanHandle.h" |
| 3 | +#include "DataFormats/TrackReco/interface/Track.h" |
| 4 | +#include "DataFormats/TrackReco/interface/TrackExtra.h" |
| 5 | +#include "DataFormats/TrackReco/interface/TrackFwd.h" |
| 6 | +#include "DataFormats/VertexReco/interface/Vertex.h" |
| 7 | +#include "DataFormats/VertexReco/interface/VertexFwd.h" |
| 8 | +#include "DataFormats/VertexSoA/interface/ZVertexHost.h" |
| 9 | +#include "FWCore/Framework/interface/Event.h" |
| 10 | +#include "FWCore/Framework/interface/EventSetup.h" |
| 11 | +#include "FWCore/Framework/interface/MakerMacros.h" |
| 12 | +#include "FWCore/Framework/interface/global/EDProducer.h" |
| 13 | +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" |
| 14 | +#include "FWCore/ParameterSet/interface/ParameterSet.h" |
| 15 | +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" |
| 16 | +#include "FWCore/PluginManager/interface/ModuleDef.h" |
| 17 | +#include "FWCore/Utilities/interface/EDGetToken.h" |
| 18 | +#include "FWCore/Utilities/interface/InputTag.h" |
| 19 | +#include "Geometry/Records/interface/TrackerTopologyRcd.h" |
| 20 | +#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h" |
| 21 | + |
| 22 | +class PixelVertexProducerFromSoAAlpaka : public edm::global::EDProducer<> { |
| 23 | +public: |
| 24 | + using IndToEdm = std::vector<uint32_t>; |
| 25 | + |
| 26 | + explicit PixelVertexProducerFromSoAAlpaka(const edm::ParameterSet &iConfig); |
| 27 | + ~PixelVertexProducerFromSoAAlpaka() override = default; |
| 28 | + |
| 29 | + static void fillDescriptions(edm::ConfigurationDescriptions &descriptions); |
| 30 | + |
| 31 | +private: |
| 32 | + void produce(edm::StreamID streamID, edm::Event &iEvent, const edm::EventSetup &iSetup) const override; |
| 33 | + |
| 34 | + edm::EDGetTokenT<ZVertexHost> tokenVertex_; |
| 35 | + edm::EDGetTokenT<reco::BeamSpot> tokenBeamSpot_; |
| 36 | + edm::EDGetTokenT<reco::TrackCollection> tokenTracks_; |
| 37 | + edm::EDGetTokenT<IndToEdm> tokenIndToEdm_; |
| 38 | +}; |
| 39 | + |
| 40 | +PixelVertexProducerFromSoAAlpaka::PixelVertexProducerFromSoAAlpaka(const edm::ParameterSet &conf) |
| 41 | + : tokenVertex_(consumes(conf.getParameter<edm::InputTag>("src"))), |
| 42 | + tokenBeamSpot_(consumes(conf.getParameter<edm::InputTag>("beamSpot"))), |
| 43 | + tokenTracks_(consumes(conf.getParameter<edm::InputTag>("TrackCollection"))), |
| 44 | + tokenIndToEdm_(consumes(conf.getParameter<edm::InputTag>("TrackCollection"))) { |
| 45 | + produces<reco::VertexCollection>(); |
| 46 | +} |
| 47 | + |
| 48 | +void PixelVertexProducerFromSoAAlpaka::fillDescriptions(edm::ConfigurationDescriptions &descriptions) { |
| 49 | + edm::ParameterSetDescription desc; |
| 50 | + |
| 51 | + desc.add<edm::InputTag>("TrackCollection", edm::InputTag("pixelTracks")); |
| 52 | + desc.add<edm::InputTag>("beamSpot", edm::InputTag("offlineBeamSpot")); |
| 53 | + desc.add<edm::InputTag>("src", edm::InputTag("pixelVerticesAlpaka")); |
| 54 | + |
| 55 | + descriptions.add("pixelVertexFromSoAAlpaka", desc); |
| 56 | +} |
| 57 | + |
| 58 | +void PixelVertexProducerFromSoAAlpaka::produce(edm::StreamID streamID, |
| 59 | + edm::Event &iEvent, |
| 60 | + const edm::EventSetup &) const { |
| 61 | + auto vertexes = std::make_unique<reco::VertexCollection>(); |
| 62 | + |
| 63 | + auto tracksHandle = iEvent.getHandle(tokenTracks_); |
| 64 | + auto tracksSize = tracksHandle->size(); |
| 65 | + auto const &indToEdm = iEvent.get(tokenIndToEdm_); |
| 66 | + auto bsHandle = iEvent.getHandle(tokenBeamSpot_); |
| 67 | + |
| 68 | + float x0 = 0, y0 = 0, z0 = 0, dxdz = 0, dydz = 0; |
| 69 | + std::vector<int32_t> itrk; |
| 70 | + itrk.reserve(64); // avoid first relocations |
| 71 | + if (!bsHandle.isValid()) { |
| 72 | + edm::LogWarning("PixelVertexProducer") << "No beamspot found. returning vertexes with (0,0,Z) "; |
| 73 | + } else { |
| 74 | + const reco::BeamSpot &bs = *bsHandle; |
| 75 | + x0 = bs.x0(); |
| 76 | + y0 = bs.y0(); |
| 77 | + z0 = bs.z0(); |
| 78 | + dxdz = bs.dxdz(); |
| 79 | + dydz = bs.dydz(); |
| 80 | + } |
| 81 | + |
| 82 | + auto const &soa = iEvent.get(tokenVertex_); |
| 83 | + |
| 84 | + int nv = soa.view().nvFinal(); |
| 85 | + |
| 86 | +#ifdef PIXVERTEX_DEBUG_PRODUCE |
| 87 | + std::cout << "converting " << nv << " vertices " |
| 88 | + << " from " << indToEdm.size() << " tracks" << std::endl; |
| 89 | +#endif // PIXVERTEX_DEBUG_PRODUCE |
| 90 | + |
| 91 | + std::set<uint32_t> uind; // for verifing index consistency |
| 92 | + for (int j = nv - 1; j >= 0; --j) { |
| 93 | + auto i = soa.view()[j].sortInd(); // on gpu sorted in ascending order.... |
| 94 | + assert(i < nv); |
| 95 | + uind.insert(i); |
| 96 | + assert(itrk.empty()); |
| 97 | + auto z = soa.view()[i].zv(); |
| 98 | + auto x = x0 + dxdz * z; |
| 99 | + auto y = y0 + dydz * z; |
| 100 | + z += z0; |
| 101 | + reco::Vertex::Error err; |
| 102 | + err(2, 2) = 1.f / soa.view()[i].wv(); |
| 103 | + err(2, 2) *= 2.; // artifically inflate error |
| 104 | + //Copy also the tracks (no intention to be efficient....) |
| 105 | + for (auto k = 0U; k < indToEdm.size(); ++k) { |
| 106 | + if (soa.view()[k].idv() == int16_t(i)) |
| 107 | + itrk.push_back(k); |
| 108 | + } |
| 109 | + auto nt = itrk.size(); |
| 110 | + if (nt == 0) { |
| 111 | +#ifdef PIXVERTEX_DEBUG_PRODUCE |
| 112 | + std::cout << "vertex " << i << " with no tracks..." << std::endl; |
| 113 | +#endif // PIXVERTEX_DEBUG_PRODUCE |
| 114 | + continue; |
| 115 | + } |
| 116 | + if (nt < 2) { |
| 117 | + itrk.clear(); |
| 118 | + continue; |
| 119 | + } // remove outliers |
| 120 | + (*vertexes).emplace_back(reco::Vertex::Point(x, y, z), err, soa.view()[i].chi2(), soa.view()[i].ndof(), nt); |
| 121 | + auto &v = (*vertexes).back(); |
| 122 | + v.reserve(itrk.size()); |
| 123 | + for (auto it : itrk) { |
| 124 | + assert(it < int(indToEdm.size())); |
| 125 | + auto k = indToEdm[it]; |
| 126 | + if (k > tracksSize) { |
| 127 | + edm::LogWarning("PixelVertexProducer") << "oops track " << it << " does not exists on CPU " << k; |
| 128 | + continue; |
| 129 | + } |
| 130 | + auto tk = reco::TrackRef(tracksHandle, k); |
| 131 | + v.add(tk); |
| 132 | + } |
| 133 | + itrk.clear(); |
| 134 | + } |
| 135 | + |
| 136 | + LogDebug("PixelVertexProducer") << ": Found " << vertexes->size() << " vertexes\n"; |
| 137 | + for (unsigned int i = 0; i < vertexes->size(); ++i) { |
| 138 | + LogDebug("PixelVertexProducer") << "Vertex number " << i << " has " << (*vertexes)[i].tracksSize() |
| 139 | + << " tracks with a position of " << (*vertexes)[i].z() << " +- " |
| 140 | + << std::sqrt((*vertexes)[i].covariance(2, 2)); |
| 141 | + } |
| 142 | + |
| 143 | + // legacy logic.... |
| 144 | + if (vertexes->empty() && bsHandle.isValid()) { |
| 145 | + const reco::BeamSpot &bs = *bsHandle; |
| 146 | + |
| 147 | + GlobalError bse(bs.rotatedCovariance3D()); |
| 148 | + if ((bse.cxx() <= 0.) || (bse.cyy() <= 0.) || (bse.czz() <= 0.)) { |
| 149 | + AlgebraicSymMatrix33 we; |
| 150 | + we(0, 0) = 10000; |
| 151 | + we(1, 1) = 10000; |
| 152 | + we(2, 2) = 10000; |
| 153 | + vertexes->push_back(reco::Vertex(bs.position(), we, 0., 0., 0)); |
| 154 | + |
| 155 | + edm::LogInfo("PixelVertexProducer") << "No vertices found. Beamspot with invalid errors " << bse.matrix() |
| 156 | + << "\nWill put Vertex derived from dummy-fake BeamSpot into Event.\n" |
| 157 | + << (*vertexes)[0].x() << "\n" |
| 158 | + << (*vertexes)[0].y() << "\n" |
| 159 | + << (*vertexes)[0].z() << "\n"; |
| 160 | + } else { |
| 161 | + vertexes->push_back(reco::Vertex(bs.position(), bs.rotatedCovariance3D(), 0., 0., 0)); |
| 162 | + |
| 163 | + edm::LogInfo("PixelVertexProducer") << "No vertices found. Will put Vertex derived from BeamSpot into Event:\n" |
| 164 | + << (*vertexes)[0].x() << "\n" |
| 165 | + << (*vertexes)[0].y() << "\n" |
| 166 | + << (*vertexes)[0].z() << "\n"; |
| 167 | + } |
| 168 | + } else if (vertexes->empty() && !bsHandle.isValid()) { |
| 169 | + edm::LogWarning("PixelVertexProducer") << "No beamspot and no vertex found. No vertex returned."; |
| 170 | + } |
| 171 | + |
| 172 | + iEvent.put(std::move(vertexes)); |
| 173 | +} |
| 174 | + |
| 175 | +DEFINE_FWK_MODULE(PixelVertexProducerFromSoAAlpaka); |
0 commit comments