File indexing completed on 2025-03-23 16:00:19
0001
0002
0003
0004
0005
0006
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 "SimDataFormats/GeneratorProducts/interface/HepMC3Product.h"
0019 #include "PhysicsTools/HepMCCandAlgos/interface/MCTruthHelper.h"
0020 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0021
0022 #include <vector>
0023 #include <string>
0024 #include <unordered_map>
0025
0026 namespace edm {
0027 class ParameterSet;
0028 }
0029 namespace HepMC {
0030 class GenParticle;
0031 class GenEvent;
0032 }
0033
0034 namespace HepMC3 {
0035 class GenParticle;
0036 class GenVertex;
0037 class GenEvent;
0038 }
0039
0040 static constexpr int PDGCacheMax = 32768;
0041
0042 namespace {
0043 struct IDto3Charge {
0044 IDto3Charge(HepPDT::ParticleDataTable const&, bool abortOnUnknownPDGCode);
0045
0046 int chargeTimesThree(int) const;
0047
0048 private:
0049 std::vector<int> chargeP_, chargeM_;
0050 std::unordered_map<int, int> chargeMap_;
0051 bool abortOnUnknownPDGCode_;
0052 };
0053
0054 IDto3Charge::IDto3Charge(HepPDT::ParticleDataTable const& iTable, bool iAbortOnUnknownPDGCode)
0055 : chargeP_(PDGCacheMax, 0), chargeM_(PDGCacheMax, 0), abortOnUnknownPDGCode_(iAbortOnUnknownPDGCode) {
0056 for (auto const& p : iTable) {
0057 const HepPDT::ParticleID& id = p.first;
0058 int pdgId = id.pid(), apdgId = std::abs(pdgId);
0059 int q3 = id.threeCharge();
0060 if (apdgId < PDGCacheMax && pdgId > 0) {
0061 chargeP_[apdgId] = q3;
0062 chargeM_[apdgId] = -q3;
0063 } else if (apdgId < PDGCacheMax) {
0064 chargeP_[apdgId] = -q3;
0065 chargeM_[apdgId] = q3;
0066 } else {
0067 chargeMap_.emplace(pdgId, q3);
0068 chargeMap_.emplace(-pdgId, -q3);
0069 }
0070 }
0071 }
0072
0073 int IDto3Charge::chargeTimesThree(int id) const {
0074 if (std::abs(id) < PDGCacheMax)
0075 return id > 0 ? chargeP_[id] : chargeM_[-id];
0076 auto f = chargeMap_.find(id);
0077 if (f == chargeMap_.end()) {
0078 if (abortOnUnknownPDGCode_)
0079 throw edm::Exception(edm::errors::LogicError) << "invalid PDG id: " << id;
0080 else
0081 return HepPDT::ParticleID(id).threeCharge();
0082 }
0083 return f->second;
0084 }
0085
0086 }
0087
0088 class GenParticleProducer : public edm::global::EDProducer<edm::RunCache<IDto3Charge>> {
0089 public:
0090
0091 GenParticleProducer(const edm::ParameterSet&);
0092
0093 ~GenParticleProducer() override;
0094
0095
0096 void produce(edm::StreamID, edm::Event& e, const edm::EventSetup&) const override;
0097 std::shared_ptr<IDto3Charge> globalBeginRun(const edm::Run&, const edm::EventSetup&) const override;
0098 void globalEndRun(edm::Run const&, edm::EventSetup const&) const override {}
0099
0100 bool convertParticle(reco::GenParticle& cand, const HepMC::GenParticle* part, const IDto3Charge& id2Charge) const;
0101 bool fillIndices(const HepMC::GenEvent* mc,
0102 std::vector<const HepMC::GenParticle*>& particles,
0103 std::vector<int>& barCodeVector,
0104 int offset,
0105 std::unordered_map<int, size_t>& barcodes) const;
0106 bool fillDaughters(reco::GenParticleCollection& cand,
0107 const HepMC::GenParticle* part,
0108 reco::GenParticleRefProd const& ref,
0109 size_t index,
0110 std::unordered_map<int, size_t>& barcodes) const;
0111
0112 bool convertParticle3(reco::GenParticle& cand, const HepMC3::GenParticle* part, const IDto3Charge& id2Charge) const;
0113 bool fillIndices3(const HepMC3::GenEvent* mc,
0114 std::vector<const HepMC3::GenParticle*>& particles,
0115 std::vector<int>& barCodeVector,
0116 int offset,
0117 std::unordered_map<int, size_t>& barcodes) const;
0118 bool fillDaughters3(reco::GenParticleCollection& cand,
0119 const HepMC3::GenParticle* part,
0120 reco::GenParticleRefProd const& ref,
0121 size_t index,
0122 std::unordered_map<int, size_t>& barcodes) const;
0123
0124 private:
0125
0126 edm::EDGetTokenT<edm::HepMCProduct> srcToken_;
0127 edm::EDGetTokenT<edm::HepMC3Product> srcToken3_;
0128 std::vector<edm::EDGetTokenT<edm::HepMCProduct>> vectorSrcTokens_;
0129 edm::EDGetTokenT<CrossingFrame<edm::HepMCProduct>> mixToken_;
0130 edm::ESGetToken<HepPDT::ParticleDataTable, edm::DefaultRecord> particleTableToken_;
0131
0132
0133 bool abortOnUnknownPDGCode_;
0134
0135 bool saveBarCodes_;
0136
0137
0138 bool doSubEvent_;
0139 bool useCF_;
0140
0141 MCTruthHelper<reco::GenParticle> mcTruthHelperGenParts_;
0142 MCTruthHelper<HepMC::GenParticle> mcTruthHelper_;
0143 MCTruthHelper<HepMC3::GenParticle> mcTruthHelper3_;
0144 };
0145
0146 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0147
0148
0149 #include "SimDataFormats/CrossingFrame/interface/MixCollection.h"
0150
0151 #include "DataFormats/Common/interface/Handle.h"
0152 #include "FWCore/Framework/interface/ESHandle.h"
0153 #include "FWCore/Framework/interface/Event.h"
0154 #include "FWCore/Framework/interface/Run.h"
0155 #include "FWCore/Framework/interface/EventSetup.h"
0156 #include "FWCore/Utilities/interface/EDMException.h"
0157 #include "FWCore/Utilities/interface/transform.h"
0158 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0159
0160 #include "HepMC3/GenEvent.h"
0161 #include "HepMC3/GenParticle.h"
0162 #include "HepMC3/GenVertex.h"
0163
0164 #include <fstream>
0165 #include <algorithm>
0166 using namespace edm;
0167 using namespace reco;
0168 using namespace std;
0169 using namespace HepMC;
0170
0171 static const double mmToCm = 0.1;
0172 static const double mmToNs = 1.0 / 299792458e-6;
0173
0174 GenParticleProducer::GenParticleProducer(const ParameterSet& cfg)
0175 : abortOnUnknownPDGCode_(false),
0176 saveBarCodes_(cfg.getUntrackedParameter<bool>("saveBarCodes", false)),
0177 doSubEvent_(cfg.getUntrackedParameter<bool>("doSubEvent", false)),
0178 useCF_(cfg.getUntrackedParameter<bool>("useCrossingFrame", false)) {
0179 particleTableToken_ = esConsumes<HepPDT::ParticleDataTable, edm::DefaultRecord, edm::Transition::BeginRun>();
0180 produces<GenParticleCollection>();
0181 produces<math::XYZPointF>("xyz0");
0182 produces<float>("t0");
0183 if (saveBarCodes_) {
0184 std::string alias(cfg.getParameter<std::string>("@module_label"));
0185 produces<vector<int>>().setBranchAlias(alias + "BarCodes");
0186 }
0187
0188 if (useCF_)
0189 mixToken_ =
0190 mayConsume<CrossingFrame<HepMCProduct>>(InputTag(cfg.getParameter<std::string>("mix"), "generatorSmeared"));
0191 else {
0192 srcToken_ = mayConsume<HepMCProduct>(cfg.getParameter<InputTag>("src"));
0193 srcToken3_ = mayConsume<HepMC3Product>(cfg.getParameter<InputTag>("src"));
0194 }
0195 }
0196
0197 GenParticleProducer::~GenParticleProducer() {}
0198
0199 std::shared_ptr<IDto3Charge> GenParticleProducer::globalBeginRun(const Run&, const EventSetup& es) const {
0200 ESHandle<HepPDT::ParticleDataTable> pdt = es.getHandle(particleTableToken_);
0201 return std::make_shared<IDto3Charge>(*pdt, abortOnUnknownPDGCode_);
0202 }
0203
0204 void GenParticleProducer::produce(StreamID, Event& evt, const EventSetup& es) const {
0205 std::unordered_map<int, size_t> barcodes;
0206
0207 size_t totalSize = 0;
0208 const GenEvent* mc = nullptr;
0209 HepMC3::GenEvent* mc3 = nullptr;
0210 MixCollection<HepMCProduct>* cfhepmcprod = nullptr;
0211 size_t npiles = 1;
0212
0213 int ivhepmc = 0;
0214 if (useCF_) {
0215 Handle<CrossingFrame<HepMCProduct>> cf;
0216 evt.getByToken(mixToken_, cf);
0217 cfhepmcprod = new MixCollection<HepMCProduct>(cf.product());
0218 npiles = cfhepmcprod->size();
0219 LogDebug("GenParticleProducer") << "npiles : " << npiles << endl;
0220 for (unsigned int icf = 0; icf < npiles; ++icf) {
0221 LogDebug("GenParticleProducer") << "subSize : " << cfhepmcprod->getObject(icf).GetEvent()->particles_size()
0222 << endl;
0223 totalSize += cfhepmcprod->getObject(icf).GetEvent()->particles_size();
0224 }
0225 LogDebug("GenParticleProducer") << "totalSize : " << totalSize << endl;
0226 } else {
0227 Handle<HepMCProduct> mcp;
0228 bool found = evt.getByToken(srcToken_, mcp);
0229 if (found) {
0230 mc = mcp->GetEvent();
0231 if (mc == nullptr)
0232 throw edm::Exception(edm::errors::InvalidReference) << "HepMC has null pointer to GenEvent" << endl;
0233 ivhepmc = 2;
0234 totalSize = mc->particles_size();
0235 } else {
0236 Handle<HepMC3Product> mcp3;
0237 bool found = evt.getByToken(srcToken3_, mcp3);
0238 if (!found)
0239 throw cms::Exception("ProductAbsent") << "No HepMCProduct, tried to get HepMC3Product, but it is also absent.";
0240 ivhepmc = 3;
0241 mc3 = new HepMC3::GenEvent();
0242 mc3->read_data(*mcp3->GetEvent());
0243 totalSize = (mc3->particles()).size();
0244 }
0245 }
0246
0247
0248 const size_t size = totalSize;
0249 vector<const HepMC::GenParticle*> particles(size);
0250 vector<const HepMC3::GenParticle*> particles3(size);
0251 auto candsPtr = std::make_unique<GenParticleCollection>(size);
0252 auto barCodeVector = std::make_unique<vector<int>>(size);
0253 std::unique_ptr<math::XYZPointF> xyz0Ptr(new math::XYZPointF(0., 0., 0.));
0254 std::unique_ptr<float> t0Ptr(new float(0.f));
0255 reco::GenParticleRefProd ref = evt.getRefBeforePut<GenParticleCollection>();
0256 GenParticleCollection& cands = *candsPtr;
0257 size_t offset = 0;
0258 size_t suboffset = 0;
0259
0260 IDto3Charge const& id2Charge = *runCache(evt.getRun().index());
0261
0262 if (doSubEvent_ || useCF_) {
0263 for (size_t ipile = 0; ipile < npiles; ++ipile) {
0264 LogDebug("GenParticleProducer") << "mixed object ipile : " << ipile << endl;
0265 barcodes.clear();
0266 if (useCF_)
0267 mc = cfhepmcprod->getObject(ipile).GetEvent();
0268
0269
0270 bool isHI = false;
0271 const HepMC::HeavyIon* hi = mc->heavy_ion();
0272 if (hi && hi->Ncoll_hard() > 1)
0273 isHI = true;
0274 size_t num_particles = mc->particles_size();
0275 LogDebug("GenParticleProducer") << "num_particles : " << num_particles << endl;
0276 if (ipile == 0) {
0277 auto origin = (*mc->vertices_begin())->position();
0278 xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0279 *t0Ptr = origin.t() * mmToNs;
0280 }
0281 fillIndices(mc, particles, *barCodeVector, offset, barcodes);
0282
0283 for (size_t ipar = offset; ipar < offset + num_particles; ++ipar) {
0284 const HepMC::GenParticle* part = particles[ipar];
0285 reco::GenParticle& cand = cands[ipar];
0286
0287 convertParticle(cand, part, id2Charge);
0288 cand.resetDaughters(ref.id());
0289 }
0290
0291 for (size_t d = offset; d < offset + num_particles; ++d) {
0292 const HepMC::GenParticle* part = particles[d];
0293 const GenVertex* productionVertex = part->production_vertex();
0294 int sub_id = 0;
0295 if (productionVertex != nullptr) {
0296 sub_id = productionVertex->id();
0297 if (!isHI)
0298 sub_id = 0;
0299
0300 fillDaughters(cands, part, ref, d, barcodes);
0301 } else {
0302 const GenVertex* endVertex = part->end_vertex();
0303 if (endVertex != nullptr)
0304 sub_id = endVertex->id();
0305 else
0306 throw cms::Exception("SubEventID")
0307 << "SubEvent not determined. Particle has no production and no end vertex!" << endl;
0308 }
0309 if (sub_id < 0)
0310 sub_id = 0;
0311 int new_id = sub_id + suboffset;
0312 GenParticleRef dref(ref, d);
0313 cands[d].setCollisionId(new_id);
0314 LogDebug("VertexId") << "SubEvent offset 3 : " << suboffset;
0315 }
0316 int nsub = -2;
0317 if (isHI) {
0318 nsub = hi->Ncoll_hard() + 1;
0319 suboffset += nsub;
0320 } else {
0321 suboffset += 1;
0322 }
0323 offset += num_particles;
0324 }
0325 } else {
0326
0327 if (ivhepmc == 2) {
0328 if (totalSize) {
0329 auto origin = (*mc->vertices_begin())->position();
0330 xyz0Ptr->SetXYZ(origin.x() * mmToCm, origin.y() * mmToCm, origin.z() * mmToCm);
0331 *t0Ptr = origin.t() * mmToNs;
0332 fillIndices(mc, particles, *barCodeVector, 0, barcodes);
0333
0334
0335 for (size_t i = 0; i < particles.size(); ++i) {
0336 const HepMC::GenParticle* part = particles[i];
0337 reco::GenParticle& cand = cands[i];
0338
0339 convertParticle(cand, part, id2Charge);
0340 cand.resetDaughters(ref.id());
0341 }
0342
0343
0344 for (size_t d = 0; d < cands.size(); ++d) {
0345 const HepMC::GenParticle* part = particles[d];
0346 const GenVertex* productionVertex = part->production_vertex();
0347
0348 if (productionVertex != nullptr) {
0349 fillDaughters(cands, part, ref, d, barcodes);
0350 cands[d].setCollisionId(0);
0351 }
0352 }
0353 }
0354 }
0355
0356 if (ivhepmc == 3) {
0357 if (totalSize) {
0358 for (const HepMC3::GenVertexPtr& v : mc3->vertices()) {
0359 xyz0Ptr->SetXYZ((v->position()).x() * mmToCm, (v->position()).y() * mmToCm, (v->position()).z() * mmToCm);
0360 *t0Ptr = (v->position()).t() * mmToNs;
0361 break;
0362 }
0363 fillIndices3(mc3, particles3, *barCodeVector, 0, barcodes);
0364
0365
0366 for (size_t i = 0; i < particles3.size(); ++i) {
0367 const HepMC3::GenParticle* part = particles3[i];
0368 reco::GenParticle& cand = cands[i];
0369
0370 convertParticle3(cand, part, id2Charge);
0371 cand.resetDaughters(ref.id());
0372 }
0373
0374
0375 for (size_t d = 0; d < cands.size(); ++d) {
0376 const HepMC3::GenParticle* part = particles3[d];
0377
0378 if (part->production_vertex() != nullptr) {
0379 fillDaughters3(cands, part, ref, d, barcodes);
0380 cands[d].setCollisionId(0);
0381 }
0382 }
0383 }
0384 }
0385 }
0386
0387 if (ivhepmc == 3)
0388 delete mc3;
0389 evt.put(std::move(candsPtr));
0390 if (saveBarCodes_)
0391 evt.put(std::move(barCodeVector));
0392 if (cfhepmcprod)
0393 delete cfhepmcprod;
0394 evt.put(std::move(xyz0Ptr), "xyz0");
0395 evt.put(std::move(t0Ptr), "t0");
0396 }
0397
0398 bool GenParticleProducer::convertParticle(reco::GenParticle& cand,
0399 const HepMC::GenParticle* part,
0400 IDto3Charge const& id2Charge) const {
0401 Candidate::LorentzVector p4(part->momentum());
0402 int pdgId = part->pdg_id();
0403 cand.setThreeCharge(id2Charge.chargeTimesThree(pdgId));
0404 cand.setPdgId(pdgId);
0405 cand.setStatus(part->status());
0406 cand.setP4(p4);
0407 cand.setCollisionId(0);
0408 const GenVertex* v = part->production_vertex();
0409 if (v != nullptr) {
0410 ThreeVector vtx = v->point3d();
0411 Candidate::Point vertex(vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm);
0412 cand.setVertex(vertex);
0413 } else {
0414 cand.setVertex(Candidate::Point(0, 0, 0));
0415 }
0416 mcTruthHelper_.fillGenStatusFlags(*part, cand.statusFlags());
0417 return true;
0418 }
0419
0420 bool GenParticleProducer::convertParticle3(reco::GenParticle& cand,
0421 const HepMC3::GenParticle* part,
0422 IDto3Charge const& id2Charge) const {
0423 Candidate::LorentzVector p4(part->momentum());
0424 int pdgId = part->pdg_id();
0425 cand.setThreeCharge(id2Charge.chargeTimesThree(pdgId));
0426 cand.setPdgId(pdgId);
0427 cand.setStatus(part->status());
0428 cand.setP4(p4);
0429 cand.setCollisionId(0);
0430 if (part->production_vertex() != nullptr) {
0431 ThreeVector vtx((part->production_vertex()->position()).x(),
0432 (part->production_vertex()->position()).y(),
0433 (part->production_vertex()->position()).z());
0434 Candidate::Point vertex(vtx.x() * mmToCm, vtx.y() * mmToCm, vtx.z() * mmToCm);
0435 cand.setVertex(vertex);
0436 } else {
0437 cand.setVertex(Candidate::Point(0, 0, 0));
0438 }
0439 mcTruthHelper3_.fillGenStatusFlags(*part, cand.statusFlags());
0440 return true;
0441 }
0442
0443 bool GenParticleProducer::fillDaughters(reco::GenParticleCollection& cands,
0444 const HepMC::GenParticle* part,
0445 reco::GenParticleRefProd const& ref,
0446 size_t index,
0447 std::unordered_map<int, size_t>& barcodes) const {
0448 const GenVertex* productionVertex = part->production_vertex();
0449 size_t numberOfMothers = productionVertex->particles_in_size();
0450 if (numberOfMothers > 0) {
0451 GenVertex::particles_in_const_iterator motherIt = productionVertex->particles_in_const_begin();
0452 for (; motherIt != productionVertex->particles_in_const_end(); motherIt++) {
0453 const HepMC::GenParticle* mother = *motherIt;
0454 size_t m = barcodes.find(mother->barcode())->second;
0455 cands[m].addDaughter(GenParticleRef(ref, index));
0456 cands[index].addMother(GenParticleRef(ref, m));
0457 }
0458 }
0459 return true;
0460 }
0461
0462 bool GenParticleProducer::fillDaughters3(reco::GenParticleCollection& cands,
0463 const HepMC3::GenParticle* part,
0464 reco::GenParticleRefProd const& ref,
0465 size_t index,
0466 std::unordered_map<int, size_t>& barcodes) const {
0467 size_t numberOfMothers = (part->production_vertex())->particles_in_size();
0468 if (numberOfMothers > 0) {
0469 for (const HepMC3::ConstGenParticlePtr& mother : (part->production_vertex())->particles_in()) {
0470 size_t m = barcodes.find(mother->id())->second;
0471 cands[m].addDaughter(GenParticleRef(ref, index));
0472 cands[index].addMother(GenParticleRef(ref, m));
0473 }
0474 }
0475 return true;
0476 }
0477
0478 bool GenParticleProducer::fillIndices(const HepMC::GenEvent* mc,
0479 vector<const HepMC::GenParticle*>& particles,
0480 vector<int>& barCodeVector,
0481 int offset,
0482 std::unordered_map<int, size_t>& barcodes) const {
0483 size_t idx = offset;
0484 HepMC::GenEvent::particle_const_iterator begin = mc->particles_begin(), end = mc->particles_end();
0485 for (HepMC::GenEvent::particle_const_iterator p = begin; p != end; ++p) {
0486 const HepMC::GenParticle* particle = *p;
0487 size_t barCode_this_event = particle->barcode();
0488 size_t barCode = barCode_this_event + offset;
0489 if (barcodes.find(barCode) != barcodes.end())
0490 throw cms::Exception("WrongReference") << "barcodes are duplicated! " << endl;
0491 particles[idx] = particle;
0492 barCodeVector[idx] = barCode;
0493 barcodes.insert(make_pair(barCode_this_event, idx++));
0494 }
0495 return true;
0496 }
0497
0498 bool GenParticleProducer::fillIndices3(const HepMC3::GenEvent* mc,
0499 vector<const HepMC3::GenParticle*>& particles,
0500 vector<int>& barCodeVector,
0501 int offset,
0502 std::unordered_map<int, size_t>& barcodes) const {
0503 size_t idx = offset;
0504 for (const auto& p : mc->particles()) {
0505 const HepMC3::GenParticle* particle = p.get();
0506 size_t barCode_this_event = particle->id();
0507 size_t barCode = barCode_this_event + offset;
0508 if (barcodes.find(barCode) != barcodes.end())
0509 throw cms::Exception("WrongReference") << "barcodes are duplicated! " << endl;
0510 particles[idx] = particle;
0511 barCodeVector[idx] = barCode;
0512 barcodes.insert(make_pair(barCode_this_event, idx++));
0513 }
0514 return true;
0515 }
0516
0517 #include "FWCore/Framework/interface/MakerMacros.h"
0518
0519 DEFINE_FWK_MODULE(GenParticleProducer);