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