TrackAssociation/plugins/PackedCandidateGenAssociationProducer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    MyTemporarySubSystem/PackedCandidateGenAssociationProducer
0004 // Class:      PackedCandidateGenAssociationProducer
0005 //
0006 /**\class PackedCandidateGenAssociationProducer PackedCandidateGenAssociationProducer.cc MyTemporarySubSystem/PackedCandidateGenAssociationProducer/plugins/PackedCandidateGenAssociationProducer.cc
0007
0008  Description: [one line class summary]
0009
0010  Implementation:
0011      [Notes on implementation]
0012 */
0013 //
0014 // Original Author:  Enrico Lusiani
0015 //         Created:  Mon, 03 May 2021 13:40:39 GMT
0016 //
0017 //
0018
0019 // system include files
0020 #include <memory>
0021
0022 // user include files
0023 #include "FWCore/Framework/interface/Frameworkfwd.h"
0024 #include "FWCore/Framework/interface/global/EDProducer.h"
0025
0026 #include "FWCore/Framework/interface/Event.h"
0027 #include "FWCore/Framework/interface/MakerMacros.h"
0028
0029 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0030 #include "FWCore/Utilities/interface/StreamID.h"
0031
0032 #include "DataFormats/Common/interface/Association.h"
0033 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0034 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0035 #include "DataFormats/TrackReco/interface/Track.h"
0036
0037 //
0038 // class declaration
0039 //
0040
0041 class PackedCandidateGenAssociationProducer : public edm::global::EDProducer<> {
0042 public:
0043   explicit PackedCandidateGenAssociationProducer(const edm::ParameterSet&);
0044
0045   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0046
0047 private:
0048   void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;
0049
0050   edm::EDGetTokenT<edm::Association<reco::GenParticleCollection>> trackToGenToken_;
0051   edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>> trackToPcToken_;
0052   edm::EDGetTokenT<edm::Association<reco::GenParticleCollection>> genToPrunedToken_;
0053   edm::EDGetTokenT<edm::Association<reco::GenParticleCollection>> genToPrunedWSOToken_;
0054   edm::EDGetTokenT<edm::View<reco::Track>> tracksToken_;
0055 };
0056 PackedCandidateGenAssociationProducer::PackedCandidateGenAssociationProducer(const edm::ParameterSet& iConfig)
0057     : trackToGenToken_(consumes<edm::Association<reco::GenParticleCollection>>(
0058           iConfig.getParameter<edm::InputTag>("trackToGenAssoc"))),
0059       trackToPcToken_(consumes<edm::Association<pat::PackedCandidateCollection>>(
0060           iConfig.getParameter<edm::InputTag>("trackToPackedCandidatesAssoc"))),
0061       genToPrunedToken_(consumes<edm::Association<reco::GenParticleCollection>>(
0062           iConfig.getParameter<edm::InputTag>("genToPrunedAssoc"))),
0063       genToPrunedWSOToken_(consumes<edm::Association<reco::GenParticleCollection>>(
0064           iConfig.getParameter<edm::InputTag>("genToPrunedAssocWithStatusOne"))),
0065       tracksToken_(consumes<edm::View<reco::Track>>(iConfig.getParameter<edm::InputTag>("tracks"))) {
0066   produces<edm::Association<reco::GenParticleCollection>>();
0067 }
0068
0069 void PackedCandidateGenAssociationProducer::produce(edm::StreamID,
0070                                                     edm::Event& iEvent,
0071                                                     const edm::EventSetup& iSetup) const {
0072   using namespace edm;
0073
0074   const auto& trackToPackedCandidatesAssoc = iEvent.get(trackToPcToken_);
0075   auto pcCollection = trackToPackedCandidatesAssoc.ref();
0076
0077   const auto& genToPrunedAssoc = iEvent.get(genToPrunedToken_);
0078   auto prunedCollection = genToPrunedAssoc.ref();
0079
0080   const auto& genToPrunedAssocWSO = iEvent.get(genToPrunedWSOToken_);
0081
0082   auto trackHandle = iEvent.getHandle(tracksToken_);
0083   const auto& tracks = *trackHandle;
0084
0085   auto out = std::make_unique<edm::Association<reco::GenParticleCollection>>(prunedCollection);
0086
0087   auto trackToGenAssocHandle = iEvent.getHandle(trackToGenToken_);
0088   if (not trackToGenAssocHandle.isValid() or not trackToGenAssocHandle->contains(trackHandle.id())) {
0089     // not track to gen association available, possibly an old AODSIM, or a missing RECOSIM step
0090     // alternatively, the track association may not contain our tracks, as in the case of RECOSIM run on an old RAWSIM
0091     // early exit with an empty collection to avoid crash
0092     iEvent.put(std::move(out));
0093     return;
0094   }
0095
0096   const auto& trackToGenAssoc = *trackToGenAssocHandle;
0097
0098   Association<reco::GenParticleCollection>::Filler filler(*out);
0099
0100   std::vector<int> indices(pcCollection->size(), -1);
0101
0102   for (size_t i = 0; i < tracks.size(); i++) {
0103     auto track = tracks.refAt(i);
0104
0105     auto pc = trackToPackedCandidatesAssoc[track];
0106     if (pc.isNull()) {
0107       continue;
0108     }
0109
0110     auto gen = trackToGenAssoc[track];
0111     if (gen.isNull()) {
0112       continue;
0113     }
0114
0115     auto newGenWSO = genToPrunedAssocWSO[gen];
0116     if (newGenWSO.isNull()) {
0117       continue;
0118     }
0119
0120     auto newGen = genToPrunedAssoc[newGenWSO];
0121     if (newGen.isNull()) {
0122       continue;
0123     }
0124
0125     indices[pc.index()] = newGen.index();
0126   }
0127   filler.insert(pcCollection, indices.begin(), indices.end());
0128   filler.fill();
0129   iEvent.put(std::move(out));
0130 }
0131
0132 void PackedCandidateGenAssociationProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0133   edm::ParameterSetDescription desc;
0134   desc.add<edm::InputTag>("genToPrunedAssoc", edm::InputTag("prunedGenParticles"));
0135   desc.add<edm::InputTag>("genToPrunedAssocWithStatusOne", edm::InputTag("prunedGenParticlesWithStatusOne"));
0136   desc.add<edm::InputTag>("trackToPackedCandidatesAssoc", edm::InputTag("packedPFCandidates"));
0137   desc.add<edm::InputTag>("trackToGenAssoc");
0138   desc.add<edm::InputTag>("tracks", edm::InputTag("generalTracks"));
0139
0140   descriptions.add("packedCandidatesGenAssociationDefault", desc);
0141 }
0142
0143 DEFINE_FWK_MODULE(PackedCandidateGenAssociationProducer);