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 { class GenEvent; }
0030
0031 using namespace std;
0032
0033 static boost::iostreams::filtering_streambuf<boost::iostreams::output> zoutbuffer;
0034
0035 Fun4AllHepMCOutputManager::Fun4AllHepMCOutputManager(const string &myname,
0036                                                      const string &filename)
0037   : Fun4AllOutputManager(myname)
0038   , outfilename(filename)
0039   , comment_written(0)
0040   , filestream(nullptr)
0041   , zipstream(nullptr)
0042   , _embedding_id(0)
0043 {
0044   TString tstr(filename);
0045   TPRegexp bzip_ext(".bz2$");
0046   TPRegexp gzip_ext(".gz$");
0047
0048   if (tstr.Contains(bzip_ext))
0049   {
0050     // use boost iosteam library to compress to bz2 file on the fly
0051     filestream = new ofstream(filename.c_str(), std::ios::out | std::ios::binary);
0052     zoutbuffer.push(boost::iostreams::bzip2_compressor(9));
0053     zoutbuffer.push(*filestream);
0054     zipstream = new ostream(&zoutbuffer);
0055     ascii_out = new HepMC::IO_GenEvent(*zipstream);
0056   }
0057   else if (tstr.Contains(gzip_ext))
0058   {
0059     // use boost iosream to compress to gzip file on the fly
0060     filestream = new ofstream(filename.c_str(), std::ios::out | std::ios::binary);
0061     zoutbuffer.push(boost::iostreams::gzip_compressor(9));
0062     zoutbuffer.push(*filestream);
0063     zipstream = new ostream(&zoutbuffer);
0064     ascii_out = new HepMC::IO_GenEvent(*zipstream);
0065   }
0066   else
0067   {
0068     // produces normal ascii hepmc file
0069     ascii_out = new HepMC::IO_GenEvent(filename, std::ios::out);
0070   }
0071
0072   if (ascii_out->rdstate())
0073   {
0074     cout << "error opening " << outfilename << " exiting " << endl;
0075     exit(1);
0076   }
0077   return;
0078 }
0079
0080 Fun4AllHepMCOutputManager::~Fun4AllHepMCOutputManager()
0081 {
0082   if (ascii_out)
0083   {
0084     if (!comment_written)
0085     {
0086       if (comment.size())
0087       {
0088         ascii_out->write_comment(comment);
0089       }
0090     }
0091     ascii_out->clear();
0092   }
0093
0094   delete ascii_out;
0095
0096   if (zoutbuffer.size() > 0) zoutbuffer.reset();
0097
0098   if (zipstream) delete zipstream;
0099   if (filestream) delete filestream;
0100
0101   return;
0102 }
0103
0104 void Fun4AllHepMCOutputManager::Print(const string &what) const
0105 {
0106   cout << Name() << " writes " << outfilename << endl;
0107   if (comment.size())
0108   {
0109     cout << "comment : " << comment << endl;
0110   }
0111   // base class print method
0112   Fun4AllOutputManager::Print(what);
0113
0114   return;
0115 }
0116
0117 int Fun4AllHepMCOutputManager::Write(PHCompositeNode *topNode)
0118 {
0119   if (!comment_written)
0120   {
0121     if (comment.size())
0122     {
0123       ascii_out->write_comment(comment);
0124     }
0125     comment_written = 1;
0126   }
0127
0128   PHHepMCGenEventMap *geneventmap = findNode::getClass<PHHepMCGenEventMap>(topNode, "PHHepMCGenEventMap");
0129
0130   if (!geneventmap)
0131   {
0132     cout << "Fun4AllHepMCOutputManager::Write - Fatal Error - missing source node PHHepMCGenEventMap" << endl;
0133     return Fun4AllReturnCodes::ABORTRUN;
0134   }
0135   assert(geneventmap);
0136
0137   PHHepMCGenEvent *genevt = geneventmap->get(_embedding_id);
0138   if (!genevt)
0139   {
0140     cout << "Fun4AllHepMCOutputManager::Write - Warning - missing sub-event with embedding ID" << _embedding_id << " on node PHHepMCGenEventMap" << endl;
0141     return Fun4AllReturnCodes::DISCARDEVENT;
0142   }
0143   assert(genevt);
0144
0145   HepMC::GenEvent *evt = genevt->getEvent();
0146   if (!evt)
0147   {
0148     cout << PHWHERE << "0 HepMC Pointer" << endl;
0149     return Fun4AllReturnCodes::ABORTRUN;
0150   }
0151   assert(evt);
0152
0153   IncrementEvents(1);
0154   ascii_out->write_event(evt);
0155   return Fun4AllReturnCodes::EVENT_OK;
0156 }
0157
0158 int Fun4AllHepMCOutputManager::AddComment(const std::string &text)
0159 {
0160   comment = text;
0161   comment_written = 0;
0162   return 0;
0163 }