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 edm::EDGetTokenT<reco::DeDxHitInfoAss> dedxHitAssToken_;
0033   const edm::EDGetTokenT<edm::ValueMap<std::vector<float>>> dedxHitMomToken_;
0034   const std::map<std::string, edm::EDGetTokenT<edm::ValueMap<reco::DeDxData>>> dedxEstimatorsTokens_;
0035   const std::map<std::string, edm::EDGetTokenT<pat::PackedCandidateCollection>> packedCandidatesTokens_;
0036   const std::map<std::string, edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>>> trk2pcTokens_;
0037 };
0038
0039 void DeDxEstimatorRekeyer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0040   edm::ParameterSetDescription desc;
0041   desc.add<edm::InputTag>("tracks", {"generalTracks"});
0042   desc.add<edm::InputTag>("dedxHits", {"dedxHitInfo"});
0043   desc.add<edm::InputTag>("dedxMomentum", {"dedxHitInfo:momentumAtHit"});
0044   desc.add<std::vector<edm::InputTag>>(
0045       "packedCandidates",
0046       {edm::InputTag("packedPFCandidates"), edm::InputTag("lostTracks"), edm::InputTag("lostTracks:eleTracks")});
0047   desc.add<std::vector<edm::InputTag>>("dedxEstimators",
0048                                        {edm::InputTag("dedxHarmonic2"), edm::InputTag("dedxPixelHarmonic2")});
0049   descriptions.addWithDefaultLabel(desc);
0050 }
0051
0052 DeDxEstimatorRekeyer::DeDxEstimatorRekeyer(const edm::ParameterSet& iConfig)
0053     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracks"))),
0054       dedxHitAssToken_(consumes<reco::DeDxHitInfoAss>(iConfig.getParameter<edm::InputTag>("dedxHits"))),
0055       dedxHitMomToken_(
0056           consumes<edm::ValueMap<std::vector<float>>>(iConfig.getParameter<edm::InputTag>("dedxMomentum"))),
0057       dedxEstimatorsTokens_(
0058           getTokens<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<std::vector<edm::InputTag>>("dedxEstimators"))),
0059       packedCandidatesTokens_(getTokens<pat::PackedCandidateCollection>(
0060           iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))),
0061       trk2pcTokens_(getTokens<edm::Association<pat::PackedCandidateCollection>>(
0062           iConfig.getParameter<std::vector<edm::InputTag>>("packedCandidates"))) {
0063   for (const auto& d : dedxEstimatorsTokens_)
0064     produces<edm::ValueMap<reco::DeDxData>>(d.first);
0065   produces<reco::DeDxHitInfoCollection>();
0066   produces<reco::DeDxHitInfoAss>();
0067   produces<edm::ValueMap<std::vector<float>>>("momentumAtHit");
0068 }
0069
0070 void DeDxEstimatorRekeyer::produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
0071   // Get input collections
0072   const auto& tracks = iEvent.getHandle(tracksToken_);
0073   std::vector<reco::TrackRef> trackRefs;
0074   trackRefs.reserve(tracks->size());
0075   for (auto track = tracks->begin(); track != tracks->end(); track++)
0076     trackRefs.emplace_back(tracks, track - tracks->begin());
0077
0078   typedef std::map<pat::PackedCandidateRef, reco::TrackRef> PCTrkMap;
0079   std::vector<std::pair<edm::Handle<pat::PackedCandidateCollection>, PCTrkMap>> pcTrkMap;
0080   pcTrkMap.reserve(packedCandidatesTokens_.size());
0081   for (const auto& p : packedCandidatesTokens_) {
0082     PCTrkMap map;
0083     const auto& trk2pc = iEvent.get(trk2pcTokens_.at(p.first));
0084     for (const auto& track : trackRefs) {
0085       const auto& pc = trk2pc[track];
0086       if (pc.isNonnull())
0087         map.emplace(pc, track);
0088     }
0089     pcTrkMap.emplace_back(iEvent.getHandle(p.second), map);
0090   }
0091
0092   // Rekey dEdx estimators
0093   for (const auto& d : dedxEstimatorsTokens_) {
0094     const auto& dedxEstimators = iEvent.get(d.second);
0095     auto trackDeDxValueMap = std::make_unique<edm::ValueMap<reco::DeDxData>>();
0096     edm::ValueMap<reco::DeDxData>::Filler filler(*trackDeDxValueMap);
0097     // Loop over packed candidates
0098     for (const auto& h : pcTrkMap) {
0099       std::vector<reco::DeDxData> dedxEstimate(h.first->size());
0100       for (const auto& p : h.second)
0101         dedxEstimate[p.first.key()] = dedxEstimators[p.second];
0102       filler.insert(h.first, dedxEstimate.begin(), dedxEstimate.end());
0103     }
0104     // Fill the value map and put it into the event
0105     filler.fill();
0106     iEvent.put(std::move(trackDeDxValueMap), d.first);
0107   }
0108
0109   // Rekey dEdx hit info
0110   const auto& dedxHitMom = iEvent.get(dedxHitMomToken_);
0111   const auto& dedxHitAss = iEvent.get(dedxHitAssToken_);
0112   const auto& dedxHitInfoHandle = iEvent.getRefBeforePut<reco::DeDxHitInfoCollection>();
0113   auto dedxHitInfoAssociation = std::make_unique<reco::DeDxHitInfoAss>(dedxHitInfoHandle);
0114   reco::DeDxHitInfoAss::Filler filler(*dedxHitInfoAssociation);
0115   auto resultdedxHitColl = std::make_unique<reco::DeDxHitInfoCollection>();
0116   resultdedxHitColl->reserve(!pcTrkMap.empty() ? pcTrkMap.size() * pcTrkMap[0].second.size() : 0);
0117   std::vector<std::vector<float>> momenta;
0118   momenta.reserve(resultdedxHitColl->capacity());
0119   // Loop over packed candidates
0120   for (const auto& h : pcTrkMap) {
0121     std::vector<int> indices(h.first->size(), -1);
0122     for (const auto& p : h.second) {
0123       const auto& dedxHit = dedxHitAss[p.second];
0124       if (dedxHit.isNull())
0125         continue;
0126       indices[p.first.key()] = resultdedxHitColl->size();
0127       resultdedxHitColl->emplace_back(*dedxHit);
0128       momenta.emplace_back(dedxHitMom[dedxHit]);
0129     }
0130     filler.insert(h.first, indices.begin(), indices.end());
0131   }
0132   const auto& dedxHitCollHandle = iEvent.put(std::move(resultdedxHitColl));
0133   // Fill the association map and put it into the event
0134   filler.fill();
0135   iEvent.put(std::move(dedxHitInfoAssociation));
0136   // Fill the value map and put it into the event
0137   auto dedxMomenta = std::make_unique<edm::ValueMap<std::vector<float>>>();
0138   edm::ValueMap<std::vector<float>>::Filler mfiller(*dedxMomenta);
0139   mfiller.insert(dedxHitCollHandle, momenta.begin(), momenta.end());
0140   mfiller.fill();
0141   iEvent.put(std::move(dedxMomenta), "momentumAtHit");
0142 }
0143
0144 //define this as a plug-in
0145 #include "FWCore/Framework/interface/MakerMacros.h"
0146 DEFINE_FWK_MODULE(DeDxEstimatorRekeyer);