Tracking/GenFitTrackProp/GenFitTrackProp.cc

0001 /*!
0002  *  \file       GenFitTrackProp.cc
0003  *  \brief      Example module to extrapolate SvtxTrack with GenFit swiming
0004  *  \author     Haiwang Yu <yuhw@nmsu.edu>
0005  */
0006
0007 #include "GenFitTrackProp.h"
0008
0009 //#include <phgenfit/Tools.h>
0010 #include <phgenfit/Fitter.h>
0011 #include <phgenfit/PlanarMeasurement.h>
0012 #include <phgenfit/Track.h>
0013 #include <phgenfit/SpacepointMeasurement.h>
0014
0015 #include <GenFit/RKTrackRep.h>
0016
0017 #include <phool/phool.h>
0018 #include <phool/getClass.h>
0019 #include <phgeom/PHGeomUtility.h>
0020 #include <phfield/PHFieldUtility.h>
0021 #include <fun4all/Fun4AllReturnCodes.h>
0022 #include <g4main/PHG4HitContainer.h>
0023 #include <g4main/PHG4TruthInfoContainer.h>
0024 #include <g4main/PHG4Particle.h>
0025 #include <g4main/PHG4Hit.h>
0026 #include <g4main/PHG4VtxPoint.h>
0027 #include <fun4all/PHTFileServer.h>
0028 #include <fun4all/Fun4AllServer.h>
0029
0030 #include <g4hough/SvtxVertexMap.h>
0031 #include <g4hough/SvtxVertex.h>
0032 #include <g4hough/SvtxTrackMap.h>
0033 #include <g4hough/SvtxTrack.h>
0034 #include <g4hough/SvtxTrack_FastSim.h>
0035 #include <g4hough/SvtxClusterMap.h>
0036 #include <g4hough/SvtxCluster.h>
0037 #include <g4hough/SvtxHitMap.h>
0038 #include <g4hough/SvtxHit.h>
0039
0040 #include <g4eval/SvtxEvalStack.h>
0041 #include <g4eval/SvtxTrackEval.h>
0042 #include <g4eval/SvtxClusterEval.h>
0043 #include <g4eval/SvtxTruthEval.h>
0044 #include <g4eval/SvtxVertexEval.h>
0045 #include <g4eval/SvtxHitEval.h>
0046
0047 #include <TTree.h>
0048 #include <TH2D.h>
0049 #include <TVector3.h>
0050 #include <TDatabasePDG.h>
0051
0052 #include <memory>
0053 #include <iostream>
0054
0055 #define LogDebug(exp)       std::cout<<"DEBUG: "    <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
0056 #define LogError(exp)       std::cout<<"ERROR: "    <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
0057 #define LogWarning(exp) std::cout<<"WARNING: "  <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
0058
0059 using namespace std;
0060
0061 GenFitTrackProp::GenFitTrackProp(const string &name, const int pid_guess) :
0062         SubsysReco(name),
0063
0064         _fitter(nullptr),
0065         _track_fitting_alg_name("DafRef"),
0066         _do_evt_display(false),
0067
0068         _pid_guess(pid_guess),
0069
0070         _outfile_name("GenFitTrackProp.root"),
0071         _eval_tree_tracks( NULL)
0072 {
0073     //initialize
0074     _event = 0;
0075 }
0076
0077
0078 int GenFitTrackProp::Init(PHCompositeNode *topNode) {
0079
0080     cout << PHWHERE << " Openning file " << _outfile_name << endl;
0081     PHTFileServer::get().open(_outfile_name, "RECREATE");
0082
0083     // create TTree
0084     _eval_tree_tracks = new TTree("T", " => tracks");
0085
0086     _eval_tree_tracks->Branch("event", &event, "event/I");
0087     _eval_tree_tracks->Branch("gtrackID", &gtrackID, "gtrackID/I");
0088     _eval_tree_tracks->Branch("gflavor", &gflavor, "gflavor/I");
0089
0090     _eval_tree_tracks->Branch("gpx", &gpx, "gpx/D");
0091     _eval_tree_tracks->Branch("gpy", &gpy, "gpy/D");
0092     _eval_tree_tracks->Branch("gpz", &gpz, "gpz/D");
0093     _eval_tree_tracks->Branch("gpt", &gpt, "gpt/D");
0094
0095     _eval_tree_tracks->Branch("gvx", &gvx, "gvx/D");
0096     _eval_tree_tracks->Branch("gvy", &gvy, "gvy/D");
0097     _eval_tree_tracks->Branch("gvz", &gvz, "gvz/D");
0098
0099     _eval_tree_tracks->Branch("trackID", &trackID, "trackID/I");
0100     _eval_tree_tracks->Branch("charge", &charge, "charge/I");
0101     _eval_tree_tracks->Branch("nhits", &nhits, "nhits/I");
0102
0103     _eval_tree_tracks->Branch("px", &px, "px/D");
0104     _eval_tree_tracks->Branch("py", &py, "py/D");
0105     _eval_tree_tracks->Branch("pz", &pz, "pz/D");
0106     _eval_tree_tracks->Branch("pt", &pt, "pt/D");
0107
0108     _eval_tree_tracks->Branch("dca2d", &dca2d, "dca2d/D");
0109
0110     _eval_tree_tracks->Branch("radius80", &radius80, "radius80/D");
0111     _eval_tree_tracks->Branch("pathlength80", &pathlength80, "pathlength80/D");
0112     _eval_tree_tracks->Branch("pathlength85", &pathlength85, "pathlength85/D");
0113
0114     return Fun4AllReturnCodes::EVENT_OK;
0115 }
0116
0117 int GenFitTrackProp::End(PHCompositeNode *topNode) {
0118
0119     PHTFileServer::get().cd(_outfile_name);
0120
0121     _eval_tree_tracks->Write();
0122
0123     //PHTFileServer::get().close();
0124
0125     delete _svtxevalstack;
0126
0127     delete _fitter;
0128
0129     return Fun4AllReturnCodes::EVENT_OK;
0130 }
0131
0132 int GenFitTrackProp::InitRun(PHCompositeNode *topNode) {
0133
0134     TGeoManager* tgeo_manager = PHGeomUtility::GetTGeoManager(topNode);
0135   PHField * field = PHFieldUtility::GetFieldMapNode(nullptr, topNode);
0136
0137     _fitter = PHGenFit::Fitter::getInstance(tgeo_manager,
0138         field, _track_fitting_alg_name,
0139             "RKTrackRep", _do_evt_display);
0140
0141     _fitter->set_verbosity(verbosity);
0142
0143     if (!_fitter) {
0144         cerr << PHWHERE << endl;
0145         return Fun4AllReturnCodes::ABORTRUN;
0146     }
0147
0148     return Fun4AllReturnCodes::EVENT_OK;
0149 }
0150
0151 int GenFitTrackProp::process_event(PHCompositeNode *topNode) {
0152     _event++;
0153     if (verbosity >= 2 and _event % 1 == 0)
0154         cout << PHWHERE << "Events processed: " << _event << endl;
0155
0156       if (!_svtxevalstack) {
0157         _svtxevalstack = new SvtxEvalStack(topNode);
0158         _svtxevalstack->set_verbosity(verbosity+1);
0159       } else {
0160         _svtxevalstack->next_event(topNode);
0161       }
0162
0163     GetNodes(topNode);
0164
0165     fill_tree(topNode);
0166
0167     return Fun4AllReturnCodes::EVENT_OK;
0168 }
0169
0170
0171 /*!
0172  * Fill the trees with truth, track fit, and cluster information
0173  */
0174 void GenFitTrackProp::fill_tree(PHCompositeNode *topNode) {
0175
0176     // Make sure to reset all the TTree variables before trying to set them.
0177     reset_variables();
0178
0179     event = _event;
0180
0181     if (!_truth_container) {
0182         LogError("_truth_container not found!");
0183         return;
0184     }
0185
0186     if (!_trackmap) {
0187         LogError("_trackmap not found!");
0188         return;
0189     }
0190
0191     SvtxTrackEval* trackeval = _svtxevalstack->get_track_eval();
0192
0193     for (auto track_itr = _trackmap->begin(); track_itr != _trackmap->end();
0194             track_itr++) {
0195
0196         SvtxTrack* track = track_itr->second;
0197
0198         if(!track) continue;
0199
0200         trackID = track->get_id();
0201         charge = track->get_charge();
0202         nhits = track->size_clusters();
0203
0204         px = track->get_px();
0205         py = track->get_py();
0206         pz = track->get_pz();
0207
0208         pt = sqrt(px*px+py*py);
0209
0210         dca2d = track->get_dca2d();
0211
0212         auto last_state_iter = --track->end_states();
0213
0214         SvtxTrackState * trackstate = last_state_iter->second;
0215
0216         if(!trackstate) continue;
0217
0218 //      cout
0219 //      <<__LINE__
0220 //      <<": track->size_states(): " << track->size_states()
0221 //      <<": track->size_clusters(): " << track->size_clusters()
0222 //      <<endl;
0223
0224         genfit::MeasuredStateOnPlane* msop80 = nullptr;
0225
0226         auto pdg = unique_ptr<TDatabasePDG> (TDatabasePDG::Instance());
0227         int reco_charge = track->get_charge();
0228         int gues_charge = pdg->GetParticle(_pid_guess)->Charge();
0229         if(reco_charge*gues_charge<0) _pid_guess *= -1;
0230
0231         genfit::AbsTrackRep* rep = new genfit::RKTrackRep(_pid_guess);
0232
0233         trackstate->get_x();
0234         {
0235             TVector3 pos(trackstate->get_x(), trackstate->get_y(),
0236                     trackstate->get_z());
0237
0238             TVector3 mom(trackstate->get_px(), trackstate->get_py(),
0239                     trackstate->get_pz());
0240             TMatrixDSym cov(6);
0241             for (int i = 0; i < 6; ++i) {
0242                 for (int j = 0; j < 6; ++j) {
0243                     cov[i][j] = trackstate->get_error(i, j);
0244                 }
0245             }
0246
0247             msop80 = new genfit::MeasuredStateOnPlane(rep);
0248             msop80->setPosMomCov(pos, mom, cov);
0249
0250             radius80 = pos.Perp();
0251         }
0252
0253         double radius = 85;
0254         TVector3 line_point(0,0,0);
0255         TVector3 line_direction(0,0,1);
0256
0257         pathlength80 = last_state_iter->first;
0258         genfit::MeasuredStateOnPlane* msop85 = new genfit::MeasuredStateOnPlane(*msop80);
0259         rep->extrapolateToCylinder(*msop85, radius, line_point, line_direction);
0260         //pathlength85 = pathlength80 + rep->extrapolateToCylinder(*msop85, radius, line_point, line_direction);
0261
0262         TVector3 tof_hit_pos(msop85->getPos());
0263         TVector3 tof_hit_norm(msop85->getPos().X(),msop85->getPos().Y(),0);
0264         genfit::SharedPlanePtr tof_module_plane (new genfit::DetPlane(tof_hit_pos,tof_hit_norm));
0265
0266         genfit::MeasuredStateOnPlane* msop_tof_module = new genfit::MeasuredStateOnPlane(*msop80);
0267         pathlength85 = pathlength80 + rep->extrapolateToPlane(*msop_tof_module, tof_module_plane);
0268
0269         //! Truth information
0270         PHG4Particle* g4particle = trackeval->max_truth_particle_by_nclusters(
0271                 track);
0272
0273         if(!g4particle) continue;
0274
0275         gtrackID = g4particle->get_track_id();
0276         gflavor = g4particle->get_pid();
0277
0278         gpx = g4particle->get_px();
0279         gpy = g4particle->get_py();
0280         gpz = g4particle->get_pz();
0281
0282         gpt = sqrt(gpx*gpx+gpy*gpy);
0283
0284         _eval_tree_tracks->Fill();
0285     }
0286
0287     return;
0288 }
0289
0290 /*!
0291  * Reset all the tree variables to their default values.
0292  * Needs to be called at the start of every event
0293  */
0294 void GenFitTrackProp::reset_variables() {
0295     event = -9999;
0296
0297     //-- truth
0298     gtrackID = -9999;
0299     gflavor = -9999;
0300     gpx = -9999;
0301     gpy = -9999;
0302     gpz = -9999;
0303     gvx = -9999;
0304     gvy = -9999;
0305     gvz = -9999;
0306
0307     //-- reco
0308     trackID = -9999;
0309     charge = -9999;
0310     nhits = -9999;
0311     px = -9999;
0312     py = -9999;
0313     pz = -9999;
0314     dca2d = -9999;
0315 }
0316
0317 int GenFitTrackProp::GetNodes(PHCompositeNode * topNode) {
0318     //DST objects
0319     //Truth container
0320     _truth_container = findNode::getClass<PHG4TruthInfoContainer>(topNode,
0321             "G4TruthInfo");
0322     if (!_truth_container && _event < 2) {
0323         cout << PHWHERE << " PHG4TruthInfoContainer node not found on node tree"
0324                 << endl;
0325         return Fun4AllReturnCodes::ABORTEVENT;
0326     }
0327
0328     _trackmap = findNode::getClass<SvtxTrackMap>(topNode,
0329             "SvtxTrackMap");
0330     if (!_trackmap && _event < 2) {
0331         cout << PHWHERE << "SvtxTrackMap node not found on node tree"
0332                 << endl;
0333         return Fun4AllReturnCodes::ABORTEVENT;
0334     }
0335
0336     return Fun4AllReturnCodes::EVENT_OK;
0337 }
0338