g4simulation/g4main/ReadEICFiles.cc

0001 #include "ReadEICFiles.h"
0002
0003 #include "EicEventHeader.h"
0004 #include "EicEventHeaderv1.h"
0005
0006 #include <fun4all/Fun4AllReturnCodes.h>
0007
0008 #include <phool/PHCompositeNode.h>
0009 #include <phool/PHIODataNode.h>
0010 #include <phool/PHNode.h>  // for PHNode
0011 #include <phool/PHNodeIterator.h>
0012 #include <phool/PHObject.h>  // for PHObject
0013 #include <phool/getClass.h>
0014 #include <phool/phool.h>  // for PHWHERE
0015
0016 #include <HepMC/GenEvent.h>
0017 #include <HepMC/GenParticle.h>  // for GenParticle
0018 #include <HepMC/GenVertex.h>
0019 #include <HepMC/PdfInfo.h>       // for PdfInfo
0020 #include <HepMC/SimpleVector.h>  // for FourVector
0021 #include <HepMC/Units.h>         // for GEV, MM
0022
0023 // eicsmear classes
0024 #include <eicsmear/erhic/EventMC.h>
0025 #include <eicsmear/erhic/EventMilou.h>
0026 #include <eicsmear/erhic/ParticleMC.h>  // for ParticleMC
0027 #include <eicsmear/erhic/Pid.h>         // for Pid
0028
0029 #include <eicsmear/erhic/EventDEMP.h>
0030
0031 // General Root and C++ classes
0032 #include <TBranch.h>  // for TBranch
0033 #include <TChain.h>
0034 #include <TVector3.h>  // for TVector3
0035
0036 #include <iostream>  // for operator<<, endl, basic_ostream
0037 #include <vector>    // for vector
0038
0039 class PHHepMCGenEvent;
0040
0041 ///////////////////////////////////////////////////////////////////
0042
0043 ReadEICFiles::ReadEICFiles(const std::string &name)
0044   : SubsysReco(name)
0045   , Tin(nullptr)
0046   , nEntries(0)
0047   , entry(0)
0048   , m_EvtGenId(EvtGen::Unknown)
0049   , GenEvent(nullptr)
0050   , _node_name("PHHepMCGenEvent")
0051 {
0052   PHHepMCGenHelper::set_embedding_id(1);  // default embedding ID to 1
0053   return;
0054 }
0055
0056 ///////////////////////////////////////////////////////////////////
0057
0058 ReadEICFiles::~ReadEICFiles()
0059 {
0060   delete Tin;
0061 }
0062
0063 ///////////////////////////////////////////////////////////////////
0064
0065 bool ReadEICFiles::OpenInputFile(const std::string &name)
0066 {
0067   filename = name;
0068   Tin = new TChain("EICTree", "EICTree");
0069   Tin->Add(name.c_str());
0070   GetTree();
0071   return true;
0072 }
0073
0074 ///////////////////////////////////////////////////////////////////
0075
0076 void ReadEICFiles::GetTree()
0077 {
0078   /* Print the actual class of the event branch record,
0079      i.e. erhic::EventMilou or other */
0080   std::string EventClass(Tin->GetBranch("event")->GetClassName());
0081   if (Verbosity() > 0)
0082   {
0083     std::cout << "ReadEICFiles: Input Branch Event Class = "
0084               << EventClass << std::endl;
0085   }
0086   if (EventClass.find("Milou") != std::string::npos)
0087   {
0088     m_EvtGenId = EvtGen::Milou;
0089   }
0090
0091   if (EventClass.find("DEMP") != std::string::npos)
0092   {
0093     m_EvtGenId = EvtGen::DEMP;
0094   }
0095
0096   Tin->SetBranchAddress("event", &GenEvent);
0097   nEntries = Tin->GetEntries();
0098 }
0099
0100 int ReadEICFiles::Init(PHCompositeNode *topNode)
0101 {
0102   /* Create node tree */
0103   CreateNodeTree(topNode);
0104   return 0;
0105 }
0106
0107 int ReadEICFiles::process_event(PHCompositeNode *topNode)
0108 {
0109   /* Check if there is an unused event left in input file */
0110   if (entry >= nEntries)
0111   {
0112     if (!filename.empty())
0113     {
0114       std::cout << "input file " << filename << " exhausted" << std::endl;
0115     }
0116     else
0117     {
0118       std::cout << "no file opened" << std::endl;
0119     }
0120     return Fun4AllReturnCodes::ABORTRUN;
0121   }
0122
0123   /* Get event record from input file */
0124   Tin->GetEntry(entry);
0125
0126   EicEventHeader *evthead = findNode::getClass<EicEventHeader>(topNode, "EicEventHeader");
0127   switch (m_EvtGenId)
0128   {
0129   case EvtGen::Milou:
0130   {
0131     erhic::EventMilou *gen = dynamic_cast<erhic::EventMilou *>(GenEvent);
0132     evthead->set_eventgenerator_type(EicEventHeader::EvtGen::Milou);
0133     evthead->set_milou_weight(gen->weight);
0134     evthead->set_milou_trueX(gen->trueX);
0135     evthead->set_milou_trueQ2(gen->trueQ2);
0136   }
0137   break;
0138   case EvtGen::DEMP:
0139   {
0140     erhic::EventDEMP *gen = dynamic_cast<erhic::EventDEMP *>(GenEvent);
0141     evthead->set_eventgenerator_type(EicEventHeader::EvtGen::DEMP);
0142     evthead->set_demp_weight(gen->weight);
0143   }
0144   break;
0145   case EvtGen::Unknown:
0146     std::cout << "unknown event generator" << std::endl;
0147     break;
0148   default:
0149     std::cout << "what is this " << m_EvtGenId << " ????" << std::endl;
0150     break;
0151   }
0152
0153   /* Create GenEvent */
0154   HepMC::GenEvent *evt = new HepMC::GenEvent();
0155
0156   /* define the units (Pythia uses GeV and mm) */
0157   evt->use_units(HepMC::Units::GEV, HepMC::Units::MM);
0158
0159   /* add global information to the event */
0160   evt->set_event_number(entry);
0161
0162   /* process ID from pythia */
0163   evt->set_signal_process_id(GenEvent->GetProcess());
0164
0165   /* Set the PDF information */
0166   HepMC::PdfInfo pdfinfo;
0167   pdfinfo.set_x1(1);
0168   pdfinfo.set_x2(GenEvent->GetX());
0169   pdfinfo.set_scalePDF(GenEvent->GetQ2());
0170   evt->set_pdf_info(pdfinfo);
0171
0172   /* Loop over all particles for this event in input file and fill
0173    * vector with HepMC particles */
0174   std::vector<HepMC::GenParticle *> hepmc_particles;
0175   std::vector<unsigned> origin_index;
0176
0177   /* save pointers to beam particles */
0178   HepMC::GenParticle *hepmc_beam1 = nullptr;
0179   HepMC::GenParticle *hepmc_beam2 = nullptr;
0180
0181   for (unsigned ii = 0; ii < GenEvent->GetNTracks(); ii++)
0182   {
0183     /* Get particle / track from event records.
0184      * Class documentation for erhic::VirtualParticle at
0185      * http://www.star.bnl.gov/~tpb/eic-smear/classerhic_1_1_virtual_particle.html */
0186     erhic::ParticleMC *track_ii = GenEvent->GetTrack(ii);
0187
0188     if (Verbosity() > 1)
0189     {
0190       std::cout << __PRETTY_FUNCTION__ << " : " << __LINE__ << std::endl;
0191       std::cout << "\ttrack " << ii << std::endl;
0192       std::cout << "\t4mom\t" << track_ii->GetPx() << "\t" << track_ii->GetPy() << "\t" << track_ii->GetPz() << "\t" << track_ii->GetE() << std::endl;
0193       std::cout << "\tstatus= " << track_ii->GetStatus() << "\tindex= " << track_ii->GetIndex() << "\tmass= " << track_ii->GetM() << std::endl;
0194     }
0195
0196     /* Create HepMC particle record */
0197     HepMC::GenParticle *hepmcpart = new HepMC::GenParticle(HepMC::FourVector(track_ii->GetPx(),
0198                                                                              track_ii->GetPy(),
0199                                                                              track_ii->GetPz(),
0200                                                                              track_ii->GetE()),
0201                                                            track_ii->Id());
0202
0203     /* translate eic-smear status codes to HepMC status codes */
0204     switch (track_ii->GetStatus())
0205     {
0206     case 1:
0207       hepmcpart->set_status(1);
0208       break;  // 'stable particle'
0209
0210     case 21:
0211       hepmcpart->set_status(3);
0212       break;  // 'documentation line'
0213
0214     default:
0215       hepmcpart->set_status(0);
0216       break;  // 'null entry'
0217     }
0218
0219     /* assume the first two particles are the beam particles (which getsHepMC status 4)*/
0220     if (ii < 2)
0221     {
0222       hepmcpart->set_status(4);
0223     }
0224
0225     /* add particle information */
0226     hepmcpart->setGeneratedMass(track_ii->GetM());
0227
0228     /* append particle to vector */
0229     hepmc_particles.push_back(hepmcpart);
0230     origin_index.push_back(track_ii->GetIndex());
0231
0232     /* if first particle, call this the first beam particle */
0233     if (ii == 0)
0234     {
0235       hepmc_beam1 = hepmcpart;
0236     }
0237
0238     /* if second particle, call this the second beam particle */
0239     if (ii == 1)
0240     {
0241       hepmc_beam2 = hepmcpart;
0242     }
0243   }
0244
0245   /* Check if hepmc_particles and origin_index vectors are the same size */
0246   if (hepmc_particles.size() != origin_index.size())
0247   {
0248     std::cout << "ReadEICFiles::process_event - Lengths of HepMC particles and Origin index vectors do not match!" << std::endl;
0249
0250     delete evt;
0251     return Fun4AllReturnCodes::ABORTRUN;
0252   }
0253
0254   /* add HepMC particles to Hep MC vertices; skip first two particles
0255    * in loop, assuming that they are the beam particles */
0256   std::vector<HepMC::GenVertex *> hepmc_vertices;
0257
0258   for (unsigned p = 2; p < hepmc_particles.size(); p++)
0259   {
0260     HepMC::GenParticle *pp = hepmc_particles.at(p);
0261
0262     /* continue if vertices for particle are already set */
0263     if (pp->production_vertex() && pp->end_vertex())
0264     {
0265       continue;
0266     }
0267
0268     /* access mother particle vertex */
0269     erhic::ParticleMC *track_pp = GenEvent->GetTrack(p);
0270
0271     unsigned parent_index = track_pp->GetParentIndex();
0272
0273     HepMC::GenParticle *pmother = nullptr;
0274     for (unsigned m = 0; m < hepmc_particles.size(); m++)
0275     {
0276       if (origin_index.at(m) == parent_index)
0277       {
0278         pmother = hepmc_particles.at(m);
0279         break;
0280       }
0281     }
0282
0283     /* if mother does not exist: create new vertex and add this particle as outgoing to vertex */
0284     if (!pmother)
0285     {
0286       HepMC::GenVertex *hepmcvtx = new HepMC::GenVertex(HepMC::FourVector(track_pp->GetVertex()[0],
0287                                                                           track_pp->GetVertex()[1],
0288                                                                           track_pp->GetVertex()[2],
0289                                                                           0));
0290       hepmc_vertices.push_back(hepmcvtx);
0291       hepmcvtx->add_particle_out(pp);
0292       continue;
0293     }
0294     /* if mother exists and has end vertex: add this particle as outgoing to the mother's end vertex */
0295     if (pmother->end_vertex())
0296     {
0297       pmother->end_vertex()->add_particle_out(pp);
0298     }
0299     /* if mother exists and has no end vertex: create new vertex */
0300     else
0301     {
0302       HepMC::GenVertex *hepmcvtx = new HepMC::GenVertex(HepMC::FourVector(track_pp->GetVertex()[0],
0303                                                                           track_pp->GetVertex()[1],
0304                                                                           track_pp->GetVertex()[2],
0305                                                                           0));
0306       hepmc_vertices.push_back(hepmcvtx);
0307       hepmcvtx->add_particle_in(pmother);
0308       pmother->end_vertex()->add_particle_out(pp);
0309     }
0310   }
0311
0312   /* Add end vertex to beam particles if they don't have one yet */
0313   for (unsigned p = 0; p < 2; p++)
0314   {
0315     HepMC::GenParticle *pp = hepmc_particles.at(p);
0316
0317     if (!pp->end_vertex())
0318     {
0319       /* create collision vertex */
0320       HepMC::GenVertex *hepmcvtx = new HepMC::GenVertex(HepMC::FourVector(0,
0321                                                                           0,
0322                                                                           0,
0323                                                                           0));
0324       hepmc_vertices.push_back(hepmcvtx);
0325       hepmcvtx->add_particle_in(pp);
0326     }
0327   }
0328
0329   /* Check that all particles (except beam particles) have a production vertex */
0330   for (unsigned p = 2; p < hepmc_particles.size(); p++)
0331   {
0332     if (!hepmc_particles.at(p)->production_vertex())
0333     {
0334       std::cout << "ReadEICFiles::process_event - Missing production vertex for one or more non-beam particles!" << std::endl;
0335       delete evt;
0336       return Fun4AllReturnCodes::ABORTRUN;
0337     }
0338   }
0339
0340   /* Add HepMC vertices to event */
0341   for (auto &hepmc_vertice : hepmc_vertices)
0342   {
0343     evt->add_vertex(hepmc_vertice);
0344   }
0345
0346   /* set beam particles */
0347   evt->set_beam_particles(hepmc_beam1, hepmc_beam2);
0348
0349   /* pass HepMC to PHNode*/
0350   PHHepMCGenEvent *success = PHHepMCGenHelper::insert_event(evt);
0351   if (Verbosity() > 1)
0352   {
0353     std::cout << __PRETTY_FUNCTION__ << " : " << __LINE__ << std::endl;
0354     evt->print();
0355     std::cout << std::endl
0356               << std::endl;
0357   }
0358
0359   if (!success)
0360   {
0361     std::cout << "ReadEICFiles::process_event - Failed to add event to HepMC record!" << std::endl;
0362     return Fun4AllReturnCodes::ABORTRUN;
0363   }
0364   /* Count up number of 'used' events from input file */
0365   entry++;
0366
0367   /* Done */
0368   return 0;
0369 }
0370
0371 int ReadEICFiles::CreateNodeTree(PHCompositeNode *topNode)
0372 {
0373   PHHepMCGenHelper::create_node_tree(topNode);
0374
0375   PHNodeIterator iter(topNode);
0376   // Looking for the DST node
0377   PHCompositeNode *dstNode = dynamic_cast<PHCompositeNode *>(iter.findFirst("PHCompositeNode", "DST"));
0378   if (!dstNode)
0379   {
0380     std::cout << PHWHERE << "DST Node missing, doing nothing." << std::endl;
0381     return Fun4AllReturnCodes::ABORTRUN;
0382   }
0383   EicEventHeader *evthead = findNode::getClass<EicEventHeader>(topNode, "EicEventHeader");
0384   if (!evthead)
0385   {
0386     evthead = new EicEventHeaderv1();
0387     PHIODataNode<PHObject> *HeaderNode = new PHIODataNode<PHObject>(evthead, "EicEventHeader", "PHObject");
0388     dstNode->addNode(HeaderNode);
0389   }
0390   return Fun4AllReturnCodes::EVENT_OK;
0391 }