HepMCCandAlgos/plugins/GenParticleDecaySelector.cc

0001 /* \class GenParticleDecaySelector
0002  *
0003  * \author Luca Lista, INFN
0004  *
0005  *
0006  */
0007
0008 #include "FWCore/Framework/interface/stream/EDProducer.h"
0009 #include "FWCore/Framework/interface/EventSetup.h"
0010 #include "FWCore/Utilities/interface/InputTag.h"
0011 #include "FWCore/Utilities/interface/ESGetToken.h"
0012 #include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
0013 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0014 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0015
0016 class GenParticleDecaySelector : public edm::stream::EDProducer<> {
0017 public:
0018   /// constructor
0019   GenParticleDecaySelector(const edm::ParameterSet&);
0020
0021   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0022
0023 private:
0024   /// process one event
0025   void produce(edm::Event& e, const edm::EventSetup&) override;
0026   bool firstEvent_;
0027   /// source collection name
0028   edm::EDGetTokenT<reco::GenParticleCollection> srcToken_;
0029   /// particle type
0030   edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> tableToken_;
0031   PdtEntry particle_;
0032   /// particle status
0033   int status_;
0034   /// recursively add a new particle to the output collection
0035   std::pair<reco::GenParticleRef, reco::GenParticle*> add(reco::GenParticleCollection&,
0036                                                           const reco::GenParticle&,
0037                                                           reco::GenParticleRefProd);
0038 };
0039
0040 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0041 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0042 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0043 #include "DataFormats/Common/interface/Handle.h"
0044 #include "FWCore/Framework/interface/Event.h"
0045 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0046
0047 using namespace edm;
0048 using namespace reco;
0049 using namespace std;
0050
0051 GenParticleDecaySelector::GenParticleDecaySelector(const edm::ParameterSet& cfg)
0052     : firstEvent_(true),
0053       srcToken_(consumes<GenParticleCollection>(cfg.getParameter<InputTag>("src"))),
0054       tableToken_(esConsumes()),
0055       particle_(cfg.getParameter<PdtEntry>("particle")),
0056       status_(cfg.getParameter<int>("status")) {
0057   produces<GenParticleCollection>();
0058 }
0059
0060 void GenParticleDecaySelector::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0061   edm::ParameterSetDescription desc;
0062   desc.add<edm::InputTag>("src");
0063   desc.addNode(edm::ParameterDescription<int>("particle", true) xor
0064                edm::ParameterDescription<std::string>("particle", true))
0065       ->setComment("the PdtEntry can be specified as either an 'int' or via its name using a 'string'");
0066   desc.add<int>("status");
0067
0068   descriptions.addDefault(desc);
0069 }
0070
0071 void GenParticleDecaySelector::produce(edm::Event& evt, const edm::EventSetup& es) {
0072   if (firstEvent_) {
0073     auto const& pdt = es.getData(tableToken_);
0074     particle_.setup(pdt);
0075     firstEvent_ = false;
0076   }
0077
0078   Handle<GenParticleCollection> genParticles;
0079   evt.getByToken(srcToken_, genParticles);
0080   auto decay = std::make_unique<GenParticleCollection>();
0081   const GenParticleRefProd ref = evt.getRefBeforePut<GenParticleCollection>();
0082   for (GenParticleCollection::const_iterator g = genParticles->begin(); g != genParticles->end(); ++g)
0083     if (g->pdgId() == particle_.pdgId() && g->status() == status_)
0084       add(*decay, *g, ref);
0085   evt.put(std::move(decay));
0086 }
0087
0088 pair<GenParticleRef, GenParticle*> GenParticleDecaySelector::add(GenParticleCollection& decay,
0089                                                                  const GenParticle& p,
0090                                                                  GenParticleRefProd ref) {
0091   size_t idx = decay.size();
0092   GenParticleRef r(ref, idx);
0093   decay.resize(idx + 1);
0094   const LeafCandidate& part = p;
0095   GenParticle g(part);
0096   size_t n = p.numberOfDaughters();
0097   for (size_t i = 0; i < n; ++i) {
0098     pair<GenParticleRef, GenParticle*> d = add(decay, *p.daughterRef(i), ref);
0099     d.second->addMother(r);
0100     g.addDaughter(d.first);
0101   }
0102   GenParticle& gp = decay[idx] = g;
0103   return make_pair(r, &gp);
0104 }
0105
0106 #include "FWCore/Framework/interface/MakerMacros.h"
0107
0108 DEFINE_FWK_MODULE(GenParticleDecaySelector);