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     bool found = evt.getByToken(srcToken_, mcp);
0200     if (found) {
0201       mc = mcp->GetEvent();
0202       if (mc == nullptr)
0203         throw edm::Exception(edm::errors::InvalidReference) << "HepMC has null pointer to GenEvent" << endl;
0204       totalSize = mc->particles_size();
0205     } else {
0206       totalSize = 0;
0207     }
0208   }
0209
0210   // initialise containers
0211   const size_t size = totalSize;
0212   vector<const HepMC::GenParticle*> particles(size);
0213   auto candsPtr = std::make_unique<GenParticleCollection>(size);
0214   auto barCodeVector = std::make_unique<vector<int>>(size);
0215   std::unique_ptr<math::XYZPointF> xyz0Ptr(new math::XYZPointF(0., 0., 0.));
0216   std::unique_ptr<float> t0Ptr(new float(0.f));
0217   reco::GenParticleRefProd ref = evt.getRefBeforePut<GenParticleCollection>();
0218   GenParticleCollection& cands = *candsPtr;
0219   size_t offset = 0;
0220   size_t suboffset = 0;
0221
0222   IDto3Charge const& id2Charge = *runCache(evt.getRun().index());
0223   /// fill indices
0224   if (doSubEvent_ || useCF_) {
0225     for (size_t ipile = 0; ipile < npiles; ++ipile) {
0226       LogDebug("GenParticleProducer") << "mixed object ipile : " << ipile << endl;
0227       barcodes.clear();
0228       if (useCF_)
0229         mc = cfhepmcprod->getObject(ipile).GetEvent();
0230
0231       //Look whether heavy ion/signal event
0232       bool isHI = false;
0233       const HepMC::HeavyIon* hi = mc->heavy_ion();
0234       if (hi && hi->Ncoll_hard() > 1)
0235         isHI = true;
0236       size_t num_particles = mc->particles_size();
0237       LogDebug("GenParticleProducer") << "num_particles : " << num_particles << endl;
0238       if (ipile == 0) {
0239         auto origin = (*mc->vertices_begin())->position();
0240         xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0241         *t0Ptr = origin.t() * mmToNs;
0242       }
0243       fillIndices(mc, particles, *barCodeVector, offset, barcodes);
0244       // fill output collection and save association
0245       for (size_t ipar = offset; ipar < offset + num_particles; ++ipar) {
0246         const HepMC::GenParticle* part = particles[ipar];
0247         reco::GenParticle& cand = cands[ipar];
0248         // convert HepMC::GenParticle to new reco::GenParticle
0249         convertParticle(cand, part, id2Charge);
0250         cand.resetDaughters(ref.id());
0251       }
0252
0253       for (size_t d = offset; d < offset + num_particles; ++d) {
0254         const HepMC::GenParticle* part = particles[d];
0255         const GenVertex* productionVertex = part->production_vertex();
0256         int sub_id = 0;
0257         if (productionVertex != nullptr) {
0258           sub_id = productionVertex->id();
0259           if (!isHI)
0260             sub_id = 0;
0261           // search barcode map and attach daughters
0262           fillDaughters(cands, part, ref, d, barcodes);
0263         } else {
0264           const GenVertex* endVertex = part->end_vertex();
0265           if (endVertex != nullptr)
0266             sub_id = endVertex->id();
0267           else
0268             throw cms::Exception("SubEventID")
0269                 << "SubEvent not determined. Particle has no production and no end vertex!" << endl;
0270         }
0271         if (sub_id < 0)
0272           sub_id = 0;
0273         int new_id = sub_id + suboffset;
0274         GenParticleRef dref(ref, d);
0275         cands[d].setCollisionId(new_id);  // For new GenParticle
0276         LogDebug("VertexId") << "SubEvent offset 3 : " << suboffset;
0277       }
0278       int nsub = -2;
0279       if (isHI) {
0280         nsub = hi->Ncoll_hard() + 1;
0281         suboffset += nsub;
0282       } else {
0283         suboffset += 1;
0284       }
0285       offset += num_particles;
0286     }
0287   } else {
0288     if (totalSize) {
0289       auto origin = (*mc->vertices_begin())->position();
0290       xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0291       *t0Ptr = origin.t() * mmToNs;
0292       fillIndices(mc, particles, *barCodeVector, 0, barcodes);
0293
0294       // fill output collection and save association
0295       for (size_t i = 0; i < particles.size(); ++i) {
0296         const HepMC::GenParticle* part = particles[i];
0297         reco::GenParticle& cand = cands[i];
0298         // convert HepMC::GenParticle to new reco::GenParticle
0299         convertParticle(cand, part, id2Charge);
0300         cand.resetDaughters(ref.id());
0301       }
0302
0303       // fill references to daughters
0304       for (size_t d = 0; d < cands.size(); ++d) {
0305         const HepMC::GenParticle* part = particles[d];
0306         const GenVertex* productionVertex = part->production_vertex();
0307         // search barcode map and attach daughters
0308         if (productionVertex != nullptr)
0309           fillDaughters(cands, part, ref, d, barcodes);
0310         cands[d].setCollisionId(0);
0311       }
0312     }
0313   }
0314
0315   evt.put(std::move(candsPtr));
0316   if (saveBarCodes_)
0317     evt.put(std::move(barCodeVector));
0318   if (cfhepmcprod)
0319     delete cfhepmcprod;
0320   evt.put(std::move(xyz0Ptr), "xyz0");
0321   evt.put(std::move(t0Ptr), "t0");
0322 }
0323
0324 bool GenParticleProducer::convertParticle(reco::GenParticle& cand,
0325                                           const HepMC::GenParticle* part,
0326                                           IDto3Charge const& id2Charge) const {
0327   Candidate::LorentzVector p4(part->momentum());
0328   int pdgId = part->pdg_id();
0329   cand.setThreeCharge(id2Charge.chargeTimesThree(pdgId));
0330   cand.setPdgId(pdgId);
0331   cand.setStatus(part->status());
0332   cand.setP4(p4);
0333   cand.setCollisionId(0);
0334   const GenVertex* v = part->production_vertex();
0335   if (v != nullptr) {
0336     ThreeVector vtx = v->point3d();
0337     Candidate::Point vertex(vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm);
0338     cand.setVertex(vertex);
0339   } else {
0340     cand.setVertex(Candidate::Point(0, 0, 0));
0341   }
0342   mcTruthHelper_.fillGenStatusFlags(*part, cand.statusFlags());
0343   return true;
0344 }
0345
0346 bool GenParticleProducer::fillDaughters(reco::GenParticleCollection& cands,
0347                                         const HepMC::GenParticle* part,
0348                                         reco::GenParticleRefProd const& ref,
0349                                         size_t index,
0350                                         std::unordered_map<int, size_t>& barcodes) const {
0351   const GenVertex* productionVertex = part->production_vertex();
0352   size_t numberOfMothers = productionVertex->particles_in_size();
0353   if (numberOfMothers > 0) {
0354     GenVertex::particles_in_const_iterator motherIt = productionVertex->particles_in_const_begin();
0355     for (; motherIt != productionVertex->particles_in_const_end(); motherIt++) {
0356       const HepMC::GenParticle* mother = *motherIt;
0357       size_t m = barcodes.find(mother->barcode())->second;
0358       cands[m].addDaughter(GenParticleRef(ref, index));
0359       cands[index].addMother(GenParticleRef(ref, m));
0360     }
0361   }
0362   return true;
0363 }
0364
0365 bool GenParticleProducer::fillIndices(const HepMC::GenEvent* mc,
0366                                       vector<const HepMC::GenParticle*>& particles,
0367                                       vector<int>& barCodeVector,
0368                                       int offset,
0369                                       std::unordered_map<int, size_t>& barcodes) const {
0370   size_t idx = offset;
0371   HepMC::GenEvent::particle_const_iterator begin = mc->particles_begin(), end = mc->particles_end();
0372   for (HepMC::GenEvent::particle_const_iterator p = begin; p != end; ++p) {
0373     const HepMC::GenParticle* particle = *p;
0374     size_t barCode_this_event = particle->barcode();
0375     size_t barCode = barCode_this_event + offset;
0376     if (barcodes.find(barCode) != barcodes.end())
0377       throw cms::Exception("WrongReference") << "barcodes are duplicated! " << endl;
0378     particles[idx] = particle;
0379     barCodeVector[idx] = barCode;
0380     barcodes.insert(make_pair(barCode_this_event, idx++));
0381   }
0382   return true;
0383 }
0384
0385 #include "FWCore/Framework/interface/MakerMacros.h"
0386
0387 DEFINE_FWK_MODULE(GenParticleProducer);