packages/trackreco/PHRaveVertexing.cc

0001 /*!
0002  *  \file       PHRaveVertexing.C
0003  *  \brief      Refit SvtxTracks with PHGenFit.
0004  *  \details    Refit SvtxTracks with PHGenFit.
0005  *  \author     Haiwang Yu <yuhw@nmsu.edu>
0006  */
0007
0008 #include "PHRaveVertexing.h"
0009
0010 #include <trackbase_historic/SvtxTrack.h>
0011 #include <trackbase_historic/SvtxTrackMap.h>
0012 #include <trackbase_historic/SvtxTrackState.h>    // for SvtxTrackState
0013 #include <globalvertex/SvtxVertexMap_v1.h>
0014 #include <globalvertex/SvtxVertex_v1.h>
0015
0016 #include <g4main/PHG4TruthInfoContainer.h>
0017
0018 #include <phgenfit/Fitter.h>
0019
0020 #include <phfield/PHFieldUtility.h>
0021
0022 #include <phgeom/PHGeomUtility.h>
0023
0024 #include <fun4all/Fun4AllReturnCodes.h>
0025 #include <fun4all/SubsysReco.h>                   // for SubsysReco
0026
0027 #include <phool/PHCompositeNode.h>
0028 #include <phool/PHIODataNode.h>
0029 #include <phool/PHNode.h>                         // for PHNode
0030 #include <phool/PHNodeIterator.h>
0031 #include <phool/PHObject.h>                       // for PHObject
0032 #include <phool/PHTimer.h>
0033 #include <phool/getClass.h>
0034 #include <phool/phool.h>
0035
0036 #include <GenFit/FitStatus.h>                     // for FitStatus
0037 #include <GenFit/GFRaveTrackParameters.h>         // for GFRaveTrackParameters
0038 #include <GenFit/GFRaveVertex.h>
0039 #include <GenFit/GFRaveVertexFactory.h>
0040 #include <GenFit/KalmanFittedStateOnPlane.h>      // for KalmanFittedStateOn...
0041 #include <GenFit/KalmanFitterInfo.h>
0042 #include <GenFit/MeasuredStateOnPlane.h>
0043 #include <GenFit/RKTrackRep.h>
0044 #include <GenFit/Track.h>
0045 #include <GenFit/TrackPoint.h>                    // for TrackPoint
0046
0047 #include <TMatrixDSymfwd.h>                       // for TMatrixDSym
0048 #include <TMatrixTSym.h>                          // for TMatrixTSym
0049 #include <TMatrixTUtils.h>                        // for TMatrixTRow
0050 #include <TVector3.h>
0051
0052 #include <iostream>
0053 #include <map>
0054 #include <memory>
0055 #include <utility>
0056 #include <vector>
0057
0058 class PHField;
0059 class TGeoManager;
0060 namespace genfit { class AbsTrackRep; }
0061
0062 #define LogDebug(exp) std::cout << "DEBUG: " << __FILE__ << ": " << __LINE__ << ": " << exp << std::endl
0063 #define LogError(exp) std::cout << "ERROR: " << __FILE__ << ": " << __LINE__ << ": " << exp << std::endl
0064 #define LogWarning(exp) std::cout << "WARNING: " << __FILE__ << ": " << __LINE__ << ": " << exp << std::endl
0065
0066 //#define _DEBUG_
0067
0068 using namespace std;
0069
0070 /*
0071  * Constructor
0072  */
0073 PHRaveVertexing::PHRaveVertexing(const string& name)
0074   : SubsysReco(name)
0075   , _over_write_svtxvertexmap(false)
0076   , _svtxvertexmaprefit_node_name("SvtxVertexMapRefit")
0077   , _fitter(nullptr)
0078   , _primary_pid_guess(211)
0079   , _vertex_min_ndf(20)
0080   , _vertex_finder(nullptr)
0081   , _vertexing_method("avf-smoothing:1")
0082   , _trackmap(nullptr)
0083   , _vertexmap(nullptr)
0084   , _vertexmap_refit(nullptr)
0085   , _t_translate(nullptr)
0086   , _t_rave(nullptr)
0087 {
0088   Verbosity(0);
0089
0090   _event = 0;
0091 }
0092
0093 /*
0094  * Init
0095  */
0096 int PHRaveVertexing::Init(PHCompositeNode* /*topNode*/)
0097 {
0098   return Fun4AllReturnCodes::EVENT_OK;
0099 }
0100
0101 /*
0102  * Init run
0103  */
0104 int PHRaveVertexing::InitRun(PHCompositeNode* topNode)
0105 {
0106   CreateNodes(topNode);
0107
0108   TGeoManager* tgeo_manager = PHGeomUtility::GetTGeoManager(topNode);
0109   PHField* field = PHFieldUtility::GetFieldMapNode(nullptr, topNode);
0110
0111   //_fitter = new PHGenFit::Fitter("sPHENIX_Geo.root","sPHENIX.2d.root", 1.4 / 1.5);
0112   _fitter = PHGenFit::Fitter::getInstance(tgeo_manager,
0113                                           field, "DafRef",
0114                                           "RKTrackRep", false);
0115   if (!_fitter)
0116   {
0117     std::cout << PHWHERE << " PHGenFit::Fitter::getInstance returned nullptr"
0118           << std::endl;
0119     return Fun4AllReturnCodes::ABORTRUN;
0120   }
0121   _fitter->set_verbosity(Verbosity());
0122
0123   if (!_fitter)
0124   {
0125     cerr << PHWHERE << endl;
0126     return Fun4AllReturnCodes::ABORTRUN;
0127   }
0128
0129   _vertex_finder = new genfit::GFRaveVertexFactory(Verbosity());
0130   if (!_vertex_finder)
0131   {
0132     std::cout << PHWHERE << " genfit::GFRaveVertexFactory returned null ptr" << std::endl;
0133     return Fun4AllReturnCodes::ABORTRUN;
0134   }
0135   _vertex_finder->setMethod(_vertexing_method.data());
0136   //_vertex_finder->setBeamspot();
0137
0138
0139   _t_translate = new PHTimer("_t_translate");
0140   _t_translate->stop();
0141
0142   _t_rave = new PHTimer("_t_rave");
0143   _t_rave->stop();
0144
0145   return Fun4AllReturnCodes::EVENT_OK;
0146 }
0147 /*!
0148  * process_event():
0149  *  Call user instructions for every event.
0150  *  This function contains the analysis structure.
0151  *
0152  */
0153 int PHRaveVertexing::process_event(PHCompositeNode* topNode)
0154 {
0155   _event++;
0156
0157   if (Verbosity() > 1)
0158     std::cout << PHWHERE << "Events processed: " << _event << std::endl;
0159
0160   GetNodes(topNode);
0161
0162   //! stands for Refit_GenFit_Tracks
0163   GenFitTrackMap gf_track_map;
0164   vector<genfit::Track*> gf_tracks;
0165   if (Verbosity() > 1) _t_translate->restart();
0166   for (SvtxTrackMap::Iter iter = _trackmap->begin(); iter != _trackmap->end();
0167        ++iter)
0168   {
0169     SvtxTrack* svtx_track = iter->second;
0170     if (!svtx_track)
0171       continue;
0172
0173     if (!(svtx_track->get_ndf() >= _vertex_min_ndf))
0174       continue;
0175
0176     // require MVTX association
0177     if(_nmvtx_required > 0)
0178       {
0179     unsigned int nmvtx = 0;
0180     for(auto clusit = svtx_track->begin_cluster_keys(); clusit != svtx_track->end_cluster_keys(); ++clusit)
0181       {
0182         if(TrkrDefs::getTrkrId(*clusit) == TrkrDefs::mvtxId )
0183           {
0184         nmvtx++;
0185           }
0186         if(nmvtx >=  _nmvtx_required) break;
0187       }
0188     if(nmvtx < _nmvtx_required) continue;
0189     if(Verbosity() > 1) std::cout << " track " << iter->first << "  has nmvtx at least " << nmvtx << std::endl;
0190       }
0191
0192     //auto genfit_track = shared_ptr<genfit::Track> (TranslateSvtxToGenFitTrack(svtx_track));
0193     auto genfit_track = TranslateSvtxToGenFitTrack(svtx_track);
0194     if (!genfit_track)
0195       continue;
0196     gf_track_map.insert({genfit_track, iter->first});
0197     gf_tracks.push_back(const_cast<genfit::Track*>(genfit_track));
0198   }
0199   if (Verbosity() > 1) _t_translate->stop();
0200
0201   if (Verbosity() > 1) _t_rave->restart();
0202   vector<genfit::GFRaveVertex*> rave_vertices;
0203   if (gf_tracks.size() >= 2)
0204   {
0205     try
0206     {
0207       _vertex_finder->findVertices(&rave_vertices, gf_tracks);
0208     }
0209     catch (...)
0210     {
0211       if (Verbosity() > 1)
0212         std::cout << PHWHERE << "GFRaveVertexFactory::findVertices failed!";
0213     }
0214   }
0215   if (Verbosity() > 1) _t_rave->stop();
0216   FillSvtxVertexMap(rave_vertices, gf_track_map);
0217
0218   for (auto iter : gf_track_map) delete iter.first;
0219
0220   if (Verbosity() > 1)
0221   {
0222     _vertexmap->identify();
0223
0224     std::cout << "=============== Timers: ===============" << std::endl;
0225     std::cout << "Event: " << _event << std::endl;
0226     std::cout << "Translate:                " << _t_translate->get_accumulated_time() / 1000. << " sec" << std::endl;
0227     std::cout << "RAVE:                     " << _t_rave->get_accumulated_time() / 1000. << " sec" << std::endl;
0228     std::cout << "=======================================" << std::endl;
0229   }
0230
0231   return Fun4AllReturnCodes::EVENT_OK;
0232 }
0233
0234 /*
0235  * End
0236  */
0237 int PHRaveVertexing::End(PHCompositeNode* /*topNode*/)
0238 {
0239   return Fun4AllReturnCodes::EVENT_OK;
0240 }
0241
0242 /*
0243  * dtor
0244  */
0245 PHRaveVertexing::~PHRaveVertexing()
0246 {
0247   delete _fitter;
0248   delete _vertex_finder;
0249 }
0250
0251 int PHRaveVertexing::CreateNodes(PHCompositeNode* topNode)
0252 {
0253   // create nodes...
0254   PHNodeIterator iter(topNode);
0255
0256   PHCompositeNode* dstNode = static_cast<PHCompositeNode*>(iter.findFirst(
0257       "PHCompositeNode", "DST"));
0258   if (!dstNode)
0259   {
0260     cerr << PHWHERE << "DST Node missing, doing nothing." << endl;
0261     return Fun4AllReturnCodes::ABORTEVENT;
0262   }
0263   PHNodeIterator iter_dst(dstNode);
0264
0265   // Create the SVTX node
0266   PHCompositeNode* tb_node = dynamic_cast<PHCompositeNode*>(iter_dst.findFirst(
0267       "PHCompositeNode", "SVTX"));
0268   if (!tb_node)
0269   {
0270     tb_node = new PHCompositeNode("SVTX");
0271     dstNode->addNode(tb_node);
0272     if (Verbosity() > 0)
0273       cout << "SVTX node added" << endl;
0274   }
0275
0276   if (!(_over_write_svtxvertexmap))
0277   {
0278     _vertexmap_refit = new SvtxVertexMap_v1;
0279     PHIODataNode<PHObject>* vertexes_node = new PHIODataNode<PHObject>(
0280         _vertexmap_refit, _svtxvertexmaprefit_node_name.c_str(), "PHObject");
0281     tb_node->addNode(vertexes_node);
0282     if (Verbosity() > 0)
0283       cout << "Svtx/SvtxVertexMapRefit node added" << endl;
0284   }
0285   else if (!findNode::getClass<SvtxVertexMap>(topNode, "SvtxVertexMap"))
0286   {
0287     _vertexmap = new SvtxVertexMap_v1;
0288     PHIODataNode<PHObject>* vertexes_node = new PHIODataNode<PHObject>(
0289         _vertexmap, "SvtxVertexMap", "PHObject");
0290     tb_node->addNode(vertexes_node);
0291     if (Verbosity() > 0)
0292       cout << "Svtx/SvtxVertexMap node added" << endl;
0293   }
0294
0295   return Fun4AllReturnCodes::EVENT_OK;
0296 }
0297
0298 /*
0299  * GetNodes():
0300  *  Get all the all the required nodes off the node tree
0301  */
0302 int PHRaveVertexing::GetNodes(PHCompositeNode* topNode)
0303 {
0304   //DST objects
0305   // Input Svtx Tracks
0306   _trackmap = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
0307   if (!_trackmap && _event < 2)
0308   {
0309     cout << PHWHERE << " SvtxTrackMap node not found on node tree"
0310          << endl;
0311     return Fun4AllReturnCodes::ABORTEVENT;
0312   }
0313
0314   // Input Svtx Vertices
0315   _vertexmap = findNode::getClass<SvtxVertexMap>(topNode, "SvtxVertexMap");
0316   if (!_vertexmap && _event < 2)
0317   {
0318     cout << PHWHERE << " SvtxVertexrMap node not found on node tree"
0319          << endl;
0320     return Fun4AllReturnCodes::ABORTEVENT;
0321   }
0322
0323   // Output Svtx Vertices
0324   if (!(_over_write_svtxvertexmap))
0325   {
0326     _vertexmap_refit = findNode::getClass<SvtxVertexMap>(topNode,
0327                                                          _svtxvertexmaprefit_node_name.c_str());
0328     if (!_vertexmap_refit && _event < 2)
0329     {
0330       cout << PHWHERE << " SvtxVertexMapRefit node not found on node tree"
0331            << endl;
0332       return Fun4AllReturnCodes::ABORTEVENT;
0333     }
0334   }
0335
0336   return Fun4AllReturnCodes::EVENT_OK;
0337 }
0338
0339 /*
0340  * Fill SvtxVertexMap from GFRaveVertexes and Tracks
0341  */
0342 bool PHRaveVertexing::FillSvtxVertexMap(
0343     const std::vector<genfit::GFRaveVertex*>& rave_vertices,
0344     const GenFitTrackMap& gf_track_map)
0345 {
0346   if (_over_write_svtxvertexmap)
0347   {
0348     _vertexmap->clear();
0349   }
0350
0351   //    for (unsigned int ivtx = 0; ivtx < rave_vertices.size(); ++ivtx) {
0352   //        genfit::GFRaveVertex* rave_vtx = rave_vertices[ivtx];
0353
0354   for (genfit::GFRaveVertex* rave_vtx : rave_vertices)
0355   {
0356     if (!rave_vtx)
0357     {
0358       cerr << PHWHERE << endl;
0359       return false;
0360     }
0361
0362     std::shared_ptr<SvtxVertex> svtx_vtx(new SvtxVertex_v1());
0363
0364     svtx_vtx->set_chisq(rave_vtx->getChi2());
0365     svtx_vtx->set_ndof(rave_vtx->getNdf());
0366     svtx_vtx->set_position(0, rave_vtx->getPos().X());
0367     svtx_vtx->set_position(1, rave_vtx->getPos().Y());
0368     svtx_vtx->set_position(2, rave_vtx->getPos().Z());
0369
0370     for (int i = 0; i < 3; i++)
0371       for (int j = 0; j < 3; j++)
0372         svtx_vtx->set_error(i, j, rave_vtx->getCov()[i][j]);
0373
0374     for (unsigned int i = 0; i < rave_vtx->getNTracks(); i++)
0375     {
0376       //TODO improve speed
0377       const genfit::Track* rave_track =
0378           rave_vtx->getParameters(i)->getTrack();
0379       //            for(auto iter : gf_track_map) {
0380       //                if (iter.second == rave_track)
0381       //                    svtx_vtx->insert_track(iter.first);
0382       //            }
0383       auto iter = gf_track_map.find(rave_track);
0384       if (iter != gf_track_map.end())
0385         svtx_vtx->insert_track(iter->second);
0386     }
0387
0388     if (_over_write_svtxvertexmap)
0389     {
0390       if (_vertexmap)
0391       {
0392         _vertexmap->insert_clone(svtx_vtx.get());
0393       }
0394       else
0395       {
0396         LogError("!_vertexmap");
0397       }
0398     }
0399     else
0400     {
0401       if (_vertexmap_refit)
0402       {
0403         _vertexmap_refit->insert_clone(svtx_vtx.get());
0404       }
0405       else
0406       {
0407         LogError("!_vertexmap_refit");
0408       }
0409     }
0410
0411 #ifdef _DEBUG_
0412     cout << __LINE__ << endl;
0413     svtx_vtx->identify();
0414 #endif
0415
0416   }  //loop over RAVE vertices
0417
0418   return true;
0419 }
0420
0421 genfit::Track* PHRaveVertexing::TranslateSvtxToGenFitTrack(SvtxTrack* svtx_track)
0422 {
0423   try
0424   {
0425     // The first state is extracted to PCA, second one is the one with measurement
0426     SvtxTrackState* svtx_state(nullptr);
0427     //SvtxTrackState* svtx_state = (svtx_track->begin_states())->second;
0428
0429     if (svtx_track->begin_states() == svtx_track->end_states())
0430     {
0431       LogDebug("TranslateSvtxToGenFitTrack no state in track!");
0432       return nullptr;
0433     }
0434     else if (++(svtx_track->begin_states()) == svtx_track->end_states())
0435     {
0436       // only one state in track
0437       svtx_state = (svtx_track->begin_states())->second;
0438     }
0439     else
0440     {
0441       // multiple state in track
0442       // The first state is extracted to PCA, second one is the one with measurement
0443       svtx_state = (++(svtx_track->begin_states()))->second;
0444     }
0445
0446     if (!svtx_state)
0447     {
0448       LogDebug("TranslateSvtxToGenFitTrack invalid state found on track!");
0449       return nullptr;
0450     }
0451
0452     TVector3 pos(svtx_state->get_x(), svtx_state->get_y(), svtx_state->get_z());
0453     TVector3 mom(svtx_state->get_px(), svtx_state->get_py(), svtx_state->get_pz());
0454     TMatrixDSym cov(6);
0455     for (int i = 0; i < 6; ++i)
0456     {
0457       for (int j = 0; j < 6; ++j)
0458       {
0459         cov[i][j] = svtx_state->get_error(i, j);
0460       }
0461     }
0462
0463 #ifdef _DEBUG_
0464     {
0465       cout << "DEBUG" << __LINE__ << endl;
0466       cout << "path length:      " << svtx_state->get_pathlength() << endl;
0467       cout << "radius:           " << pos.Perp() << endl;
0468       cout << "DEBUG: " << __LINE__ << endl;
0469       for (int i = 0; i < 6; ++i)
0470       {
0471         for (int j = 0; j < 6; ++j)
0472         {
0473           cout << svtx_state->get_error(i, j) << "\t";
0474         }
0475         cout << endl;
0476       }
0477
0478       cov.Print();
0479     }
0480
0481 #endif
0482
0483     genfit::AbsTrackRep* rep = new genfit::RKTrackRep(_primary_pid_guess);
0484     genfit::Track* genfit_track = new genfit::Track(rep, TVector3(0, 0, 0), TVector3(0, 0, 0));
0485
0486     genfit::FitStatus* fs = new genfit::FitStatus();
0487     fs->setCharge(svtx_track->get_charge());
0488     fs->setChi2(svtx_track->get_chisq());
0489     fs->setNdf(svtx_track->get_ndf());
0490     fs->setIsFitted(true);
0491     fs->setIsFitConvergedFully(true);
0492
0493     genfit_track->setFitStatus(fs, rep);
0494
0495     genfit::TrackPoint* tp = new genfit::TrackPoint(genfit_track);
0496
0497     genfit::KalmanFitterInfo* fi = new genfit::KalmanFitterInfo(tp, rep);
0498     tp->setFitterInfo(fi);
0499
0500     genfit::MeasuredStateOnPlane* ms = new genfit::MeasuredStateOnPlane(rep);
0501     ms->setPosMomCov(pos, mom, cov);
0502 #ifdef _DEBUG_
0503     {
0504       cout << "DEBUG: " << __LINE__ << endl;
0505       ms->Print();
0506       cout << "Orig: " << __LINE__ << endl;
0507       cov.Print();
0508       cout << "Translate: " << __LINE__ << endl;
0509       ms->get6DCov().Print();
0510     }
0511 #endif
0512     genfit::KalmanFittedStateOnPlane* kfs = new genfit::KalmanFittedStateOnPlane(*ms, 1., 1.);
0513
0514     //< Acording to the special order of using the stored states
0515     fi->setForwardUpdate(kfs);
0516
0517     genfit_track->insertPoint(tp);
0518
0519 #ifdef _DEBUG_
0520 //      {
0521 //          cout << "DEBUG" << __LINE__ << endl;
0522 //          TVector3 pos, mom;
0523 //          TMatrixDSym cov;
0524 //          genfit_track->getFittedState().getPosMomCov(pos, mom, cov);
0525 //          pos.Print();
0526 //          mom.Print();
0527 //          cov.Print();
0528 //      }
0529 #endif
0530
0531     return genfit_track;
0532   }
0533   catch (...)
0534   {
0535     LogDebug("TranslateSvtxToGenFitTrack failed!");
0536   }
0537
0538   return nullptr;
0539 }