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