analysis/ForwardTracking/PHG4HitKalmanFitter.cc

0001 /*!
0002  *  \file       PHG4HitKalmanFitter.cc
0003  *  \brief      Kalman Filter based on smeared truth PHG4Hit
0004  *  \details    Kalman Filter based on smeared truth PHG4Hit
0005  *  \author     Haiwang Yu <yuhw@nmsu.edu>
0006  */
0007
0008 #include <cmath>
0009 #include <map>
0010
0011 #include <GenFit/RKTrackRep.h>
0012 #include <fun4all/Fun4AllReturnCodes.h>
0013 #include <phool/PHCompositeNode.h>
0014 #include <phool/PHIODataNode.h>
0015 #include <phool/PHNodeIterator.h>
0016 #include <phool/getClass.h>
0017 #include <phool/phool.h>
0018
0019 #include <GenFit/StateOnPlane.h>
0020
0021 #include <TMath.h>
0022 #include <TRandom.h>
0023 #include <TString.h>
0024
0025 #include <g4hough/SvtxTrackMap.h>
0026 #include <g4hough/SvtxTrackMap_v1.h>
0027 #include <g4hough/SvtxTrack_v1.h>
0028 #include <g4main/PHG4Hit.h>
0029 #include <g4main/PHG4Particle.h>
0030 #include <g4main/PHG4TruthInfoContainer.h>
0031 #include <phgenfit/Fitter.h>
0032 #include <phgenfit/PlanarMeasurement.h>
0033 #include <phgenfit/Track.h>
0034
0035 #include <phfield/PHFieldUtility.h>
0036 #include <phgeom/PHGeomUtility.h>
0037
0038 #include "PHG4HitKalmanFitter.h"
0039
0040 #define LogDebug(exp) std::cout << "DEBUG: " << __FILE__ << ": " << __LINE__ << ": " << exp << "\n"
0041 #define LogError(exp) std::cout << "ERROR: " << __FILE__ << ": " << __LINE__ << ": " << exp << "\n"
0042 #define LogWarning(exp) std::cout << "WARNING: " << __FILE__ << ": " << __LINE__ << ": " << exp << "\n"
0043
0044 #define _N_DETECTOR_LAYER 5
0045
0046 using namespace std;
0047
0048 PHG4TrackFastSim::PHG4TrackFastSim(const std::string& name)
0049   : SubsysReco(name)
0050   , _truth_container(NULL)
0051   , _trackmap_out(NULL)
0052   , _fitter(NULL)
0053   , _fit_alg_name("KalmanFitterRefTrack")
0054   , _do_evt_display(false)
0055   , _phi_resolution(50E-4)
0056   ,  //100um
0057   _r_resolution(1.)
0058   , _pat_rec_hit_finding_eff(1.)
0059   , _pat_rec_nosise_prob(0.)
0060
0061 {
0062   _event = 0;
0063
0064   for (int i = 0; i < _N_DETECTOR_LAYER; i++)
0065   {
0066     _phg4hits_names.push_back(Form("G4HIT_FGEM_%1d", i));
0067   }
0068 }
0069
0070 PHG4TrackFastSim::~PHG4TrackFastSim()
0071 {
0072   delete _fitter;
0073 }
0074
0075 /*
0076  * Init
0077  */
0078 int PHG4TrackFastSim::Init(PHCompositeNode* topNode)
0079 {
0080   return Fun4AllReturnCodes::EVENT_OK;
0081 }
0082
0083 int PHG4TrackFastSim::InitRun(PHCompositeNode* topNode)
0084 {
0085   CreateNodes(topNode);
0086
0087   TGeoManager* tgeo_manager = PHGeomUtility::GetTGeoManager(topNode);
0088   PHField* field = PHFieldUtility::GetFieldMapNode(nullptr, topNode);
0089   //_fitter = new PHGenFit::Fitter("sPHENIX_Geo.root","sPHENIX.2d.root", 1.4 / 1.5);
0090   _fitter = PHGenFit::Fitter::getInstance(tgeo_manager,
0091                                           field, "KalmanFitterRefTrack", "RKTrackRep",
0092                                           _do_evt_display);
0093
0094   if (!_fitter)
0095   {
0096     cerr << PHWHERE << endl;
0097     return Fun4AllReturnCodes::ABORTRUN;
0098   }
0099
0100   return Fun4AllReturnCodes::EVENT_OK;
0101 }
0102
0103 int PHG4TrackFastSim::End(PHCompositeNode* topNode)
0104 {
0105   return Fun4AllReturnCodes::EVENT_OK;
0106 }
0107
0108 int PHG4TrackFastSim::process_event(PHCompositeNode* topNode)
0109 {
0110   GetNodes(topNode);
0111
0112   if (_trackmap_out)
0113     _trackmap_out->empty();
0114   else
0115   {
0116     LogError("_trackmap_out not found!");
0117     return Fun4AllReturnCodes::ABORTRUN;
0118   }
0119
0120   for (PHG4TruthInfoContainer::ConstIterator itr =
0121            _truth_container->GetPrimaryParticleRange().first;
0122        itr != _truth_container->GetPrimaryParticleRange().second; ++itr)
0123   {
0124     PHG4Particle* particle = itr->second;
0125
0126     TVector3 seed_pos(0, 0, 0);
0127     TVector3 seed_mom(0, 0, 0);
0128     TMatrixDSym seed_cov(6);
0129
0130     //! Create measurements
0131     std::vector<PHGenFit::Measurement*> measurements;
0132
0133     const bool _use_vertex_in_fitting = true;
0134
0135     PHGenFit::Measurement* vtx_meas = NULL;
0136
0137     if (_use_vertex_in_fitting)
0138     {
0139       vtx_meas = VertexMeasurement(
0140           TVector3(0, 0, 0), 0.0050, 0.0050);
0141       measurements.push_back(vtx_meas);
0142     }
0143
0144     PseudoPatternRecognition(particle, measurements, seed_pos, seed_mom, seed_cov);
0145
0146     if (measurements.size() < 3)
0147     {
0148       if (verbosity >= 2)
0149         LogWarning("measurements.size() < 3");
0150       continue;
0151     }
0152
0153     //! Build TrackRep from particle assumption
0154     /*!
0155          * mu+: -13
0156          * mu-: 13
0157          * pi+: 211
0158          * pi-: -211
0159          * e-:  11
0160          * e+:  -11
0161          */
0162     int pid = 13;  //
0163     //SMART(genfit::AbsTrackRep) rep = NEW(genfit::RKTrackRep)(pid);
0164     genfit::AbsTrackRep* rep = new genfit::RKTrackRep(pid);
0165
0166     //! Initiallize track with seed from pattern recognition
0167     PHGenFit::Track* track = new PHGenFit::Track(rep, seed_pos, seed_mom,
0168                                                  seed_cov);
0169
0170     //LogDEBUG;
0171     //! Add measurements to track
0172     track->addMeasurements(measurements);
0173
0174     //LogDEBUG;
0175     //! Fit the track
0176     int fitting_err = _fitter->processTrack(track, _do_evt_display);
0177
0178     if (fitting_err != 0)
0179     {
0180       //            LogDEBUG;
0181       //            std::cout<<"event: "<<ientry<<"\n";
0182       continue;
0183     }
0184
0185     SvtxTrack* svtx_track_out = MakeSvtxTrack(track);
0186
0187     _trackmap_out->insert(svtx_track_out);
0188   }
0189
0190   return Fun4AllReturnCodes::EVENT_OK;
0191 }
0192
0193 int PHG4TrackFastSim::CreateNodes(PHCompositeNode* topNode)
0194 {
0195   // create nodes...
0196   PHNodeIterator iter(topNode);
0197
0198   PHCompositeNode* dstNode = static_cast<PHCompositeNode*>(iter.findFirst(
0199       "PHCompositeNode", "DST"));
0200   if (!dstNode)
0201   {
0202     cerr << PHWHERE << "DST Node missing, doing nothing." << endl;
0203     return Fun4AllReturnCodes::ABORTEVENT;
0204   }
0205   PHNodeIterator iter_dst(dstNode);
0206
0207   // Create the FGEM node
0208   PHCompositeNode* tb_node = dynamic_cast<PHCompositeNode*>(iter_dst.findFirst(
0209       "PHCompositeNode", "FGEM"));
0210   if (!tb_node)
0211   {
0212     tb_node = new PHCompositeNode("FGEM");
0213     dstNode->addNode(tb_node);
0214     if (verbosity > 0)
0215       cout << "FGEM node added" << endl;
0216   }
0217
0218   _trackmap_out = new SvtxTrackMap_v1;
0219
0220   PHIODataNode<PHObject>* tracks_node = new PHIODataNode<PHObject>(
0221       _trackmap_out, "ForwardTrackMap", "PHObject");
0222   tb_node->addNode(tracks_node);
0223   if (verbosity > 0)
0224     cout << "FGEM/ForwardTrackMap node added" << endl;
0225
0226   return Fun4AllReturnCodes::EVENT_OK;
0227 }
0228
0229 int PHG4TrackFastSim::GetNodes(PHCompositeNode* topNode)
0230 {
0231   //DST objects
0232   //Truth container
0233   _truth_container = findNode::getClass<PHG4TruthInfoContainer>(topNode,
0234                                                                 "G4TruthInfo");
0235   if (!_truth_container && _event < 2)
0236   {
0237     cout << PHWHERE << " PHG4TruthInfoContainer node not found on node tree"
0238          << endl;
0239     return Fun4AllReturnCodes::ABORTEVENT;
0240   }
0241
0242   for (int i = 0; i < _N_DETECTOR_LAYER; i++)
0243   {
0244     PHG4HitContainer* phg4hit = findNode::getClass<PHG4HitContainer>(topNode,
0245                                                                      _phg4hits_names[i].c_str());
0246     if (!phg4hit && _event < 2)
0247     {
0248       cout << PHWHERE << _phg4hits_names[i].c_str() << " node not found on node tree"
0249            << endl;
0250       return Fun4AllReturnCodes::ABORTEVENT;
0251     }
0252
0253     _phg4hits.push_back(phg4hit);
0254   }
0255
0256   _trackmap_out = findNode::getClass<SvtxTrackMap>(topNode,
0257                                                    "ForwardTrackMap");
0258   if (!_trackmap_out && _event < 2)
0259   {
0260     cout << PHWHERE << " ForwardTrackMap node not found on node tree"
0261          << endl;
0262     return Fun4AllReturnCodes::ABORTEVENT;
0263   }
0264
0265   return Fun4AllReturnCodes::EVENT_OK;
0266 }
0267
0268 int PHG4TrackFastSim::PseudoPatternRecognition(
0269     const PHG4Particle* particle,
0270     std::vector<PHGenFit::Measurement*>& meas_out, TVector3& seed_pos,
0271     TVector3& seed_mom, TMatrixDSym& seed_cov, const bool do_smearing)
0272 {
0273   seed_pos.SetXYZ(0, 0, 0);
0274   seed_mom.SetXYZ(0, 0, 10);
0275   seed_cov.ResizeTo(6, 6);
0276
0277   for (int i = 0; i < 3; i++)
0278   {
0279     seed_cov[i][i] = _phi_resolution * _phi_resolution;
0280   }
0281
0282   for (int i = 3; i < 6; i++)
0283   {
0284     seed_cov[i][i] = 10;
0285   }
0286
0287   if (particle)
0288   {
0289     TVector3 True_mom(particle->get_px(), particle->get_py(),
0290                       particle->get_pz());
0291
0292     seed_mom.SetXYZ(particle->get_px(), particle->get_py(),
0293                     particle->get_pz());
0294     if (do_smearing)
0295     {
0296       const double momSmear = 3. / 180. * TMath::Pi();  // rad
0297       const double momMagSmear = 0.1;                   // relative
0298
0299       seed_mom.SetPhi(gRandom->Gaus(True_mom.Phi(), momSmear));
0300       seed_mom.SetTheta(gRandom->Gaus(True_mom.Theta(), momSmear));
0301       seed_mom.SetMag(
0302           gRandom->Gaus(True_mom.Mag(),
0303                         momMagSmear * True_mom.Mag()));
0304     }
0305   }
0306
0307   meas_out.clear();
0308
0309   for (int ilayer = 0; ilayer < _N_DETECTOR_LAYER; ilayer++)
0310   {
0311     if (!_phg4hits[ilayer])
0312     {
0313       LogError("No _phg4hits[i] found!");
0314       continue;
0315     }
0316
0317     for (PHG4HitContainer::ConstIterator itr =
0318              _phg4hits[ilayer]->getHits().first;
0319          itr != _phg4hits[ilayer]->getHits().second; ++itr)
0320     {
0321       PHG4Hit* hit = itr->second;
0322       if (!hit)
0323       {
0324         LogDebug("No PHG4Hit Found!");
0325         continue;
0326       }
0327       if (hit->get_trkid() == particle->get_track_id() || gRandom->Uniform(0, 1) < _pat_rec_nosise_prob)
0328       {
0329         PHGenFit::Measurement* meas = PHG4HitToMeasurementVerticalPlane(hit);
0330         if (gRandom->Uniform(0, 1) <= _pat_rec_hit_finding_eff)
0331           meas_out.push_back(meas);
0332       }
0333     }
0334
0335   } /*Loop detector layers*/
0336
0337   return Fun4AllReturnCodes::EVENT_OK;
0338 }
0339
0340 SvtxTrack* PHG4TrackFastSim::MakeSvtxTrack(
0341     const PHGenFit::Track* phgf_track)
0342 {
0343   double chi2 = phgf_track->get_chi2();
0344   double ndf = phgf_track->get_ndf();
0345
0346   genfit::MeasuredStateOnPlane* gf_state = phgf_track->extrapolateToPlane(TVector3(0., 0., 0.), TVector3(0., 0., 1.));
0347   TVector3 mom = gf_state->getMom();
0348   TVector3 pos = gf_state->getPos();
0349   TMatrixDSym cov = gf_state->get6DCov();
0350
0351   //    SvtxTrack_v1* out_track = new SvtxTrack_v1(*static_cast<const SvtxTrack_v1*> (svtx_track));
0352   SvtxTrack_v1* out_track = new SvtxTrack_v1();
0353
0354   /*!
0355      * FIXME: check the definition
0356      *  1/p, u'/z', v'/z', u, v
0357      *  u is defined as mom X beam line at POCA
0358      *  so u is the dca2d direction
0359      */
0360   double dca2d = gf_state->getState()[3];
0361   out_track->set_dca2d(dca2d);
0362   out_track->set_dca2d_error(gf_state->getCov()[3][3]);
0363   double dca3d = sqrt(
0364       dca2d * dca2d +
0365       gf_state->getState()[4] * gf_state->getState()[4]);
0366   out_track->set_dca(dca3d);
0367
0368   out_track->set_chisq(chi2);
0369   out_track->set_ndf(ndf);
0370   out_track->set_charge(phgf_track->get_charge());
0371
0372   out_track->set_px(mom.Px());
0373   out_track->set_py(mom.Py());
0374   out_track->set_pz(mom.Pz());
0375
0376   out_track->set_x(pos.X());
0377   out_track->set_y(pos.Y());
0378   out_track->set_z(pos.Z());
0379
0380   for (int i = 0; i < 6; i++)
0381   {
0382     for (int j = i; j < 6; j++)
0383     {
0384       out_track->set_error(i, j, cov[i][j]);
0385     }
0386   }
0387
0388   return out_track;
0389 }
0390
0391 PHGenFit::PlanarMeasurement* PHG4TrackFastSim::PHG4HitToMeasurementVerticalPlane(const PHG4Hit* g4hit)
0392 {
0393   PHGenFit::PlanarMeasurement* meas = NULL;
0394
0395   TVector3 pos(
0396       g4hit->get_avg_x(),
0397       g4hit->get_avg_y(),
0398       g4hit->get_avg_z());
0399
0400   TVector3 v(pos.X(), pos.Y(), 0);
0401   v = 1 / v.Mag() * v;
0402
0403   TVector3 u = v.Cross(TVector3(0, 0, 1));
0404   u = 1 / u.Mag() * u;
0405
0406   double u_smear = gRandom->Gaus(0, _phi_resolution);
0407   double v_smear = gRandom->Gaus(0, _r_resolution);
0408   pos.SetX(g4hit->get_avg_x() + u_smear * u.X() + v_smear * v.X());
0409   pos.SetY(g4hit->get_avg_y() + u_smear * u.Y() + v_smear * v.Y());
0410
0411   meas = new PHGenFit::PlanarMeasurement(pos, u, v, _phi_resolution, _r_resolution);
0412
0413   //    std::cout<<"------------\n";
0414   //    pos.Print();
0415   //    std::cout<<"at "<<istation<<" station, "<<ioctant << " octant \n";
0416   //    u.Print();
0417   //    v.Print();
0418
0419   //dynamic_cast<PHGenFit::PlanarMeasurement*> (meas)->getMeasurement()->Print();
0420
0421   return meas;
0422 }
0423
0424 PHGenFit::PlanarMeasurement* PHG4TrackFastSim::VertexMeasurement(const TVector3& vtx, const double dr,
0425                                                                  const double dphi)
0426 {
0427   PHGenFit::PlanarMeasurement* meas = NULL;
0428
0429   TVector3 u(1, 0, 0);
0430   TVector3 v(0, 1, 0);
0431
0432   TVector3 pos = vtx;
0433
0434   double u_smear = gRandom->Gaus(0, dphi);
0435   double v_smear = gRandom->Gaus(0, dr);
0436   pos.SetX(vtx.X() + u_smear * u.X() + v_smear * v.X());
0437   pos.SetY(vtx.Y() + u_smear * u.Y() + v_smear * v.Y());
0438
0439   meas = new PHGenFit::PlanarMeasurement(pos, u, v, dr, dphi);
0440
0441   return meas;
0442 }