Merge pull request #48989 from sihyunjeon/cmssw_merge/hepmc_writer

add hepmc writer

Author: cmsbuild

GeneratorInterface/RivetInterface/plugins/GenParticles2HepMCConverter.cc

@@ -22,6 +22,7 @@
 #include "HepMC3/GenParticle.h"
 #include "HepMC3/GenVertex.h"
 #include "HepMC3/Print.h"
+#include "HepMC3/WriterAscii.h"
 
 #include <iostream>
 #include <map>
@@ -31,7 +32,7 @@ using namespace std;
 class GenParticles2HepMCConverter : public edm::stream::EDProducer<> {
 public:
   explicit GenParticles2HepMCConverter(const edm::ParameterSet& pset);
-  ~GenParticles2HepMCConverter() override {}
+  ~GenParticles2HepMCConverter() override;
 
   void beginRun(edm::Run const& iRun, edm::EventSetup const&) override;
   void produce(edm::Event& event, const edm::EventSetup& eventSetup) override;
@@ -43,7 +44,10 @@ class GenParticles2HepMCConverter : public edm::stream::EDProducer<> {
   edm::ESGetToken<HepPDT::ParticleDataTable, PDTRecord> pTable_;
 
   const double cmEnergy_;
+  const bool writeHepMC_;
+  const std::string outputFile_;
   HepMC3::GenCrossSectionPtr xsec_;
+  std::shared_ptr<HepMC3::Writer> writer_;
 
 private:
   inline HepMC3::FourVector FourVector(const reco::Candidate::Point& point) {
@@ -57,14 +61,27 @@ class GenParticles2HepMCConverter : public edm::stream::EDProducer<> {
 };
 
 GenParticles2HepMCConverter::GenParticles2HepMCConverter(const edm::ParameterSet& pset)
-    // dummy value to set incident proton pz for particle gun samples
-    : cmEnergy_(pset.getUntrackedParameter<double>("cmEnergy", 13000)) {
+    : cmEnergy_(pset.getUntrackedParameter<double>("cmEnergy", 13000)),
+      writeHepMC_(pset.getUntrackedParameter<bool>("writeHepMC", false)),
+      outputFile_(pset.getUntrackedParameter<std::string>("outputFile", "events.hepmc")) {
   genParticlesToken_ = consumes<reco::CandidateView>(pset.getParameter<edm::InputTag>("genParticles"));
   genEventInfoToken_ = consumes<GenEventInfoProduct>(pset.getParameter<edm::InputTag>("genEventInfo"));
   genRunInfoToken_ = consumes<GenRunInfoProduct, edm::InRun>(pset.getParameter<edm::InputTag>("genEventInfo"));
   pTable_ = esConsumes<HepPDT::ParticleDataTable, PDTRecord>();
 
   produces<edm::HepMC3Product>("unsmeared");
+
+  // Initialize HepMC3 writer
+  if (writeHepMC_) {
+    writer_ = std::make_shared<HepMC3::WriterAscii>(outputFile_);
+  }
+}
+
+GenParticles2HepMCConverter::~GenParticles2HepMCConverter() {
+  if (writeHepMC_) {
+    writer_->close();
+    writer_.reset();
+  }
 }
 
 void GenParticles2HepMCConverter::beginRun(edm::Run const& iRun, edm::EventSetup const&) {
@@ -208,6 +225,11 @@ void GenParticles2HepMCConverter::produce(edm::Event& event, const edm::EventSet
     }
   }
 
+  // Write event to HepMC file
+  if (writeHepMC_) {
+    writer_->write_event(hepmc_event);
+  }
+
   // Finalize HepMC event record
   auto hepmc_product = std::make_unique<edm::HepMC3Product>(&hepmc_event);
   event.put(std::move(hepmc_product), "unsmeared");

GeneratorInterface/RivetInterface/test/genParticles2HepMC_cfg.py

@@ -26,5 +26,8 @@
 process.load("GeneratorInterface.RivetInterface.genParticles2HepMC_cfi")
 process.genParticles2HepMC.genParticles = cms.InputTag("mergedGenParticles")
 
+# Turn on to write HepMC file
+#process.genParticles2HepMC.writeHepMC = cms.untracked.bool(True)
+
 process.path = cms.Path(process.mergedGenParticles*process.genParticles2HepMC)