ECCE @ EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
MomentumEvaluator.cc
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file MomentumEvaluator.cc
1 #include "MomentumEvaluator.h"
2 
5 
6 #include <g4main/PHG4Hit.h>
8 #include <g4main/PHG4Particle.h>
10 #include <g4main/PHG4VtxPoint.h>
11 
13 #include <phool/getClass.h>
14 
15 #include <TFile.h>
16 #include <TNtuple.h>
17 
18 #include <cstdlib>
19 #include <iostream>
20 #include <map>
21 #include <memory>
22 #include <vector>
23 #include <utility>
24 
25 using namespace std;
26 
28 {
29  public:
30  float px, py, pz;
31  float dcax, dcay, dcaz;
32  float quality;
33  TrivialTrack(float x, float y, float z, float dx, float dy, float dz, float qual = 0.)
34  : px(x)
35  , py(y)
36  , pz(z)
37  , dcax(dx)
38  , dcay(dy)
39  , dcaz(dz)
40  , quality(qual)
41  {
42  }
44 };
45 
47 {
48  protected:
49  float px_lo, px_hi;
50  float py_lo, py_hi;
51  float pz_lo, pz_hi;
52 
53  int level;
54  int maxlevel;
55 
56  unsigned int x_pos, y_pos, z_pos;
57 
58  RecursiveMomentumContainer* containers[2][2][2];
59 
60  public:
61  RecursiveMomentumContainer(float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI, int MLEV, int LEV = 0)
62  : px_lo(PX_LO)
63  , px_hi(PX_HI)
64  , py_lo(PY_LO)
65  , py_hi(PY_HI)
66  , pz_lo(PZ_LO)
67  , pz_hi(PZ_HI)
68  , level(LEV)
69  , maxlevel(MLEV)
70  , x_pos(0)
71  , y_pos(0)
72  , z_pos(0)
73  {
74  for (unsigned int i = 0; i < 2; ++i)
75  {
76  for (unsigned int j = 0; j < 2; ++j)
77  {
78  for (unsigned int k = 0; k < 2; ++k)
79  {
80  containers[i][j][k] = nullptr;
81  }
82  }
83  }
84  }
86  {
87  for (unsigned int i = 0; i < 2; ++i)
88  {
89  for (unsigned int j = 0; j < 2; ++j)
90  {
91  for (unsigned int k = 0; k < 2; ++k)
92  {
93  if (containers[i][j][k] != nullptr)
94  {
95  delete containers[i][j][k];
96  }
97  }
98  }
99  }
100  }
101 
102  virtual bool insert(TrivialTrack& track);
103 
104  virtual TrivialTrack* begin()
105  {
106  x_pos = 0;
107  y_pos = 0;
108  z_pos = 0;
109  while (true)
110  {
111  if (containers[x_pos][y_pos][z_pos] == nullptr)
112  {
113  if (z_pos == 0)
114  {
115  z_pos = 1;
116  continue;
117  }
118  else
119  {
120  if (y_pos == 0)
121  {
122  z_pos = 0;
123  y_pos = 1;
124  continue;
125  }
126  else
127  {
128  if (x_pos == 0)
129  {
130  z_pos = 0;
131  y_pos = 0;
132  x_pos = 1;
133  continue;
134  }
135  else
136  {
137  return nullptr;
138  }
139  }
140  }
141  }
142  else
143  {
144  return containers[x_pos][y_pos][z_pos]->begin();
145  }
146  }
147  }
148 
149  virtual TrivialTrack* next()
150  {
151  bool block_changed = false;
152  while (true)
153  {
154  if (containers[x_pos][y_pos][z_pos] == nullptr)
155  {
156  block_changed = true;
157  if (z_pos == 0)
158  {
159  z_pos = 1;
160  continue;
161  }
162  else
163  {
164  if (y_pos == 0)
165  {
166  z_pos = 0;
167  y_pos = 1;
168  continue;
169  }
170  else
171  {
172  if (x_pos == 0)
173  {
174  z_pos = 0;
175  y_pos = 0;
176  x_pos = 1;
177  continue;
178  }
179  else
180  {
181  return nullptr;
182  }
183  }
184  }
185  }
186  TrivialTrack* val = nullptr;
187  if (block_changed == true)
188  {
189  val = containers[x_pos][y_pos][z_pos]->begin();
190  }
191  else
192  {
193  val = containers[x_pos][y_pos][z_pos]->next();
194  }
195 
196  if (val == nullptr)
197  {
198  block_changed = true;
199  if (z_pos == 0)
200  {
201  z_pos = 1;
202  continue;
203  }
204  else
205  {
206  if (y_pos == 0)
207  {
208  z_pos = 0;
209  y_pos = 1;
210  continue;
211  }
212  else
213  {
214  if (x_pos == 0)
215  {
216  z_pos = 0;
217  y_pos = 0;
218  x_pos = 1;
219  continue;
220  }
221  else
222  {
223  return nullptr;
224  }
225  }
226  }
227  }
228  else
229  {
230  return val;
231  }
232  }
233  }
234 
235  virtual void append_list(vector<TrivialTrack*>& track_list, float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI)
236  {
237  for (unsigned int i = 0; i < 2; ++i)
238  {
239  for (unsigned int j = 0; j < 2; ++j)
240  {
241  for (unsigned int k = 0; k < 2; ++k)
242  {
243  if (containers[i][j][k] == nullptr)
244  {
245  continue;
246  }
247 
248  if ((containers[i][j][k]->px_hi < PX_LO) || (containers[i][j][k]->px_lo > PX_HI) || (containers[i][j][k]->py_hi < PY_LO) || (containers[i][j][k]->py_lo > PY_HI) || (containers[i][j][k]->pz_hi < PZ_LO) || (containers[i][j][k]->pz_lo > PZ_HI))
249  {
250  continue;
251  }
252 
253  containers[i][j][k]->append_list(track_list, PX_LO, PX_HI, PY_LO, PY_HI, PZ_LO, PZ_HI);
254  }
255  }
256  }
257  }
258 };
259 
261 {
262  public:
263  RecursiveMomentumContainerEnd(float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI, int MLEV, int LEV = 0)
264  : RecursiveMomentumContainer(PX_LO, PX_HI, PY_LO, PY_HI, PZ_LO, PZ_HI, MLEV, LEV)
265  {
266  }
267 
269  {
270  }
271 
272  bool insert(TrivialTrack& track) override
273  {
274  tracks.push_back(track);
275  return true;
276  }
277 
278  TrivialTrack* begin() override
279  {
280  x_pos = 0;
281  return (&(tracks.at(0)));
282  }
283 
284  TrivialTrack* next() override
285  {
286  if (x_pos >= (tracks.size() - 1))
287  {
288  return nullptr;
289  }
290  else
291  {
292  x_pos += 1;
293  return (&(tracks[x_pos]));
294  }
295  }
296 
297  void append_list(vector<TrivialTrack*>& track_list, float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI) override
298  {
299  for (unsigned int i = 0; i < tracks.size(); ++i)
300  {
301  if ((tracks[i].px < PX_LO) || (tracks[i].px > PX_HI) || (tracks[i].py < PY_LO) || (tracks[i].py > PY_HI) || (tracks[i].pz < PZ_LO) || (tracks[i].pz > PZ_HI))
302  {
303  continue;
304  }
305  track_list.push_back(&(tracks[i]));
306  }
307  }
308 
309  protected:
310  vector<TrivialTrack> tracks;
311 };
312 
314 {
315  if ((track.px < px_lo) || (track.py < py_lo) || (track.pz < pz_lo) || (track.px > px_hi) || (track.py > py_hi) || (track.pz > pz_hi))
316  {
317  return false;
318  }
319 
320  int x_ind = 0;
321  if (track.px > (px_lo + 0.5 * (px_hi - px_lo)))
322  {
323  x_ind = 1;
324  }
325  int y_ind = 0;
326  if (track.py > (py_lo + 0.5 * (py_hi - py_lo)))
327  {
328  y_ind = 1;
329  }
330  int z_ind = 0;
331  if (track.pz > (pz_lo + 0.5 * (pz_hi - pz_lo)))
332  {
333  z_ind = 1;
334  }
335 
336  if (containers[x_ind][y_ind][z_ind] == nullptr)
337  {
338  float px_lo_new = px_lo + (float(x_ind)) * 0.5 * (px_hi - px_lo);
339  float px_hi_new = px_lo_new + 0.5 * (px_hi - px_lo);
340 
341  float py_lo_new = py_lo + (float(y_ind)) * 0.5 * (py_hi - py_lo);
342  float py_hi_new = py_lo_new + 0.5 * (py_hi - py_lo);
343 
344  float pz_lo_new = pz_lo + (float(z_ind)) * 0.5 * (pz_hi - pz_lo);
345  float pz_hi_new = pz_lo_new + 0.5 * (pz_hi - pz_lo);
346 
347  if (level < maxlevel)
348  {
349  containers[x_ind][y_ind][z_ind] = new RecursiveMomentumContainer(px_lo_new, px_hi_new, py_lo_new, py_hi_new, pz_lo_new, pz_hi_new, maxlevel, level + 1);
350  }
351  else
352  {
353  containers[x_ind][y_ind][z_ind] = new RecursiveMomentumContainerEnd(px_lo_new, px_hi_new, py_lo_new, py_hi_new, pz_lo_new, pz_hi_new, maxlevel, level + 1);
354  }
355  }
356  return containers[x_ind][y_ind][z_ind]->insert(track);
357 }
358 
359 MomentumEvaluator::MomentumEvaluator(const std::string& fname, float pt_s, float pz_s, unsigned int /*n_l*/, unsigned int n_i, unsigned int n_r, float i_z, float o_z)
360  : ntp_true(nullptr)
361  , ntp_reco(nullptr)
362  , pt_search_scale(pt_s)
363  , pz_search_scale(pz_s)
364  , event_counter(0)
365  , file_name(fname)
366  , n_inner_layers(n_i)
367  , n_required_layers(n_r)
368  , inner_z_length(i_z)
369  , outer_z_length(o_z)
370 {
371 }
373 {
374  if (ntp_true != nullptr)
375  {
376  delete ntp_true;
377  }
378  if (ntp_reco != nullptr)
379  {
380  delete ntp_reco;
381  }
382 }
383 
385 {
386  if (ntp_true != nullptr)
387  {
388  delete ntp_true;
389  }
390  if (ntp_reco != nullptr)
391  {
392  delete ntp_reco;
393  }
394 
395  ntp_true = new TNtuple("ntp_true", "true simulated tracks", "event:px:py:pz:dcax:dcay:dcaz:r_px:r_py:r_pz:r_dcax:r_dcay:r_dcaz:quality");
396  ntp_reco = new TNtuple("ntp_reco", "reconstructed tracks", "event:px:py:pz:dcax:dcay:dcaz:t_px:t_py:t_pz:t_dcax:t_dcay:t_dcaz:quality");
397  event_counter = 0;
398 
400 }
401 
403 {
404  PHG4TruthInfoContainer* truthinfo = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
405 
406  PHG4HitContainer* g4hits = findNode::getClass<PHG4HitContainer>(topNode, "G4HIT_SVTX");
407  if (g4hits == nullptr)
408  {
409  cout << "can't find PHG4HitContainer" << endl;
410  exit(1);
411  }
412  PHG4HitContainer::ConstRange g4range = g4hits->getHits();
413 
414  // set<int> trkids;
415  map<int, pair<unsigned int, unsigned int> > trkids;
416 
417  for (PHG4HitContainer::ConstIterator iter = g4range.first; iter != g4range.second; ++iter)
418  {
419  PHG4Hit* hit = iter->second;
420 
421  int layer = hit->get_layer();
422  float length = outer_z_length;
423  if (((unsigned int) layer) < n_inner_layers)
424  {
425  length = inner_z_length;
426  }
427  if (fabs(hit->get_z(0)) > length)
428  {
429  continue;
430  }
431 
432  int trk_id = hit->get_trkid();
433  if (trkids.find(trk_id) == trkids.end())
434  {
435  trkids[trk_id].first = 0;
436  trkids[trk_id].second = 0;
437  }
438  if (hit->get_layer() < 32)
439  {
440  trkids[trk_id].first = (trkids[trk_id].first | (1 << (hit->get_layer())));
441  }
442  else
443  {
444  trkids[trk_id].second = (trkids[trk_id].second | (1 << (hit->get_layer() - 32)));
445  }
446 
447  // cout<<"trk_id = "<<trk_id<<endl;
448  // cout<<"layer = "<<hit->get_layer()<<endl;
449  // cout<<"nlayer = "<<__builtin_popcount(trkids[trk_id].first)+__builtin_popcount(trkids[trk_id].second)<<endl<<endl;
450  // trkids.insert(trk_id);
451  }
452 
453  SvtxTrackMap* trackmap = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
454 
455  PHG4VtxPoint* gvertex = truthinfo->GetPrimaryVtx(truthinfo->GetPrimaryVertexIndex());
456  float gvx = gvertex->get_x();
457  float gvy = gvertex->get_y();
458  float gvz = gvertex->get_z();
459 
460  RecursiveMomentumContainer true_sorted(-20., 20., -20., 20., -20., 20., 10);
461 
462  // PHG4TruthInfoContainer::Map primarymap = truthinfo->GetPrimaryMap();
463  PHG4TruthInfoContainer::Map primarymap = truthinfo->GetMap();
464  for (PHG4TruthInfoContainer::Iterator iter = primarymap.begin(); iter != primarymap.end(); ++iter)
465  {
466  PHG4Particle* particle = iter->second;
467 
468  float vx = truthinfo->GetVtx(particle->get_vtx_id())->get_x();
469  float vy = truthinfo->GetVtx(particle->get_vtx_id())->get_y();
470  float vz = truthinfo->GetVtx(particle->get_vtx_id())->get_z();
471 
472  TrivialTrack track(particle->get_px(), particle->get_py(), particle->get_pz(), vx - gvx, vy - gvy, vz - gvz);
473 
474  if (((track.px * track.px) + (track.py * track.py)) < (0.1 * 0.1))
475  {
476  continue;
477  }
478 
479  if (trkids.find(particle->get_track_id()) == trkids.end())
480  {
481  continue;
482  }
483 
484  // cout<<"trk, nhits = "<<particle->get_track_id()<<" "<<__builtin_popcount(trkids[particle->get_track_id()].first)+__builtin_popcount(trkids[particle->get_track_id()].second)<<endl;
485 
486  if (__builtin_popcount(trkids[particle->get_track_id()].first) + __builtin_popcount(trkids[particle->get_track_id()].second) < (int) n_required_layers)
487  {
488  continue;
489  }
490 
491  true_sorted.insert(track);
492  }
493 
494  RecursiveMomentumContainer reco_sorted(-20., 20., -20., 20., -20., 20., 10);
495  for (SvtxTrackMap::Iter iter = trackmap->begin(); iter != trackmap->end(); ++iter)
496  {
497  SvtxTrack* track = iter->second;
498 
499  TrivialTrack ttrack(track->get_px(), track->get_py(), track->get_pz(), track->get_x() - gvx, track->get_y() - gvy, track->get_z() - gvz, track->get_quality());
500  reco_sorted.insert(ttrack);
501  }
502 
503  TrivialTrack* t_track = true_sorted.begin();
504  vector<TrivialTrack*> pointer_list;
505  while (t_track != nullptr)
506  {
507  pointer_list.clear();
508 
509  float pt = sqrt((t_track->px * t_track->px) + (t_track->py * t_track->py));
510  float pt_diff = pt * pt_search_scale;
511  float px_lo = t_track->px - pt_diff;
512  float px_hi = t_track->px + pt_diff;
513  float py_lo = t_track->py - pt_diff;
514  float py_hi = t_track->py + pt_diff;
515  float pz_diff = fabs(t_track->pz) * pz_search_scale;
516  float pz_lo = t_track->pz - pz_diff;
517  float pz_hi = t_track->pz + pz_diff;
518 
519  reco_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
520 
521  if (pointer_list.size() > 0)
522  {
523  float mom_true = sqrt(pt * pt + (t_track->pz) * (t_track->pz));
524  float best_ind = 0;
525  float mom_reco = sqrt((pointer_list[0]->px) * (pointer_list[0]->px) + (pointer_list[0]->py) * (pointer_list[0]->py) + (pointer_list[0]->pz) * (pointer_list[0]->pz));
526  float best_mom = mom_reco;
527  for (unsigned int i = 1; i < pointer_list.size(); ++i)
528  {
529  mom_reco = sqrt((pointer_list[i]->px) * (pointer_list[i]->px) + (pointer_list[i]->py) * (pointer_list[i]->py) + (pointer_list[i]->pz) * (pointer_list[i]->pz));
530  if (fabs(mom_true - mom_reco) < fabs(mom_true - best_mom))
531  {
532  best_mom = mom_reco;
533  best_ind = i;
534  }
535  }
536 
537  float ntp_data[14] = {(float) event_counter, t_track->px, t_track->py, t_track->pz, t_track->dcax, t_track->dcay, t_track->dcaz, pointer_list[best_ind]->px, pointer_list[best_ind]->py, pointer_list[best_ind]->pz, pointer_list[best_ind]->dcax, pointer_list[best_ind]->dcay, pointer_list[best_ind]->dcaz, pointer_list[best_ind]->quality};
538  ntp_true->Fill(ntp_data);
539  }
540  else
541  {
542  float ntp_data[14] = {(float) event_counter, t_track->px, t_track->py, t_track->pz, t_track->dcax, t_track->dcay, t_track->dcaz, -9999., -9999., -9999., -9999., -9999., -9999., -9999.};
543  ntp_true->Fill(ntp_data);
544  }
545 
546  t_track = true_sorted.next();
547  }
548 
549  TrivialTrack* r_track = reco_sorted.begin();
550  while (r_track != nullptr)
551  {
552  pointer_list.clear();
553 
554  float pt = sqrt((r_track->px * r_track->px) + (r_track->py * r_track->py));
555  float pt_diff = pt * pt_search_scale;
556  float px_lo = r_track->px - pt_diff;
557  float px_hi = r_track->px + pt_diff;
558  float py_lo = r_track->py - pt_diff;
559  float py_hi = r_track->py + pt_diff;
560  float pz_diff = fabs(r_track->pz) * pz_search_scale;
561  float pz_lo = r_track->pz - pz_diff;
562  float pz_hi = r_track->pz + pz_diff;
563 
564  true_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
565 
566  if (pointer_list.size() > 0)
567  {
568  float mom_reco = sqrt(pt * pt + (r_track->pz) * (r_track->pz));
569  float best_ind = 0;
570  float mom_true = sqrt((pointer_list[0]->px) * (pointer_list[0]->px) + (pointer_list[0]->py) * (pointer_list[0]->py) + (pointer_list[0]->pz) * (pointer_list[0]->pz));
571  float best_mom = mom_true;
572  for (unsigned int i = 1; i < pointer_list.size(); ++i)
573  {
574  mom_true = sqrt((pointer_list[i]->px) * (pointer_list[i]->px) + (pointer_list[i]->py) * (pointer_list[i]->py) + (pointer_list[i]->pz) * (pointer_list[i]->pz));
575  if (fabs(mom_reco - mom_true) < fabs(mom_reco - best_mom))
576  {
577  best_mom = mom_true;
578  best_ind = i;
579  }
580  }
581 
582  float ntp_data[14] = {(float) event_counter, r_track->px, r_track->py, r_track->pz, r_track->dcax, r_track->dcay, r_track->dcaz, pointer_list[best_ind]->px, pointer_list[best_ind]->py, pointer_list[best_ind]->pz, pointer_list[best_ind]->dcax, pointer_list[best_ind]->dcay, pointer_list[best_ind]->dcaz, r_track->quality};
583  ntp_reco->Fill(ntp_data);
584  }
585  else
586  {
587  float ntp_data[14] = {(float) event_counter, r_track->px, r_track->py, r_track->pz, r_track->dcax, r_track->dcay, r_track->dcaz, -9999., -9999., -9999., -9999., -9999., -9999., r_track->quality};
588  ntp_reco->Fill(ntp_data);
589  }
590 
591  r_track = reco_sorted.next();
592  }
593 
594  event_counter += 1;
596 }
597 
599 {
600  TFile outfile(file_name.c_str(), "recreate");
601  outfile.cd();
602  ntp_true->Write();
603  ntp_reco->Write();
604  outfile.Close();
605 
607 }