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