HepMCCandAlgos/plugins/GenPlusSimParticleProducer.cc

0001 /**
0002   \class    pat::GenPlusSimParticleProducer GenPlusSimParticleProducer.h "PhysicsTools/PatAlgos/interface/GenPlusSimParticleProducer.h"
0003   \brief    Produces reco::GenParticle from SimTracks
0004
0005 The GenPlusSimParticleProducer produces GenParticles from SimTracks,
0006 so they can be used for MC matching. A copy of the original genParticle list
0007 is made and the genParticles created from the GEANT/FAMOS particles are added
0008 to the list including all ancestors and correct mother/daughter references
0009
0010 Sample useage in cfg.py file:
0011 process.genParticlePlusGEANT = cms.EDProducer("GenPlusSimParticleProducer",
0012         src           = cms.InputTag("g4SimHits"), # use "fastSimProducer" for FastSim
0013         setStatus     = cms.int32(8),             # set status = 8 for GEANT GPs
0014         particleTypes = cms.vstring("pi+"),       # also picks pi- (optional)
0015         filter        = cms.vstring("pt > 0.0"),  # just for testing
0016         genParticles   = cms.InputTag("genParticles") # original genParticle list
0017 )
0018
0019   \author   Jordan Tucker (original module), Keith Ulmer (generalization)
0020 */
0021
0022 #include "FWCore/Framework/interface/Frameworkfwd.h"
0023 #include "FWCore/Framework/interface/stream/EDProducer.h"
0024 #include "FWCore/Framework/interface/Event.h"
0025 #include "FWCore/Framework/interface/EventSetup.h"
0026 #include "FWCore/Framework/interface/MakerMacros.h"
0027 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0028
0029 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0030 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0031 #include "SimDataFormats/Track/interface/SimTrackContainer.h"
0032 #include "SimDataFormats/Vertex/interface/SimVertexContainer.h"
0033
0034 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0035 #include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
0036 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0037
0038 #include <ext/algorithm>
0039 #include <memory>
0040
0041 namespace pat {
0042   class GenPlusSimParticleProducer : public edm::stream::EDProducer<> {
0043   public:
0044     explicit GenPlusSimParticleProducer(const edm::ParameterSet &);
0045
0046   private:
0047     void produce(edm::Event &, const edm::EventSetup &) override;
0048
0049     bool firstEvent_;
0050     edm::EDGetTokenT<edm::SimTrackContainer> simtracksToken_;
0051     edm::EDGetTokenT<edm::SimVertexContainer> simverticesToken_;
0052     int setStatus_;
0053     std::set<int> pdgIds_;        // these are the ones we really use
0054     std::vector<PdtEntry> pdts_;  // these are needed before we get the EventSetup
0055
0056     typedef StringCutObjectSelector<reco::GenParticle> StrFilter;
0057     std::unique_ptr<StrFilter> filter_;
0058
0059     /// Collection of GenParticles I need to make refs to. It must also have its associated vector<int> of barcodes, aligned with them.
0060     edm::EDGetTokenT<reco::GenParticleCollection> gensToken_;
0061     edm::EDGetTokenT<std::vector<int>> genBarcodesToken_;
0062
0063     edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
0064     /// Try to link the GEANT particle to the generator particle it came from
0065     /** Arguments:
0066    * -- Specific --
0067    *    gp: GenParticle made from the GEANT particle
0068    *    st: The GEANT simTrack for which we create a genParticle
0069    *
0070    * -- GEANT related --
0071    *    simtks:  A list of GEANT tracks, sorted by track id
0072    *    simvtxs: The list of GEANT vertices, in their natural order (skimtks have pointers into this vector!)
0073    *
0074    * -- GenParticle related --
0075    *    gens             : Handle to the GenParticles, to make the ref to
0076    *    genBarcodes      : Barcodes for each GenParticle, in a vector aligned to the GenParticleCollection.
0077    *    barcodesAreSorted: true if the barcodes are sorted (which means I can use binary_search on them) */
0078     void addGenParticle(const SimTrack &stMom,
0079                         const SimTrack &stDau,
0080                         unsigned int momidx,
0081                         const edm::SimTrackContainer &simtks,
0082                         const edm::SimVertexContainer &simvtxs,
0083                         reco::GenParticleCollection &mergedGens,
0084                         const reco::GenParticleRefProd ref,
0085                         std::vector<int> &genBarcodes,
0086                         bool &barcodesAreSorted) const;
0087     struct LessById {
0088       bool operator()(const SimTrack &tk1, const SimTrack &tk2) const { return tk1.trackId() < tk2.trackId(); }
0089       bool operator()(const SimTrack &tk1, unsigned int id) const { return tk1.trackId() < id; }
0090       bool operator()(unsigned int id, const SimTrack &tk2) const { return id < tk2.trackId(); }
0091     };
0092   };
0093 }  // namespace pat
0094
0095 using namespace std;
0096 using namespace edm;
0097 using namespace reco;
0098 using pat::GenPlusSimParticleProducer;
0099
0100 GenPlusSimParticleProducer::GenPlusSimParticleProducer(const ParameterSet &cfg)
0101     : firstEvent_(true),
0102       simtracksToken_(consumes<SimTrackContainer>(cfg.getParameter<InputTag>("src"))),  // source sim tracks & vertices
0103       simverticesToken_(
0104           consumes<SimVertexContainer>(cfg.getParameter<InputTag>("src"))),  // source sim tracks & vertices
0105       setStatus_(cfg.getParameter<int32_t>("setStatus")),                    // set status of GenParticle to this code
0106       gensToken_(consumes<GenParticleCollection>(
0107           cfg.getParameter<InputTag>("genParticles"))),  // get the genParticles to add GEANT particles to
0108       genBarcodesToken_(consumes<std::vector<int>>(
0109           cfg.getParameter<InputTag>("genParticles")))  // get the genParticles to add GEANT particles to
0110 {
0111   // Possibly allow a list of particle types
0112   if (cfg.exists("particleTypes")) {
0113     pdts_ = cfg.getParameter<vector<PdtEntry>>("particleTypes");
0114     if (!pdts_.empty()) {
0115       tableToken_ = esConsumes();
0116     }
0117   }
0118
0119   // Possibly allow a string cut
0120   if (cfg.existsAs<string>("filter")) {
0121     string filter = cfg.getParameter<string>("filter");
0122     if (!filter.empty()) {
0123       filter_ = std::make_unique<StrFilter>(filter);
0124     }
0125   }
0126   produces<GenParticleCollection>();
0127   produces<vector<int>>();
0128 }
0129
0130 void GenPlusSimParticleProducer::addGenParticle(const SimTrack &stMom,
0131                                                 const SimTrack &stDau,
0132                                                 unsigned int momidx,
0133                                                 const SimTrackContainer &simtracksSorted,
0134                                                 const SimVertexContainer &simvertices,
0135                                                 reco::GenParticleCollection &mergedGens,
0136                                                 const GenParticleRefProd ref,
0137                                                 std::vector<int> &genBarcodes,
0138                                                 bool &barcodesAreSorted) const {
0139   // Make the genParticle for stDau and add it to the new collection and update the parent-child relationship
0140   // Make up a GenParticleCandidate from the GEANT track info.
0141   int charge = static_cast<int>(stDau.charge());
0142   const Particle::LorentzVector &p4 = stDau.momentum();
0143   Particle::Point vtx;  // = (0,0,0) by default
0144   if (!stDau.noVertex())
0145     vtx = simvertices[stDau.vertIndex()].position();
0146   GenParticle genp(charge, p4, vtx, stDau.type(), setStatus_, true);
0147
0148   // Maybe apply filter on the particle
0149   if (filter_.get() != nullptr) {
0150     if (!(*filter_)(genp))
0151       return;
0152   }
0153
0154   reco::GenParticleRef parentRef(ref, momidx);
0155   genp.addMother(parentRef);
0156   mergedGens.push_back(genp);
0157   // get the index for the daughter just added
0158   unsigned int dauidx = mergedGens.size() - 1;
0159
0160   // update add daughter relationship
0161   reco::GenParticle &cand = mergedGens[momidx];
0162   cand.addDaughter(GenParticleRef(ref, dauidx));
0163
0164   //look for simtrack daughters of stDau to see if we need to recur further down the chain
0165
0166   for (SimTrackContainer::const_iterator isimtrk = simtracksSorted.begin(); isimtrk != simtracksSorted.end();
0167        ++isimtrk) {
0168     if (!isimtrk->noVertex()) {
0169       // Pick the vertex (isimtrk.vertIndex() is really an index)
0170       const SimVertex &vtx = simvertices[isimtrk->vertIndex()];
0171
0172       // Check if the vertex has a parent track (otherwise, we're lost)
0173       if (!vtx.noParent()) {
0174         // Now note that vtx.parentIndex() is NOT an index, it's a track id, so I have to search for it
0175         unsigned int idx = vtx.parentIndex();
0176         SimTrackContainer::const_iterator it =
0177             std::lower_bound(simtracksSorted.begin(), simtracksSorted.end(), idx, LessById());
0178         if ((it != simtracksSorted.end()) && (it->trackId() == idx)) {
0179           if (it->trackId() == stDau.trackId()) {
0180             //need the genparticle index of stDau which is dauidx
0181             addGenParticle(
0182                 stDau, *isimtrk, dauidx, simtracksSorted, simvertices, mergedGens, ref, genBarcodes, barcodesAreSorted);
0183           }
0184         }
0185       }
0186     }
0187   }
0188 }
0189
0190 void GenPlusSimParticleProducer::produce(Event &event, const EventSetup &iSetup) {
0191   if (firstEvent_) {
0192     if (!pdts_.empty()) {
0193       auto const &pdt = iSetup.getData(tableToken_);
0194       pdgIds_.clear();
0195       for (vector<PdtEntry>::iterator itp = pdts_.begin(), edp = pdts_.end(); itp != edp; ++itp) {
0196         itp->setup(pdt);  // decode string->pdgId and vice-versa
0197         pdgIds_.insert(std::abs(itp->pdgId()));
0198       }
0199       pdts_.clear();
0200     }
0201     firstEvent_ = false;
0202   }
0203
0204   // Simulated tracks (i.e. GEANT particles).
0205   Handle<SimTrackContainer> simtracks;
0206   event.getByToken(simtracksToken_, simtracks);
0207
0208   // Need to check that SimTrackContainer is sorted; otherwise, copy and sort :-(
0209   std::unique_ptr<SimTrackContainer> simtracksTmp;
0210   const SimTrackContainer *simtracksSorted = &*simtracks;
0211   if (!__gnu_cxx::is_sorted(simtracks->begin(), simtracks->end(), LessById())) {
0212     simtracksTmp = std::make_unique<SimTrackContainer>(*simtracks);
0213     std::sort(simtracksTmp->begin(), simtracksTmp->end(), LessById());
0214     simtracksSorted = &*simtracksTmp;
0215   }
0216
0217   // Get the associated vertices
0218   Handle<SimVertexContainer> simvertices;
0219   event.getByToken(simverticesToken_, simvertices);
0220
0221   // Get the GenParticles and barcodes, if needed to set mother and daughter links
0222   Handle<GenParticleCollection> gens;
0223   Handle<std::vector<int>> genBarcodes;
0224   bool barcodesAreSorted = true;
0225   event.getByToken(gensToken_, gens);
0226   event.getByToken(genBarcodesToken_, genBarcodes);
0227   if (gens->size() != genBarcodes->size())
0228     throw cms::Exception("Corrupt data") << "Barcodes not of the same size as GenParticles!\n";
0229
0230   // make the output collection
0231   auto candsPtr = std::make_unique<GenParticleCollection>();
0232   GenParticleCollection &cands = *candsPtr;
0233
0234   const GenParticleRefProd ref = event.getRefBeforePut<GenParticleCollection>();
0235
0236   // add the original genParticles to the merged output list
0237   for (size_t i = 0; i < gens->size(); ++i) {
0238     reco::GenParticle cand((*gens)[i]);
0239     cands.push_back(cand);
0240   }
0241
0242   // make new barcodes vector and fill it with the original list
0243   auto newGenBarcodes = std::make_unique<vector<int>>();
0244   for (unsigned int i = 0; i < genBarcodes->size(); ++i) {
0245     newGenBarcodes->push_back((*genBarcodes)[i]);
0246   }
0247   barcodesAreSorted = __gnu_cxx::is_sorted(newGenBarcodes->begin(), newGenBarcodes->end());
0248
0249   for (size_t i = 0; i < cands.size(); ++i) {
0250     reco::GenParticle &cand = cands[i];
0251     size_t nDaus = cand.numberOfDaughters();
0252     GenParticleRefVector daus = cand.daughterRefVector();
0253     cand.resetDaughters(ref.id());
0254     for (size_t d = 0; d < nDaus; ++d) {
0255       cand.addDaughter(GenParticleRef(ref, daus[d].key()));
0256     }
0257
0258     size_t nMoms = cand.numberOfMothers();
0259     GenParticleRefVector moms = cand.motherRefVector();
0260     cand.resetMothers(ref.id());
0261     for (size_t m = 0; m < nMoms; ++m) {
0262       cand.addMother(GenParticleRef(ref, moms[m].key()));
0263     }
0264   }
0265
0266   for (SimTrackContainer::const_iterator isimtrk = simtracks->begin(); isimtrk != simtracks->end(); ++isimtrk) {
0267     // Skip PYTHIA tracks.
0268     if (isimtrk->genpartIndex() != -1)
0269       continue;
0270
0271     // Maybe apply the PdgId filter
0272     if (!pdgIds_.empty()) {  // if we have a filter on pdg ids
0273       if (pdgIds_.find(std::abs(isimtrk->type())) == pdgIds_.end())
0274         continue;
0275     }
0276
0277     // find simtrack that has a genParticle match to its parent
0278     // Look at the production vertex. If there is no vertex, I can do nothing...
0279     if (!isimtrk->noVertex()) {
0280       // Pick the vertex (isimtrk.vertIndex() is really an index)
0281       const SimVertex &vtx = (*simvertices)[isimtrk->vertIndex()];
0282
0283       // Check if the vertex has a parent track (otherwise, we're lost)
0284       if (!vtx.noParent()) {
0285         // Now note that vtx.parentIndex() is NOT an index, it's a track id, so I have to search for it
0286         unsigned int idx = vtx.parentIndex();
0287         SimTrackContainer::const_iterator it =
0288             std::lower_bound(simtracksSorted->begin(), simtracksSorted->end(), idx, LessById());
0289         if ((it != simtracksSorted->end()) && (it->trackId() == idx)) {  //it is the parent sim track
0290           if (it->genpartIndex() != -1) {
0291             std::vector<int>::const_iterator itIndex;
0292             if (barcodesAreSorted) {
0293               itIndex = std::lower_bound(genBarcodes->begin(), genBarcodes->end(), it->genpartIndex());
0294             } else {
0295               itIndex = std::find(genBarcodes->begin(), genBarcodes->end(), it->genpartIndex());
0296             }
0297
0298             // Check that I found something
0299             // I need to check '*itIndex == it->genpartIndex()' because lower_bound just finds the right spot for an item in a sorted list, not the item
0300             if ((itIndex != genBarcodes->end()) && (*itIndex == it->genpartIndex())) {
0301               // Ok, I'll make the genParticle for st and add it to the new collection updating the map and parent-child relationship
0302               // pass the mother and daughter sim tracks and the mother genParticle to method to create the daughter genParticle and recur
0303               unsigned int momidx = itIndex - genBarcodes->begin();
0304               addGenParticle(*it,
0305                              *isimtrk,
0306                              momidx,
0307                              *simtracksSorted,
0308                              *simvertices,
0309                              *candsPtr,
0310                              ref,
0311                              *newGenBarcodes,
0312                              barcodesAreSorted);
0313             }
0314           }
0315         }
0316       }
0317     }
0318   }
0319
0320   event.put(std::move(candsPtr));
0321   event.put(std::move(newGenBarcodes));
0322 }
0323
0324 DEFINE_FWK_MODULE(GenPlusSimParticleProducer);