source: svn/trunk/src/HepMCConverter.cc@ 606

/***********************************************************************
**                                                                    **
**          /----------------------------------------------\          **
**         |  Delphes, a framework for the fast simulation  |         **
**         |       of a  generic collider experiment        |         **
**          \-------------  arXiv:0903.2225v1  ------------/          **
**                                                                    **
**                                                                    **
**   This package uses:                                               **
**   ------------------                                               **
**   ROOT: Nucl. Inst. & Meth. in Phys. Res. A389 (1997) 81-86        **
**   FastJet algorithm: Phys. Lett. B641 (2006) [hep-ph/0512210]      **
**   Hector: JINST 2:P09005 (2007) [physics.acc-ph:0707.1198v2]       **
**   FROG: [hep-ex/0901.2718v1]                                       **
**   HepMC: Comput. Phys. Commun.134 (2001) 41                        **
**                                                                    **
** ------------------------------------------------------------------ **
**                                                                    **
**   Main authors:                                                    **
**   -------------                                                    **
**                                                                    **
**                Severine Ovyn                Xavier Rouby           **
**          severine.ovyn@uclouvain.be      xavier.rouby@cern         **
**                                                                    **
**         Center for Particle Physics and Phenomenology (CP3)        **
**                Universite catholique de Louvain (UCL)              **
**                     Louvain-la-Neuve, Belgium                      **
**                                                                    **
**                      Copyright (C) 2008-2009,                      **
**                        All rights reserved.                        **
**                                                                    **
***********************************************************************/

#include <iostream>
#include <fstream>
#include "TLorentzVector.h"
#include "BlockClasses.h"

#include "PdgParticle.h"
#include "ExRootTreeWriter.h"
#include "ExRootTreeBranch.h"
#include "HepMCConverter.h"

#include "GenParticle.h"
#include "GenVertex.h"
#include "IO_AsciiParticles.h"
#include "IO_GenEvent.h"

//-------------------------------------------------------------------------
int HepMCConverter::find_in_map( const std::map<HepMC::GenParticle*,int>& m, HepMC::GenParticle *p) const
{
  std::map<HepMC::GenParticle*,int>::const_iterator iter = m.find(p);
  return (iter == m.end()) ? 0 : iter->second;
}

//--------------------------------------------------------------------------
void  HepMCConverter::ReadStats()
{

 unsigned int particle_counter=0;
 index_to_particle.clear(); 
 particle_to_index.clear();
 index_to_particle.reserve(evt->particles_size()); 
 index_to_particle[0] = 0; 
 HepMC::GenEvent::vertex_const_iterator v;
 for (v = evt->vertices_begin(); v != evt->vertices_end(); ++v ) 
   {
     // making a list of incoming particles of the vertices
     // so that the mother indices in HEPEVT can be filled properly
     HepMC::GenVertex::particles_out_const_iterator p1;
     for (p1 = (*v)->particles_in_const_begin();p1 != (*v)->particles_in_const_end(); ++p1 ) 
       {

         ++particle_counter;
         //particle_counter can be very large for heavy ions
         if(particle_counter >= index_to_particle.size() ) 
           {
             //make it large enough to hold up to this index
             index_to_particle.resize(particle_counter+1);
           }                 
         index_to_particle[particle_counter] = *p1;
         particle_to_index[*p1] = particle_counter;
       }   
     // daughters are entered only if they aren't a mother of
     // another vertex
   HepMC::GenVertex::particles_out_const_iterator p2;
     for (p2 = (*v)->particles_out_const_begin();p2 != (*v)->particles_out_const_end(); ++p2) 
       {
         if (!(*p2)->end_vertex()) 
           {
             ++particle_counter;
             //particle_counter can be very large for heavy ions
             if(particle_counter  >= index_to_particle.size() ) 
               {        
                 //make it large enough to hold up to this index
                 index_to_particle.resize(particle_counter+1);
               }                   
             index_to_particle[particle_counter] = *p2;
             particle_to_index[*p2] = particle_counter;
           }
       }

   } 
}


//-------------------------------------------------------------------------
void HepMCConverter::getStatsFromTuple(int &mo1, int &mo2, int &da1, int &da2, int &status, int &pid, int j) const
{
  if (!evt) 
    {
      cout <<  "HepMCFileReader: Got no event :-(  Game over already  ?" <<endl; 
    } 
  else 
    {
      status =  index_to_particle[j]->status();
      pid = index_to_particle[j]->pdg_id();
      if ( index_to_particle[j]->production_vertex() ) 
        {
          int num_mothers = index_to_particle[j]->production_vertex()->particles_in_size();
          if (num_mothers ==0) { 
            mo1 = 0; 
            mo2 = 0; 
          }
          else {
            int first_mother = find_in_map( particle_to_index,*(index_to_particle[j]->production_vertex()->particles_in_const_begin()));
            int last_mother = first_mother + num_mothers - 1;
            if ( first_mother == 0 ) last_mother = 0;
            mo1=first_mother;
            mo2=last_mother;
          } // if num_mothers !=0
        } 
      else // no data on production_vertex
        {
          mo1 =0;
          mo2 =0;
        }   
       if (index_to_particle[j]->end_vertex())
        {
          // make sure first and last daughter are indeed the first and last
          int first_daughter = find_in_map( particle_to_index,*(index_to_particle[j]->end_vertex()->particles_begin(HepMC::children)));
          int last_daughter = 0;

          HepMC::GenVertex::particle_iterator ic;
          for (ic = index_to_particle[j]->end_vertex()->particles_begin(HepMC::children);ic != index_to_particle[j]->end_vertex()->particles_end(HepMC::children); ++ic) {
            int current_daughter = find_in_map( particle_to_index,*ic) ; 
            if (current_daughter < first_daughter)
              first_daughter = current_daughter;
            if (current_daughter > last_daughter)
              last_daughter = current_daughter;
          }
          if (first_daughter== 0) last_daughter = 0;
          da1=first_daughter;
          da2=last_daughter;

        } 
      else 
        {
          da1=0;
          da2=0;
        }
    } 
} 



//---------------------------------------------------------------------------

void HepMCConverter::AnalyseEvent(ExRootTreeBranch *branch,HepMC::GenEvent& evt,const Long64_t eventNumber)
{
  TRootLHEFEvent *element;
  
  element = static_cast<TRootLHEFEvent*>(branch->NewEntry());
  element->Number = eventNumber;
  element->Nparticles = evt.particles_size();
  element->ProcessID = evt.signal_process_id();
  //  element->Weight = hepeup.XWGTUP;
  element->ScalePDF = evt.event_scale();
  element->CouplingQED = evt.alphaQED();
  element->CouplingQCD = evt.alphaQCD();
  
}

//---------------------------------------------------------------------------

void HepMCConverter::AnalyseParticles(ExRootTreeBranch *branch, const HepMC::GenEvent& evt)
{
  TRootC::GenParticle *element;

  TLorentzVector momentum;
  Double_t signPz;

  ReadStats();
  for(int n=1; n<=evt.particles_size(); n++)
    {
      getStatsFromTuple( mo1,mo2,da1,da2,status,pid,n);
    
      element = static_cast<TRootC::GenParticle*>(branch->NewEntry());

      element->PID = pid;
      element->Status = status;
      element->M1 = mo1 - 1; // added -1 as the numbering in the tree starts from 0
      element->M2 = mo2 - 1;
      element->D1 = da1 - 1;
      element->D2 = da2 - 1;

      element->E = index_to_particle[n]->momentum().e();
      element->Px = index_to_particle[n]->momentum().px();
      element->Py = index_to_particle[n]->momentum().py();
      element->Pz = index_to_particle[n]->momentum().pz();


      //cout << "element->PID = " << pid << "\t"; 
      //PdgParticle pdg_part(PdgID[pid]);
      //element->M =  index_to_particle[n]->momentum().m(); // this is the particle virtuality, not its rest mass
      //element->M      = pdg_part.mass();
      //element->Charge = pdg_part.charge();
      //cout << "element->M = " << element->M << " \t element->Charge = " << element->Charge << endl;

      element->PT = sqrt(pow(element->Px,2)+pow(element->Py,2));

      momentum.SetPxPyPzE(element->Px, element->Py, element->Pz, element->E);
      signPz = (element->Pz >= 0.0) ? 1.0 : -1.0;
      //element->Eta = element->PT == 0.0 ? signPz*999.9 : momentum.Eta(); to avoid a warning from ROOT, replace the "==0" by "< 1e-6"
      element->Eta = element->PT < 1e-6 ? signPz*999.9 : momentum.Eta();
      element->Phi = index_to_particle[n]->momentum().phi();
 
    //In particle at vertex
    HepMC::GenVertex* vrtI = (index_to_particle[n])->production_vertex();
    HepMC::GenVertex::particles_in_const_iterator partI;

    if(vrtI)
      {
        element->T = vrtI->position().t();
        element->X = vrtI->position().x();
        element->Y = vrtI->position().y();
        element->Z = vrtI->position().z();
      }
    else
     {
        element->T = 0.;
        element->X = 0.;
        element->Y = 0.;
        element->Z = 0.;
     }  
  }
}

//

//------------------------------------------------------------------------------

HepMCConverter::~HepMCConverter()
{
   delete evt;

 /*cout << "delete index to particle" << endl;
 std::vector<HepMC::GenParticle*>::iterator i;
 / for (i=index_to_particle.begin();i != index_to_particle.end(); i++) delete (*i);
 cout << "done" << endl;

 std::map<HepMC::GenParticle*,int>::iterator j;
 for (j=particle_to_index.begin(); j != particle_to_index.end(); j++) delete j->first;
 */

}

//------------------------------------------------------------------------------

HepMCConverter::HepMCConverter(const string& inputFileList, const string& outputFileName, const PdgTable& pdg, const int& Nevents) : DataConverter(pdg,Nevents)
{
  
  ExRootTreeWriter *treeWriter = new ExRootTreeWriter(outputFileName, "GEN");
  // information about generated event
  ExRootTreeBranch *branchGenEvent = treeWriter->NewBranch("Event", TRootLHEFEvent::Class());
  // generated particles from HEPEVT
  ExRootTreeBranch *branchGenParticle = treeWriter->NewBranch("Particle", TRootC::GenParticle::Class());
 
  int nevt_already_processed=0;
 
  // Open a stream connected to an event file:
  ifstream infile(inputFileList.c_str());
  string filename;
  if(!infile.is_open()) {
     cerr << left  << setw(30) <<"**   ERROR: Can't open     "<<""
          << left  << setw(20) << inputFileList                        <<""
          << right << setw(19) <<"for input **"<<""<<endl;
     exit(1);
  }
 
  Long64_t entry = 0;
  while(1)
    {
      infile >> filename;
      if(!infile.good()) break; // end of listfile reached
      if (Nevt>0 && nevt_already_processed >=Nevt) break; // enough events already processed   

      ifstream checking_the_file(filename.c_str());
      if(!checking_the_file.good()) 
        { 
          cerr << left  << setw(30) <<"**   ERROR: Can't find file  "<<""
               << left  << setw(36) << filename                      <<""
               << right << setw(3) <<" **"<<""<<endl;
          continue;
        }
      else checking_the_file.close();
      
      //Open the file
      HepMC::IO_GenEvent ascii_in(filename,std::ios::in);
      // get the first event

      evt = ascii_in.read_next_event();
      
      while ( evt ) 
        {
          if(Nevt>0 && entry>=Nevt) break;
          if(Nevt>0 && nevt_already_processed >=Nevt) break; // enough events already processed

          treeWriter->Clear();
          AnalyseEvent(branchGenEvent, *evt,entry+1);
          AnalyseParticles(branchGenParticle, *evt);      
          delete evt;
          // read the next event
          ascii_in >> evt;
          treeWriter->Fill();
          ++entry;
          ++nevt_already_processed;
        }
    }
  
  treeWriter->Write();
  delete treeWriter;
  
}