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 using namespace std;
0024
0025 PHHepMCParticleSelectorDecayProductChain::PHHepMCParticleSelectorDecayProductChain(const string& name)
0026   : SubsysReco(name)
0027   , _embedding_id(0)
0028 {
0029   _theParticle = 11;
0030   return;
0031 }
0032
0033 int PHHepMCParticleSelectorDecayProductChain::InitRun(PHCompositeNode* /*topNode*/)
0034 {
0035   return Fun4AllReturnCodes::EVENT_OK;
0036 }
0037
0038 HepMC::GenParticle* PHHepMCParticleSelectorDecayProductChain::GetParent(HepMC::GenParticle* p, HepMC::GenEvent* /*event*/)
0039 {
0040   HepMC::GenParticle* parent = nullptr;
0041   if (!p->production_vertex()) return parent;
0042
0043   for (HepMC::GenVertex::particle_iterator mother = p->production_vertex()->particles_begin(HepMC::ancestors);
0044        mother != p->production_vertex()->particles_end(HepMC::ancestors);
0045        ++mother)
0046   {
0047     for (unsigned int i = 0; i < _theAncestors.size(); i++)
0048     {
0049       if (abs((*mother)->pdg_id()) == _theAncestors[i])
0050       {
0051         parent = *mother;
0052         break;
0053       }
0054     }
0055     if (parent != nullptr) break;
0056   }
0057
0058   return parent;
0059 }
0060
0061 int PHHepMCParticleSelectorDecayProductChain::process_event(PHCompositeNode* topNode)
0062 {
0063   if (_theParticle == 0 && _theDaughters.size() == 0)
0064   {
0065     cout << PHWHERE << "Doing nothing." << endl;
0066     return Fun4AllReturnCodes::EVENT_OK;
0067   }
0068
0069   PHHepMCGenEventMap* geneventmap = findNode::getClass<PHHepMCGenEventMap>(topNode, "PHHepMCGenEventMap");
0070   if (!geneventmap)
0071   {
0072     cerr << "ERROR: PHHepMCGenEventMap node not found!" << endl;
0073     return Fun4AllReturnCodes::ABORTEVENT;
0074   }
0075   PHHepMCGenEvent* inEvent = geneventmap->get(_embedding_id);
0076   if (!inEvent)
0077   {
0078     cerr << "PHHepMCParticleSelectorDecayProductChain::process_event - WARNING: PHHepMCGenEvent with embedding ID " << _embedding_id << " is not found! Move on." << endl;
0079     return Fun4AllReturnCodes::DISCARDEVENT;
0080   }
0081
0082   HepMC::GenEvent* event = inEvent->getEvent();
0083   int npart = event->particles_size();
0084   int nvert = event->vertices_size();
0085   if (Verbosity() > 0) cout << "=========== Event " << event->event_number() << " contains " << npart << " particles and " << nvert << " vertices." << endl;
0086
0087   // list of vertices to keep
0088   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.size() > 0)
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.size() > 0)
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 (unsigned int i = 0; i < _theDaughters.size(); i++)
0121               {
0122                 if (abs((*des)->pdg_id()) == _theDaughters[i])
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 (unsigned int ip = 0; ip < _theDaughters.size(); ip++)
0140       {
0141         if (abs((*p)->pdg_id()) == _theDaughters[ip])
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 (unsigned int i = 0; i < vkeep.size(); i++)
0154     {
0155       HepMC::GenVertex tmp1 = (*(*v));
0156       HepMC::GenVertex tmp2 = vkeep[i];
0157       if (tmp1 == tmp2)
0158       {
0159         goodvertex = true;
0160       }
0161     }
0162     if (!goodvertex)
0163     {
0164       bool tmp = event->remove_vertex((*v));
0165       if (Verbosity() > 10 && tmp)
0166       {
0167         cout << PHWHERE << " Erasing empty vertex." << 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 (unsigned int j = 0; j < _theDaughters.size(); j++)
0187       {
0188         if (abs((*itpart)->pdg_id()) == _theDaughters[j] && (*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 (unsigned int j = 0; j < _theDaughters.size(); j++)
0209       {
0210         if (abs((*itpart)->pdg_id()) == _theDaughters[j] && (*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 (unsigned int k = 0; k < removep.size(); k++)
0222     {
0223       HepMC::GenParticle* tmp = (*v)->remove_particle(removep[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     cout << "FINAL Event " << event->event_number() << " contains " << event->particles_size() << " particles and " << event->vertices_size() << " vertices." << endl;
0235     cout << "FINAL LIST OF PARTICLES:" << 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       cout << pid << " " << mass << " " << status << " " << pt << " " << pz << " " << eta << " " << (*p)->production_vertex() << " " << (*p)->end_vertex() << 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       cout << "EVENT ABORTED: No particles to write out." << endl;
0258     }
0259     return Fun4AllReturnCodes::ABORTEVENT;
0260   }
0261   else
0262   {
0263     return Fun4AllReturnCodes::EVENT_OK;
0264   }
0265 }
0266
0267 void PHHepMCParticleSelectorDecayProductChain::SetParticle(const int pid)
0268 {
0269   _theParticle = pid;
0270   return;
0271 }
0272
0273 void PHHepMCParticleSelectorDecayProductChain::AddAncestor(const int pid)
0274 {
0275   _theAncestors.push_back(pid);
0276   return;
0277 }
0278
0279 void PHHepMCParticleSelectorDecayProductChain::AddDaughter(const int pid)
0280 {
0281   _theDaughters.push_back(pid);
0282   return;
0283 }