Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
GenFitTrackProp.cc
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file GenFitTrackProp.cc
1 
7 #include "GenFitTrackProp.h"
8 
9 //#include <phgenfit/Tools.h>
10 #include <phgenfit/Fitter.h>
11 #include <phgenfit/PlanarMeasurement.h>
12 #include <phgenfit/Track.h>
13 #include <phgenfit/SpacepointMeasurement.h>
14 
15 #include <GenFit/RKTrackRep.h>
16 
17 #include <phool/phool.h>
18 #include <phool/getClass.h>
19 #include <phgeom/PHGeomUtility.h>
20 #include <phfield/PHFieldUtility.h>
21 #include <fun4all/Fun4AllReturnCodes.h>
22 #include <g4main/PHG4HitContainer.h>
23 #include <g4main/PHG4TruthInfoContainer.h>
24 #include <g4main/PHG4Particle.h>
25 #include <g4main/PHG4Hit.h>
26 #include <g4main/PHG4VtxPoint.h>
27 #include <fun4all/PHTFileServer.h>
28 #include <fun4all/Fun4AllServer.h>
29 
30 #include <g4hough/SvtxVertexMap.h>
31 #include <g4hough/SvtxVertex.h>
32 #include <g4hough/SvtxTrackMap.h>
33 #include <g4hough/SvtxTrack.h>
34 #include <g4hough/SvtxTrack_FastSim.h>
35 #include <g4hough/SvtxClusterMap.h>
36 #include <g4hough/SvtxCluster.h>
37 #include <g4hough/SvtxHitMap.h>
38 #include <g4hough/SvtxHit.h>
39 
40 #include <g4eval/SvtxEvalStack.h>
41 #include <g4eval/SvtxTrackEval.h>
42 #include <g4eval/SvtxClusterEval.h>
43 #include <g4eval/SvtxTruthEval.h>
44 #include <g4eval/SvtxVertexEval.h>
45 #include <g4eval/SvtxHitEval.h>
46 
47 #include <TTree.h>
48 #include <TH2D.h>
49 #include <TVector3.h>
50 #include <TDatabasePDG.h>
51 
52 #include <memory>
53 #include <iostream>
54 
55 #define LogDebug(exp) std::cout<<"DEBUG: " <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
56 #define LogError(exp) std::cout<<"ERROR: " <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
57 #define LogWarning(exp) std::cout<<"WARNING: " <<__FILE__<<": "<<__LINE__<<": "<< exp <<"\n"
58 
59 using namespace std;
60 
61 GenFitTrackProp::GenFitTrackProp(const string &name, const int pid_guess) :
62  SubsysReco(name),
63 
64  _fitter(nullptr),
65  _track_fitting_alg_name("DafRef"),
66  _do_evt_display(false),
67 
68  _pid_guess(pid_guess),
69 
70  _outfile_name("GenFitTrackProp.root"),
71  _eval_tree_tracks( NULL)
72 {
73  //initialize
74  _event = 0;
75 }
76 
77 
78 int GenFitTrackProp::Init(PHCompositeNode *topNode) {
79 
80  cout << PHWHERE << " Openning file " << _outfile_name << endl;
81  PHTFileServer::get().open(_outfile_name, "RECREATE");
82 
83  // create TTree
84  _eval_tree_tracks = new TTree("T", " => tracks");
85 
86  _eval_tree_tracks->Branch("event", &event, "event/I");
87  _eval_tree_tracks->Branch("gtrackID", &gtrackID, "gtrackID/I");
88  _eval_tree_tracks->Branch("gflavor", &gflavor, "gflavor/I");
89 
90  _eval_tree_tracks->Branch("gpx", &gpx, "gpx/D");
91  _eval_tree_tracks->Branch("gpy", &gpy, "gpy/D");
92  _eval_tree_tracks->Branch("gpz", &gpz, "gpz/D");
93  _eval_tree_tracks->Branch("gpt", &gpt, "gpt/D");
94 
95  _eval_tree_tracks->Branch("gvx", &gvx, "gvx/D");
96  _eval_tree_tracks->Branch("gvy", &gvy, "gvy/D");
97  _eval_tree_tracks->Branch("gvz", &gvz, "gvz/D");
98 
99  _eval_tree_tracks->Branch("trackID", &trackID, "trackID/I");
100  _eval_tree_tracks->Branch("charge", &charge, "charge/I");
101  _eval_tree_tracks->Branch("nhits", &nhits, "nhits/I");
102 
103  _eval_tree_tracks->Branch("px", &px, "px/D");
104  _eval_tree_tracks->Branch("py", &py, "py/D");
105  _eval_tree_tracks->Branch("pz", &pz, "pz/D");
106  _eval_tree_tracks->Branch("pt", &pt, "pt/D");
107 
108  _eval_tree_tracks->Branch("dca2d", &dca2d, "dca2d/D");
109 
110  _eval_tree_tracks->Branch("radius80", &radius80, "radius80/D");
111  _eval_tree_tracks->Branch("pathlength80", &pathlength80, "pathlength80/D");
112  _eval_tree_tracks->Branch("pathlength85", &pathlength85, "pathlength85/D");
113 
114  return Fun4AllReturnCodes::EVENT_OK;
115 }
116 
117 int GenFitTrackProp::End(PHCompositeNode *topNode) {
118 
119  PHTFileServer::get().cd(_outfile_name);
120 
121  _eval_tree_tracks->Write();
122 
123  //PHTFileServer::get().close();
124 
125  delete _svtxevalstack;
126 
127  delete _fitter;
128 
129  return Fun4AllReturnCodes::EVENT_OK;
130 }
131 
132 int GenFitTrackProp::InitRun(PHCompositeNode *topNode) {
133 
134  TGeoManager* tgeo_manager = PHGeomUtility::GetTGeoManager(topNode);
135  PHField * field = PHFieldUtility::GetFieldMapNode(nullptr, topNode);
136 
137  _fitter = PHGenFit::Fitter::getInstance(tgeo_manager,
138  field, _track_fitting_alg_name,
139  "RKTrackRep", _do_evt_display);
140 
141  _fitter->set_verbosity(verbosity);
142 
143  if (!_fitter) {
144  cerr << PHWHERE << endl;
145  return Fun4AllReturnCodes::ABORTRUN;
146  }
147 
148  return Fun4AllReturnCodes::EVENT_OK;
149 }
150 
151 int GenFitTrackProp::process_event(PHCompositeNode *topNode) {
152  _event++;
153  if (verbosity >= 2 and _event % 1 == 0)
154  cout << PHWHERE << "Events processed: " << _event << endl;
155 
156  if (!_svtxevalstack) {
157  _svtxevalstack = new SvtxEvalStack(topNode);
158  _svtxevalstack->set_verbosity(verbosity+1);
159  } else {
160  _svtxevalstack->next_event(topNode);
161  }
162 
163  GetNodes(topNode);
164 
165  fill_tree(topNode);
166 
167  return Fun4AllReturnCodes::EVENT_OK;
168 }
169 
170 
174 void GenFitTrackProp::fill_tree(PHCompositeNode *topNode) {
175 
176  // Make sure to reset all the TTree variables before trying to set them.
177  reset_variables();
178 
179  event = _event;
180 
181  if (!_truth_container) {
182  LogError("_truth_container not found!");
183  return;
184  }
185 
186  if (!_trackmap) {
187  LogError("_trackmap not found!");
188  return;
189  }
190 
191  SvtxTrackEval* trackeval = _svtxevalstack->get_track_eval();
192 
193  for (auto track_itr = _trackmap->begin(); track_itr != _trackmap->end();
194  track_itr++) {
195 
196  SvtxTrack* track = track_itr->second;
197 
198  if(!track) continue;
199 
200  trackID = track->get_id();
201  charge = track->get_charge();
202  nhits = track->size_clusters();
203 
204  px = track->get_px();
205  py = track->get_py();
206  pz = track->get_pz();
207 
208  pt = sqrt(px*px+py*py);
209 
210  dca2d = track->get_dca2d();
211 
212  auto last_state_iter = --track->end_states();
213 
214  SvtxTrackState * trackstate = last_state_iter->second;
215 
216  if(!trackstate) continue;
217 
218 // cout
219 // <<__LINE__
220 // <<": track->size_states(): " << track->size_states()
221 // <<": track->size_clusters(): " << track->size_clusters()
222 // <<endl;
223 
224  genfit::MeasuredStateOnPlane* msop80 = nullptr;
225 
226  auto pdg = unique_ptr<TDatabasePDG> (TDatabasePDG::Instance());
227  int reco_charge = track->get_charge();
228  int gues_charge = pdg->GetParticle(_pid_guess)->Charge();
229  if(reco_charge*gues_charge<0) _pid_guess *= -1;
230 
231  genfit::AbsTrackRep* rep = new genfit::RKTrackRep(_pid_guess);
232 
233  trackstate->get_x();
234  {
235  TVector3 pos(trackstate->get_x(), trackstate->get_y(),
236  trackstate->get_z());
237 
238  TVector3 mom(trackstate->get_px(), trackstate->get_py(),
239  trackstate->get_pz());
240  TMatrixDSym cov(6);
241  for (int i = 0; i < 6; ++i) {
242  for (int j = 0; j < 6; ++j) {
243  cov[i][j] = trackstate->get_error(i, j);
244  }
245  }
246 
247  msop80 = new genfit::MeasuredStateOnPlane(rep);
248  msop80->setPosMomCov(pos, mom, cov);
249 
250  radius80 = pos.Perp();
251  }
252 
253  double radius = 85;
254  TVector3 line_point(0,0,0);
255  TVector3 line_direction(0,0,1);
256 
257  pathlength80 = last_state_iter->first;
258  genfit::MeasuredStateOnPlane* msop85 = new genfit::MeasuredStateOnPlane(*msop80);
259  rep->extrapolateToCylinder(*msop85, radius, line_point, line_direction);
260  //pathlength85 = pathlength80 + rep->extrapolateToCylinder(*msop85, radius, line_point, line_direction);
261 
262  TVector3 tof_hit_pos(msop85->getPos());
263  TVector3 tof_hit_norm(msop85->getPos().X(),msop85->getPos().Y(),0);
264  genfit::SharedPlanePtr tof_module_plane (new genfit::DetPlane(tof_hit_pos,tof_hit_norm));
265 
266  genfit::MeasuredStateOnPlane* msop_tof_module = new genfit::MeasuredStateOnPlane(*msop80);
267  pathlength85 = pathlength80 + rep->extrapolateToPlane(*msop_tof_module, tof_module_plane);
268 
270  PHG4Particle* g4particle = trackeval->max_truth_particle_by_nclusters(
271  track);
272 
273  if(!g4particle) continue;
274 
275  gtrackID = g4particle->get_track_id();
276  gflavor = g4particle->get_pid();
277 
278  gpx = g4particle->get_px();
279  gpy = g4particle->get_py();
280  gpz = g4particle->get_pz();
281 
282  gpt = sqrt(gpx*gpx+gpy*gpy);
283 
284  _eval_tree_tracks->Fill();
285  }
286 
287  return;
288 }
289 
294 void GenFitTrackProp::reset_variables() {
295  event = -9999;
296 
297  //-- truth
298  gtrackID = -9999;
299  gflavor = -9999;
300  gpx = -9999;
301  gpy = -9999;
302  gpz = -9999;
303  gvx = -9999;
304  gvy = -9999;
305  gvz = -9999;
306 
307  //-- reco
308  trackID = -9999;
309  charge = -9999;
310  nhits = -9999;
311  px = -9999;
312  py = -9999;
313  pz = -9999;
314  dca2d = -9999;
315 }
316 
317 int GenFitTrackProp::GetNodes(PHCompositeNode * topNode) {
318  //DST objects
319  //Truth container
320  _truth_container = findNode::getClass<PHG4TruthInfoContainer>(topNode,
321  "G4TruthInfo");
322  if (!_truth_container && _event < 2) {
323  cout << PHWHERE << " PHG4TruthInfoContainer node not found on node tree"
324  << endl;
325  return Fun4AllReturnCodes::ABORTEVENT;
326  }
327 
328  _trackmap = findNode::getClass<SvtxTrackMap>(topNode,
329  "SvtxTrackMap");
330  if (!_trackmap && _event < 2) {
331  cout << PHWHERE << "SvtxTrackMap node not found on node tree"
332  << endl;
333  return Fun4AllReturnCodes::ABORTEVENT;
334  }
335 
336  return Fun4AllReturnCodes::EVENT_OK;
337 }
338