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