HepMCCandAlgos/plugins/GenParticlePruner.cc

0001 #include "FWCore/Framework/interface/stream/EDProducer.h"
0002 #include "FWCore/Utilities/interface/InputTag.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "DataFormats/Common/interface/Handle.h"
0006 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0007 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0008 #include "FWCore/Utilities/interface/EDMException.h"
0009 #include "PhysicsTools/HepMCCandAlgos/interface/pdgEntryReplace.h"
0010 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0011 #include "DataFormats/Common/interface/Association.h"
0012
0013 namespace helper {
0014   struct SelectCode {
0015     enum KeepOrDrop { kKeep, kDrop };
0016     enum FlagDepth { kNone, kFirst, kAll };
0017     KeepOrDrop keepOrDrop_;
0018     FlagDepth daughtersDepth_, mothersDepth_;
0019     bool all_;
0020   };
0021 }  // namespace helper
0022
0023 class GenParticlePruner : public edm::stream::EDProducer<> {
0024 public:
0025   GenParticlePruner(const edm::ParameterSet &);
0026
0027 private:
0028   void produce(edm::Event &, const edm::EventSetup &) override;
0029   bool firstEvent_;
0030   edm::EDGetTokenT<reco::GenParticleCollection> srcToken_;
0031   edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
0032   int keepOrDropAll_;
0033   std::vector<std::string> selection_;
0034   std::vector<std::pair<StringCutObjectSelector<reco::GenParticle>, helper::SelectCode>> select_;
0035   std::vector<int> flags_;
0036   std::vector<size_t> indices_;
0037   void parse(const std::string &selection, helper::SelectCode &code, std::string &cut) const;
0038   void flagDaughters(const reco::GenParticle &, int);
0039   void flagMothers(const reco::GenParticle &, int);
0040   void recursiveFlagDaughters(size_t, const reco::GenParticleCollection &, int, std::vector<size_t> &);
0041   void recursiveFlagMothers(size_t, const reco::GenParticleCollection &, int, std::vector<size_t> &);
0042   void getDaughterKeys(std::vector<size_t> &, std::vector<size_t> &, const reco::GenParticleRefVector &) const;
0043   void getMotherKeys(std::vector<size_t> &, std::vector<size_t> &, const reco::GenParticleRefVector &) const;
0044 };
0045
0046 using namespace edm;
0047 using namespace std;
0048 using namespace reco;
0049
0050 const int keep = 1, drop = -1;
0051
0052 void GenParticlePruner::parse(const std::string &selection, ::helper::SelectCode &code, std::string &cut) const {
0053   using namespace ::helper;
0054   size_t f = selection.find_first_not_of(' ');
0055   size_t n = selection.size();
0056   string command;
0057   char c;
0058   for (; (c = selection[f]) != ' ' && f < n; ++f) {
0059     command.push_back(c);
0060   }
0061   if (command[0] == '+') {
0062     command.erase(0, 1);
0063     if (command[0] == '+') {
0064       command.erase(0, 1);
0065       code.mothersDepth_ = SelectCode::kAll;
0066     } else {
0067       code.mothersDepth_ = SelectCode::kFirst;
0068     }
0069   } else
0070     code.mothersDepth_ = SelectCode::kNone;
0071
0072   if (command[command.size() - 1] == '+') {
0073     command.erase(command.size() - 1);
0074     if (command[command.size() - 1] == '+') {
0075       command.erase(command.size() - 1);
0076       code.daughtersDepth_ = SelectCode::kAll;
0077     } else {
0078       code.daughtersDepth_ = SelectCode::kFirst;
0079     }
0080   } else
0081     code.daughtersDepth_ = SelectCode::kNone;
0082
0083   if (command == "keep")
0084     code.keepOrDrop_ = SelectCode::kKeep;
0085   else if (command == "drop")
0086     code.keepOrDrop_ = SelectCode::kDrop;
0087   else {
0088     throw Exception(errors::Configuration) << "invalid selection command: " << command << "\n" << endl;
0089   }
0090   for (; f < n; ++f) {
0091     if (selection[f] != ' ')
0092       break;
0093   }
0094   cut = string(selection, f);
0095   if (cut[0] == '*')
0096     cut = string(cut, 0, cut.find_first_of(' '));
0097   code.all_ = cut == "*";
0098 }
0099
0100 GenParticlePruner::GenParticlePruner(const ParameterSet &cfg)
0101     : firstEvent_(true),
0102       srcToken_(consumes<GenParticleCollection>(cfg.getParameter<InputTag>("src"))),
0103       tableToken_(esConsumes()),
0104       keepOrDropAll_(drop),
0105       selection_(cfg.getParameter<vector<string>>("select")) {
0106   using namespace ::helper;
0107   produces<GenParticleCollection>();
0108   produces<edm::Association<reco::GenParticleCollection>>();
0109 }
0110
0111 void GenParticlePruner::flagDaughters(const reco::GenParticle &gen, int keepOrDrop) {
0112   const GenParticleRefVector &daughters = gen.daughterRefVector();
0113   for (GenParticleRefVector::const_iterator i = daughters.begin(); i != daughters.end(); ++i)
0114     flags_[i->key()] = keepOrDrop;
0115 }
0116
0117 void GenParticlePruner::flagMothers(const reco::GenParticle &gen, int keepOrDrop) {
0118   const GenParticleRefVector &mothers = gen.motherRefVector();
0119   for (GenParticleRefVector::const_iterator i = mothers.begin(); i != mothers.end(); ++i)
0120     flags_[i->key()] = keepOrDrop;
0121 }
0122
0123 void GenParticlePruner::recursiveFlagDaughters(size_t index,
0124                                                const reco::GenParticleCollection &src,
0125                                                int keepOrDrop,
0126                                                std::vector<size_t> &allIndices) {
0127   GenParticleRefVector daughters = src[index].daughterRefVector();
0128   // avoid infinite recursion if the daughters are set to "this" particle.
0129   size_t cachedIndex = index;
0130   for (GenParticleRefVector::const_iterator i = daughters.begin(); i != daughters.end(); ++i) {
0131     index = i->key();
0132     // To also avoid infinite recursion if a "loop" is found in the daughter list,
0133     // check to make sure the index hasn't already been added.
0134     if (find(allIndices.begin(), allIndices.end(), index) == allIndices.end()) {
0135       allIndices.push_back(index);
0136       if (cachedIndex != index) {
0137         flags_[index] = keepOrDrop;
0138         recursiveFlagDaughters(index, src, keepOrDrop, allIndices);
0139       }
0140     }
0141   }
0142 }
0143
0144 void GenParticlePruner::recursiveFlagMothers(size_t index,
0145                                              const reco::GenParticleCollection &src,
0146                                              int keepOrDrop,
0147                                              std::vector<size_t> &allIndices) {
0148   GenParticleRefVector mothers = src[index].motherRefVector();
0149   // avoid infinite recursion if the mothers are set to "this" particle.
0150   size_t cachedIndex = index;
0151   for (GenParticleRefVector::const_iterator i = mothers.begin(); i != mothers.end(); ++i) {
0152     index = i->key();
0153     // To also avoid infinite recursion if a "loop" is found in the daughter list,
0154     // check to make sure the index hasn't already been added.
0155     if (find(allIndices.begin(), allIndices.end(), index) == allIndices.end()) {
0156       allIndices.push_back(index);
0157       if (cachedIndex != index) {
0158         flags_[index] = keepOrDrop;
0159         recursiveFlagMothers(index, src, keepOrDrop, allIndices);
0160       }
0161     }
0162   }
0163 }
0164
0165 void GenParticlePruner::produce(Event &evt, const EventSetup &es) {
0166   if (firstEvent_) {
0167     auto const &pdt = es.getData(tableToken_);
0168
0169     for (vector<string>::const_iterator i = selection_.begin(); i != selection_.end(); ++i) {
0170       string cut;
0171       ::helper::SelectCode code;
0172       parse(*i, code, cut);
0173       if (code.all_) {
0174         if (i != selection_.begin())
0175           throw Exception(errors::Configuration)
0176               << "selections \"keep *\" and \"drop *\" can be used only as first options. Here used in position # "
0177               << (i - selection_.begin()) + 1 << "\n"
0178               << endl;
0179         switch (code.keepOrDrop_) {
0180           case ::helper::SelectCode::kDrop:
0181             keepOrDropAll_ = drop;
0182             break;
0183           case ::helper::SelectCode::kKeep:
0184             keepOrDropAll_ = keep;
0185         };
0186       } else {
0187         cut = pdgEntryReplace(cut, pdt);
0188         select_.push_back(make_pair(StringCutObjectSelector<GenParticle>(cut), code));
0189       }
0190     }
0191     firstEvent_ = false;
0192   }
0193
0194   using namespace ::helper;
0195   Handle<GenParticleCollection> src;
0196   evt.getByToken(srcToken_, src);
0197   const size_t n = src->size();
0198   flags_.clear();
0199   flags_.resize(n, keepOrDropAll_);
0200   for (size_t j = 0; j < select_.size(); ++j) {
0201     const pair<StringCutObjectSelector<GenParticle>, SelectCode> &sel = select_[j];
0202     SelectCode code = sel.second;
0203     const StringCutObjectSelector<GenParticle> &cut = sel.first;
0204     for (size_t i = 0; i < n; ++i) {
0205       const GenParticle &p = (*src)[i];
0206       if (cut(p)) {
0207         int keepOrDrop = keep;
0208         switch (code.keepOrDrop_) {
0209           case SelectCode::kKeep:
0210             keepOrDrop = keep;
0211             break;
0212           case SelectCode::kDrop:
0213             keepOrDrop = drop;
0214         };
0215         flags_[i] = keepOrDrop;
0216         std::vector<size_t> allIndicesDa;
0217         std::vector<size_t> allIndicesMo;
0218         switch (code.daughtersDepth_) {
0219           case SelectCode::kAll:
0220             recursiveFlagDaughters(i, *src, keepOrDrop, allIndicesDa);
0221             break;
0222           case SelectCode::kFirst:
0223             flagDaughters(p, keepOrDrop);
0224             break;
0225           case SelectCode::kNone:;
0226         };
0227         switch (code.mothersDepth_) {
0228           case SelectCode::kAll:
0229             recursiveFlagMothers(i, *src, keepOrDrop, allIndicesMo);
0230             break;
0231           case SelectCode::kFirst:
0232             flagMothers(p, keepOrDrop);
0233             break;
0234           case SelectCode::kNone:;
0235         };
0236       }
0237     }
0238   }
0239   indices_.clear();
0240   int counter = 0;
0241   for (size_t i = 0; i < n; ++i) {
0242     if (flags_[i] == keep) {
0243       indices_.push_back(i);
0244       flags_[i] = counter++;
0245     } else {
0246       flags_[i] = -1;  //set to invalid ref
0247     }
0248   }
0249
0250   auto out = std::make_unique<GenParticleCollection>();
0251   GenParticleRefProd outRef = evt.getRefBeforePut<GenParticleCollection>();
0252   out->reserve(counter);
0253
0254   for (vector<size_t>::const_iterator i = indices_.begin(); i != indices_.end(); ++i) {
0255     size_t index = *i;
0256     const GenParticle &gen = (*src)[index];
0257     const LeafCandidate &part = gen;
0258     out->push_back(GenParticle(part));
0259     GenParticle &newGen = out->back();
0260     //fill status flags
0261     newGen.statusFlags() = gen.statusFlags();
0262     newGen.setCollisionId(gen.collisionId());
0263     // The "daIndxs" and "moIndxs" keep a list of the keys for the mother/daughter
0264     // parentage/descendency. In some cases, a circular referencing is encountered,
0265     // which would result in an infinite loop. The list is checked to
0266     // avoid this.
0267     vector<size_t> daIndxs, daNewIndxs;
0268     getDaughterKeys(daIndxs, daNewIndxs, gen.daughterRefVector());
0269     std::sort(daNewIndxs.begin(), daNewIndxs.end());
0270     for (size_t i = 0; i < daNewIndxs.size(); ++i)
0271       newGen.addDaughter(GenParticleRef(outRef, daNewIndxs[i]));
0272
0273     vector<size_t> moIndxs, moNewIndxs;
0274     getMotherKeys(moIndxs, moNewIndxs, gen.motherRefVector());
0275     std::sort(moNewIndxs.begin(), moNewIndxs.end());
0276     for (size_t i = 0; i < moNewIndxs.size(); ++i)
0277       newGen.addMother(GenParticleRef(outRef, moNewIndxs[i]));
0278   }
0279
0280   edm::OrphanHandle<reco::GenParticleCollection> oh = evt.put(std::move(out));
0281   auto orig2new = std::make_unique<edm::Association<reco::GenParticleCollection>>(oh);
0282   edm::Association<reco::GenParticleCollection>::Filler orig2newFiller(*orig2new);
0283   orig2newFiller.insert(src, flags_.begin(), flags_.end());
0284   orig2newFiller.fill();
0285   evt.put(std::move(orig2new));
0286 }
0287
0288 void GenParticlePruner::getDaughterKeys(vector<size_t> &daIndxs,
0289                                         vector<size_t> &daNewIndxs,
0290                                         const GenParticleRefVector &daughters) const {
0291   for (GenParticleRefVector::const_iterator j = daughters.begin(); j != daughters.end(); ++j) {
0292     GenParticleRef dau = *j;
0293     if (find(daIndxs.begin(), daIndxs.end(), dau.key()) == daIndxs.end()) {
0294       daIndxs.push_back(dau.key());
0295       int idx = flags_[dau.key()];
0296       if (idx > 0) {
0297         daNewIndxs.push_back(idx);
0298       } else {
0299         const GenParticleRefVector &daus = dau->daughterRefVector();
0300         if (!daus.empty())
0301           getDaughterKeys(daIndxs, daNewIndxs, daus);
0302       }
0303     }
0304   }
0305 }
0306
0307 void GenParticlePruner::getMotherKeys(vector<size_t> &moIndxs,
0308                                       vector<size_t> &moNewIndxs,
0309                                       const GenParticleRefVector &mothers) const {
0310   for (GenParticleRefVector::const_iterator j = mothers.begin(); j != mothers.end(); ++j) {
0311     GenParticleRef mom = *j;
0312     if (find(moIndxs.begin(), moIndxs.end(), mom.key()) == moIndxs.end()) {
0313       moIndxs.push_back(mom.key());
0314       int idx = flags_[mom.key()];
0315       if (idx >= 0) {
0316         moNewIndxs.push_back(idx);
0317       } else {
0318         const GenParticleRefVector &moms = mom->motherRefVector();
0319         if (!moms.empty())
0320           getMotherKeys(moIndxs, moNewIndxs, moms);
0321       }
0322     }
0323   }
0324 }
0325
0326 #include "FWCore/Framework/interface/MakerMacros.h"
0327
0328 DEFINE_FWK_MODULE(GenParticlePruner);