PatAlgos/plugins/DeDxEstimatorRekeyer.cc

0001 #include "FWCore/Framework/interface/global/EDProducer.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0004 #include "DataFormats/TrackReco/interface/DeDxData.h"
0005 #include "DataFormats/TrackReco/interface/DeDxHitInfo.h"
0006 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0007 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0008
0009 //
0010 // class declaration
0011 //
0012
0013 class DeDxEstimatorRekeyer : public edm::global::EDProducer<> {
0014 public:
0015   explicit DeDxEstimatorRekeyer(const edm::ParameterSet&);
0016   ~DeDxEstimatorRekeyer() override {}
0017   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0018
0019 private:
0020   void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0021
0022   template <typename T>
0023   std::map<std::string, edm::EDGetTokenT<T>> getTokens(const std::vector<edm::InputTag>& tags) {
0024     std::map<std::string, edm::EDGetTokenT<T>> tokens;
0025     for (const auto& tag : tags)
0026       tokens.emplace(tag.label(), consumes<T>(tag));
0027     return tokens;
0028   };
0029
0030   // ----------member data ---------------------------
0031   const edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
0032   const bool rekey_dedxHits_;
0033   const edm::EDGetTokenT<reco::DeDxHitInfoAss> dedxHitAssToken_;
0034   const edm::EDGetTokenT<edm::ValueMap<std::vector<float>>> dedxHitMomToken_;
0035   const std::map<std::string, edm::EDGetTokenT<edm::ValueMap<reco::DeDxData>>> dedxEstimatorsTokens_;
0036   const std::map<std::string, edm::EDGetTokenT<pat::PackedCandidateCollection>> packedCandidatesTokens_;
0037   const std::map<std::string, edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>>> trk2pcTokens_;
0038 };
0039
0040 void DeDxEstimatorRekeyer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0041   edm::ParameterSetDescription desc;
0042   desc.add<edm::InputTag>("tracks", {"generalTracks"});
0043   desc.add<edm::InputTag>("dedxHits", {});
0044   desc.add<edm::InputTag>("dedxMomentum", {});
0045   desc.add<std::vector<edm::InputTag>>(
0046       "packedCandidates",
0047       {edm::InputTag("packedPFCandidates"), edm::InputTag("lostTracks"), edm::InputTag("lostTracks:eleTracks")});
0048   desc.add<std::vector<edm::InputTag>>("dedxEstimators",
0049                                        {edm::InputTag("dedxHarmonic2"), edm::InputTag("dedxPixelHarmonic2")});
0050   descriptions.addWithDefaultLabel(desc);
0051 }
0052
0053 DeDxEstimatorRekeyer::DeDxEstimatorRekeyer(const edm::ParameterSet& iConfig)
0054     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracks"))),
0055       rekey_dedxHits_(not iConfig.getParameter<edm::InputTag>("dedxHits").label().empty()),
0056       dedxHitAssToken_(consumes<reco::DeDxHitInfoAss>(iConfig.getParameter<edm::InputTag>("dedxHits"))),
0057       dedxHitMomToken_(
0058           consumes<edm::ValueMap<std::vector<float>>>(iConfig.getParameter<edm::InputTag>("dedxMomentum"))),
0059       dedxEstimatorsTokens_(
0060           getTokens<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<std::vector<edm::InputTag>>("dedxEstimators"))),
0061       packedCandidatesTokens_(getTokens<pat::PackedCandidateCollection>(
0062           iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))),
0063       trk2pcTokens_(getTokens<edm::Association<pat::PackedCandidateCollection>>(
0064           iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))) {
0065   for (const auto& d : dedxEstimatorsTokens_)
0066     produces<edm::ValueMap<reco::DeDxData>>(d.first);
0067   if (rekey_dedxHits_) {
0068     produces<reco::DeDxHitInfoCollection>();
0069     produces<reco::DeDxHitInfoAss>();
0070     produces<edm::ValueMap<std::vector<float>>>("momentumAtHit");
0071   }
0072 }
0073
0074 void DeDxEstimatorRekeyer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0075   // Get input collections
0076   const auto& tracks = iEvent.getHandle(tracksToken_);
0077   std::vector<reco::TrackRef> trackRefs;
0078   trackRefs.reserve(tracks->size());
0079   for (auto track = tracks->begin(); track != tracks->end(); track++)
0080     trackRefs.emplace_back(tracks, track - tracks->begin());
0081
0082   typedef std::map<pat::PackedCandidateRef, reco::TrackRef> PCTrkMap;
0083   std::vector<std::pair<edm::Handle<pat::PackedCandidateCollection>, PCTrkMap>> pcTrkMap;
0084   pcTrkMap.reserve(packedCandidatesTokens_.size());
0085   for (const auto& p : packedCandidatesTokens_) {
0086     PCTrkMap map;
0087     const auto& trk2pc = iEvent.get(trk2pcTokens_.at(p.first));
0088     for (const auto& track : trackRefs) {
0089       const auto& pc = trk2pc[track];
0090       if (pc.isNonnull())
0091         map.emplace(pc, track);
0092     }
0093     pcTrkMap.emplace_back(iEvent.getHandle(p.second), map);
0094   }
0095
0096   // Rekey dEdx estimators
0097   for (const auto& d : dedxEstimatorsTokens_) {
0098     const auto& dedxEstimators = iEvent.get(d.second);
0099     auto trackDeDxValueMap = std::make_unique<edm::ValueMap<reco::DeDxData>>();
0100     edm::ValueMap<reco::DeDxData>::Filler filler(*trackDeDxValueMap);
0101     // Loop over packed candidates
0102     for (const auto& h : pcTrkMap) {
0103       std::vector<reco::DeDxData> dedxEstimate(h.first->size());
0104       for (const auto& p : h.second)
0105         dedxEstimate[p.first.key()] = dedxEstimators[p.second];
0106       filler.insert(h.first, dedxEstimate.begin(), dedxEstimate.end());
0107     }
0108     // Fill the value map and put it into the event
0109     filler.fill();
0110     iEvent.put(std::move(trackDeDxValueMap), d.first);
0111   }
0112
0113   // Rekey dEdx hit info
0114   if (not rekey_dedxHits_)
0115     return;
0116   const auto& dedxHitMom = iEvent.get(dedxHitMomToken_);
0117   const auto& dedxHitAss = iEvent.get(dedxHitAssToken_);
0118   const auto& dedxHitInfoHandle = iEvent.getRefBeforePut<reco::DeDxHitInfoCollection>();
0119   auto dedxHitInfoAssociation = std::make_unique<reco::DeDxHitInfoAss>(dedxHitInfoHandle);
0120   reco::DeDxHitInfoAss::Filler filler(*dedxHitInfoAssociation);
0121   auto resultdedxHitColl = std::make_unique<reco::DeDxHitInfoCollection>();
0122   resultdedxHitColl->reserve(!pcTrkMap.empty() ? pcTrkMap.size() * pcTrkMap[0].second.size() : 0);
0123   std::vector<std::vector<float>> momenta;
0124   momenta.reserve(resultdedxHitColl->capacity());
0125   // Loop over packed candidates
0126   for (const auto& h : pcTrkMap) {
0127     std::vector<int> indices(h.first->size(), -1);
0128     for (const auto& p : h.second) {
0129       const auto& dedxHit = dedxHitAss[p.second];
0130       if (dedxHit.isNull())
0131         continue;
0132       indices[p.first.key()] = resultdedxHitColl->size();
0133       resultdedxHitColl->emplace_back(*dedxHit);
0134       momenta.emplace_back(dedxHitMom[dedxHit]);
0135     }
0136     filler.insert(h.first, indices.begin(), indices.end());
0137   }
0138   const auto& dedxHitCollHandle = iEvent.put(std::move(resultdedxHitColl));
0139   // Fill the association map and put it into the event
0140   filler.fill();
0141   iEvent.put(std::move(dedxHitInfoAssociation));
0142   // Fill the value map and put it into the event
0143   auto dedxMomenta = std::make_unique<edm::ValueMap<std::vector<float>>>();
0144   edm::ValueMap<std::vector<float>>::Filler mfiller(*dedxMomenta);
0145   mfiller.insert(dedxHitCollHandle, momenta.begin(), momenta.end());
0146   mfiller.fill();
0147   iEvent.put(std::move(dedxMomenta), "momentumAtHit");
0148 }
0149
0150 //define this as a plug-in
0151 #include "FWCore/Framework/interface/MakerMacros.h"
0152 DEFINE_FWK_MODULE(DeDxEstimatorRekeyer);