Analysis Software
Documentation for sPHENIX simulation 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 
14 #include <phool/getClass.h>
15 
16 #include <TFile.h>
17 #include <TNtuple.h>
18 #include <TSystem.h>
19 
20 #include <cstdlib>
21 #include <iostream>
22 #include <map>
23 #include <memory>
24 #include <utility>
25 #include <vector>
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  }
43  ~TrivialTrack() = default;
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 
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 (auto& container : containers)
75  {
76  for (auto& j : container)
77  {
78  for (unsigned int k = 0; k < 2; ++k)
79  {
80  j[k] = nullptr;
81  }
82  }
83  }
84  }
86  {
87  for (auto& container : containers)
88  {
89  for (auto& j : container)
90  {
91  for (unsigned int k = 0; k < 2; ++k)
92  {
93  if (j[k] != nullptr)
94  {
95  delete 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(std::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 (auto& container : containers)
238  {
239  for (auto& j : container)
240  {
241  for (unsigned int k = 0; k < 2; ++k)
242  {
243  if (j[k] == nullptr)
244  {
245  continue;
246  }
247 
248  if ((j[k]->px_hi < PX_LO) || (j[k]->px_lo > PX_HI) || (j[k]->py_hi < PY_LO) || (j[k]->py_lo > PY_HI) || (j[k]->pz_hi < PZ_LO) || (j[k]->pz_lo > PZ_HI))
249  {
250  continue;
251  }
252 
253  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 
268  ~RecursiveMomentumContainerEnd() override = default;
269 
270  bool insert(TrivialTrack& track) override
271  {
272  tracks.push_back(track);
273  return true;
274  }
275 
276  TrivialTrack* begin() override
277  {
278  x_pos = 0;
279  return (&(tracks.at(0)));
280  }
281 
282  TrivialTrack* next() override
283  {
284  if (x_pos >= (tracks.size() - 1))
285  {
286  return nullptr;
287  }
288  else
289  {
290  x_pos += 1;
291  return (&(tracks[x_pos]));
292  }
293  }
294 
295  void append_list(std::vector<TrivialTrack*>& track_list, float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI) override
296  {
297  for (auto& track : tracks)
298  {
299  if ((track.px < PX_LO) || (track.px > PX_HI) || (track.py < PY_LO) || (track.py > PY_HI) || (track.pz < PZ_LO) || (track.pz > PZ_HI))
300  {
301  continue;
302  }
303  track_list.push_back(&track);
304  }
305  }
306 
307  protected:
308  std::vector<TrivialTrack> tracks;
309 };
310 
312 {
313  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))
314  {
315  return false;
316  }
317 
318  int x_ind = 0;
319  if (track.px > (px_lo + 0.5 * (px_hi - px_lo)))
320  {
321  x_ind = 1;
322  }
323  int y_ind = 0;
324  if (track.py > (py_lo + 0.5 * (py_hi - py_lo)))
325  {
326  y_ind = 1;
327  }
328  int z_ind = 0;
329  if (track.pz > (pz_lo + 0.5 * (pz_hi - pz_lo)))
330  {
331  z_ind = 1;
332  }
333 
334  if (containers[x_ind][y_ind][z_ind] == nullptr)
335  {
336  float px_lo_new = px_lo + (float(x_ind)) * 0.5 * (px_hi - px_lo);
337  float px_hi_new = px_lo_new + 0.5 * (px_hi - px_lo);
338 
339  float py_lo_new = py_lo + (float(y_ind)) * 0.5 * (py_hi - py_lo);
340  float py_hi_new = py_lo_new + 0.5 * (py_hi - py_lo);
341 
342  float pz_lo_new = pz_lo + (float(z_ind)) * 0.5 * (pz_hi - pz_lo);
343  float pz_hi_new = pz_lo_new + 0.5 * (pz_hi - pz_lo);
344 
345  if (level < maxlevel)
346  {
347  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);
348  }
349  else
350  {
351  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);
352  }
353  }
354  return containers[x_ind][y_ind][z_ind]->insert(track);
355 }
356 
357 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)
358  : ntp_true(nullptr)
359  , ntp_reco(nullptr)
360  , pt_search_scale(pt_s)
361  , pz_search_scale(pz_s)
362  , event_counter(0)
363  , file_name(fname)
364  , n_inner_layers(n_i)
365  , n_required_layers(n_r)
366  , inner_z_length(i_z)
367  , outer_z_length(o_z)
368 {
369 }
371 {
372  if (ntp_true != nullptr)
373  {
374  delete ntp_true;
375  }
376  if (ntp_reco != nullptr)
377  {
378  delete ntp_reco;
379  }
380 }
381 
383 {
384  if (ntp_true != nullptr)
385  {
386  delete ntp_true;
387  }
388  if (ntp_reco != nullptr)
389  {
390  delete ntp_reco;
391  }
392 
393  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");
394  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");
395  event_counter = 0;
396 
398 }
399 
401 {
402  PHG4TruthInfoContainer* truthinfo = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
403 
404  PHG4HitContainer* g4hits = findNode::getClass<PHG4HitContainer>(topNode, "G4HIT_SVTX");
405  if (g4hits == nullptr)
406  {
407  std::cout << "can't find PHG4HitContainer" << std::endl;
408  exit(1);
409  }
410  PHG4HitContainer::ConstRange g4range = g4hits->getHits();
411 
412  // set<int> trkids;
413  std::map<int, std::pair<unsigned int, unsigned int> > trkids;
414 
415  for (PHG4HitContainer::ConstIterator iter = g4range.first; iter != g4range.second; ++iter)
416  {
417  PHG4Hit* hit = iter->second;
418 
419  int layer = hit->get_layer();
420  float length = outer_z_length;
421  if (((unsigned int) layer) < n_inner_layers)
422  {
423  length = inner_z_length;
424  }
425  if (std::fabs(hit->get_z(0)) > length)
426  {
427  continue;
428  }
429 
430  int trk_id = hit->get_trkid();
431  if (trkids.find(trk_id) == trkids.end())
432  {
433  trkids[trk_id].first = 0;
434  trkids[trk_id].second = 0;
435  }
436  if (hit->get_layer() < 32)
437  {
438  trkids[trk_id].first = (trkids[trk_id].first | (1 << (hit->get_layer())));
439  }
440  else if (hit->get_layer() < 64)
441  {
442  trkids[trk_id].second = (trkids[trk_id].second | (1 << (hit->get_layer() - 32)));
443  }
444  else
445  {
446  std::cout << PHWHERE << "hit layer out of bounds (0-63) " << hit->get_layer() << std::endl;
447  gSystem->Exit(1);
448  exit(1);
449  }
450 
451  // std::cout<<"trk_id = "<<trk_id<<std::endl;
452  // std::cout<<"layer = "<<hit->get_layer()<<std::endl;
453  // std::cout<<"nlayer = "<<__builtin_popcount(trkids[trk_id].first)+__builtin_popcount(trkids[trk_id].second)<<std::endl<<std::endl;
454  // trkids.insert(trk_id);
455  }
456 
457  SvtxTrackMap* trackmap = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
458 
459  PHG4VtxPoint* gvertex = truthinfo->GetPrimaryVtx(truthinfo->GetPrimaryVertexIndex());
460  float gvx = gvertex->get_x();
461  float gvy = gvertex->get_y();
462  float gvz = gvertex->get_z();
463 
464  RecursiveMomentumContainer true_sorted(-20., 20., -20., 20., -20., 20., 10);
465 
466  // PHG4TruthInfoContainer::Map primarymap = truthinfo->GetPrimaryMap();
467  PHG4TruthInfoContainer::Map primarymap = truthinfo->GetMap();
468  for (auto& iter : primarymap)
469  {
470  PHG4Particle* particle = iter.second;
471 
472  float vx = truthinfo->GetVtx(particle->get_vtx_id())->get_x();
473  float vy = truthinfo->GetVtx(particle->get_vtx_id())->get_y();
474  float vz = truthinfo->GetVtx(particle->get_vtx_id())->get_z();
475 
476  TrivialTrack track(particle->get_px(), particle->get_py(), particle->get_pz(), vx - gvx, vy - gvy, vz - gvz);
477 
478  if (((track.px * track.px) + (track.py * track.py)) < (0.1 * 0.1))
479  {
480  continue;
481  }
482 
483  if (trkids.find(particle->get_track_id()) == trkids.end())
484  {
485  continue;
486  }
487 
488  // std::cout<<"trk, nhits = "<<particle->get_track_id()<<" "<<__builtin_popcount(trkids[particle->get_track_id()].first)+__builtin_popcount(trkids[particle->get_track_id()].second)<<std::endl;
489 
490  if (__builtin_popcount(trkids[particle->get_track_id()].first) + __builtin_popcount(trkids[particle->get_track_id()].second) < (int) n_required_layers)
491  {
492  continue;
493  }
494 
495  true_sorted.insert(track);
496  }
497 
498  RecursiveMomentumContainer reco_sorted(-20., 20., -20., 20., -20., 20., 10);
499  for (auto& iter : *trackmap)
500  {
501  SvtxTrack* track = iter.second;
502 
503  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());
504  reco_sorted.insert(ttrack);
505  }
506 
507  TrivialTrack* t_track = true_sorted.begin();
508  std::vector<TrivialTrack*> pointer_list;
509  while (t_track != nullptr)
510  {
511  pointer_list.clear();
512 
513  float pt = std::sqrt((t_track->px * t_track->px) + (t_track->py * t_track->py));
514  float pt_diff = pt * pt_search_scale;
515  float px_lo = t_track->px - pt_diff;
516  float px_hi = t_track->px + pt_diff;
517  float py_lo = t_track->py - pt_diff;
518  float py_hi = t_track->py + pt_diff;
519  float pz_diff = std::fabs(t_track->pz) * pz_search_scale;
520  float pz_lo = t_track->pz - pz_diff;
521  float pz_hi = t_track->pz + pz_diff;
522 
523  reco_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
524 
525  if (pointer_list.size() > 0)
526  {
527  float mom_true = std::sqrt(pt * pt + (t_track->pz) * (t_track->pz));
528  float best_ind = 0;
529  float mom_reco = std::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));
530  float best_mom = mom_reco;
531  for (unsigned int i = 1; i < pointer_list.size(); ++i)
532  {
533  mom_reco = std::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));
534  if (std::fabs(mom_true - mom_reco) < std::fabs(mom_true - best_mom))
535  {
536  best_mom = mom_reco;
537  best_ind = i;
538  }
539  }
540 
541  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};
542  ntp_true->Fill(ntp_data);
543  }
544  else
545  {
546  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.};
547  ntp_true->Fill(ntp_data);
548  }
549 
550  t_track = true_sorted.next();
551  }
552 
553  TrivialTrack* r_track = reco_sorted.begin();
554  while (r_track != nullptr)
555  {
556  pointer_list.clear();
557 
558  float pt = std::sqrt((r_track->px * r_track->px) + (r_track->py * r_track->py));
559  float pt_diff = pt * pt_search_scale;
560  float px_lo = r_track->px - pt_diff;
561  float px_hi = r_track->px + pt_diff;
562  float py_lo = r_track->py - pt_diff;
563  float py_hi = r_track->py + pt_diff;
564  float pz_diff = std::fabs(r_track->pz) * pz_search_scale;
565  float pz_lo = r_track->pz - pz_diff;
566  float pz_hi = r_track->pz + pz_diff;
567 
568  true_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
569 
570  if (pointer_list.size() > 0)
571  {
572  float mom_reco = std::sqrt(pt * pt + (r_track->pz) * (r_track->pz));
573  float best_ind = 0;
574  float mom_true = std::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));
575  float best_mom = mom_true;
576  for (unsigned int i = 1; i < pointer_list.size(); ++i)
577  {
578  mom_true = std::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));
579  if (std::fabs(mom_reco - mom_true) < std::fabs(mom_reco - best_mom))
580  {
581  best_mom = mom_true;
582  best_ind = i;
583  }
584  }
585 
586  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};
587  ntp_reco->Fill(ntp_data);
588  }
589  else
590  {
591  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};
592  ntp_reco->Fill(ntp_data);
593  }
594 
595  r_track = reco_sorted.next();
596  }
597 
598  event_counter += 1;
600 }
601 
603 {
604  TFile outfile(file_name.c_str(), "recreate");
605  outfile.cd();
606  ntp_true->Write();
607  ntp_reco->Write();
608  outfile.Close();
609 
611 }