PileupInformation/plugins/GenPUProtonProducer.cc

0001 /* \class GenPUProtonProducer
0002  *
0003  * Modification of GenParticleProducer.
0004  * Saves final state protons from HepMC events in Crossing Frame, in the generator-particle format.
0005  *
0006  * Note: Use the option USER_CXXFLAGS=-DEDM_ML_DEBUG with SCRAM in order to enable debug messages.
0007  *
0008  *    March  9, 2017   : Initial version.
0009  *    March 14, 2017   : Updated debug messages.
0010  *     July 27, 2017   : Removed extra loop initially inherited from GenParticleProducer.
0011  *   August 17, 2017   : Replaced std::auto_ptr with std::unique_ptr.
0012  * September 6, 2017   : Updated module to edm::global::EDProducer with ConvertParticle as RunCache following GenParticleProducer.
0013  *
0014  */
0015
0016 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0017 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0018 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0019 #include "FWCore/Utilities/interface/EDMException.h"
0020 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0021
0022 #include <vector>
0023 #include <string>
0024 #include <unordered_map>
0025
0026 namespace {
0027
0028   class ConvertParticle {
0029   public:
0030     static constexpr int PDGCacheMax = 32768;
0031     static constexpr double mmToCm = 0.1;
0032
0033     ConvertParticle()
0034         : abortOnUnknownPDGCode_(true), initialized_(false), chargeP_(PDGCacheMax, 0), chargeM_(PDGCacheMax, 0) {}
0035
0036     ConvertParticle(bool abortOnUnknownPDGCode)
0037         : abortOnUnknownPDGCode_(abortOnUnknownPDGCode),
0038           initialized_(false),
0039           chargeP_(PDGCacheMax, 0),
0040           chargeM_(PDGCacheMax, 0) {}
0041
0042     ~ConvertParticle() {}
0043
0044     bool initialized() const { return initialized_; }
0045
0046     void init(HepPDT::ParticleDataTable const& pdt) {
0047       if (!initialized_) {
0048         for (HepPDT::ParticleDataTable::const_iterator p = pdt.begin(); p != pdt.end(); ++p) {
0049           HepPDT::ParticleID const& id = p->first;
0050           int pdgId = id.pid(), apdgId = std::abs(pdgId);
0051           int q3 = id.threeCharge();
0052           if (apdgId < PDGCacheMax && pdgId > 0) {
0053             chargeP_[apdgId] = q3;
0054             chargeM_[apdgId] = -q3;
0055           } else if (apdgId < PDGCacheMax) {
0056             chargeP_[apdgId] = -q3;
0057             chargeM_[apdgId] = q3;
0058           } else {
0059             chargeMap_.emplace(pdgId, q3);
0060             chargeMap_.emplace(-pdgId, -q3);
0061           }
0062         }
0063         initialized_ = true;
0064       }
0065     }
0066
0067     bool operator()(reco::GenParticle& cand, HepMC::GenParticle const* part) const {
0068       reco::Candidate::LorentzVector p4(part->momentum());
0069       int pdgId = part->pdg_id();
0070       cand.setThreeCharge(chargeTimesThree(pdgId));
0071       cand.setPdgId(pdgId);
0072       cand.setStatus(part->status());
0073       cand.setP4(p4);
0074       cand.setCollisionId(0);
0075       HepMC::GenVertex const* v = part->production_vertex();
0076       if (v != nullptr) {
0077         HepMC::ThreeVector vtx = v->point3d();
0078         reco::Candidate::Point vertex(vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm);
0079         cand.setVertex(vertex);
0080       } else {
0081         cand.setVertex(reco::Candidate::Point(0, 0, 0));
0082       }
0083       return true;
0084     }
0085
0086   private:
0087     bool abortOnUnknownPDGCode_;
0088     bool initialized_;
0089     std::vector<int> chargeP_, chargeM_;
0090     std::unordered_map<int, int> chargeMap_;
0091
0092     int chargeTimesThree(int id) const {
0093       if (std::abs(id) < PDGCacheMax)
0094         return id > 0 ? chargeP_[id] : chargeM_[-id];
0095
0096       auto f = chargeMap_.find(id);
0097       if (f == chargeMap_.end()) {
0098         if (abortOnUnknownPDGCode_)
0099           throw edm::Exception(edm::errors::LogicError) << "invalid PDG id: " << id << std::endl;
0100         else
0101           return HepPDT::ParticleID(id).threeCharge();
0102       }
0103       return f->second;
0104     }
0105   };
0106
0107   class SelectProton {
0108   public:
0109     bool operator()(HepMC::GenParticle const* part, double minPz) const {
0110       bool selection = ((!part->end_vertex() && part->status() == 1) && (part->pdg_id() == 2212) &&
0111                         (TMath::Abs(part->momentum().pz()) >= minPz));
0112       return selection;
0113     }
0114   };
0115
0116 }  // Anonymous namespace
0117
0118 #include "FWCore/Framework/interface/global/EDProducer.h"
0119 #include "FWCore/Utilities/interface/InputTag.h"
0120 #include "FWCore/Utilities/interface/ESGetToken.h"
0121 #include "DataFormats/Candidate/interface/CandidateFwd.h"
0122 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0123
0124 #include "SimDataFormats/CrossingFrame/interface/CrossingFrame.h"
0125 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0126 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0127 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0128 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0129
0130 namespace edm {
0131   class ParameterSet;
0132 }
0133
0134 class GenPUProtonProducer : public edm::global::EDProducer<edm::RunCache<ConvertParticle> > {
0135 public:
0136   GenPUProtonProducer(const edm::ParameterSet&);
0137   ~GenPUProtonProducer() override;
0138
0139   void produce(edm::StreamID, edm::Event& e, const edm::EventSetup&) const override;
0140   std::shared_ptr<ConvertParticle> globalBeginRun(const edm::Run&, const edm::EventSetup&) const override;
0141   void globalEndRun(edm::Run const&, edm::EventSetup const&) const override {}
0142
0143 private:
0144   edm::EDGetTokenT<CrossingFrame<edm::HepMCProduct> > mixToken_;
0145   edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> pdtToken_;
0146   bool abortOnUnknownPDGCode_;
0147   std::vector<int> bunchList_;
0148   double minPz_;
0149   SelectProton select_;
0150 };
0151
0152 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0153 #include "DataFormats/Common/interface/Handle.h"
0154 #include "FWCore/Framework/interface/ESHandle.h"
0155 #include "FWCore/Framework/interface/Event.h"
0156 #include "FWCore/Framework/interface/Run.h"
0157 #include "FWCore/Framework/interface/EventSetup.h"
0158 #include "FWCore/Utilities/interface/EDMException.h"
0159
0160 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0161 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0162
0163 #include <algorithm>
0164
0165 using namespace edm;
0166 using namespace reco;
0167 using namespace std;
0168 using namespace HepMC;
0169
0170 GenPUProtonProducer::GenPUProtonProducer(const ParameterSet& cfg)
0171     : abortOnUnknownPDGCode_(cfg.getUntrackedParameter<bool>("abortOnUnknownPDGCode", true)),
0172       bunchList_(cfg.getParameter<vector<int> >("bunchCrossingList")),
0173       minPz_(cfg.getParameter<double>("minPz")) {
0174   produces<GenParticleCollection>();
0175   mixToken_ =
0176       consumes<CrossingFrame<HepMCProduct> >(InputTag(cfg.getParameter<std::string>("mix"), "generatorSmeared"));
0177   pdtToken_ = esConsumes<HepPDT::ParticleDataTable, edm::DefaultRecord, edm::Transition::BeginRun>();
0178 }
0179
0180 GenPUProtonProducer::~GenPUProtonProducer() {}
0181
0182 std::shared_ptr<ConvertParticle> GenPUProtonProducer::globalBeginRun(const Run&, const EventSetup& es) const {
0183   ESHandle<HepPDT::ParticleDataTable> pdt = es.getHandle(pdtToken_);
0184   auto convert_ptr = std::make_shared<ConvertParticle>(abortOnUnknownPDGCode_);
0185   if (!convert_ptr->initialized())
0186     convert_ptr->init(*pdt);
0187
0188   return convert_ptr;
0189 }
0190
0191 void GenPUProtonProducer::produce(StreamID, Event& evt, const EventSetup& es) const {
0192   size_t totalSize = 0;
0193   size_t npiles = 1;
0194
0195   Handle<CrossingFrame<HepMCProduct> > cf;
0196   evt.getByToken(mixToken_, cf);
0197   std::unique_ptr<MixCollection<HepMCProduct> > cfhepmcprod(new MixCollection<HepMCProduct>(cf.product()));
0198   npiles = cfhepmcprod->size();
0199
0200   LogDebug("GenPUProtonProducer") << " Number of pile-up events : " << npiles << endl;
0201
0202   for (size_t icf = 0; icf < npiles; ++icf) {
0203     LogDebug("GenPUProtonProducer") << "CF " << icf
0204                                     << " size : " << cfhepmcprod->getObject(icf).GetEvent()->particles_size() << endl;
0205     totalSize += cfhepmcprod->getObject(icf).GetEvent()->particles_size();
0206   }
0207   LogDebug("GenPUProtonProducer") << "Total size : " << totalSize << endl;
0208
0209   // Initialise containers
0210   auto candsPtr = std::make_unique<GenParticleCollection>();
0211   GenParticleCollection& cands = *candsPtr;
0212
0213   // Loop over pile-up events
0214   ConvertParticle const& convertParticle_ = *runCache(evt.getRun().index());
0215
0216   MixCollection<HepMCProduct>::MixItr mixHepMC_itr;
0217   unsigned int total_number_of_protons = 0;
0218   size_t idx_mix = 0;
0219   // Fill collection
0220   for (mixHepMC_itr = cfhepmcprod->begin(); mixHepMC_itr != cfhepmcprod->end(); ++mixHepMC_itr, ++idx_mix) {
0221     int bunch = mixHepMC_itr.bunch();
0222
0223     if (find(bunchList_.begin(), bunchList_.end(), bunch) != bunchList_.end()) {
0224       auto event = (*mixHepMC_itr).GetEvent();
0225
0226       // Fill output collection
0227       unsigned int number_of_protons = 0;
0228       for (auto p = event->particles_begin(); p != event->particles_end(); ++p) {
0229         HepMC::GenParticle const* part = *p;
0230         if (select_(part, minPz_)) {
0231           reco::GenParticle cand;
0232           convertParticle_(cand, part);
0233           ++number_of_protons;
0234           cands.push_back(cand);
0235         }
0236       }
0237       LogDebug("GenPUProtonProducer") << "Idx : " << idx_mix << " Bunch : " << bunch
0238                                       << " Number of particles : " << event->particles_size()
0239                                       << " Number of protons : " << number_of_protons << endl;
0240
0241       total_number_of_protons += number_of_protons;
0242     }
0243   }
0244   LogDebug("GenPUProtonProducer") << "Total number of protons : " << total_number_of_protons << endl;
0245   LogDebug("GenPUProtonProducer") << "Output collection size : " << cands.size() << endl;
0246
0247   evt.put(std::move(candsPtr));
0248 }
0249
0250 #include "FWCore/Framework/interface/MakerMacros.h"
0251
0252 DEFINE_FWK_MODULE(GenPUProtonProducer);