generators/phhepmc/Fun4AllHepMCOutputManager.cc

0001 #include "Fun4AllHepMCOutputManager.h"
0002
0003 #include "PHHepMCGenEvent.h"
0004 #include "PHHepMCGenEventMap.h"
0005
0006 #include <fun4all/Fun4AllOutputManager.h>  // for Fun4AllOutp...
0007 #include <fun4all/Fun4AllReturnCodes.h>
0008
0009 #include <phool/getClass.h>
0010 #include <phool/phool.h>  // for PHWHERE
0011
0012 #include <HepMC/IO_GenEvent.h>
0013
0014 #include <TPRegexp.h>
0015 #include <TString.h>
0016
0017 #include <boost/iostreams/filter/bzip2.hpp>
0018 #include <boost/iostreams/filter/gzip.hpp>
0019 #include <boost/iostreams/filtering_streambuf.hpp>
0020
0021 #include <cassert>
0022 #include <cstdlib>  // for exit
0023 #include <fstream>
0024 #include <iostream>
0025 #include <ostream>
0026 #include <string>
0027 #include <utility>  // for swap
0028
0029 namespace HepMC
0030 {
0031   class GenEvent;
0032 }
0033
0034 namespace
0035 {
0036   boost::iostreams::filtering_streambuf<boost::iostreams::output> zoutbuffer;
0037 }
0038
0039 Fun4AllHepMCOutputManager::Fun4AllHepMCOutputManager(const std::string &myname,
0040                                                      const std::string &filename)
0041   : Fun4AllOutputManager(myname)
0042   , outfilename(filename)
0043   , comment_written(0)
0044   , filestream(nullptr)
0045   , zipstream(nullptr)
0046   , _embedding_id(0)
0047 {
0048   TString tstr(filename);
0049   TPRegexp bzip_ext(".bz2$");
0050   TPRegexp gzip_ext(".gz$");
0051
0052   if (tstr.Contains(bzip_ext))
0053   {
0054     // use boost iosteam library to compress to bz2 file on the fly
0055     filestream = new std::ofstream(filename.c_str(), std::ios::out | std::ios::binary);
0056     zoutbuffer.push(boost::iostreams::bzip2_compressor(9));
0057     zoutbuffer.push(*filestream);
0058     zipstream = new std::ostream(&zoutbuffer);
0059     ascii_out = new HepMC::IO_GenEvent(*zipstream);
0060   }
0061   else if (tstr.Contains(gzip_ext))
0062   {
0063     // use boost iosream to compress to gzip file on the fly
0064     filestream = new std::ofstream(filename.c_str(), std::ios::out | std::ios::binary);
0065     zoutbuffer.push(boost::iostreams::gzip_compressor(9));
0066     zoutbuffer.push(*filestream);
0067     zipstream = new std::ostream(&zoutbuffer);
0068     ascii_out = new HepMC::IO_GenEvent(*zipstream);
0069   }
0070   else
0071   {
0072     // produces normal ascii hepmc file
0073     ascii_out = new HepMC::IO_GenEvent(filename, std::ios::out);
0074   }
0075
0076   if (ascii_out->rdstate())
0077   {
0078     std::cout << "error opening " << outfilename << " exiting " << std::endl;
0079     exit(1);
0080   }
0081   return;
0082 }
0083
0084 Fun4AllHepMCOutputManager::~Fun4AllHepMCOutputManager()
0085 {
0086   try
0087   {
0088     if (ascii_out)
0089     {
0090       if (!comment_written && !comment.empty())
0091       {
0092         ascii_out->write_comment(comment);
0093       }
0094       ascii_out->clear();
0095       delete ascii_out;
0096       ascii_out = nullptr;
0097     }
0098
0099     if (!zoutbuffer.empty())
0100     {
0101       zoutbuffer.reset();
0102     }
0103
0104     delete zipstream;
0105     zipstream = nullptr;
0106
0107     delete filestream;
0108     filestream = nullptr;
0109   }
0110   catch (const std::exception &e)
0111   {
0112     std::cout << "Exception caught in ~Fun4AllHepMCOutputManager: "
0113               << e.what() << std::endl;
0114   }
0115   catch (...)
0116   {
0117     std::cout << "Unknown exception caught in ~Fun4AllHepMCOutputManager"
0118               << std::endl;
0119   }
0120
0121   return;
0122 }
0123
0124 void Fun4AllHepMCOutputManager::Print(const std::string &what) const
0125 {
0126   std::cout << Name() << " writes " << outfilename << std::endl;
0127   if (!comment.empty())
0128   {
0129     std::cout << "comment : " << comment << std::endl;
0130   }
0131   // base class print method
0132   Fun4AllOutputManager::Print(what);
0133
0134   return;
0135 }
0136
0137 int Fun4AllHepMCOutputManager::Write(PHCompositeNode *topNode)
0138 {
0139   if (!comment_written)
0140   {
0141     if (!comment.empty())
0142     {
0143       ascii_out->write_comment(comment);
0144     }
0145     comment_written = 1;
0146   }
0147
0148   PHHepMCGenEventMap *geneventmap = findNode::getClass<PHHepMCGenEventMap>(topNode, "PHHepMCGenEventMap");
0149
0150   if (!geneventmap)
0151   {
0152     std::cout << "Fun4AllHepMCOutputManager::Write - Fatal Error - missing source node PHHepMCGenEventMap" << std::endl;
0153     return Fun4AllReturnCodes::ABORTRUN;
0154   }
0155   assert(geneventmap);
0156
0157   PHHepMCGenEvent *genevt = geneventmap->get(_embedding_id);
0158   if (!genevt)
0159   {
0160     std::cout << "Fun4AllHepMCOutputManager::Write - Warning - missing sub-event with embedding ID" << _embedding_id << " on node PHHepMCGenEventMap" << std::endl;
0161     return Fun4AllReturnCodes::DISCARDEVENT;
0162   }
0163   assert(genevt);
0164
0165   HepMC::GenEvent *evt = genevt->getEvent();
0166   if (!evt)
0167   {
0168     std::cout << PHWHERE << "0 HepMC Pointer" << std::endl;
0169     return Fun4AllReturnCodes::ABORTRUN;
0170   }
0171   assert(evt);
0172
0173   IncrementEvents(1);
0174   ascii_out->write_event(evt);
0175   return Fun4AllReturnCodes::EVENT_OK;
0176 }
0177
0178 int Fun4AllHepMCOutputManager::AddComment(const std::string &text)
0179 {
0180   comment = text;
0181   comment_written = 0;
0182   return 0;
0183 }