generators/phhepmc/Fun4AllOscarInputManager.cc

0001 #include "Fun4AllOscarInputManager.h"
0002
0003 #include "PHHepMCGenEvent.h"  // for PHHepMCGenE...
0004 #include "PHHepMCGenEventMap.h"
0005
0006 #include <frog/FROG.h>
0007
0008 #include <fun4all/Fun4AllInputManager.h>  // for Fun4AllInpu...
0009 #include <fun4all/Fun4AllReturnCodes.h>
0010 #include <fun4all/Fun4AllServer.h>
0011 #include <fun4all/Fun4AllSyncManager.h>
0012
0013 #include <phool/PHCompositeNode.h>
0014 #include <phool/PHIODataNode.h>    // for PHIODataNode
0015 #include <phool/PHNodeIterator.h>  // for PHNodeIterator
0016 #include <phool/PHObject.h>
0017 #include <phool/getClass.h>
0018 #include <phool/recoConsts.h>
0019
0020 #include <HepMC/GenEvent.h>
0021 #include <HepMC/GenParticle.h>   // for GenParticle
0022 #include <HepMC/GenVertex.h>     // for GenVertex
0023 #include <HepMC/SimpleVector.h>  // for FourVector
0024 #include <HepMC/Units.h>         // for GEV, MM
0025
0026 #include <TPRegexp.h>
0027 #include <TString.h>
0028
0029 #include <boost/iostreams/filter/bzip2.hpp>
0030 #include <boost/iostreams/filter/gzip.hpp>
0031 #include <boost/iostreams/filtering_streambuf.hpp>
0032
0033 #include <cstdlib>
0034 #include <fstream>
0035 #include <iostream>
0036 #include <map>  // for _Rb_tree_it...
0037 #include <sstream>
0038 #include <utility>  // for swap, pair
0039 #include <vector>   // for vector
0040
0041 using namespace std;
0042
0043 static boost::iostreams::filtering_streambuf<boost::iostreams::input> zinbuffer;
0044 static const double toMM = 1.e-12;
0045 typedef PHIODataNode<PHObject> PHObjectNode_t;
0046
0047 Fun4AllOscarInputManager::Fun4AllOscarInputManager(const string &name, const string &topnodename)
0048   : Fun4AllInputManager(name, "")
0049   , events_total(0)
0050   , events_thisfile(0)
0051   , topNodeName(topnodename)
0052   , evt(nullptr)
0053   , skipEvents(0)
0054   , skippedEvents(0)
0055   , filestream(nullptr)
0056   , unzipstream(nullptr)
0057   , isCompressed(false)
0058 {
0059   Fun4AllServer *se = Fun4AllServer::instance();
0060   topNode = se->topNode(topNodeName.c_str());
0061   PHNodeIterator iter(topNode);
0062   PHCompositeNode *dstNode = se->getNode(InputNode(), topNodeName);
0063
0064   PHHepMCGenEventMap *geneventmap = findNode::getClass<PHHepMCGenEventMap>(topNode, "PHHepMCGenEventMap");
0065   if (!geneventmap)
0066   {
0067     geneventmap = new PHHepMCGenEventMap();
0068     PHIODataNode<PHObject> *newmapnode = new PHIODataNode<PHObject>(geneventmap, "PHHepMCGenEventMap", "PHObject");
0069     dstNode->addNode(newmapnode);
0070   }
0071
0072   PHHepMCGenHelper::set_geneventmap(geneventmap);
0073 }
0074
0075 Fun4AllOscarInputManager::~Fun4AllOscarInputManager()
0076 {
0077   fileclose();
0078   delete filestream;
0079   delete unzipstream;
0080 }
0081
0082 int Fun4AllOscarInputManager::fileopen(const string &filenam)
0083 {
0084   if (!MySyncManager())
0085   {
0086     cout << "Call fileopen only after you registered your Input Manager " << Name() << " with the Fun4AllServer" << endl;
0087     exit(1);
0088   }
0089   if (IsOpen())
0090   {
0091     cout << "Closing currently open file "
0092          << filename
0093          << " and opening " << filenam << endl;
0094     fileclose();
0095   }
0096   filename = filenam;
0097   FROG frog;
0098   string fname(frog.location(filename.c_str()));
0099   if (Verbosity() > 0)
0100   {
0101     cout << Name() << ": opening file " << fname << endl;
0102   }
0103
0104   TString tstr(fname);
0105   TPRegexp bzip_ext(".bz2$");
0106   TPRegexp gzip_ext(".gz$");
0107   if (tstr.Contains(bzip_ext))
0108   {
0109     // use boost iosteam library to decompress bz2 on the fly
0110     filestream = new ifstream(fname.c_str(), std::ios::in | std::ios::binary);
0111     zinbuffer.push(boost::iostreams::bzip2_decompressor());
0112     zinbuffer.push(*filestream);
0113     unzipstream = new istream(&zinbuffer);
0114     isCompressed = true;
0115   }
0116   else if (tstr.Contains(gzip_ext))
0117   {
0118     // use boost iosream to decompress the gzip file on the fly
0119     filestream = new ifstream(fname.c_str(), std::ios::in | std::ios::binary);
0120     zinbuffer.push(boost::iostreams::gzip_decompressor());
0121     zinbuffer.push(*filestream);
0122     unzipstream = new istream(&zinbuffer);
0123     isCompressed = true;
0124   }
0125   else
0126   {
0127     theOscarFile.open(fname.c_str());
0128   }
0129
0130   recoConsts *rc = recoConsts::instance();
0131   static bool run_number_forced = rc->FlagExist("RUNNUMBER");
0132   if (run_number_forced)
0133   {
0134     MySyncManager()->CurrentRun(rc->get_IntFlag("RUNNUMBER"));
0135   }
0136   else
0137   {
0138     MySyncManager()->CurrentRun(-1);
0139   }
0140   events_thisfile = 0;
0141   IsOpen(1);
0142   AddToFileOpened(fname);  // add file to the list of files which were opened
0143   return 0;
0144 }
0145
0146 int Fun4AllOscarInputManager::run(const int nevents)
0147 {
0148 readagain:
0149   if (!IsOpen())
0150   {
0151     if (FileListEmpty())
0152     {
0153       if (Verbosity() > 0)
0154       {
0155         cout << Name() << ": No Input file open" << endl;
0156       }
0157       return 0;
0158     }
0159     else
0160     {
0161       if (OpenNextFile())
0162       {
0163         cout << Name() << ": No Input file from filelist opened" << endl;
0164         return 0;
0165       }
0166     }
0167   }
0168
0169   //Read oscar
0170   evt = new HepMC::GenEvent(HepMC::Units::GEV, HepMC::Units::MM);
0171   int code = ConvertFromOscar();
0172   /*
0173   if(skippedEvents < skipEvents || code == 3)
0174     {
0175       goto readagain;
0176     }
0177   */
0178   if (code == 1 || evt == nullptr)
0179   {
0180     if (Verbosity() > 1) cout << "Finished file!" << endl;
0181     fileclose();
0182     goto readagain;
0183   }
0184
0185   //  if(Verbosity() > 4) cout << "SIZE: " << phhepmcgenevt->size() << endl;
0186   //MySyncManager()->CurrentEvent(evt->event_number());
0187   events_total++;
0188   events_thisfile++;
0189
0190   // check if the local SubsysReco discards this event
0191   if (RejectEvent() != Fun4AllReturnCodes::EVENT_OK)
0192   {
0193     //ResetEvent();
0194     goto readagain;
0195   }
0196
0197   if (events_total < nevents)
0198   {
0199     //ResetEvent();
0200     goto readagain;
0201   }
0202
0203   return 0;
0204 }
0205
0206 int Fun4AllOscarInputManager::fileclose()
0207 {
0208   if (!IsOpen())
0209   {
0210     cout << Name() << ": fileclose: No Input file open" << endl;
0211     return -1;
0212   }
0213   if (isCompressed)
0214   {
0215     filestream->close();
0216   }
0217   else
0218   {
0219     theOscarFile.close();
0220   }
0221   IsOpen(0);
0222   // if we have a file list, move next entry to top of the list
0223   // or repeat the same entry again
0224   UpdateFileList();
0225   return 0;
0226 }
0227
0228 void Fun4AllOscarInputManager::Print(const string &what) const
0229 {
0230   Fun4AllInputManager::Print(what);
0231   return;
0232 }
0233
0234 int Fun4AllOscarInputManager::ResetEvent()
0235 {
0236   //delete evt;
0237   //evt = nullptr;
0238   return 0;
0239 }
0240
0241 int Fun4AllOscarInputManager::PushBackEvents(const int i)
0242 {
0243   int counter = 0;
0244
0245   while (counter < i)
0246   {
0247     std::string theLine;
0248     while (getline(theOscarFile, theLine))
0249     {
0250       if (theLine.compare(0, 1, "#") == 0) continue;
0251       vector<double> theInfo;
0252       double number = NAN;
0253       for (istringstream numbers_iss(theLine); numbers_iss >> number;)
0254       {
0255         theInfo.push_back(number);
0256       }
0257
0258       if (theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] == 0)
0259       {
0260         counter++;
0261         skippedEvents++;
0262         break;
0263       }
0264       else if (theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] > 0)
0265       {
0266         continue;
0267       }
0268     }
0269   }
0270
0271   if (theOscarFile.eof()) return -1;
0272
0273   return 0;
0274 }
0275
0276 int Fun4AllOscarInputManager::ConvertFromOscar()
0277 {
0278   delete evt;
0279   evt = nullptr;
0280
0281   if (theOscarFile.eof())  // if the file is exhausted bail out during this next read
0282   {
0283     cout << "Oscar EOF" << endl;
0284     return 1;
0285   }
0286   evt = new HepMC::GenEvent(HepMC::Units::GEV, HepMC::Units::MM);
0287
0288   if (Verbosity() > 1) cout << "Reading Oscar Event " << events_total + skippedEvents + 1 << endl;
0289   //Grab New Event From Oscar
0290   string theLine;
0291   vector<vector<double> > theEventVec;
0292   //  vector<HepMC::FourVector> theVtxVec;
0293   if (isCompressed)
0294   {
0295     // while(getline(unzipstream, theLine))
0296     //  {
0297     //    if(theLine.find("#") == 0) continue;
0298     //    vector<double> theInfo; //format: N,pid,px,py,pz,E,mass,xvtx,yvtx,zvtx,?
0299     //    double number;
0300     //    for(istringstream numbers_iss(theLine); numbers_iss >> number; )
0301     //      {
0302     //        theInfo.push_back(number);
0303     //      }
0304
0305     //    if(theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] == 0)
0306     //      {
0307     //        break;
0308     //      }
0309     //    else if (theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] > 0)
0310     //      {
0311     //        continue;
0312     //      }
0313     //    else
0314     //      {
0315     //        theEventVec.push_back(theInfo);
0316     //        HepMC::FourVector vert(theInfo[8]*toMM, theInfo[9]*toMM, theInfo[10]*toMM, theInfo[11]);
0317     //        theVtxVec.push_back(vert);
0318     //      }
0319
0320     //  }//while(getline)
0321   }
0322   else
0323   {
0324     while (getline(theOscarFile, theLine))
0325     {
0326       if (theLine.compare(0, 1, "#") == 0) continue;
0327       vector<double> theInfo;  //format: N,pid,px,py,pz,E,mass,xvtx,yvtx,zvtx,?
0328       double number = NAN;
0329       for (istringstream numbers_iss(theLine); numbers_iss >> number;)
0330       {
0331         theInfo.push_back(number);
0332       }
0333
0334       if (theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] == 0)
0335       {
0336         break;
0337       }
0338       else if (theInfo.size() == 2 && theInfo[0] == 0 && theInfo[1] > 0)
0339       {
0340         continue;
0341       }
0342       else
0343       {
0344         theEventVec.push_back(theInfo);
0345       }
0346
0347     }  //while(getline)
0348   }
0349
0350   /*
0351   if(skippedEvents < skipEvents)
0352     {
0353       skippedEvents++;
0354       if (Verbosity() > 5) cout << "Skipping event " << skippedEvents << endl;
0355       return 2;
0356     }
0357   */
0358
0359   //Set Event Number
0360   evt->set_event_number(events_total + 1);
0361
0362   //Loop Over One Event, Fill particles
0363   std::vector<HepMC::GenParticle *> hepevt_particles(theEventVec.size());
0364   for (unsigned int i = 0; i < theEventVec.size(); i++)
0365   {
0366     //int N = (int)theEventVec[i][0];
0367     int pid = (int) theEventVec[i][1];
0368     double px = theEventVec[i][3];
0369     double py = theEventVec[i][4];
0370     double pz = theEventVec[i][5];
0371     double E = theEventVec[i][6];
0372     double m = theEventVec[i][7];
0373     int status = 1;  //oscar only writes final state particles
0374
0375     hepevt_particles[i] = new HepMC::GenParticle(HepMC::FourVector(px, py, pz, E), pid, status);
0376     hepevt_particles[i]->setGeneratedMass(m);
0377     hepevt_particles[i]->suggest_barcode(i + 1);
0378   }
0379
0380   for (unsigned int i = 0; i < theEventVec.size(); i++)
0381   {
0382     HepMC::GenParticle *p = hepevt_particles[i];
0383     HepMC::GenVertex *prod_vtx = p->production_vertex();
0384     if (prod_vtx) prod_vtx->add_particle_out(p);
0385
0386     bool found = false;
0387     HepMC::FourVector prod_pos(theEventVec[i][8] * toMM, theEventVec[i][9] * toMM, theEventVec[i][10] * toMM, theEventVec[i][11]);
0388     if (!prod_vtx)
0389     {
0390       //See if the vertex is already in the event
0391       for (HepMC::GenEvent::vertex_iterator v = evt->vertices_begin(); v != evt->vertices_end(); ++v)
0392       {
0393         HepMC::GenVertex *theV = *v;
0394         if (theV->position().x() != prod_pos.x()) continue;
0395         if (theV->position().y() != prod_pos.y()) continue;
0396         if (theV->position().z() != prod_pos.z()) continue;
0397         found = true;
0398         theV->add_particle_out(p);
0399       }
0400       if (!found)
0401       {
0402         //Didn't find vertex, add it
0403         prod_vtx = new HepMC::GenVertex(prod_pos);
0404         prod_vtx->add_particle_out(p);
0405         evt->add_vertex(prod_vtx);
0406       }
0407     }
0408
0409     // If prod_vtx doesn't already have position specified, fill it.
0410     if (!found && prod_vtx && prod_vtx->position() == HepMC::FourVector()) prod_vtx->set_position(prod_pos);
0411   }
0412
0413   evt->print();
0414   if (Verbosity() > 5) evt->print();
0415   if (Verbosity() > 3) cout << "Adding Event to phhepmcgenevt" << endl;
0416
0417   PHHepMCGenEventMap::Iter ievt =
0418       PHHepMCGenHelper::get_geneventmap()->find(PHHepMCGenHelper::get_embedding_id());
0419   if (ievt != PHHepMCGenHelper::get_geneventmap()->end())
0420   {
0421     // override existing event
0422     ievt->second->addEvent(evt);
0423   }
0424   else
0425     PHHepMCGenHelper::insert_event(evt);
0426   return 0;
0427 }