PrimaryVertex.cxx

Go to the documentation of this file.

#include "CLHEP/Geometry/Point3D.h"
#include "CLHEP/Vector/LorentzVector.h"
#ifndef ENABLE_BACKWARDS_COMPATIBILITY
typedef HepGeom::Point3D<double> HepPoint3D;
#endif
#include "EventModel/Event.h"
#include "EventModel/EventHeader.h"
#include "EventModel/EventModel.h"
#include "EvtRecEvent/EvtRecEvent.h"
#include "EvtRecEvent/EvtRecPrimaryVertex.h"
#include "EvtRecEvent/EvtRecTrack.h"
#include "GaudiKernel/IDataManagerSvc.h"
#include "GaudiKernel/MsgStream.h"
#include "GaudiKernel/SmartDataPtr.h"
#include "ParticleID/ParticleID.h"
#include "PrimaryVertex.h"
#include "VertexFit/BField.h"
#include "VertexFit/FastVertexFit.h"
#include "VertexFit/HTrackParameter.h"
#include "VertexFit/KalmanVertexFit.h"
#include "VertexFit/KinematicFit.h"
#include "VertexFit/VertexFit.h"
 
#include "TF1.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TMath.h"
#include <assert.h>
#include <fstream>
#include <iostream>
#include <map>
 
using CLHEP::HepLorentzVector;
using namespace std;
 
const double xmass[5] = { 0.000511, 0.105658, 0.139570, 0.493677, 0.938272 }; // GeV
const double ecms     = 3.097;
const double mpi0     = 0.134976;
const double momega   = 0.78265;
 
typedef std::vector<int>              Vint;
typedef std::vector<HepLorentzVector> Vp4;
 
DECLARE_COMPONENT( PrimaryVertex )
//*******************************************************************************
PrimaryVertex::PrimaryVertex( const std::string& name, ISvcLocator* pSvcLocator )
    : Algorithm( name, pSvcLocator ) {
  // Declare the properties
  declareProperty( "Output", m_output = 0 );
  declareProperty( "TrackNumberCut", m_trackNumberCut = 1 );
  declareProperty( "CosThetaCut", m_cosThetaCut = 0.93 );
  declareProperty( "Vz0Cut", m_vz0Cut = 40.0 );
  // fit Method
  declareProperty( "FitMethod", m_fitMethod = 0 );
  // for global method
  declareProperty( "Chi2CutOfGlobal", m_globalChisqCut = 200.0 );
 
  // for Kalman method
  declareProperty( "ChisqCut", m_chisqCut = 200.0 );
  declareProperty( "TrackIteration", m_trackIteration = 5 );
  declareProperty( "VertexIteration", m_vertexIteration = 5 );
  declareProperty( "Chi2CutforTrkIter", m_chi2CutforTrkIter = 0.1 );
  declareProperty( "FreedomCut", m_freedomCut = 1 );
  declareProperty( "Chi2CutforSmooth", m_chi2CutforSmooth = 200.0 );
  // declareProperty("MinDistance", m_minDistance = 100.0);
  // declareProperty("MinPointX", m_minPointX = 0.1);
  // declareProperty("MinPointY", m_minPointY = 0.1);
}
 
//*******************************************************************************
StatusCode PrimaryVertex::initialize() {
  MsgStream log( msgSvc(), name() );
  log << MSG::INFO << "in initialize()" << endmsg;
 
  for ( int i = 0; i < 15; i++ ) { m_sel_number[i] = 0; }
 
  StatusCode status;
 
  if ( m_output == 1 )
  {
    NTuplePtr nt1( ntupleSvc(), "FILE999/glbvtx" );
    if ( nt1 ) m_tuple1 = nt1;
    else
    {
      m_tuple1 = ntupleSvc()->book( "FILE999/glbvtx", CLID_ColumnWiseTuple, "global vertex" );
      if ( m_tuple1 )
      {
        status = m_tuple1->addItem( "gvx", m_gvx );
        status = m_tuple1->addItem( "gvy", m_gvy );
        status = m_tuple1->addItem( "gvz", m_gvz );
        status = m_tuple1->addItem( "chig", m_chig );
        status = m_tuple1->addItem( "ndofg", m_ndofg );
        status = m_tuple1->addItem( "probg", m_probg );
      }
      else
      {
        log << MSG::ERROR << "Cannot book N-tuple:" << long( m_tuple1 ) << endmsg;
        return StatusCode::FAILURE;
      }
    }
 
    NTuplePtr nt2( ntupleSvc(), "FILE999/chisq" ); // chisq of Kalman filter
    if ( nt2 ) m_tuple2 = nt2;
    else
    {
      m_tuple2 = ntupleSvc()->book( "FILE999/chisq", CLID_ColumnWiseTuple, "chi-square of " );
      if ( m_tuple2 )
      {
        status = m_tuple2->addItem( "chis", m_chis );
        status = m_tuple2->addItem( "probs", m_probs );
        status = m_tuple2->addItem( "chif", m_chif );
        status = m_tuple2->addItem( "probf", m_probf );
      }
      else
      {
        log << MSG::ERROR << "Cannot book N-tuple:" << long( m_tuple2 ) << endmsg;
        return StatusCode::FAILURE;
      }
    }
 
    NTuplePtr nt3( ntupleSvc(), "FILE999/Pull" );
    if ( nt3 ) m_tuple3 = nt3;
    else
    {
      m_tuple3 = ntupleSvc()->book( "FILE999/Pull", CLID_ColumnWiseTuple, "Pull" );
      if ( m_tuple3 )
      {
        status = m_tuple3->addItem( "pull_drho", m_pull_drho );
        status = m_tuple3->addItem( "pull_phi", m_pull_phi );
        status = m_tuple3->addItem( "pull_kapha", m_pull_kapha );
        status = m_tuple3->addItem( "pull_dz", m_pull_dz );
        status = m_tuple3->addItem( "pull_lamb", m_pull_lamb );
        status = m_tuple3->addItem( "pull_momentum", m_pull_momentum );
      }
      else
      {
        log << MSG::ERROR << "Cannot book N-tuple:" << long( m_tuple3 ) << endmsg;
        return StatusCode::FAILURE;
      }
    }
 
    NTuplePtr nt4( ntupleSvc(), "FILE999/kalvtx" );
    if ( nt4 ) m_tuple4 = nt4;
    else
    {
      m_tuple4 = ntupleSvc()->book( "FILE999/kalvtx", CLID_ColumnWiseTuple, "kalman vertex" );
      if ( m_tuple4 )
      {
        status = m_tuple4->addItem( "kvx", m_kvx );
        status = m_tuple4->addItem( "kvy", m_kvy );
        status = m_tuple4->addItem( "kvz", m_kvz );
        status = m_tuple4->addItem( "chik", m_chik );
        status = m_tuple4->addItem( "ndofk", m_ndofk );
        status = m_tuple4->addItem( "probk", m_probk );
      }
      else
      {
        log << MSG::ERROR << "Cannot book N-tuple:" << long( m_tuple4 ) << endmsg;
        return StatusCode::FAILURE;
      }
    }
  } // end if output
 
  log << MSG::INFO << "successfully return from initialize()" << endmsg;
  return StatusCode::SUCCESS;
}
//*******************************************************************************
 
StatusCode
PrimaryVertex::RegisterEvtRecPrimaryVertex( EvtRecPrimaryVertex* aNewEvtRecPrimaryVertex,
                                            MsgStream&           log ) {
  DataObject* aEvtRecEvent;
  eventSvc()->findObject( "/Event/EvtRec", aEvtRecEvent );
  if ( aEvtRecEvent == NULL )
  {
    aEvtRecEvent  = new EvtRecEvent();
    StatusCode sc = eventSvc()->registerObject( "/Event/EvtRec", aEvtRecEvent );
    if ( sc != StatusCode::SUCCESS )
    {
      log << MSG::FATAL << "Could not register EvtRecEvent" << endmsg;
      return StatusCode::FAILURE;
    }
  }
  SmartIF<IDataManagerSvc> dataManSvc( eventSvc() );
  DataObject*              aEvtRecPrimaryVertex;
  eventSvc()->findObject( "/Event/EvtRec/EvtRecPrimaryVertex", aEvtRecPrimaryVertex );
  if ( aEvtRecPrimaryVertex != NULL )
  {
    // IDataManagerSvc *dataManSvc = dynamic_cast<IDataManagerSvc*> (eventSvc());
    dataManSvc->clearSubTree( "/Event/EvtRec/EvtRecPrimaryVertex" );
    eventSvc()->unregisterObject( "/Event/EvtRec/EvtRecPrimaryVertex" );
  }
  StatusCode sc = eventSvc()->registerObject( "/Event/EvtRec/EvtRecPrimaryVertex",
                                              aNewEvtRecPrimaryVertex );
  if ( sc != StatusCode::SUCCESS )
  {
    log << MSG::FATAL << "Could not register EvtRecPrimaryVertex in TDS!" << endmsg;
    return StatusCode::FAILURE;
  }
 
  return StatusCode::SUCCESS;
}
 
///////////////////////////////////////////////////
//  Select good charged tracks in MDC ( no PID )
//////////////////////////////////////////////////
void PrimaryVertex::SelectGoodChargedTracks( SmartDataPtr<EvtRecEvent>&    recEvent,
                                             SmartDataPtr<EvtRecTrackCol>& recTrackCol,
                                             Vint& icp, Vint& icm, Vint& iGood ) {
  assert( recEvent->totalCharged() <= recTrackCol->size() );
  for ( unsigned int i = 0; i < recEvent->totalCharged(); i++ )
  {
    EvtRecTrackIterator itTrk = recTrackCol->begin() + i;
    if ( !( *itTrk )->isMdcTrackValid() ) continue;
    if ( !( *itTrk )->isMdcKalTrackValid() ) continue;
    RecMdcTrack* mdcTrk = ( *itTrk )->mdcTrack();
    if ( fabs( cos( mdcTrk->theta() ) ) >= m_cosThetaCut ) continue;
    if ( fabs( mdcTrk->z() ) >= m_vz0Cut ) continue;
    iGood.push_back( ( *itTrk )->trackId() );
    if ( mdcTrk->charge() > 0 ) { icp.push_back( ( *itTrk )->trackId() ); }
    else if ( mdcTrk->charge() < 0 ) { icm.push_back( ( *itTrk )->trackId() ); }
  }
}
 
void InitVertexParameter( VertexParameter& vxpar ) {
  HepPoint3D   vx( 0., 0., 0. );
  HepSymMatrix Evx( 3, 0 );
  double       bx = 1E+6;
  double       by = 1E+6;
  double       bz = 1E+6;
  Evx[0][0]       = bx * bx;
  Evx[1][1]       = by * by;
  Evx[2][2]       = bz * bz;
  vxpar.setVx( vx );
  vxpar.setEvx( Evx );
}
 
//***************************************************************************
StatusCode PrimaryVertex::execute() {
  MsgStream log( msgSvc(), name() );
  log << MSG::INFO << "in execute()" << endmsg;
 
  cout.precision( 20 );
  ///////////////////////////////////////////
  //  Read REC information
  //////////////////////////////////////////
  SmartDataPtr<Event::EventHeader> eventHeader( eventSvc(), "/Event/EventHeader" );
  SmartDataPtr<EvtRecEvent>        recEvent( eventSvc(), EventModel::EvtRec::EvtRecEvent );
  SmartDataPtr<EvtRecTrackCol> recTrackCol( eventSvc(), EventModel::EvtRec::EvtRecTrackCol );
  log << MSG::DEBUG << "run and event = " << eventHeader->runNumber() << " "
      << eventHeader->eventNumber() << endmsg;
  log << MSG::DEBUG << "ncharg, nneu, tottks = " << recEvent->totalCharged() << " , "
      << recEvent->totalNeutral() << " , " << recEvent->totalTracks() << endmsg;
  m_sel_number[0]++;
  /*
  if (eventHeader->eventNumber() % 1000 == 0) {
    cout << "Event number = " << eventHeader->eventNumber() << endl;
  }*/
 
  EvtRecPrimaryVertex* aNewEvtRecPrimaryVertex = new EvtRecPrimaryVertex;
 
  StatusCode sc = RegisterEvtRecPrimaryVertex( aNewEvtRecPrimaryVertex, log );
  if ( sc != StatusCode::SUCCESS ) { return sc; }
 
  // Cut 1 : for anything sample, at least 3 charged tracks
  if ( recEvent->totalCharged() < m_trackNumberCut ) { return StatusCode::SUCCESS; }
  m_sel_number[1]++;
 
  Vint icp, icm, iGood;
  SelectGoodChargedTracks( recEvent, recTrackCol, icp, icm, iGood );
 
  // Cut 2 : for anything sample, at least 3 good charged tracks
  if ( iGood.size() < m_trackNumberCut ) { return StatusCode::SUCCESS; }
  m_sel_number[2]++;
 
  //  Vertex Fit
  VertexParameter vxpar;
  InitVertexParameter( vxpar );
 
  // For Global Vertex Fitting
  if ( m_fitMethod == 0 )
  {
    VertexFit* vtxfit = VertexFit::instance();
    vtxfit->init();
    Vint tlis, trkidlis;
    for ( int i = 0; i < iGood.size(); i++ )
    {
      int                 trk_id = iGood[i];
      EvtRecTrackIterator itTrk  = recTrackCol->begin() + trk_id;
      RecMdcKalTrack*     kalTrk = ( *itTrk )->mdcKalTrack();
      RecMdcKalTrack::setPidType( RecMdcKalTrack::pion );
      WTrackParameter wtrk( xmass[2], kalTrk->helix(), kalTrk->err() );
      vtxfit->AddTrack( i, wtrk );
      tlis.push_back( i );
      trkidlis.push_back( trk_id );
    }
    vtxfit->AddVertex( 0, vxpar, tlis );
    if ( !vtxfit->Fit( 0 ) ) return StatusCode::SUCCESS;                     // FIXME
    if ( vtxfit->chisq( 0 ) > m_globalChisqCut ) return StatusCode::SUCCESS; // FIXME
    if ( m_output == 1 )
    {
      m_chig          = vtxfit->chisq( 0 );
      m_ndofg         = 2 * iGood.size() - 3;
      m_probg         = TMath::Prob( m_chig, m_ndofg );
      HepVector glvtx = vtxfit->Vx( 0 );
      m_gvx           = glvtx[0];
      m_gvy           = glvtx[1];
      m_gvz           = glvtx[2];
      m_tuple1->write();
    }
    /*
        cout << "========  After global vertex fitting =============" << endl;
        cout << "Event number = " << eventHeader->eventNumber() << endl;
        cout << "NTracks: " << iGood.size() << endl;
        cout << "Chisq: " << vtxfit->chisq(0) << endl;
        cout << "Ndof: " << 2 * iGood.size() - 3 << endl;
        cout << "FitMethod: "<< " 0 " << endl;
        cout << "Vertex position: " << vtxfit->Vx(0)<< endl;
        cout << "Vertex resolution: " << vtxfit->Evx(0) << endl;
        cout << "track id list : " << endl;
        for (int j = 0; j < trkidlis.size(); j++) {
          cout << trkidlis[j] << "  ";
        }
        cout << " " << endl; */
 
    aNewEvtRecPrimaryVertex->setNTracks( iGood.size() );
    aNewEvtRecPrimaryVertex->setTrackIdList( trkidlis );
    aNewEvtRecPrimaryVertex->setChi2( vtxfit->chisq( 0 ) );
    aNewEvtRecPrimaryVertex->setNdof( 2 * iGood.size() - 3 );
    aNewEvtRecPrimaryVertex->setFitMethod( 0 );
    aNewEvtRecPrimaryVertex->setVertex( vtxfit->Vx( 0 ) );
    aNewEvtRecPrimaryVertex->setErrorVertex( vtxfit->Evx( 0 ) );
    aNewEvtRecPrimaryVertex->setIsValid( true );
  }
  else if ( m_fitMethod == 1 )
  {
    // For Kalman Vertex Fitting
    KalmanVertexFit* kalvtx = KalmanVertexFit::instance();
    kalvtx->init();
    kalvtx->initVertex( vxpar );
    kalvtx->setChisqCut( m_chisqCut );
    kalvtx->setTrackIteration( m_trackIteration );
    kalvtx->setVertexIteration( m_vertexIteration );
    kalvtx->setTrackIterationCut( m_chi2CutforTrkIter );
    for ( int i = 0; i < iGood.size(); i++ )
    {
      int                 trk_id = iGood[i];
      EvtRecTrackIterator itTrk  = recTrackCol->begin() + trk_id;
      RecMdcKalTrack*     kalTrk = ( *itTrk )->mdcKalTrack();
      RecMdcKalTrack::setPidType( RecMdcKalTrack::pion );
      HTrackParameter htrk( kalTrk->helix(), kalTrk->err(), trk_id, 2 );
      kalvtx->addTrack( htrk );
    }
    kalvtx->filter();
 
    // int freedomCut = -3 + 2*m_trackNumberCut;
    if ( kalvtx->ndof() >= m_freedomCut )
    { // Cut : add 2 when add every track
      // for(int i = 0; i < kalvtx->numTrack(); i++) {
      for ( int i = 0; i < iGood.size(); i++ )
      {
        if ( m_output == 1 )
        {
          m_chif  = kalvtx->chiF( i );
          m_chis  = kalvtx->chiS( i );
          m_probs = TMath::Prob( m_chis, 2 );
          m_probf = TMath::Prob( m_chif, 2 );
          m_tuple2->write();
        }
        if ( kalvtx->chiS( i ) > m_chi2CutforSmooth ) kalvtx->remove( i );
      }
    }
    if ( kalvtx->ndof() >= m_freedomCut )
    { // FIXME  to Rhopi is 0 , others are 1
      for ( int i = 0; i < iGood.size(); i++ )
      {
        kalvtx->smooth( i );
 
        HepVector Pull( 5, 0 );
        Pull = kalvtx->pull( i );
        if ( m_output == 1 )
        {
          m_pull_drho     = Pull[0];
          m_pull_phi      = Pull[1];
          m_pull_kapha    = Pull[2];
          m_pull_dz       = Pull[3];
          m_pull_lamb     = Pull[4];
          m_pull_momentum = kalvtx->pullmomentum( i );
          m_tuple3->write();
        }
      }
      if ( m_output == 1 )
      {
        m_chik  = kalvtx->chisq();
        m_ndofk = kalvtx->ndof();
        m_probk = TMath::Prob( m_chik, m_ndofk );
        HepVector xp( 3, 0 );
        xp    = kalvtx->x();
        m_kvx = xp[0];
        m_kvy = xp[1];
        m_kvz = xp[2];
        m_tuple4->write();
      }
 
      m_sel_number[3]++;
      /*
        cout << "========  After Kalman vertex fitting =============" << endl;
        cout << "NTracks: " << kalvtx->numTrack() << endl;
        cout << "Chisq: " << kalvtx->chisq() << endl;
        cout << "Ndof: " <<  kalvtx->ndof() << endl;
        cout << "FitMethod: "<< " 1 " << endl;
        cout << "Vertex position: " << kalvtx->x() << endl;
        cout << "Vertex resolution: " << kalvtx->Ex() << endl;
        for (int j = 0; j < (kalvtx->trackIDList()).size(); j++) {
          cout << kalvtx->trackIDList()[j] << "  ";
        }
        cout << " " << endl;*/
 
      aNewEvtRecPrimaryVertex->setNTracks( kalvtx->numTrack() );
      aNewEvtRecPrimaryVertex->setTrackIdList( kalvtx->trackIDList() );
      aNewEvtRecPrimaryVertex->setChi2( kalvtx->chisq() );
      aNewEvtRecPrimaryVertex->setNdof( kalvtx->ndof() );
      aNewEvtRecPrimaryVertex->setFitMethod( 1 );
      aNewEvtRecPrimaryVertex->setVertex( kalvtx->x() );
      aNewEvtRecPrimaryVertex->setErrorVertex( kalvtx->Ex() );
      aNewEvtRecPrimaryVertex->setIsValid( true );
    }
  }
  return StatusCode::SUCCESS;
}
 
//***************************************************************************
 
StatusCode PrimaryVertex::finalize() {
  MsgStream log( msgSvc(), name() );
  log << MSG::INFO << "in finalize()" << endmsg;
 
  log << MSG::ALWAYS << "==================================" << endmsg;
  log << MSG::ALWAYS << "survived event :";
  for ( int i = 0; i < 10; i++ ) { log << MSG::ALWAYS << m_sel_number[i] << " "; }
  log << MSG::ALWAYS << endmsg;
  log << MSG::ALWAYS << "==================================" << endmsg;
  return StatusCode::SUCCESS;
}