generators/phhepmc/PHHepMCParticleSelectorDecayProductChain.cc

0001 #include "PHHepMCParticleSelectorDecayProductChain.h"
0002
0003 #include "PHHepMCGenEvent.h"
0004 #include "PHHepMCGenEventMap.h"
0005
0006 #include <fun4all/Fun4AllReturnCodes.h>
0007 #include <fun4all/SubsysReco.h>  // for SubsysReco
0008
0009 #include <phool/getClass.h>
0010 #include <phool/phool.h>  // for PHWHERE
0011
0012 #include <HepMC/GenEvent.h>       // for GenEvent::particle_const_ite...
0013 #include <HepMC/GenParticle.h>    // for GenParticle
0014 #include <HepMC/GenVertex.h>      // for GenVertex, GenVertex::partic...
0015 #include <HepMC/IteratorRange.h>  // for ancestors, children, descend...
0016 #include <HepMC/SimpleVector.h>   // for FourVector
0017
0018 #include <cstdlib>   // for abs
0019 #include <iostream>  // for operator<<, basic_ostream::o...
0020
0021 class PHCompositeNode;
0022
0023 PHHepMCParticleSelectorDecayProductChain::PHHepMCParticleSelectorDecayProductChain(const std::string& name)
0024   : SubsysReco(name)
0025 {
0026   return;
0027 }
0028
0029 HepMC::GenParticle* PHHepMCParticleSelectorDecayProductChain::GetParent(HepMC::GenParticle* p, HepMC::GenEvent* /*event*/)
0030 {
0031   HepMC::GenParticle* parent = nullptr;
0032   if (!p->production_vertex())
0033   {
0034     return parent;
0035   }
0036
0037   for (HepMC::GenVertex::particle_iterator mother = p->production_vertex()->particles_begin(HepMC::ancestors);
0038        mother != p->production_vertex()->particles_end(HepMC::ancestors);
0039        ++mother)
0040   {
0041     for (int _theAncestor : _theAncestors)
0042     {
0043       if (abs((*mother)->pdg_id()) == _theAncestor)
0044       {
0045         parent = *mother;
0046         break;
0047       }
0048     }
0049     if (parent != nullptr)
0050     {
0051       break;
0052     }
0053   }
0054
0055   return parent;
0056 }
0057
0058 int PHHepMCParticleSelectorDecayProductChain::process_event(PHCompositeNode* topNode)
0059 {
0060   if (_theParticle == 0 && _theDaughters.empty())
0061   {
0062     std::cout << PHWHERE << "Doing nothing." << std::endl;
0063     return Fun4AllReturnCodes::EVENT_OK;
0064   }
0065
0066   PHHepMCGenEventMap* geneventmap = findNode::getClass<PHHepMCGenEventMap>(topNode, "PHHepMCGenEventMap");
0067   if (!geneventmap)
0068   {
0069     std::cout << "ERROR: PHHepMCGenEventMap node not found!" << std::endl;
0070     return Fun4AllReturnCodes::ABORTEVENT;
0071   }
0072   PHHepMCGenEvent* inEvent = geneventmap->get(_embedding_id);
0073   if (!inEvent)
0074   {
0075     std::cout << "PHHepMCParticleSelectorDecayProductChain::process_event - WARNING: PHHepMCGenEvent with embedding ID " << _embedding_id << " is not found! Move on." << std::endl;
0076     return Fun4AllReturnCodes::DISCARDEVENT;
0077   }
0078
0079   HepMC::GenEvent* event = inEvent->getEvent();
0080   int npart = event->particles_size();
0081   int nvert = event->vertices_size();
0082   if (Verbosity() > 0)
0083   {
0084     std::cout << "=========== Event " << event->event_number() << " contains " << npart << " particles and " << nvert << " vertices." << std::endl;
0085   }
0086
0087   // list of vertices to keep
0088   std::vector<HepMC::GenVertex> vkeep;
0089
0090   if (_theParticle != 0)  // keep _theParticle and all its daughters (if any)
0091   {
0092     for (HepMC::GenEvent::particle_const_iterator p = event->particles_begin(); p != event->particles_end(); ++p)
0093     {
0094       // find _thePartcle
0095       if (abs((*p)->pdg_id()) == _theParticle)
0096       {
0097         // do we need to check for ancestors?
0098         if (!_theAncestors.empty())
0099         {
0100           HepMC::GenParticle* parent = GetParent(*p, event);
0101           if (parent)
0102           {
0103             vkeep.push_back(*(*p)->production_vertex());
0104           }
0105         }
0106         else
0107         {
0108           vkeep.push_back(*(*p)->production_vertex());
0109         }
0110
0111         // do we need to keep the daughters?
0112         if (!_theDaughters.empty())
0113         {
0114           if ((*p)->end_vertex())
0115           {
0116             for (HepMC::GenVertex::particle_iterator des = (*p)->end_vertex()->particles_begin(HepMC::descendants);
0117                  des != (*p)->end_vertex()->particles_end(HepMC::descendants);
0118                  ++des)
0119             {
0120               for (int _theDaughter : _theDaughters)
0121               {
0122                 if (abs((*des)->pdg_id()) == _theDaughter)
0123                 {
0124                   vkeep.push_back(*(*p)->end_vertex());
0125                   break;
0126                 }
0127               }
0128             }
0129           }
0130         }  // there are daughters
0131
0132       }  // this is _theParticle
0133     }  // end loop over particles
0134   }
0135   else  // save only particles in _theDaughters list no matter where they came from
0136   {
0137     for (HepMC::GenEvent::particle_const_iterator p = event->particles_begin(); p != event->particles_end(); ++p)
0138     {
0139       for (int _theDaughter : _theDaughters)
0140       {
0141         if (abs((*p)->pdg_id()) == _theDaughter)
0142         {
0143           vkeep.push_back(*(*p)->production_vertex());
0144         }
0145       }
0146     }
0147   }
0148
0149   // loop over vertices and keep only selected ones.
0150   for (HepMC::GenEvent::vertex_const_iterator v = event->vertices_begin(); v != event->vertices_end(); ++v)
0151   {
0152     bool goodvertex = false;
0153     for (const auto& tmp2 : vkeep)
0154     {
0155       HepMC::GenVertex tmp1 = (*(*v));
0156       if (tmp1 == tmp2)
0157       {
0158         goodvertex = true;
0159         break;
0160       }
0161     }
0162     if (!goodvertex)
0163     {
0164       bool tmp = event->remove_vertex((*v));
0165       if (Verbosity() > 10 && tmp)
0166       {
0167         std::cout << PHWHERE << " Erasing empty vertex." << std::endl;
0168       }
0169     }
0170   }
0171
0172   // clean up the vertices
0173   for (HepMC::GenEvent::vertex_const_iterator v = event->vertices_begin(); v != event->vertices_end(); ++v)
0174   {
0175     std::vector<HepMC::GenParticle*> removep;
0176
0177     for (HepMC::GenVertex::particle_iterator itpart = (*v)->particles_begin(HepMC::children);
0178          itpart != (*v)->particles_end(HepMC::children);
0179          ++itpart)
0180     {
0181       bool keepparticle = false;
0182       if (abs((*itpart)->pdg_id()) == _theParticle)
0183       {
0184         keepparticle = true;
0185       }
0186       for (int _theDaughter : _theDaughters)
0187       {
0188         if (abs((*itpart)->pdg_id()) == _theDaughter && (*itpart)->status() == 1)
0189         {
0190           keepparticle = true;
0191         }
0192       }
0193       if (!keepparticle)
0194       {
0195         removep.push_back((*itpart));
0196       }
0197     }  // end loop over particles in this vertex
0198
0199     for (HepMC::GenVertex::particle_iterator itpart = (*v)->particles_begin(HepMC::parents);
0200          itpart != (*v)->particles_end(HepMC::parents);
0201          ++itpart)
0202     {
0203       bool keepparticle = false;
0204       if (abs((*itpart)->pdg_id()) == _theParticle)
0205       {
0206         keepparticle = true;
0207       }
0208       for (int _theDaughter : _theDaughters)
0209       {
0210         if (abs((*itpart)->pdg_id()) == _theDaughter && (*itpart)->status() == 1)
0211         {
0212           keepparticle = true;
0213         }
0214       }
0215       if (!keepparticle)
0216       {
0217         removep.push_back((*itpart));
0218       }
0219     }  // end loop over particles in this vertex
0220
0221     for (auto& k : removep)
0222     {
0223       HepMC::GenParticle* tmp = (*v)->remove_particle(k);
0224       if (tmp->end_vertex())
0225       {
0226         delete tmp->end_vertex();
0227       }
0228     }
0229   }
0230
0231   int partcount = 0;
0232   if (Verbosity() > 0)
0233   {
0234     std::cout << "FINAL Event " << event->event_number() << " contains " << event->particles_size() << " particles and " << event->vertices_size() << " vertices." << std::endl;
0235     std::cout << "FINAL LIST OF PARTICLES:" << std::endl;
0236   }
0237   for (HepMC::GenEvent::particle_const_iterator p = event->particles_begin(); p != event->particles_end(); ++p)
0238   {
0239     int pid = (*p)->pdg_id();
0240     int status = (*p)->status();
0241     double pz = ((*p)->momentum()).pz();
0242     double pt = ((*p)->momentum()).perp();
0243     double eta = ((*p)->momentum()).eta();
0244     double mass = ((*p)->momentum()).m();
0245     if (Verbosity() > 0)
0246     {
0247       std::cout << pid << " " << mass << " " << status << " " << pt << " " << pz << " " << eta << " " << (*p)->production_vertex() << " " << (*p)->end_vertex() << std::endl;
0248     }
0249     partcount++;
0250   }
0251
0252   // if there is nothing to write out the code crashes
0253   if (partcount == 0)
0254   {
0255     if (Verbosity() > 0)
0256     {
0257       std::cout << "EVENT ABORTED: No particles to write out." << std::endl;
0258     }
0259     return Fun4AllReturnCodes::ABORTEVENT;
0260   }
0261
0262   return Fun4AllReturnCodes::EVENT_OK;
0263 }
0264
0265 void PHHepMCParticleSelectorDecayProductChain::SetParticle(const int pid)
0266 {
0267   _theParticle = pid;
0268   return;
0269 }
0270
0271 void PHHepMCParticleSelectorDecayProductChain::AddAncestor(const int pid)
0272 {
0273   _theAncestors.push_back(pid);
0274   return;
0275 }
0276
0277 void PHHepMCParticleSelectorDecayProductChain::AddDaughter(const int pid)
0278 {
0279   _theDaughters.push_back(pid);
0280   return;
0281 }