PatAlgos/plugins/PATGenCandsFromSimTracksProducer.cc

0001 /**
0002   \class    pat::PATGenCandsFromSimTracksProducer PATGenCandsFromSimTracksProducer.h "PhysicsTools/PatAlgos/interface/PATGenCandsFromSimTracksProducer.h"
0003   \brief    Produces reco::GenParticle from SimTracks
0004
0005    The PATGenCandsFromSimTracksProducer produces GenParticles from SimTracks, so they can be used for MC matching.
0006
0007
0008   \author   Jordan Tucker (original module), Giovanni Petrucciani (PAT integration)
0009   \version  $Id: PATGenCandsFromSimTracksProducer.cc,v 1.8 2010/10/20 23:09:25 wmtan Exp $
0010 */
0011
0012 #include <ext/algorithm>
0013 #include <memory>
0014
0015 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0016 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0017 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0018 #include "FWCore/Framework/interface/Event.h"
0019 #include "FWCore/Framework/interface/Frameworkfwd.h"
0020 #include "FWCore/Framework/interface/MakerMacros.h"
0021 #include "FWCore/Framework/interface/stream/EDProducer.h"
0022 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0023 #include "SimDataFormats/Track/interface/SimTrackContainer.h"
0024 #include "SimDataFormats/Vertex/interface/SimVertexContainer.h"
0025 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0026 #include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
0027
0028 namespace pat {
0029   class PATGenCandsFromSimTracksProducer : public edm::stream::EDProducer<> {
0030   public:
0031     explicit PATGenCandsFromSimTracksProducer(const edm::ParameterSet &);
0032     ~PATGenCandsFromSimTracksProducer() override {}
0033
0034   private:
0035     void produce(edm::Event &, const edm::EventSetup &) override;
0036
0037     bool firstEvent_;
0038     edm::EDGetTokenT<edm::SimTrackContainer> simTracksToken_;
0039     edm::EDGetTokenT<edm::SimVertexContainer> simVertexToken_;
0040     int setStatus_;
0041     std::set<int> pdgIds_;              // these are the ones we really use
0042     std::vector<PdtEntry> pdts_;        // these are needed before we get the EventSetup
0043     std::set<int> motherPdgIds_;        // these are the ones we really use
0044     std::vector<PdtEntry> motherPdts_;  // these are needed before we get the EventSetup
0045
0046     typedef StringCutObjectSelector<reco::GenParticle> StrFilter;
0047     const StrFilter filter_;
0048
0049     /// If true, I'll try to make a link from the GEANT particle to a GenParticle
0050     bool makeMotherLink_;
0051     /// If true, I'll save GenParticles corresponding to the ancestors of this GEANT particle. Common ancestors are only written once.
0052     bool writeAncestors_;
0053
0054     /// Collection of GenParticles I need to make refs to. It must also have its associated vector<int> of barcodes, aligned with them.
0055     edm::EDGetTokenT<reco::GenParticleCollection> genParticlesToken_;
0056     edm::EDGetTokenT<std::vector<int> > genBarcodesToken_;
0057
0058     edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
0059     /// Global context for all recursive methods
0060     struct GlobalContext {
0061       GlobalContext(const edm::SimTrackContainer &simtks1,
0062                     const edm::SimVertexContainer &simvtxs1,
0063                     const edm::Handle<reco::GenParticleCollection> &gens1,
0064                     const edm::Handle<std::vector<int> > &genBarcodes1,
0065                     bool barcodesAreSorted1,
0066                     reco::GenParticleCollection &output1,
0067                     const edm::RefProd<reco::GenParticleCollection> &refprod1)
0068           : simtks(simtks1),
0069             simvtxs(simvtxs1),
0070             gens(gens1),
0071             genBarcodes(genBarcodes1),
0072             barcodesAreSorted(barcodesAreSorted1),
0073             output(output1),
0074             refprod(refprod1),
0075             simTksProcessed() {}
0076       // GEANT info
0077       const edm::SimTrackContainer &simtks;
0078       const edm::SimVertexContainer &simvtxs;
0079       // PYTHIA info
0080       const edm::Handle<reco::GenParticleCollection> &gens;
0081       const edm::Handle<std::vector<int> > &genBarcodes;
0082       bool barcodesAreSorted;
0083       // MY OUTPUT info
0084       reco::GenParticleCollection &output;
0085       const edm::RefProd<reco::GenParticleCollection> &refprod;
0086       // BOOK-KEEPING
0087       std::map<unsigned int, int> simTksProcessed;  // key = sim track id;
0088       // val = 0: not processed;
0089       //       i>0:  (index+1) in my output
0090       //       i<0: -(index+1) in pythia [NOT USED]
0091     };
0092
0093     /// Find the mother of a given GEANT track (or NULL if it can't be found).
0094     const SimTrack *findGeantMother(const SimTrack &tk, const GlobalContext &g) const;
0095     /// Find the GenParticle reference for a given GEANT or PYTHIA track.
0096     ///  - if the track corresponds to a PYTHIA particle, return a ref to that particle
0097     ///  - otherwise, if this simtrack has no mother simtrack, return a null ref
0098     ///  - otherwise, if writeAncestors is true,  make a GenParticle for it and return a ref to it
0099     ///  - otherwise, if writeAncestors is false, return the ref to the GEANT mother of this track
0100     edm::Ref<reco::GenParticleCollection> findRef(const SimTrack &tk, GlobalContext &g) const;
0101
0102     /// Used by findRef if the track is a PYTHIA particle
0103     edm::Ref<reco::GenParticleCollection> generatorRef_(const SimTrack &tk, const GlobalContext &g) const;
0104     /// Make a GenParticle for this SimTrack, with a given mother
0105     reco::GenParticle makeGenParticle_(const SimTrack &tk,
0106                                        const edm::Ref<reco::GenParticleCollection> &mother,
0107                                        const GlobalContext &g) const;
0108
0109     struct LessById {
0110       bool operator()(const SimTrack &tk1, const SimTrack &tk2) const { return tk1.trackId() < tk2.trackId(); }
0111       bool operator()(const SimTrack &tk1, unsigned int id) const { return tk1.trackId() < id; }
0112       bool operator()(unsigned int id, const SimTrack &tk2) const { return id < tk2.trackId(); }
0113     };
0114   };
0115 }  // namespace pat
0116
0117 using namespace std;
0118 using namespace edm;
0119 using namespace reco;
0120 using pat::PATGenCandsFromSimTracksProducer;
0121
0122 PATGenCandsFromSimTracksProducer::PATGenCandsFromSimTracksProducer(const ParameterSet &cfg)
0123     : firstEvent_(true),
0124       simTracksToken_(consumes<SimTrackContainer>(cfg.getParameter<InputTag>("src"))),   // source sim tracks
0125       simVertexToken_(consumes<SimVertexContainer>(cfg.getParameter<InputTag>("src"))),  // source sim  vertices
0126       setStatus_(cfg.getParameter<int32_t>("setStatus")),  // set status of GenParticle to this code
0127       filter_(cfg.existsAs<string>("filter") ? cfg.getParameter<string>("filter") : std::string("1 == 1")),
0128       makeMotherLink_(cfg.existsAs<bool>("makeMotherLink") ? cfg.getParameter<bool>("makeMotherLink") : false),
0129       writeAncestors_(cfg.existsAs<bool>("writeAncestors") ? cfg.getParameter<bool>("writeAncestors") : false),
0130       genParticlesToken_(mayConsume<GenParticleCollection>(cfg.getParameter<InputTag>("genParticles"))),
0131       genBarcodesToken_(mayConsume<std::vector<int> >(cfg.getParameter<InputTag>("genParticles"))),
0132       tableToken_(esConsumes()) {
0133   // Possibly allow a list of particle types
0134   if (cfg.exists("particleTypes")) {
0135     pdts_ = cfg.getParameter<vector<PdtEntry> >("particleTypes");
0136   }
0137   if (cfg.exists("motherTypes")) {
0138     motherPdts_ = cfg.getParameter<vector<PdtEntry> >("motherTypes");
0139   }
0140
0141   if (writeAncestors_ && !makeMotherLink_) {
0142     edm::LogWarning("Configuration")
0143         << "PATGenCandsFromSimTracksProducer: "
0144         << "you have set 'writeAncestors' to 'true' and 'makeMotherLink' to false;"
0145         << "GEANT particles with generator level (e.g.PYHIA) mothers won't have mother links.\n";
0146   }
0147   produces<GenParticleCollection>();
0148 }
0149
0150 const SimTrack *PATGenCandsFromSimTracksProducer::findGeantMother(const SimTrack &tk, const GlobalContext &g) const {
0151   assert(tk.genpartIndex() == -1);  // MUST NOT be called with a PYTHIA track
0152   if (!tk.noVertex()) {
0153     const SimVertex &vtx = g.simvtxs[tk.vertIndex()];
0154     if (!vtx.noParent()) {
0155       unsigned int idx = vtx.parentIndex();
0156       SimTrackContainer::const_iterator it = std::lower_bound(g.simtks.begin(), g.simtks.end(), idx, LessById());
0157       if ((it != g.simtks.end()) && (it->trackId() == idx)) {
0158         return &*it;
0159       }
0160     }
0161   }
0162   return nullptr;
0163 }
0164
0165 edm::Ref<reco::GenParticleCollection> PATGenCandsFromSimTracksProducer::findRef(const SimTrack &tk,
0166                                                                                 GlobalContext &g) const {
0167   if (tk.genpartIndex() != -1)
0168     return makeMotherLink_ ? generatorRef_(tk, g) : edm::Ref<reco::GenParticleCollection>();
0169   const SimTrack *simMother = findGeantMother(tk, g);
0170
0171   edm::Ref<reco::GenParticleCollection> motherRef;
0172   if (simMother != nullptr)
0173     motherRef = findRef(*simMother, g);
0174
0175   if (writeAncestors_) {
0176     // If writing ancestors, I need to serialize myself, and then to return a ref to me
0177     // But first check if I've already been serialized
0178     std::map<unsigned int, int>::const_iterator it = g.simTksProcessed.find(tk.trackId());
0179     if (it != g.simTksProcessed.end()) {
0180       // just return a ref to it
0181       assert(it->second > 0);
0182       return edm::Ref<reco::GenParticleCollection>(g.refprod, (it->second) - 1);
0183     } else {
0184       // make genParticle, save, update the map, and return ref to myself
0185       reco::GenParticle p = makeGenParticle_(tk, motherRef, g);
0186       g.output.push_back(p);
0187       g.simTksProcessed[tk.trackId()] = g.output.size();
0188       return edm::Ref<reco::GenParticleCollection>(g.refprod, g.output.size() - 1);
0189     }
0190   } else {
0191     // Otherwise, I just return a ref to my mum
0192     return motherRef;
0193   }
0194 }
0195
0196 edm::Ref<reco::GenParticleCollection> PATGenCandsFromSimTracksProducer::generatorRef_(const SimTrack &st,
0197                                                                                       const GlobalContext &g) const {
0198   assert(st.genpartIndex() != -1);
0199   // Note that st.genpartIndex() is the barcode, not the index within GenParticleCollection, so I have to search the particle
0200   std::vector<int>::const_iterator it;
0201   if (g.barcodesAreSorted) {
0202     it = std::lower_bound(g.genBarcodes->begin(), g.genBarcodes->end(), st.genpartIndex());
0203   } else {
0204     it = std::find(g.genBarcodes->begin(), g.genBarcodes->end(), st.genpartIndex());
0205   }
0206
0207   // Check that I found something
0208   // I need to check '*it == st.genpartIndex()' because lower_bound just finds the right spot for an item in a sorted list, not the item
0209   if ((it != g.genBarcodes->end()) && (*it == st.genpartIndex())) {
0210     return reco::GenParticleRef(g.gens, it - g.genBarcodes->begin());
0211   } else {
0212     return reco::GenParticleRef();
0213   }
0214 }
0215
0216 reco::GenParticle PATGenCandsFromSimTracksProducer::makeGenParticle_(
0217     const SimTrack &tk, const edm::Ref<reco::GenParticleCollection> &mother, const GlobalContext &g) const {
0218   // Make up a GenParticleCandidate from the GEANT track info.
0219   int charge = static_cast<int>(tk.charge());
0220   const Particle::LorentzVector &p4 = tk.momentum();
0221   Particle::Point vtx;  // = (0,0,0) by default
0222   if (!tk.noVertex())
0223     vtx = g.simvtxs[tk.vertIndex()].position();
0224   GenParticle gp(charge, p4, vtx, tk.type(), setStatus_, true);
0225   if (mother.isNonnull())
0226     gp.addMother(mother);
0227   return gp;
0228 }
0229
0230 void PATGenCandsFromSimTracksProducer::produce(Event &event, const EventSetup &iSetup) {
0231   if (firstEvent_) {
0232     auto const &pdt = iSetup.getData(tableToken_);
0233     if (!pdts_.empty()) {
0234       pdgIds_.clear();
0235       for (vector<PdtEntry>::iterator itp = pdts_.begin(), edp = pdts_.end(); itp != edp; ++itp) {
0236         itp->setup(pdt);  // decode string->pdgId and vice-versa
0237         pdgIds_.insert(std::abs(itp->pdgId()));
0238       }
0239       pdts_.clear();
0240     }
0241     if (!motherPdts_.empty()) {
0242       motherPdgIds_.clear();
0243       for (vector<PdtEntry>::iterator itp = motherPdts_.begin(), edp = motherPdts_.end(); itp != edp; ++itp) {
0244         itp->setup(pdt);  // decode string->pdgId and vice-versa
0245         motherPdgIds_.insert(std::abs(itp->pdgId()));
0246       }
0247       motherPdts_.clear();
0248     }
0249     firstEvent_ = false;
0250   }
0251
0252   // Simulated tracks (i.e. GEANT particles).
0253   Handle<SimTrackContainer> simtracks;
0254   event.getByToken(simTracksToken_, simtracks);
0255
0256   // Need to check that SimTrackContainer is sorted; otherwise, copy and sort :-(
0257   std::unique_ptr<SimTrackContainer> simtracksTmp;
0258   const SimTrackContainer *simtracksSorted = &*simtracks;
0259   if (makeMotherLink_ || writeAncestors_) {
0260     if (!__gnu_cxx::is_sorted(simtracks->begin(), simtracks->end(), LessById())) {
0261       simtracksTmp = std::make_unique<SimTrackContainer>(*simtracks);
0262       std::sort(simtracksTmp->begin(), simtracksTmp->end(), LessById());
0263       simtracksSorted = &*simtracksTmp;
0264     }
0265   }
0266
0267   // Get the associated vertices
0268   Handle<SimVertexContainer> simvertices;
0269   event.getByToken(simVertexToken_, simvertices);
0270
0271   // Get the GenParticles and barcodes, if needed to set mother links
0272   Handle<GenParticleCollection> gens;
0273   Handle<std::vector<int> > genBarcodes;
0274   bool barcodesAreSorted = true;
0275   if (makeMotherLink_) {
0276     event.getByToken(genParticlesToken_, gens);
0277     event.getByToken(genBarcodesToken_, genBarcodes);
0278     if (gens->size() != genBarcodes->size())
0279       throw cms::Exception("Corrupt data") << "Barcodes not of the same size as GenParticles!\n";
0280     barcodesAreSorted = __gnu_cxx::is_sorted(genBarcodes->begin(), genBarcodes->end());
0281   }
0282
0283   // make the output collection
0284   auto cands = std::make_unique<GenParticleCollection>();
0285   edm::RefProd<GenParticleCollection> refprod = event.getRefBeforePut<GenParticleCollection>();
0286
0287   GlobalContext globals(*simtracksSorted, *simvertices, gens, genBarcodes, barcodesAreSorted, *cands, refprod);
0288
0289   for (SimTrackContainer::const_iterator isimtrk = simtracks->begin(); isimtrk != simtracks->end(); ++isimtrk) {
0290     // Skip PYTHIA tracks.
0291     if (isimtrk->genpartIndex() != -1)
0292       continue;
0293
0294     // Maybe apply the PdgId filter
0295     if (!pdgIds_.empty()) {  // if we have a filter on pdg ids
0296       if (pdgIds_.find(std::abs(isimtrk->type())) == pdgIds_.end())
0297         continue;
0298     }
0299
0300     GenParticle genp = makeGenParticle_(*isimtrk, Ref<GenParticleCollection>(), globals);
0301
0302     // Maybe apply filter on the particle
0303     if (!(filter_(genp)))
0304       continue;
0305
0306     if (!motherPdgIds_.empty()) {
0307       const SimTrack *motherSimTk = findGeantMother(*isimtrk, globals);
0308       if (motherSimTk == nullptr)
0309         continue;
0310       if (motherPdgIds_.find(std::abs(motherSimTk->type())) == motherPdgIds_.end())
0311         continue;
0312     }
0313
0314     if (makeMotherLink_ || writeAncestors_) {
0315       Ref<GenParticleCollection> motherRef;
0316       const SimTrack *mother = findGeantMother(*isimtrk, globals);
0317       if (mother != nullptr)
0318         motherRef = findRef(*mother, globals);
0319       if (motherRef.isNonnull())
0320         genp.addMother(motherRef);
0321     }
0322
0323     cands->push_back(genp);
0324   }
0325
0326   // Write to the Event, and get back a handle (which can be useful for debugging)
0327 #ifdef DEBUG_PATGenCandsFromSimTracksProducer
0328   edm::OrphanHandle<reco::GenParticleCollection> orphans =
0329 #endif
0330       event.put(std::move(cands));
0331
0332 #ifdef DEBUG_PATGenCandsFromSimTracksProducer
0333   std::cout << "Produced a list of " << orphans->size() << " genParticles." << std::endl;
0334   for (GenParticleCollection::const_iterator it = orphans->begin(), ed = orphans->end(); it != ed; ++it) {
0335     std::cout << "    ";
0336     std::cout << "GenParticle #" << (it - orphans->begin()) << ": pdgId " << it->pdgId() << ", pt = " << it->pt()
0337               << ", eta = " << it->eta() << ", phi = " << it->phi() << ", rho = " << it->vertex().Rho()
0338               << ", z = " << it->vertex().Z() << std::endl;
0339     edm::Ref<GenParticleCollection> mom = it->motherRef();
0340     size_t depth = 2;
0341     while (mom.isNonnull()) {
0342       if (mom.id() == orphans.id()) {
0343         // I need to re-make the ref because they are not working until this module returns.
0344         mom = edm::Ref<GenParticleCollection>(orphans, mom.key());
0345       }
0346       for (size_t i = 0; i < depth; ++i)
0347         std::cout << "    ";
0348       std::cout << "GenParticleRef [" << mom.id() << "/" << mom.key() << "]: pdgId " << mom->pdgId()
0349                 << ", status = " << mom->status() << ", pt = " << mom->pt() << ", eta = " << mom->eta()
0350                 << ", phi = " << mom->phi() << ", rho = " << mom->vertex().Rho() << ", z = " << mom->vertex().Z()
0351                 << std::endl;
0352       if (mom.id() != orphans.id())
0353         break;
0354       if ((mom->motherRef().id() == mom.id()) && (mom->motherRef().key() == mom.key())) {
0355         throw cms::Exception("Corrupt Data") << "A particle is it's own mother.\n";
0356       }
0357       mom = mom->motherRef();
0358       depth++;
0359     }
0360   }
0361   std::cout << std::endl;
0362 #endif
0363 }
0364
0365 DEFINE_FWK_MODULE(PATGenCandsFromSimTracksProducer);