RivetInterface/plugins/GenParticles2HepMCConverter.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/EventSetup.h"
0006 #include "FWCore/Framework/interface/ESHandle.h"
0007 #include "FWCore/Framework/interface/Run.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009
0010 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0011 #include "FWCore/Utilities/interface/InputTag.h"
0012 #include "DataFormats/Common/interface/Handle.h"
0013
0014 #include "DataFormats/Candidate/interface/CandidateFwd.h"
0015 #include "DataFormats/HepMCCandidate/interface/GenParticleFwd.h"
0016 #include "DataFormats/HepMCCandidate/interface/GenParticle.h"
0017 #include "SimDataFormats/GeneratorProducts/interface/HepMC3Product.h"
0018 #include "SimDataFormats/GeneratorProducts/interface/GenRunInfoProduct.h"
0019 #include "SimDataFormats/GeneratorProducts/interface/GenEventInfoProduct.h"
0020 #include "SimGeneral/HepPDTRecord/interface/ParticleDataTable.h"
0021
0022 #include "HepMC3/GenParticle.h"
0023 #include "HepMC3/GenVertex.h"
0024 #include "HepMC3/Print.h"
0025
0026 #include <iostream>
0027 #include <map>
0028
0029 using namespace std;
0030
0031 class GenParticles2HepMCConverter : public edm::stream::EDProducer<> {
0032 public:
0033   explicit GenParticles2HepMCConverter(const edm::ParameterSet& pset);
0034   ~GenParticles2HepMCConverter() override {}
0035
0036   void beginRun(edm::Run const& iRun, edm::EventSetup const&) override;
0037   void produce(edm::Event& event, const edm::EventSetup& eventSetup) override;
0038
0039 private:
0040   edm::EDGetTokenT<reco::CandidateView> genParticlesToken_;
0041   edm::EDGetTokenT<GenEventInfoProduct> genEventInfoToken_;
0042   edm::EDGetTokenT<GenRunInfoProduct> genRunInfoToken_;
0043   edm::ESGetToken<HepPDT::ParticleDataTable, PDTRecord> pTable_;
0044
0045   const double cmEnergy_;
0046   HepMC3::GenCrossSectionPtr xsec_;
0047
0048 private:
0049   inline HepMC3::FourVector FourVector(const reco::Candidate::Point& point) {
0050     return HepMC3::FourVector(10 * point.x(), 10 * point.y(), 10 * point.z(), 0);
0051   };
0052
0053   inline HepMC3::FourVector FourVector(const reco::Candidate::LorentzVector& lvec) {
0054     // Avoid negative mass, set minimum m^2 = 0
0055     return HepMC3::FourVector(lvec.px(), lvec.py(), lvec.pz(), std::hypot(lvec.P(), std::max(0., lvec.mass())));
0056   };
0057 };
0058
0059 GenParticles2HepMCConverter::GenParticles2HepMCConverter(const edm::ParameterSet& pset)
0060     // dummy value to set incident proton pz for particle gun samples
0061     : cmEnergy_(pset.getUntrackedParameter<double>("cmEnergy", 13000)) {
0062   genParticlesToken_ = consumes<reco::CandidateView>(pset.getParameter<edm::InputTag>("genParticles"));
0063   genEventInfoToken_ = consumes<GenEventInfoProduct>(pset.getParameter<edm::InputTag>("genEventInfo"));
0064   genRunInfoToken_ = consumes<GenRunInfoProduct, edm::InRun>(pset.getParameter<edm::InputTag>("genEventInfo"));
0065   pTable_ = esConsumes<HepPDT::ParticleDataTable, PDTRecord>();
0066
0067   produces<edm::HepMC3Product>("unsmeared");
0068 }
0069
0070 void GenParticles2HepMCConverter::beginRun(edm::Run const& iRun, edm::EventSetup const&) {
0071   edm::Handle<GenRunInfoProduct> genRunInfoHandle;
0072   iRun.getByToken(genRunInfoToken_, genRunInfoHandle);
0073
0074   xsec_ = make_shared<HepMC3::GenCrossSection>();
0075   if (genRunInfoHandle.isValid()) {
0076     xsec_->set_cross_section(genRunInfoHandle->internalXSec().value(), genRunInfoHandle->internalXSec().error());
0077   } else {
0078     // dummy cross section
0079     xsec_->set_cross_section(1., 0.);
0080   }
0081 }
0082
0083 void GenParticles2HepMCConverter::produce(edm::Event& event, const edm::EventSetup& eventSetup) {
0084   edm::Handle<reco::CandidateView> genParticlesHandle;
0085   event.getByToken(genParticlesToken_, genParticlesHandle);
0086
0087   edm::Handle<GenEventInfoProduct> genEventInfoHandle;
0088   event.getByToken(genEventInfoToken_, genEventInfoHandle);
0089
0090   auto const& pTableData = eventSetup.getData(pTable_);
0091
0092   HepMC3::GenEvent hepmc_event;
0093   hepmc_event.set_event_number(event.id().event());
0094   hepmc_event.add_attribute("signal_process_id",
0095                             std::make_shared<HepMC3::IntAttribute>(genEventInfoHandle->signalProcessID()));
0096   hepmc_event.add_attribute("event_scale", std::make_shared<HepMC3::DoubleAttribute>(genEventInfoHandle->qScale()));
0097   hepmc_event.add_attribute("alphaQCD", std::make_shared<HepMC3::DoubleAttribute>(genEventInfoHandle->alphaQCD()));
0098   hepmc_event.add_attribute("alphaQED", std::make_shared<HepMC3::DoubleAttribute>(genEventInfoHandle->alphaQED()));
0099
0100   hepmc_event.weights() = genEventInfoHandle->weights();
0101   // add dummy weight if necessary
0102   if (hepmc_event.weights().size() == 0) {
0103     hepmc_event.weights().push_back(1.);
0104   }
0105
0106   // resize cross section to number of weights
0107   if (xsec_->xsecs().size() < hepmc_event.weights().size()) {
0108     xsec_->set_cross_section(std::vector<double>(hepmc_event.weights().size(), xsec_->xsec(0)),
0109                              std::vector<double>(hepmc_event.weights().size(), xsec_->xsec_err(0)));
0110   }
0111   hepmc_event.set_cross_section(xsec_);
0112
0113   // Set PDF
0114   const gen::PdfInfo* pdf = genEventInfoHandle->pdf();
0115   if (pdf != nullptr) {
0116     const int pdf_id1 = pdf->id.first, pdf_id2 = pdf->id.second;
0117     const double pdf_x1 = pdf->x.first, pdf_x2 = pdf->x.second;
0118     const double pdf_scalePDF = pdf->scalePDF;
0119     const double pdf_xPDF1 = pdf->xPDF.first, pdf_xPDF2 = pdf->xPDF.second;
0120     HepMC3::GenPdfInfoPtr hepmc_pdfInfo = make_shared<HepMC3::GenPdfInfo>();
0121     hepmc_pdfInfo->set(pdf_id1, pdf_id2, pdf_x1, pdf_x2, pdf_scalePDF, pdf_xPDF1, pdf_xPDF2);
0122     hepmc_event.set_pdf_info(hepmc_pdfInfo);
0123   }
0124
0125   // Prepare list of HepMC3::GenParticles
0126   std::map<const reco::Candidate*, HepMC3::GenParticlePtr> genCandToHepMCMap;
0127   HepMC3::GenParticlePtr hepmc_parton1, hepmc_parton2;
0128   std::vector<HepMC3::GenParticlePtr> hepmc_particles;
0129   const reco::Candidate *parton1 = nullptr, *parton2 = nullptr;
0130   for (unsigned int i = 0, n = genParticlesHandle->size(); i < n; ++i) {
0131     const reco::Candidate* p = &genParticlesHandle->at(i);
0132     HepMC3::GenParticlePtr hepmc_particle =
0133         std::make_shared<HepMC3::GenParticle>(FourVector(p->p4()), p->pdgId(), p->status());
0134
0135     // Assign particle's generated mass from the standard particle data table
0136     double particleMass;
0137     if (pTableData.particle(p->pdgId()))
0138       particleMass = pTableData.particle(p->pdgId())->mass();
0139     else
0140       particleMass = p->mass();
0141
0142     hepmc_particle->set_generated_mass(particleMass);
0143
0144     hepmc_particles.push_back(hepmc_particle);
0145     genCandToHepMCMap[p] = hepmc_particle;
0146
0147     // Find incident proton pair
0148     if (p->mother() == nullptr and std::abs(p->eta()) > 5 and std::abs(p->pz()) > 1000) {
0149       if (!parton1 and p->pz() > 0) {
0150         parton1 = p;
0151         hepmc_parton1 = hepmc_particle;
0152       } else if (!parton2 and p->pz() < 0) {
0153         parton2 = p;
0154         hepmc_parton2 = hepmc_particle;
0155       }
0156     }
0157   }
0158
0159   HepMC3::GenVertexPtr vertex1;
0160   HepMC3::GenVertexPtr vertex2;
0161   if (parton1 == nullptr || parton2 == nullptr) {
0162     // Particle gun samples do not have incident partons. Put dummy incident particle and prod vertex
0163     // Note: leave parton1 and parton2 as nullptr since it is not used anymore after creating hepmc_parton1 and 2
0164     const reco::Candidate::LorentzVector nullP4(0, 0, 0, 0);
0165     const reco::Candidate::LorentzVector beamP4(0, 0, cmEnergy_ / 2, cmEnergy_ / 2);
0166     vertex1 = make_shared<HepMC3::GenVertex>(FourVector(nullP4));
0167     vertex2 = make_shared<HepMC3::GenVertex>(FourVector(nullP4));
0168     hepmc_parton1 = make_shared<HepMC3::GenParticle>(FourVector(+beamP4), 2212, 4);
0169     hepmc_parton2 = make_shared<HepMC3::GenParticle>(FourVector(-beamP4), 2212, 4);
0170   } else {
0171     // Put incident beam particles : proton -> parton vertex
0172     vertex1 = make_shared<HepMC3::GenVertex>(FourVector(parton1->vertex()));
0173     vertex2 = make_shared<HepMC3::GenVertex>(FourVector(parton2->vertex()));
0174   }
0175   hepmc_event.add_vertex(vertex1);
0176   hepmc_event.add_vertex(vertex2);
0177   vertex1->add_particle_in(hepmc_parton1);
0178   vertex2->add_particle_in(hepmc_parton2);
0179   //hepmc_event.set_beam_particles(hepmc_parton1, hepmc_parton2);
0180
0181   // Prepare vertex list
0182   typedef std::map<const reco::Candidate*, HepMC3::GenVertexPtr> ParticleToVertexMap;
0183   ParticleToVertexMap particleToVertexMap;
0184   particleToVertexMap[parton1] = vertex1;
0185   particleToVertexMap[parton2] = vertex2;
0186   for (unsigned int i = 0, n = genParticlesHandle->size(); i < n; ++i) {
0187     const reco::Candidate* p = &genParticlesHandle->at(i);
0188     if (p == parton1 or p == parton2)
0189       continue;
0190
0191     // Connect mother-daughters for the other cases
0192     for (unsigned int j = 0, nMothers = p->numberOfMothers(); j < nMothers; ++j) {
0193       // Mother-daughter hierarchy defines vertex
0194       const reco::Candidate* elder = p->mother(j)->daughter(0);
0195       HepMC3::GenVertexPtr vertex;
0196       if (particleToVertexMap.find(elder) == particleToVertexMap.end()) {
0197         vertex = make_shared<HepMC3::GenVertex>(FourVector(elder->vertex()));
0198         hepmc_event.add_vertex(vertex);
0199         particleToVertexMap[elder] = vertex;
0200       } else {
0201         vertex = particleToVertexMap[elder];
0202       }
0203
0204       // Vertex is found. Now connect each other
0205       const reco::Candidate* mother = p->mother(j);
0206       vertex->add_particle_in(genCandToHepMCMap[mother]);
0207       vertex->add_particle_out(hepmc_particles[i]);
0208     }
0209   }
0210
0211   // Finalize HepMC event record
0212   auto hepmc_product = std::make_unique<edm::HepMC3Product>(&hepmc_event);
0213   event.put(std::move(hepmc_product), "unsmeared");
0214 }
0215
0216 DEFINE_FWK_MODULE(GenParticles2HepMCConverter);