Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2024-04-06 12:31:00 

0001 // -*- C++ -*-
0002 //
0003 // Package:    SimGeneral/TrackingParticleSelectorByGen
0004 // Class:      TrackingParticleSelectorByGen
0005 //
0006 /**\class TrackingParticleSelectorByGen TrackingParticleSelectorByGen.cc SimGeneral/TrackingParticleSelectorByGen/plugins/TrackingParticleSelectorByGen.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, 26 Apr 2021 16:44:34 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/stream/EDProducer.h"
0025 
0026 #include "FWCore/Framework/interface/Event.h"
0027 #include "FWCore/Framework/interface/EventSetup.h"
0028 #include "FWCore/Framework/interface/MakerMacros.h"
0029 
0030 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0031 #include "FWCore/Utilities/interface/StreamID.h"
0032 
0033 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0034 #include "SimDataFormats/TrackingAnalysis/interface/TrackingParticle.h"
0035 #include "PhysicsTools/HepMCCandAlgos/interface/pdgEntryReplace.h"
0036 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0037 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0038 
0039 namespace helper {
0040   struct SelectCode {
0041     enum KeepOrDrop { kKeep, kDrop };
0042     enum FlagDepth { kNone, kFirst, kAll };
0043     KeepOrDrop keepOrDrop_;
0044     FlagDepth daughtersDepth_, mothersDepth_;
0045     bool all_;
0046   };
0047 }  // namespace helper
0048 
0049 class TrackingParticleSelectorByGen : public edm::stream::EDProducer<> {
0050   // a lot of this is copied from GenParticlePruner
0051   // refactor common parts in a separate class
0052 public:
0053   explicit TrackingParticleSelectorByGen(const edm::ParameterSet &);
0054 
0055   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0056 
0057 private:
0058   void produce(edm::Event &, const edm::EventSetup &) override;
0059   void parse(const std::string &selection, helper::SelectCode &code, std::string &cut) const;
0060   void flagDaughters(const reco::GenParticle &, int);
0061   void flagMothers(const reco::GenParticle &, int);
0062   void recursiveFlagDaughters(size_t, const reco::GenParticleCollection &, int, std::vector<size_t> &);
0063   void recursiveFlagMothers(size_t, const reco::GenParticleCollection &, int, std::vector<size_t> &);
0064   void getDaughterKeys(std::vector<size_t> &, std::vector<size_t> &, const reco::GenParticleRefVector &) const;
0065   void getMotherKeys(std::vector<size_t> &, std::vector<size_t> &, const reco::GenParticleRefVector &) const;
0066 
0067   // ----------member data ---------------------------
0068   edm::EDGetTokenT<TrackingParticleCollection> tpToken_;
0069   edm::EDGetTokenT<reco::GenParticleCollection> gpToken_;
0070   edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
0071   bool firstEvent_;
0072   int keepOrDropAll_;
0073   std::vector<int> flags_;
0074   std::vector<std::string> selection_;
0075   std::vector<std::pair<StringCutObjectSelector<reco::GenParticle>, helper::SelectCode>> select_;
0076 };
0077 
0078 using namespace edm;
0079 using namespace std;
0080 using namespace reco;
0081 
0082 const int keep = 1, drop = -1;
0083 
0084 void TrackingParticleSelectorByGen::parse(const std::string &selection,
0085                                           ::helper::SelectCode &code,
0086                                           std::string &cut) const {
0087   using namespace ::helper;
0088   size_t f = selection.find_first_not_of(' ');
0089   size_t n = selection.size();
0090   string command;
0091   char c;
0092   for (; (c = selection[f]) != ' ' && f < n; ++f) {
0093     command.push_back(c);
0094   }
0095   if (command[0] == '+') {
0096     command.erase(0, 1);
0097     if (command[0] == '+') {
0098       command.erase(0, 1);
0099       code.mothersDepth_ = SelectCode::kAll;
0100     } else {
0101       code.mothersDepth_ = SelectCode::kFirst;
0102     }
0103   } else
0104     code.mothersDepth_ = SelectCode::kNone;
0105 
0106   if (command[command.size() - 1] == '+') {
0107     command.erase(command.size() - 1);
0108     if (command[command.size() - 1] == '+') {
0109       command.erase(command.size() - 1);
0110       code.daughtersDepth_ = SelectCode::kAll;
0111     } else {
0112       code.daughtersDepth_ = SelectCode::kFirst;
0113     }
0114   } else
0115     code.daughtersDepth_ = SelectCode::kNone;
0116 
0117   if (command == "keep")
0118     code.keepOrDrop_ = SelectCode::kKeep;
0119   else if (command == "drop")
0120     code.keepOrDrop_ = SelectCode::kDrop;
0121   else {
0122     throw Exception(errors::Configuration) << "invalid selection command: " << command << "\n" << endl;
0123   }
0124   for (; f < n; ++f) {
0125     if (selection[f] != ' ')
0126       break;
0127   }
0128   cut = string(selection, f);
0129   if (cut[0] == '*')
0130     cut = string(cut, 0, cut.find_first_of(' '));
0131   code.all_ = cut == "*";
0132 }
0133 
0134 TrackingParticleSelectorByGen::TrackingParticleSelectorByGen(const edm::ParameterSet &iConfig)
0135     : tpToken_(consumes<TrackingParticleCollection>(iConfig.getParameter<edm::InputTag>("trackingParticles"))),
0136       gpToken_(consumes<reco::GenParticleCollection>(iConfig.getParameter<edm::InputTag>("genParticles"))),
0137       tableToken_(esConsumes()),
0138       firstEvent_(true),
0139       keepOrDropAll_(drop),
0140       selection_(iConfig.getParameter<vector<string>>("select")) {
0141   produces<TrackingParticleCollection>();
0142 }
0143 
0144 void TrackingParticleSelectorByGen::flagDaughters(const reco::GenParticle &gen, int keepOrDrop) {
0145   const GenParticleRefVector &daughters = gen.daughterRefVector();
0146   for (GenParticleRefVector::const_iterator i = daughters.begin(); i != daughters.end(); ++i)
0147     flags_[i->key()] = keepOrDrop;
0148 }
0149 
0150 void TrackingParticleSelectorByGen::flagMothers(const reco::GenParticle &gen, int keepOrDrop) {
0151   const GenParticleRefVector &mothers = gen.motherRefVector();
0152   for (GenParticleRefVector::const_iterator i = mothers.begin(); i != mothers.end(); ++i)
0153     flags_[i->key()] = keepOrDrop;
0154 }
0155 
0156 void TrackingParticleSelectorByGen::recursiveFlagDaughters(size_t index,
0157                                                            const reco::GenParticleCollection &src,
0158                                                            int keepOrDrop,
0159                                                            std::vector<size_t> &allIndices) {
0160   GenParticleRefVector daughters = src[index].daughterRefVector();
0161   // avoid infinite recursion if the daughters are set to "this" particle.
0162   size_t cachedIndex = index;
0163   for (GenParticleRefVector::const_iterator i = daughters.begin(); i != daughters.end(); ++i) {
0164     index = i->key();
0165     // To also avoid infinite recursion if a "loop" is found in the daughter list,
0166     // check to make sure the index hasn't already been added.
0167     if (find(allIndices.begin(), allIndices.end(), index) == allIndices.end()) {
0168       allIndices.push_back(index);
0169       if (cachedIndex != index) {
0170         flags_[index] = keepOrDrop;
0171         recursiveFlagDaughters(index, src, keepOrDrop, allIndices);
0172       }
0173     }
0174   }
0175 }
0176 
0177 void TrackingParticleSelectorByGen::recursiveFlagMothers(size_t index,
0178                                                          const reco::GenParticleCollection &src,
0179                                                          int keepOrDrop,
0180                                                          std::vector<size_t> &allIndices) {
0181   GenParticleRefVector mothers = src[index].motherRefVector();
0182   // avoid infinite recursion if the mothers are set to "this" particle.
0183   size_t cachedIndex = index;
0184   for (GenParticleRefVector::const_iterator i = mothers.begin(); i != mothers.end(); ++i) {
0185     index = i->key();
0186     // To also avoid infinite recursion if a "loop" is found in the daughter list,
0187     // check to make sure the index hasn't already been added.
0188     if (find(allIndices.begin(), allIndices.end(), index) == allIndices.end()) {
0189       allIndices.push_back(index);
0190       if (cachedIndex != index) {
0191         flags_[index] = keepOrDrop;
0192         recursiveFlagMothers(index, src, keepOrDrop, allIndices);
0193       }
0194     }
0195   }
0196 }
0197 
0198 void TrackingParticleSelectorByGen::getDaughterKeys(vector<size_t> &daIndxs,
0199                                                     vector<size_t> &daNewIndxs,
0200                                                     const GenParticleRefVector &daughters) const {
0201   for (GenParticleRefVector::const_iterator j = daughters.begin(); j != daughters.end(); ++j) {
0202     GenParticleRef dau = *j;
0203     if (find(daIndxs.begin(), daIndxs.end(), dau.key()) == daIndxs.end()) {
0204       daIndxs.push_back(dau.key());
0205       int idx = flags_[dau.key()];
0206       if (idx > 0) {
0207         daNewIndxs.push_back(idx);
0208       } else {
0209         const GenParticleRefVector &daus = dau->daughterRefVector();
0210         if (!daus.empty())
0211           getDaughterKeys(daIndxs, daNewIndxs, daus);
0212       }
0213     }
0214   }
0215 }
0216 
0217 void TrackingParticleSelectorByGen::getMotherKeys(vector<size_t> &moIndxs,
0218                                                   vector<size_t> &moNewIndxs,
0219                                                   const GenParticleRefVector &mothers) const {
0220   for (GenParticleRefVector::const_iterator j = mothers.begin(); j != mothers.end(); ++j) {
0221     GenParticleRef mom = *j;
0222     if (find(moIndxs.begin(), moIndxs.end(), mom.key()) == moIndxs.end()) {
0223       moIndxs.push_back(mom.key());
0224       int idx = flags_[mom.key()];
0225       if (idx >= 0) {
0226         moNewIndxs.push_back(idx);
0227       } else {
0228         const GenParticleRefVector &moms = mom->motherRefVector();
0229         if (!moms.empty())
0230           getMotherKeys(moIndxs, moNewIndxs, moms);
0231       }
0232     }
0233   }
0234 }
0235 
0236 void TrackingParticleSelectorByGen::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
0237   using namespace edm;
0238 
0239   if (firstEvent_) {
0240     auto const &pdt = iSetup.getData(tableToken_);
0241     for (vector<string>::const_iterator i = selection_.begin(); i != selection_.end(); ++i) {
0242       string cut;
0243       ::helper::SelectCode code;
0244       parse(*i, code, cut);
0245       if (code.all_) {
0246         if (i != selection_.begin())
0247           throw Exception(errors::Configuration)
0248               << "selections \"keep *\" and \"drop *\" can be used only as first options. Here used in position # "
0249               << (i - selection_.begin()) + 1 << "\n"
0250               << endl;
0251         switch (code.keepOrDrop_) {
0252           case ::helper::SelectCode::kDrop:
0253             keepOrDropAll_ = drop;
0254             break;
0255           case ::helper::SelectCode::kKeep:
0256             keepOrDropAll_ = keep;
0257         };
0258       } else {
0259         cut = pdgEntryReplace(cut, pdt);
0260         select_.push_back(make_pair(StringCutObjectSelector<GenParticle>(cut), code));
0261       }
0262     }
0263     firstEvent_ = false;
0264   }
0265 
0266   const auto &tps = iEvent.get(tpToken_);
0267 
0268   const auto &gps = iEvent.get(gpToken_);
0269 
0270   using namespace ::helper;
0271   const size_t n = gps.size();
0272   flags_.clear();
0273   flags_.resize(n, keepOrDropAll_);
0274   for (size_t j = 0; j < select_.size(); ++j) {
0275     const pair<StringCutObjectSelector<GenParticle>, SelectCode> &sel = select_[j];
0276     SelectCode code = sel.second;
0277     const StringCutObjectSelector<GenParticle> &cut = sel.first;
0278     for (size_t i = 0; i < n; ++i) {
0279       const GenParticle &p = gps[i];
0280       if (cut(p)) {
0281         int keepOrDrop = keep;
0282         switch (code.keepOrDrop_) {
0283           case SelectCode::kKeep:
0284             keepOrDrop = keep;
0285             break;
0286           case SelectCode::kDrop:
0287             keepOrDrop = drop;
0288         };
0289         flags_[i] = keepOrDrop;
0290         std::vector<size_t> allIndicesDa;
0291         std::vector<size_t> allIndicesMo;
0292         switch (code.daughtersDepth_) {
0293           case SelectCode::kAll:
0294             recursiveFlagDaughters(i, gps, keepOrDrop, allIndicesDa);
0295             break;
0296           case SelectCode::kFirst:
0297             flagDaughters(p, keepOrDrop);
0298             break;
0299           case SelectCode::kNone:;
0300         };
0301         switch (code.mothersDepth_) {
0302           case SelectCode::kAll:
0303             recursiveFlagMothers(i, gps, keepOrDrop, allIndicesMo);
0304             break;
0305           case SelectCode::kFirst:
0306             flagMothers(p, keepOrDrop);
0307             break;
0308           case SelectCode::kNone:;
0309         };
0310       }
0311     }
0312   }
0313 
0314   auto out = std::make_unique<TrackingParticleCollection>();
0315   out->reserve(n);
0316 
0317   for (auto &&tp : tps) {
0318     auto &associatedGenParticles = tp.genParticles();
0319 
0320     bool isSelected = false;
0321     for (auto &&assocGen : associatedGenParticles) {
0322       if (flags_[assocGen.index()] == keep)
0323         isSelected = true;
0324     }
0325 
0326     if (isSelected) {
0327       out->emplace_back(tp);
0328     }
0329   }
0330   iEvent.put(std::move(out));
0331 }
0332 
0333 void TrackingParticleSelectorByGen::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0334   edm::ParameterSetDescription desc;
0335   desc.add<edm::InputTag>("trackingParticles", edm::InputTag("mix", "MergedTrackTruth"));
0336   desc.add<edm::InputTag>("genParticles", edm::InputTag("genParticles"));
0337   desc.add<vector<string>>("select");
0338 
0339   descriptions.add("tpSelectorByGenDefault", desc);
0340 }
0341 
0342 //define this as a plug-in
0343 DEFINE_FWK_MODULE(TrackingParticleSelectorByGen);
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!