HepMCCandAlgos/plugins/ParticleTreeDrawer.cc

0001 /* class ParticleTreeDrawer
0002  *
0003  * \author Luca Lista, INFN
0004  */
0005 #include <sstream>
0006
0007 #include "FWCore/Framework/interface/one/EDAnalyzer.h"
0008 #include "FWCore/Utilities/interface/InputTag.h"
0009 #include "FWCore/Framework/interface/ESHandle.h"
0010 #include "FWCore/Utilities/interface/ESGetToken.h"
0011 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0012 #include "DataFormats/Common/interface/View.h"
0013 #include "DataFormats/Candidate/interface/Candidate.h"
0014
0015 class ParticleTreeDrawer : public edm::one::EDAnalyzer<> {
0016 public:
0017   ParticleTreeDrawer(const edm::ParameterSet &);
0018
0019 private:
0020   std::string getParticleName(int id) const;
0021   void analyze(const edm::Event &, const edm::EventSetup &) override;
0022   edm::EDGetTokenT<edm::View<reco::Candidate> > srcToken_;
0023   void printDecay(const reco::Candidate &, const std::string &pre) const;
0024   edm::ESHandle<ParticleDataTable> pdt_;
0025   edm::ESGetToken<ParticleDataTable, edm::DefaultRecord> pdtToken_;
0026   /// print parameters
0027   bool printP4_, printPtEtaPhi_, printVertex_, printStatus_, printIndex_;
0028   /// accepted status codes
0029   typedef std::vector<int> vint;
0030   vint status_;
0031   /// print 4 momenta
0032   void printInfo(const reco::Candidate &) const;
0033   /// accept candidate
0034   bool accept(const reco::Candidate &) const;
0035   /// has valid daughters in the chain
0036   bool hasValidDaughters(const reco::Candidate &) const;
0037   /// pointer to collection
0038   std::vector<const reco::Candidate *> cands_;
0039 };
0040
0041 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0042 #include "FWCore/Framework/interface/Event.h"
0043 #include "DataFormats/Common/interface/Handle.h"
0044 #include "FWCore/Framework/interface/EventSetup.h"
0045 #include "FWCore/Utilities/interface/EDMException.h"
0046 #include <iostream>
0047 #include <algorithm>
0048 using namespace std;
0049 using namespace edm;
0050 using namespace reco;
0051
0052 ParticleTreeDrawer::ParticleTreeDrawer(const ParameterSet &cfg)
0053     : srcToken_(consumes<edm::View<reco::Candidate> >(cfg.getParameter<InputTag>("src"))),
0054       pdtToken_(esConsumes<ParticleDataTable, edm::DefaultRecord>()),
0055       printP4_(cfg.getUntrackedParameter<bool>("printP4", false)),
0056       printPtEtaPhi_(cfg.getUntrackedParameter<bool>("printPtEtaPhi", false)),
0057       printVertex_(cfg.getUntrackedParameter<bool>("printVertex", false)),
0058       printStatus_(cfg.getUntrackedParameter<bool>("printStatus", false)),
0059       printIndex_(cfg.getUntrackedParameter<bool>("printIndex", false)),
0060       status_(cfg.getUntrackedParameter<vint>("status", vint())) {}
0061
0062 bool ParticleTreeDrawer::accept(const reco::Candidate &c) const {
0063   if (status_.empty())
0064     return true;
0065   return find(status_.begin(), status_.end(), c.status()) != status_.end();
0066 }
0067
0068 bool ParticleTreeDrawer::hasValidDaughters(const reco::Candidate &c) const {
0069   size_t ndau = c.numberOfDaughters();
0070   for (size_t i = 0; i < ndau; ++i)
0071     if (accept(*c.daughter(i)))
0072       return true;
0073   return false;
0074 }
0075
0076 std::string ParticleTreeDrawer::getParticleName(int id) const {
0077   const ParticleData *pd = pdt_->particle(id);
0078   if (!pd) {
0079     std::ostringstream ss;
0080     ss << "P" << id;
0081     return ss.str();
0082   } else
0083     return pd->name();
0084 }
0085
0086 void ParticleTreeDrawer::analyze(const Event &event, const EventSetup &es) {
0087   pdt_ = es.getHandle(pdtToken_);
0088   Handle<View<Candidate> > particles;
0089   event.getByToken(srcToken_, particles);
0090   cands_.clear();
0091   for (View<Candidate>::const_iterator p = particles->begin(); p != particles->end(); ++p) {
0092     cands_.push_back(&*p);
0093   }
0094   for (View<Candidate>::const_iterator p = particles->begin(); p != particles->end(); ++p) {
0095     if (accept(*p)) {
0096       if (p->mother() == nullptr) {
0097         cout << "-- decay tree: --" << endl;
0098         printDecay(*p, "");
0099       }
0100     }
0101   }
0102 }
0103
0104 void ParticleTreeDrawer::printInfo(const Candidate &c) const {
0105   if (printP4_)
0106     cout << " (" << c.px() << ", " << c.py() << ", " << c.pz() << "; " << c.energy() << ")";
0107   if (printPtEtaPhi_)
0108     cout << " [" << c.pt() << ": " << c.eta() << ", " << c.phi() << "]";
0109   if (printVertex_)
0110     cout << " {" << c.vx() << ", " << c.vy() << ", " << c.vz() << "}";
0111   if (printStatus_)
0112     cout << "{status: " << c.status() << "}";
0113   if (printIndex_) {
0114     int idx = -1;
0115     vector<const Candidate *>::const_iterator found = find(cands_.begin(), cands_.end(), &c);
0116     if (found != cands_.end()) {
0117       idx = found - cands_.begin();
0118       cout << " <idx: " << idx << ">";
0119     }
0120   }
0121 }
0122
0123 void ParticleTreeDrawer::printDecay(const Candidate &c, const string &pre) const {
0124   cout << getParticleName(c.pdgId());
0125   printInfo(c);
0126   cout << endl;
0127
0128   size_t ndau = c.numberOfDaughters(), validDau = 0;
0129   for (size_t i = 0; i < ndau; ++i)
0130     if (accept(*c.daughter(i)))
0131       ++validDau;
0132   if (validDau == 0)
0133     return;
0134
0135   bool lastLevel = true;
0136   for (size_t i = 0; i < ndau; ++i) {
0137     if (hasValidDaughters(*c.daughter(i))) {
0138       lastLevel = false;
0139       break;
0140     }
0141   }
0142
0143   if (lastLevel) {
0144     cout << pre << "+-> ";
0145     size_t vd = 0;
0146     for (size_t i = 0; i < ndau; ++i) {
0147       const Candidate *d = c.daughter(i);
0148       if (accept(*d)) {
0149         cout << getParticleName(d->pdgId());
0150         printInfo(*d);
0151         if (vd != validDau - 1)
0152           cout << " ";
0153         vd++;
0154       }
0155     }
0156     cout << endl;
0157     return;
0158   }
0159
0160   for (size_t i = 0; i < ndau; ++i) {
0161     const Candidate *d = c.daughter(i);
0162     assert(d != nullptr);
0163     if (accept(*d)) {
0164       cout << pre << "+-> ";
0165       string prepre(pre);
0166       if (i == ndau - 1)
0167         prepre += "    ";
0168       else
0169         prepre += "|   ";
0170       printDecay(*d, prepre);
0171     }
0172   }
0173 }
0174
0175 #include "FWCore/Framework/interface/MakerMacros.h"
0176
0177 DEFINE_FWK_MODULE(ParticleTreeDrawer);