NtupleWriter13/NtupleWriter13-00-01-00/src/JetTruthTrackSelectorTool.cxx

#include "GaudiKernel/Service.h"
#include "GaudiKernel/MsgStream.h"

#include "StoreGate/StoreGateSvc.h"

#include "NavFourMom/INavigable4MomentumCollection.h"

#include "HepMC/GenParticle.h"

#include "TruthHelper/IsGenInteracting.h"
#include "TruthHelper/GenAccessIO.h"

#include "CLHEP/Vector/LorentzVector.h"
#include "CLHEP/Units/SystemOfUnits.h"
#include "CLHEP/HepPDT/ParticleID.hh"
#include "CLHEP/Vector/ThreeVector.h"

#include "JetUtils/JetCollectionHelper.h"

#include "JetEvent/Jet.h"
#include "JetEvent/JetCollection.h"

#include "JetRec/JetAlgToolBase.h"
//#include "JetSimTools/JetTruthTrackSelectorTool.h"
//#include "JetSimTools/MyTruthTrackJetInfoContainer.h"
#include "NtupleWriter13/JetTruthTrackSelectorTool.h"
#include "NtupleWriter13/MyTruthTrackJetInfoContainer.h"

// track particles //
#include "Particle/TrackParticleContainer.h"

#include <string>
#include <vector>
#include <cmath>
#include <fstream>

using namespace Rec;
using namespace std;

JetTruthTrackSelectorTool::JetTruthTrackSelectorTool(const std::string& name,
                                     const std::string& type,
                                     const IInterface* parent)
  : JetAlgToolBase(name,type,parent)
    , m_ptMin(0.*MeV)
    , m_etaMax(5.)
    , m_includeMuons(true)
{
  declareProperty("MinPt", m_ptMin);
  declareProperty("MaxEta", m_etaMax);
  declareProperty("IncludeMuons", m_includeMuons);
  declareProperty("OutputCollectionName",m_outputCollectionName);

  m_ignore_multivertices = false;
  declareProperty("ignoreMultivertices", m_ignore_multivertices);

  m_Et_min_vert = 1000*MeV;
  declareProperty("MinEtVertex", m_Et_min_vert);

  m_vertex_cut = 0.1*mm;
  declareProperty("VertexCut", m_vertex_cut);

  m_remove_isol_leptons = false;
  declareProperty("removeIsolLeptons", m_remove_isol_leptons);

  m_electron_cone   = 0.;
  m_electron_et_cut = 0.;
  m_muon_cone       = 0.;
  m_muon_et_cut     = 0.;
  m_inner_cone      = 0.01;

  declareProperty("electronCone",  m_electron_cone);
  declareProperty("electronEtCut", m_electron_et_cut);
  declareProperty("muonCone",      m_muon_cone);
  declareProperty("muonEtCut",     m_muon_et_cut);
  declareProperty("innerCone",     m_inner_cone);

}

JetTruthTrackSelectorTool::~JetTruthTrackSelectorTool()
{ }


StatusCode JetTruthTrackSelectorTool::initialize()
{
  MsgStream report(msgSvc(),name());

  // allocate StoreGate service
  StatusCode checkOut = service("StoreGateSvc",m_storeGate);
  if ( checkOut.isFailure() )
    {
      report << MSG::ERROR
             << "cannot allocate StoreGate service"
             << endreq;
      return checkOut;
    }

  // allocate truth service
  m_truthAccess = new GenAccessIO();

  // check
  if ( m_truthAccess == 0 )
    {
      report << MSG::ERROR
             << "cannot allocate GenAccessIO service."
             << endreq;
      return StatusCode::FAILURE;
    }

  return checkOut;
}


StatusCode JetTruthTrackSelectorTool::execute(JetAlgToolBase::jetcollection_t* 
                                      theJets)
{
  // message service
  MsgStream report(msgSvc(),name());
  //static std::ofstream out_e("et_incone_truth_trkjets_e.txt");
  //static std::ofstream out_mu("et_incone_truth_trkjets_mu.txt");
  //static std::ofstream out("debug_truth.txt");
//   const TrackParticleContainer *track_particle_container(0);
//                                      // pointer to the track particles
//      StatusCode sc(m_storeGate->retrieve(track_particle_container,
//                                      "TrackParticleCandidate"));
//      if (sc.isFailure() || track_particle_container==0) {
//              report << MSG::ERROR
//                      << "Could not retrieve TrackParticleContainer"
//                      << endreq;
//              return sc;
//      }

  Hep3Vector primary_vertex(0.0, 0.0, 0.0); // primary vertex of highest
                                                  // momentum

  // retrieve truth
  std::vector<const HepMC::GenParticle*> theParticles;
  IsGenInteracting theGenInteracting ;
  StatusCode checkOut = 
    m_truthAccess->getMC(theParticles, &theGenInteracting );

  // check
  if ( theParticles.size() == 0 || checkOut.isFailure() )
    {
      report << MSG::WARNING
             << "no truth particles found in this event"
             << endreq;
      return StatusCode::SUCCESS;
    }

  // get the primary vertex with the highest Et sum
  //if (!m_ignore_multivertices) {
        primary_vertex = highest_Et_vertex(theParticles);
        //}

  // create new jet collection
  jetcollection_t* writeCollection = new jetcollection_t();

  // loop on particle
  std::vector<const HepMC::GenParticle*>::iterator 
    firstParticle = theParticles.begin();
  std::vector<const HepMC::GenParticle*>::iterator 
    lastParticle = theParticles.end();

  
  for( ; firstParticle != lastParticle; firstParticle++ )
  {
      // select particles
      HepLorentzVector particleKine = (*firstParticle)->momentum();
      HepPDT::ParticleID IDEN((*firstParticle)->pdg_id());
      
      if ( particleKine.perp() > m_ptMin                   &&
           fabs(particleKine.pseudoRapidity()) < m_etaMax  &&
           ( abs((*firstParticle)->pdg_id()) != 13 || m_includeMuons ) &&
           IDEN.threeCharge()!=0 && (*firstParticle)->production_vertex()!=0)
      {
      
          //static std::ofstream out("eta_truth_trkjets.txt");
          //out << particleKine.eta() <<"\t\t" << particleKine.et() << endl;

          //     out << particleKine.eta() << endl;
    
//      double deltaR(1.0e99);
//      Hep3Vector comp_vert;
//      for (unsigned int k=0; k<track_particle_container->size(); k++) {
//              HepLorentzVector p((*track_particle_container)[k]->hlv());
//              if (p.deltaR(particleKine)<deltaR) {
//                      deltaR = p.deltaR(particleKine);
//                      comp_vert = (*track_particle_container
//                                      )[k]->reconstructedVertex()->position()-
//                                      (*firstParticle)->production_vertex(
//                                      )->position().vect();
//              }
//      }
//      out << deltaR << "\t" << comp_vert.x() << "\t" << comp_vert.y() << "\t"
//              << comp_vert.z() << endl;

      

          //check isolation
          if ( m_remove_isol_leptons ){

              double etcone(0);

              //electrons
              if(abs((*firstParticle)->pdg_id()) == 11){
                  etcone = get_energy_in_cone( m_electron_cone, particleKine,
                                               (*firstParticle)->production_vertex()->position().vect(),
                                               theParticles
                                               ); 
                  
                  //out_e << etcone << endl;
                  if( etcone < m_electron_et_cut ){
                      continue;
                  }
                  
              }
              else if(abs((*firstParticle)->pdg_id()) == 13){
                  etcone = get_energy_in_cone( m_muon_cone, particleKine,
                                               (*firstParticle)->production_vertex()->position().vect(),
                                               theParticles
                                               );
                  //out_mu << etcone << endl;
                  if( etcone < m_muon_et_cut ){
                      continue;
                  }
                  
              }
          }

          
          // jet representing the particle
          if (m_ignore_multivertices) {
              writeCollection->push_back(new Jet(particleKine));
          } else {
              Hep3Vector delta((*firstParticle)->production_vertex(
                                        )->position().vect()-primary_vertex);
              if ((*firstParticle)->production_vertex()!=0 &&
                  //                    (*firstParticle)->production_vertex(
                  //                                    )->position().vect()==primary_vertex) {
                  delta.mag()<m_vertex_cut) {
                  writeCollection->push_back(new Jet(particleKine));
              }
          }
      }
  }
  
  // hand collection to StoreGate (explicit copy of each object!)
  if ( ! JetCollectionHelper::record_jets(m_storeGate,writeCollection,
                                          m_outputCollectionName) )
    {
      report << MSG::ERROR
             << "cannot convert initial jet list to JetCollection"
             << endreq;
      return StatusCode::FAILURE;
    }

  // clean up
  JetCollectionHelper::remove_object(writeCollection,
                                     writeCollection->begin(),
                                     writeCollection->end());
  delete writeCollection;

  // retrieve collection from StoreGate to satisfy Jet constructor
  const INavigable4MomentumCollection* readCollection;
  checkOut = m_storeGate->retrieve(readCollection,m_outputCollectionName);
  if ( checkOut.isFailure() )
    {
      report << MSG::ERROR
             << "cannot read back initial JetCollection"
             << endreq;
      return checkOut;
    }

  // build intermediate jet list
  for ( Jet::index_type theIndex=0; theIndex < readCollection->size();
        theIndex++ )
    {
      // jet with storable particle constituent
      theJets->push_back(new Jet(readCollection,theIndex));
    }

  //cout << "size: " <<  theJets->size() << endl;

  return checkOut;
}

//*****************************************************************************

//::::::::::::::::::::::::::::::
//:: METHOD highest_Et_vertex ::
//::::::::::::::::::::::::::::::

Hep3Vector JetTruthTrackSelectorTool::highest_Et_vertex(
                        std::vector<const HepMC::GenParticle*> theParticles) {

// VARIABLES //

        std::vector<Hep3Vector> vertex; // vector of vertices
        bool vertex_flag; // auxiliary vertex flag
        double max_Et; // maximum transerve vertex momentum found so far
        Hep3Vector primary_vertex(0.0, 0.0, 0.0); // auxiliary primary vertex
        std::vector<double> Et_vert; // Et(vertex)
        std::vector<unsigned> nb_part; // number of particles from the vertex
        MyTruthTrackJetInfoContainer *container(
                MyTruthTrackJetInfoContainer::getTruthTrackJetInfoContainer());
                                                // container for vertex data
        container->clear();

// COLLECT ALL VERTICES //

//      static std::ofstream out("debug_truth.txt");
        std::vector<const HepMC::GenParticle*>::iterator firstParticle(
                                                        theParticles.begin());
        std::vector<const HepMC::GenParticle*>::iterator lastParticle(
                                                        theParticles.end());
        for ( ; firstParticle != lastParticle; firstParticle++) {
                HepLorentzVector particleKine = (*firstParticle)->momentum();
                HepPDT::ParticleID IDEN((*firstParticle)->pdg_id());
                if (particleKine.perp()>m_Et_min_vert &&
                        fabs(particleKine.pseudoRapidity())<m_etaMax &&
                        (abs((*firstParticle)->pdg_id())!=13 || m_includeMuons)
                        && IDEN.threeCharge()!=0
                        && (*firstParticle)->production_vertex()!=0) {
//              out << particleKine.et() << std::endl;
                        vertex_flag = false;
                        for (unsigned int l=0; l<vertex.size(); l++) {
                                Hep3Vector delta(
                                        (*firstParticle)->production_vertex(
                                        )->position().vect()-vertex[l]);
//                              if ((*firstParticle)->production_vertex(
//                                      )->position().vect()==vertex[l]) {
                                if (delta.mag()<m_vertex_cut) {
                                        Et_vert[l] = Et_vert[l]+
                                                        particleKine.et();
                                        nb_part[l]++;
                                        vertex_flag = true;
                                        break;
                                }
                        }
                        if (!vertex_flag) {
                                vertex.push_back(
                                        (*firstParticle)->production_vertex(
                                                        )->position().vect());
                                Et_vert.push_back(particleKine.et());
                                nb_part.push_back(1);
                        }
                }
        }

        if (vertex.size()==0) { 
                return primary_vertex;
        }

// REJECT THOSE VERTICES FROM WHICH ONLY ONE PARTICLE EMERGES //

//      vector<Hep3Vector> new_vertex;
//      vector<double> new_Et_vert;
//      for (unsigned k=0; k<vertex.size(); k++) {
//              if (nb_part[k]>1) {
//                      new_vertex.push_back(vertex[k]);
//                      new_Et_vert.push_back(Et_vert[k]);
//              }
//      }

// GET THE PRIMARY VERTEX WITH THE HIGHEST TRANSVERSE MOMENTUM //


//      static ofstream outfile("debug.txt");
//      outfile << vertex.size() << "\t" << new_vertex.size() << endl;

        max_Et = Et_vert[0];
        primary_vertex = vertex[0];
        for (unsigned int k=1; k<vertex.size(); k++) {
                if (Et_vert[k]>max_Et) {
                        max_Et = Et_vert[k];
                        primary_vertex = vertex[k];
                }
        }

        container->setVertex(vertex);
        container->setEt(Et_vert);
        container->setNumberOfTracks(nb_part);

        return primary_vertex;

}
//*****************************************************************************

//:::::::::::::::::::::::::::::::
//:: METHOD get_energy_in_cone ::
//:::::::::::::::::::::::::::::::

double 
JetTruthTrackSelectorTool::get_energy_in_cone(const double & cone,
                                         const HepLorentzVector & p, 
                                         const Hep3Vector & vertex,
                                         vector<const HepMC::GenParticle*> theParticles) {

//-------------
// VARIABLES --
//-------------

    double ET(0.);   // transverse energy in the cone

//-------------------------------------------------------------
// LOOP OVER THE TRACKS AND CALCULATE THE ENERGY IN THE CONE --
//-------------------------------------------------------------
 
    for (unsigned int k=0; k<theParticles.size(); k++ ) {
     
        if(theParticles[k]->production_vertex()==NULL){
            continue;
        }
        
        if(theParticles[k]->momentum().eta()>m_etaMax){
            continue;
        }

        if (!m_ignore_multivertices){
            
            Hep3Vector delta = theParticles[k]
                ->production_vertex()->position().vect()-vertex;

            //if ( theParticles[k]
            //   ->production_vertex()->position().vect()!=vertex) {
            if (delta.mag()>m_vertex_cut){
                continue; 
            }
        }
    
        if ( p.deltaR(theParticles[k]->momentum()) < cone &&
             p.deltaR(theParticles[k]->momentum()) > m_inner_cone) {
            
            ET = ET+theParticles[k]->momentum().et();
        }
        
    }


    return ET;
    
}

---