HepMCCandAlgos/plugins/GenParticleProducer.cc

0001 /* \class GenParticleProducer
0002  *
0003  * \author Luca Lista, INFN
0004  *
0005  * Convert HepMC GenEvent format into a collection of type
0006  * CandidateCollection containing objects of type GenParticle
0007  *
0008  *
0009  */
0010 #include "FWCore/Framework/interface/global/EDProducer.h"
0011 #include "FWCore/Utilities/interface/InputTag.h"
0012 #include "FWCore/Utilities/interface/ESGetToken.h"
0013 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0014 #include "DataFormats/Candidate/interface/CandidateFwd.h"
0015 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0016 #include "SimDataFormats/CrossingFrame/interface/CrossingFrame.h"
0017 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0018 #include "PhysicsTools/HepMCCandAlgos/interface/MCTruthHelper.h"
0019 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0020
0021 #include <vector>
0022 #include <string>
0023 #include <unordered_map>
0024
0025 namespace edm {
0026   class ParameterSet;
0027 }
0028 namespace HepMC {
0029   class GenParticle;
0030   class GenEvent;
0031 }  // namespace HepMC
0032
0033 static constexpr int PDGCacheMax = 32768;
0034
0035 namespace {
0036   struct IDto3Charge {
0037     IDto3Charge(HepPDT::ParticleDataTable const&, bool abortOnUnknownPDGCode);
0038
0039     int chargeTimesThree(int) const;
0040
0041   private:
0042     std::vector<int> chargeP_, chargeM_;
0043     std::unordered_map<int, int> chargeMap_;
0044     bool abortOnUnknownPDGCode_;
0045   };
0046
0047   IDto3Charge::IDto3Charge(HepPDT::ParticleDataTable const& iTable, bool iAbortOnUnknownPDGCode)
0048       : chargeP_(PDGCacheMax, 0), chargeM_(PDGCacheMax, 0), abortOnUnknownPDGCode_(iAbortOnUnknownPDGCode) {
0049     for (auto const& p : iTable) {
0050       const HepPDT::ParticleID& id = p.first;
0051       int pdgId = id.pid(), apdgId = std::abs(pdgId);
0052       int q3 = id.threeCharge();
0053       if (apdgId < PDGCacheMax && pdgId > 0) {
0054         chargeP_[apdgId] = q3;
0055         chargeM_[apdgId] = -q3;
0056       } else if (apdgId < PDGCacheMax) {
0057         chargeP_[apdgId] = -q3;
0058         chargeM_[apdgId] = q3;
0059       } else {
0060         chargeMap_.emplace(pdgId, q3);
0061         chargeMap_.emplace(-pdgId, -q3);
0062       }
0063     }
0064   }
0065
0066   int IDto3Charge::chargeTimesThree(int id) const {
0067     if (std::abs(id) < PDGCacheMax)
0068       return id > 0 ? chargeP_[id] : chargeM_[-id];
0069     auto f = chargeMap_.find(id);
0070     if (f == chargeMap_.end()) {
0071       if (abortOnUnknownPDGCode_)
0072         throw edm::Exception(edm::errors::LogicError) << "invalid PDG id: " << id;
0073       else
0074         return HepPDT::ParticleID(id).threeCharge();
0075     }
0076     return f->second;
0077   }
0078
0079 }  // namespace
0080
0081 class GenParticleProducer : public edm::global::EDProducer<edm::RunCache<IDto3Charge>> {
0082 public:
0083   /// constructor
0084   GenParticleProducer(const edm::ParameterSet&);
0085   /// destructor
0086   ~GenParticleProducer() override;
0087
0088   /// process one event
0089   void produce(edm::StreamID, edm::Event& e, const edm::EventSetup&) const override;
0090   std::shared_ptr<IDto3Charge> globalBeginRun(const edm::Run&, const edm::EventSetup&) const override;
0091   void globalEndRun(edm::Run const&, edm::EventSetup const&) const override{};
0092
0093   bool convertParticle(reco::GenParticle& cand, const HepMC::GenParticle* part, const IDto3Charge& id2Charge) const;
0094   bool fillDaughters(reco::GenParticleCollection& cand,
0095                      const HepMC::GenParticle* part,
0096                      reco::GenParticleRefProd const& ref,
0097                      size_t index,
0098                      std::unordered_map<int, size_t>& barcodes) const;
0099   bool fillIndices(const HepMC::GenEvent* mc,
0100                    std::vector<const HepMC::GenParticle*>& particles,
0101                    std::vector<int>& barCodeVector,
0102                    int offset,
0103                    std::unordered_map<int, size_t>& barcodes) const;
0104
0105 private:
0106   /// source collection name
0107   edm::EDGetTokenT<edm::HepMCProduct> srcToken_;
0108   std::vector<edm::EDGetTokenT<edm::HepMCProduct>> vectorSrcTokens_;
0109   edm::EDGetTokenT<CrossingFrame<edm::HepMCProduct>> mixToken_;
0110   edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> particleTableToken_;
0111
0112   /// unknown code treatment flag
0113   bool abortOnUnknownPDGCode_;
0114   /// save bar-codes
0115   bool saveBarCodes_;
0116
0117   /// input & output modes
0118   bool doSubEvent_;
0119   bool useCF_;
0120
0121   MCTruthHelper<HepMC::GenParticle> mcTruthHelper_;
0122   MCTruthHelper<reco::GenParticle> mcTruthHelperGenParts_;
0123 };
0124
0125 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0126
0127 //#include "SimDataFormats/HiGenData/interface/SubEventMap.h"
0128 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0129
0130 #include "DataFormats/Common/interface/Handle.h"
0131 #include "FWCore/Framework/interface/ESHandle.h"
0132 #include "FWCore/Framework/interface/Event.h"
0133 #include "FWCore/Framework/interface/Run.h"
0134 #include "FWCore/Framework/interface/EventSetup.h"
0135 #include "FWCore/Utilities/interface/EDMException.h"
0136 #include "FWCore/Utilities/interface/transform.h"
0137 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0138
0139 #include <fstream>
0140 #include <algorithm>
0141 using namespace edm;
0142 using namespace reco;
0143 using namespace std;
0144 using namespace HepMC;
0145
0146 static const double mmToCm = 0.1;
0147 static const double mmToNs = 1.0 / 299792458e-6;
0148
0149 GenParticleProducer::GenParticleProducer(const ParameterSet& cfg)
0150     : abortOnUnknownPDGCode_(false),
0151       saveBarCodes_(cfg.getUntrackedParameter<bool>("saveBarCodes", false)),
0152       doSubEvent_(cfg.getUntrackedParameter<bool>("doSubEvent", false)),
0153       useCF_(cfg.getUntrackedParameter<bool>("useCrossingFrame", false)) {
0154   particleTableToken_ = esConsumes<HepPDT::ParticleDataTable, edm::DefaultRecord, edm::Transition::BeginRun>();
0155   produces<GenParticleCollection>();
0156   produces<math::XYZPointF>("xyz0");
0157   produces<float>("t0");
0158   if (saveBarCodes_) {
0159     std::string alias(cfg.getParameter<std::string>("@module_label"));
0160     produces<vector<int>>().setBranchAlias(alias + "BarCodes");
0161   }
0162
0163   if (useCF_)
0164     mixToken_ =
0165         mayConsume<CrossingFrame<HepMCProduct>>(InputTag(cfg.getParameter<std::string>("mix"), "generatorSmeared"));
0166   else
0167     srcToken_ = mayConsume<HepMCProduct>(cfg.getParameter<InputTag>("src"));
0168 }
0169
0170 GenParticleProducer::~GenParticleProducer() {}
0171
0172 std::shared_ptr<IDto3Charge> GenParticleProducer::globalBeginRun(const Run&, const EventSetup& es) const {
0173   ESHandle<HepPDT::ParticleDataTable> pdt = es.getHandle(particleTableToken_);
0174   return std::make_shared<IDto3Charge>(*pdt, abortOnUnknownPDGCode_);
0175 }
0176
0177 void GenParticleProducer::produce(StreamID, Event& evt, const EventSetup& es) const {
0178   std::unordered_map<int, size_t> barcodes;
0179
0180   size_t totalSize = 0;
0181   const GenEvent* mc = nullptr;
0182   MixCollection<HepMCProduct>* cfhepmcprod = nullptr;
0183   size_t npiles = 1;
0184
0185   if (useCF_) {
0186     Handle<CrossingFrame<HepMCProduct>> cf;
0187     evt.getByToken(mixToken_, cf);
0188     cfhepmcprod = new MixCollection<HepMCProduct>(cf.product());
0189     npiles = cfhepmcprod->size();
0190     LogDebug("GenParticleProducer") << "npiles : " << npiles << endl;
0191     for (unsigned int icf = 0; icf < npiles; ++icf) {
0192       LogDebug("GenParticleProducer") << "subSize : " << cfhepmcprod->getObject(icf).GetEvent()->particles_size()
0193                                       << endl;
0194       totalSize += cfhepmcprod->getObject(icf).GetEvent()->particles_size();
0195     }
0196     LogDebug("GenParticleProducer") << "totalSize : " << totalSize << endl;
0197   } else {
0198     Handle<HepMCProduct> mcp;
0199     evt.getByToken(srcToken_, mcp);
0200     mc = mcp->GetEvent();
0201     if (mc == nullptr)
0202       throw edm::Exception(edm::errors::InvalidReference) << "HepMC has null pointer to GenEvent" << endl;
0203     totalSize = mc->particles_size();
0204   }
0205
0206   // initialise containers
0207   const size_t size = totalSize;
0208   vector<const HepMC::GenParticle*> particles(size);
0209   auto candsPtr = std::make_unique<GenParticleCollection>(size);
0210   auto barCodeVector = std::make_unique<vector<int>>(size);
0211   std::unique_ptr<math::XYZPointF> xyz0Ptr(new math::XYZPointF(0., 0., 0.));
0212   std::unique_ptr<float> t0Ptr(new float(0.f));
0213   reco::GenParticleRefProd ref = evt.getRefBeforePut<GenParticleCollection>();
0214   GenParticleCollection& cands = *candsPtr;
0215   size_t offset = 0;
0216   size_t suboffset = 0;
0217
0218   IDto3Charge const& id2Charge = *runCache(evt.getRun().index());
0219   /// fill indices
0220   if (doSubEvent_ || useCF_) {
0221     for (size_t ipile = 0; ipile < npiles; ++ipile) {
0222       LogDebug("GenParticleProducer") << "mixed object ipile : " << ipile << endl;
0223       barcodes.clear();
0224       if (useCF_)
0225         mc = cfhepmcprod->getObject(ipile).GetEvent();
0226
0227       //Look whether heavy ion/signal event
0228       bool isHI = false;
0229       const HepMC::HeavyIon* hi = mc->heavy_ion();
0230       if (hi && hi->Ncoll_hard() > 1)
0231         isHI = true;
0232       size_t num_particles = mc->particles_size();
0233       LogDebug("GenParticleProducer") << "num_particles : " << num_particles << endl;
0234       if (ipile == 0) {
0235         auto origin = (*mc->vertices_begin())->position();
0236         xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0237         *t0Ptr = origin.t() * mmToNs;
0238       }
0239       fillIndices(mc, particles, *barCodeVector, offset, barcodes);
0240       // fill output collection and save association
0241       for (size_t ipar = offset; ipar < offset + num_particles; ++ipar) {
0242         const HepMC::GenParticle* part = particles[ipar];
0243         reco::GenParticle& cand = cands[ipar];
0244         // convert HepMC::GenParticle to new reco::GenParticle
0245         convertParticle(cand, part, id2Charge);
0246         cand.resetDaughters(ref.id());
0247       }
0248
0249       for (size_t d = offset; d < offset + num_particles; ++d) {
0250         const HepMC::GenParticle* part = particles[d];
0251         const GenVertex* productionVertex = part->production_vertex();
0252         int sub_id = 0;
0253         if (productionVertex != nullptr) {
0254           sub_id = productionVertex->id();
0255           if (!isHI)
0256             sub_id = 0;
0257           // search barcode map and attach daughters
0258           fillDaughters(cands, part, ref, d, barcodes);
0259         } else {
0260           const GenVertex* endVertex = part->end_vertex();
0261           if (endVertex != nullptr)
0262             sub_id = endVertex->id();
0263           else
0264             throw cms::Exception("SubEventID")
0265                 << "SubEvent not determined. Particle has no production and no end vertex!" << endl;
0266         }
0267         if (sub_id < 0)
0268           sub_id = 0;
0269         int new_id = sub_id + suboffset;
0270         GenParticleRef dref(ref, d);
0271         cands[d].setCollisionId(new_id);  // For new GenParticle
0272         LogDebug("VertexId") << "SubEvent offset 3 : " << suboffset;
0273       }
0274       int nsub = -2;
0275       if (isHI) {
0276         nsub = hi->Ncoll_hard() + 1;
0277         suboffset += nsub;
0278       } else {
0279         suboffset += 1;
0280       }
0281       offset += num_particles;
0282     }
0283   } else {
0284     auto origin = (*mc->vertices_begin())->position();
0285     xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0286     *t0Ptr = origin.t() * mmToNs;
0287     fillIndices(mc, particles, *barCodeVector, 0, barcodes);
0288
0289     // fill output collection and save association
0290     for (size_t i = 0; i < particles.size(); ++i) {
0291       const HepMC::GenParticle* part = particles[i];
0292       reco::GenParticle& cand = cands[i];
0293       // convert HepMC::GenParticle to new reco::GenParticle
0294       convertParticle(cand, part, id2Charge);
0295       cand.resetDaughters(ref.id());
0296     }
0297
0298     // fill references to daughters
0299     for (size_t d = 0; d < cands.size(); ++d) {
0300       const HepMC::GenParticle* part = particles[d];
0301       const GenVertex* productionVertex = part->production_vertex();
0302       // search barcode map and attach daughters
0303       if (productionVertex != nullptr)
0304         fillDaughters(cands, part, ref, d, barcodes);
0305       cands[d].setCollisionId(0);
0306     }
0307   }
0308
0309   evt.put(std::move(candsPtr));
0310   if (saveBarCodes_)
0311     evt.put(std::move(barCodeVector));
0312   if (cfhepmcprod)
0313     delete cfhepmcprod;
0314   evt.put(std::move(xyz0Ptr), "xyz0");
0315   evt.put(std::move(t0Ptr), "t0");
0316 }
0317
0318 bool GenParticleProducer::convertParticle(reco::GenParticle& cand,
0319                                           const HepMC::GenParticle* part,
0320                                           IDto3Charge const& id2Charge) const {
0321   Candidate::LorentzVector p4(part->momentum());
0322   int pdgId = part->pdg_id();
0323   cand.setThreeCharge(id2Charge.chargeTimesThree(pdgId));
0324   cand.setPdgId(pdgId);
0325   cand.setStatus(part->status());
0326   cand.setP4(p4);
0327   cand.setCollisionId(0);
0328   const GenVertex* v = part->production_vertex();
0329   if (v != nullptr) {
0330     ThreeVector vtx = v->point3d();
0331     Candidate::Point vertex(vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm);
0332     cand.setVertex(vertex);
0333   } else {
0334     cand.setVertex(Candidate::Point(0, 0, 0));
0335   }
0336   mcTruthHelper_.fillGenStatusFlags(*part, cand.statusFlags());
0337   return true;
0338 }
0339
0340 bool GenParticleProducer::fillDaughters(reco::GenParticleCollection& cands,
0341                                         const HepMC::GenParticle* part,
0342                                         reco::GenParticleRefProd const& ref,
0343                                         size_t index,
0344                                         std::unordered_map<int, size_t>& barcodes) const {
0345   const GenVertex* productionVertex = part->production_vertex();
0346   size_t numberOfMothers = productionVertex->particles_in_size();
0347   if (numberOfMothers > 0) {
0348     GenVertex::particles_in_const_iterator motherIt = productionVertex->particles_in_const_begin();
0349     for (; motherIt != productionVertex->particles_in_const_end(); motherIt++) {
0350       const HepMC::GenParticle* mother = *motherIt;
0351       size_t m = barcodes.find(mother->barcode())->second;
0352       cands[m].addDaughter(GenParticleRef(ref, index));
0353       cands[index].addMother(GenParticleRef(ref, m));
0354     }
0355   }
0356
0357   return true;
0358 }
0359
0360 bool GenParticleProducer::fillIndices(const HepMC::GenEvent* mc,
0361                                       vector<const HepMC::GenParticle*>& particles,
0362                                       vector<int>& barCodeVector,
0363                                       int offset,
0364                                       std::unordered_map<int, size_t>& barcodes) const {
0365   size_t idx = offset;
0366   HepMC::GenEvent::particle_const_iterator begin = mc->particles_begin(), end = mc->particles_end();
0367   for (HepMC::GenEvent::particle_const_iterator p = begin; p != end; ++p) {
0368     const HepMC::GenParticle* particle = *p;
0369     size_t barCode_this_event = particle->barcode();
0370     size_t barCode = barCode_this_event + offset;
0371     if (barcodes.find(barCode) != barcodes.end())
0372       throw cms::Exception("WrongReference") << "barcodes are duplicated! " << endl;
0373     particles[idx] = particle;
0374     barCodeVector[idx] = barCode;
0375     barcodes.insert(make_pair(barCode_this_event, idx++));
0376   }
0377   return true;
0378 }
0379
0380 #include "FWCore/Framework/interface/MakerMacros.h"
0381
0382 DEFINE_FWK_MODULE(GenParticleProducer);