File indexing completed on 2025-08-05 08:17:57
0001 #include "MomentumEvaluator.h"
0002
0003 #include <trackbase_historic/SvtxTrack.h>
0004 #include <trackbase_historic/SvtxTrackMap.h>
0005
0006 #include <g4main/PHG4Hit.h>
0007 #include <g4main/PHG4HitContainer.h>
0008 #include <g4main/PHG4Particle.h>
0009 #include <g4main/PHG4TruthInfoContainer.h>
0010 #include <g4main/PHG4VtxPoint.h>
0011
0012 #include <fun4all/Fun4AllReturnCodes.h>
0013
0014 #include <phool/getClass.h>
0015
0016 #include <TFile.h>
0017 #include <TNtuple.h>
0018 #include <TSystem.h>
0019
0020 #include <cstdlib>
0021 #include <iostream>
0022 #include <map>
0023 #include <memory>
0024 #include <utility>
0025 #include <vector>
0026
0027 class TrivialTrack
0028 {
0029 public:
0030 float px, py, pz;
0031 float dcax, dcay, dcaz;
0032 float quality;
0033 TrivialTrack(float x, float y, float z, float dx, float dy, float dz, float qual = 0.)
0034 : px(x)
0035 , py(y)
0036 , pz(z)
0037 , dcax(dx)
0038 , dcay(dy)
0039 , dcaz(dz)
0040 , quality(qual)
0041 {
0042 }
0043 ~TrivialTrack() = default;
0044 };
0045
0046 class RecursiveMomentumContainer
0047 {
0048 protected:
0049 float px_lo, px_hi;
0050 float py_lo, py_hi;
0051 float pz_lo, pz_hi;
0052
0053 int level;
0054 int maxlevel;
0055
0056 unsigned int x_pos, y_pos, z_pos;
0057
0058 RecursiveMomentumContainer* containers[2][2][2]{};
0059
0060 public:
0061 RecursiveMomentumContainer(float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI, int MLEV, int LEV = 0)
0062 : px_lo(PX_LO)
0063 , px_hi(PX_HI)
0064 , py_lo(PY_LO)
0065 , py_hi(PY_HI)
0066 , pz_lo(PZ_LO)
0067 , pz_hi(PZ_HI)
0068 , level(LEV)
0069 , maxlevel(MLEV)
0070 , x_pos(0)
0071 , y_pos(0)
0072 , z_pos(0)
0073 {
0074 for (auto& container : containers)
0075 {
0076 for (auto& j : container)
0077 {
0078 for (unsigned int k = 0; k < 2; ++k)
0079 {
0080 j[k] = nullptr;
0081 }
0082 }
0083 }
0084 }
0085 virtual ~RecursiveMomentumContainer()
0086 {
0087 for (auto& container : containers)
0088 {
0089 for (auto& j : container)
0090 {
0091 for (unsigned int k = 0; k < 2; ++k)
0092 {
0093 if (j[k] != nullptr)
0094 {
0095 delete j[k];
0096 }
0097 }
0098 }
0099 }
0100 }
0101
0102 virtual bool insert(TrivialTrack& track);
0103
0104 virtual TrivialTrack* begin()
0105 {
0106 x_pos = 0;
0107 y_pos = 0;
0108 z_pos = 0;
0109 while (true)
0110 {
0111 if (containers[x_pos][y_pos][z_pos] == nullptr)
0112 {
0113 if (z_pos == 0)
0114 {
0115 z_pos = 1;
0116 continue;
0117 }
0118 else
0119 {
0120 if (y_pos == 0)
0121 {
0122 z_pos = 0;
0123 y_pos = 1;
0124 continue;
0125 }
0126 else
0127 {
0128 if (x_pos == 0)
0129 {
0130 z_pos = 0;
0131 y_pos = 0;
0132 x_pos = 1;
0133 continue;
0134 }
0135 else
0136 {
0137 return nullptr;
0138 }
0139 }
0140 }
0141 }
0142 else
0143 {
0144 return containers[x_pos][y_pos][z_pos]->begin();
0145 }
0146 }
0147 }
0148
0149 virtual TrivialTrack* next()
0150 {
0151 bool block_changed = false;
0152 while (true)
0153 {
0154 if (containers[x_pos][y_pos][z_pos] == nullptr)
0155 {
0156 block_changed = true;
0157 if (z_pos == 0)
0158 {
0159 z_pos = 1;
0160 continue;
0161 }
0162 else
0163 {
0164 if (y_pos == 0)
0165 {
0166 z_pos = 0;
0167 y_pos = 1;
0168 continue;
0169 }
0170 else
0171 {
0172 if (x_pos == 0)
0173 {
0174 z_pos = 0;
0175 y_pos = 0;
0176 x_pos = 1;
0177 continue;
0178 }
0179 else
0180 {
0181 return nullptr;
0182 }
0183 }
0184 }
0185 }
0186 TrivialTrack* val = nullptr;
0187 if (block_changed == true)
0188 {
0189 val = containers[x_pos][y_pos][z_pos]->begin();
0190 }
0191 else
0192 {
0193 val = containers[x_pos][y_pos][z_pos]->next();
0194 }
0195
0196 if (val == nullptr)
0197 {
0198 block_changed = true;
0199 if (z_pos == 0)
0200 {
0201 z_pos = 1;
0202 continue;
0203 }
0204 else
0205 {
0206 if (y_pos == 0)
0207 {
0208 z_pos = 0;
0209 y_pos = 1;
0210 continue;
0211 }
0212 else
0213 {
0214 if (x_pos == 0)
0215 {
0216 z_pos = 0;
0217 y_pos = 0;
0218 x_pos = 1;
0219 continue;
0220 }
0221 else
0222 {
0223 return nullptr;
0224 }
0225 }
0226 }
0227 }
0228 else
0229 {
0230 return val;
0231 }
0232 }
0233 }
0234
0235 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)
0236 {
0237 for (auto& container : containers)
0238 {
0239 for (auto& j : container)
0240 {
0241 for (unsigned int k = 0; k < 2; ++k)
0242 {
0243 if (j[k] == nullptr)
0244 {
0245 continue;
0246 }
0247
0248 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))
0249 {
0250 continue;
0251 }
0252
0253 j[k]->append_list(track_list, PX_LO, PX_HI, PY_LO, PY_HI, PZ_LO, PZ_HI);
0254 }
0255 }
0256 }
0257 }
0258 };
0259
0260 class RecursiveMomentumContainerEnd : public RecursiveMomentumContainer
0261 {
0262 public:
0263 RecursiveMomentumContainerEnd(float PX_LO, float PX_HI, float PY_LO, float PY_HI, float PZ_LO, float PZ_HI, int MLEV, int LEV = 0)
0264 : RecursiveMomentumContainer(PX_LO, PX_HI, PY_LO, PY_HI, PZ_LO, PZ_HI, MLEV, LEV)
0265 {
0266 }
0267
0268 ~RecursiveMomentumContainerEnd() override = default;
0269
0270 bool insert(TrivialTrack& track) override
0271 {
0272 tracks.push_back(track);
0273 return true;
0274 }
0275
0276 TrivialTrack* begin() override
0277 {
0278 x_pos = 0;
0279 return (&(tracks.at(0)));
0280 }
0281
0282 TrivialTrack* next() override
0283 {
0284 if (x_pos >= (tracks.size() - 1))
0285 {
0286 return nullptr;
0287 }
0288 else
0289 {
0290 x_pos += 1;
0291 return (&(tracks[x_pos]));
0292 }
0293 }
0294
0295 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
0296 {
0297 for (auto& track : tracks)
0298 {
0299 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))
0300 {
0301 continue;
0302 }
0303 track_list.push_back(&track);
0304 }
0305 }
0306
0307 protected:
0308 std::vector<TrivialTrack> tracks;
0309 };
0310
0311 bool RecursiveMomentumContainer::insert(TrivialTrack& track)
0312 {
0313 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))
0314 {
0315 return false;
0316 }
0317
0318 int x_ind = 0;
0319 if (track.px > (px_lo + 0.5 * (px_hi - px_lo)))
0320 {
0321 x_ind = 1;
0322 }
0323 int y_ind = 0;
0324 if (track.py > (py_lo + 0.5 * (py_hi - py_lo)))
0325 {
0326 y_ind = 1;
0327 }
0328 int z_ind = 0;
0329 if (track.pz > (pz_lo + 0.5 * (pz_hi - pz_lo)))
0330 {
0331 z_ind = 1;
0332 }
0333
0334 if (containers[x_ind][y_ind][z_ind] == nullptr)
0335 {
0336 float px_lo_new = px_lo + (float(x_ind)) * 0.5 * (px_hi - px_lo);
0337 float px_hi_new = px_lo_new + 0.5 * (px_hi - px_lo);
0338
0339 float py_lo_new = py_lo + (float(y_ind)) * 0.5 * (py_hi - py_lo);
0340 float py_hi_new = py_lo_new + 0.5 * (py_hi - py_lo);
0341
0342 float pz_lo_new = pz_lo + (float(z_ind)) * 0.5 * (pz_hi - pz_lo);
0343 float pz_hi_new = pz_lo_new + 0.5 * (pz_hi - pz_lo);
0344
0345 if (level < maxlevel)
0346 {
0347 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);
0348 }
0349 else
0350 {
0351 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);
0352 }
0353 }
0354 return containers[x_ind][y_ind][z_ind]->insert(track);
0355 }
0356
0357 MomentumEvaluator::MomentumEvaluator(const std::string& fname, float pt_s, float pz_s, unsigned int , unsigned int n_i, unsigned int n_r, float i_z, float o_z)
0358 : ntp_true(nullptr)
0359 , ntp_reco(nullptr)
0360 , pt_search_scale(pt_s)
0361 , pz_search_scale(pz_s)
0362 , event_counter(0)
0363 , file_name(fname)
0364 , n_inner_layers(n_i)
0365 , n_required_layers(n_r)
0366 , inner_z_length(i_z)
0367 , outer_z_length(o_z)
0368 {
0369 }
0370 MomentumEvaluator::~MomentumEvaluator()
0371 {
0372 if (ntp_true != nullptr)
0373 {
0374 delete ntp_true;
0375 }
0376 if (ntp_reco != nullptr)
0377 {
0378 delete ntp_reco;
0379 }
0380 }
0381
0382 int MomentumEvaluator::Init(PHCompositeNode* )
0383 {
0384 if (ntp_true != nullptr)
0385 {
0386 delete ntp_true;
0387 }
0388 if (ntp_reco != nullptr)
0389 {
0390 delete ntp_reco;
0391 }
0392
0393 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");
0394 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");
0395 event_counter = 0;
0396
0397 return Fun4AllReturnCodes::EVENT_OK;
0398 }
0399
0400 int MomentumEvaluator::process_event(PHCompositeNode* topNode)
0401 {
0402 PHG4TruthInfoContainer* truthinfo = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
0403
0404 PHG4HitContainer* g4hits = findNode::getClass<PHG4HitContainer>(topNode, "G4HIT_SVTX");
0405 if (g4hits == nullptr)
0406 {
0407 std::cout << "can't find PHG4HitContainer" << std::endl;
0408 exit(1);
0409 }
0410 PHG4HitContainer::ConstRange g4range = g4hits->getHits();
0411
0412
0413 std::map<int, std::pair<unsigned int, unsigned int> > trkids;
0414
0415 for (PHG4HitContainer::ConstIterator iter = g4range.first; iter != g4range.second; ++iter)
0416 {
0417 PHG4Hit* hit = iter->second;
0418
0419 int layer = hit->get_layer();
0420 float length = outer_z_length;
0421 if (((unsigned int) layer) < n_inner_layers)
0422 {
0423 length = inner_z_length;
0424 }
0425 if (std::fabs(hit->get_z(0)) > length)
0426 {
0427 continue;
0428 }
0429
0430 int trk_id = hit->get_trkid();
0431 if (trkids.find(trk_id) == trkids.end())
0432 {
0433 trkids[trk_id].first = 0;
0434 trkids[trk_id].second = 0;
0435 }
0436 if (hit->get_layer() < 32)
0437 {
0438 trkids[trk_id].first = (trkids[trk_id].first | (1 << (hit->get_layer())));
0439 }
0440 else if (hit->get_layer() < 64)
0441 {
0442 trkids[trk_id].second = (trkids[trk_id].second | (1 << (hit->get_layer() - 32)));
0443 }
0444 else
0445 {
0446 std::cout << PHWHERE << "hit layer out of bounds (0-63) " << hit->get_layer() << std::endl;
0447 gSystem->Exit(1);
0448 exit(1);
0449 }
0450
0451
0452
0453
0454
0455 }
0456
0457 SvtxTrackMap* trackmap = findNode::getClass<SvtxTrackMap>(topNode, "SvtxTrackMap");
0458
0459 PHG4VtxPoint* gvertex = truthinfo->GetPrimaryVtx(truthinfo->GetPrimaryVertexIndex());
0460 float gvx = gvertex->get_x();
0461 float gvy = gvertex->get_y();
0462 float gvz = gvertex->get_z();
0463
0464 RecursiveMomentumContainer true_sorted(-20., 20., -20., 20., -20., 20., 10);
0465
0466
0467 PHG4TruthInfoContainer::Map primarymap = truthinfo->GetMap();
0468 for (auto& iter : primarymap)
0469 {
0470 PHG4Particle* particle = iter.second;
0471
0472 float vx = truthinfo->GetVtx(particle->get_vtx_id())->get_x();
0473 float vy = truthinfo->GetVtx(particle->get_vtx_id())->get_y();
0474 float vz = truthinfo->GetVtx(particle->get_vtx_id())->get_z();
0475
0476 TrivialTrack track(particle->get_px(), particle->get_py(), particle->get_pz(), vx - gvx, vy - gvy, vz - gvz);
0477
0478 if (((track.px * track.px) + (track.py * track.py)) < (0.1 * 0.1))
0479 {
0480 continue;
0481 }
0482
0483 if (trkids.find(particle->get_track_id()) == trkids.end())
0484 {
0485 continue;
0486 }
0487
0488
0489
0490 if (__builtin_popcount(trkids[particle->get_track_id()].first) + __builtin_popcount(trkids[particle->get_track_id()].second) < (int) n_required_layers)
0491 {
0492 continue;
0493 }
0494
0495 true_sorted.insert(track);
0496 }
0497
0498 RecursiveMomentumContainer reco_sorted(-20., 20., -20., 20., -20., 20., 10);
0499 for (auto& iter : *trackmap)
0500 {
0501 SvtxTrack* track = iter.second;
0502
0503 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());
0504 reco_sorted.insert(ttrack);
0505 }
0506
0507 TrivialTrack* t_track = true_sorted.begin();
0508 std::vector<TrivialTrack*> pointer_list;
0509 while (t_track != nullptr)
0510 {
0511 pointer_list.clear();
0512
0513 float pt = std::sqrt((t_track->px * t_track->px) + (t_track->py * t_track->py));
0514 float pt_diff = pt * pt_search_scale;
0515 float px_lo = t_track->px - pt_diff;
0516 float px_hi = t_track->px + pt_diff;
0517 float py_lo = t_track->py - pt_diff;
0518 float py_hi = t_track->py + pt_diff;
0519 float pz_diff = std::fabs(t_track->pz) * pz_search_scale;
0520 float pz_lo = t_track->pz - pz_diff;
0521 float pz_hi = t_track->pz + pz_diff;
0522
0523 reco_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
0524
0525 if (pointer_list.size() > 0)
0526 {
0527 float mom_true = std::sqrt(pt * pt + (t_track->pz) * (t_track->pz));
0528 float best_ind = 0;
0529 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));
0530 float best_mom = mom_reco;
0531 for (unsigned int i = 1; i < pointer_list.size(); ++i)
0532 {
0533 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));
0534 if (std::fabs(mom_true - mom_reco) < std::fabs(mom_true - best_mom))
0535 {
0536 best_mom = mom_reco;
0537 best_ind = i;
0538 }
0539 }
0540
0541 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};
0542 ntp_true->Fill(ntp_data);
0543 }
0544 else
0545 {
0546 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.};
0547 ntp_true->Fill(ntp_data);
0548 }
0549
0550 t_track = true_sorted.next();
0551 }
0552
0553 TrivialTrack* r_track = reco_sorted.begin();
0554 while (r_track != nullptr)
0555 {
0556 pointer_list.clear();
0557
0558 float pt = std::sqrt((r_track->px * r_track->px) + (r_track->py * r_track->py));
0559 float pt_diff = pt * pt_search_scale;
0560 float px_lo = r_track->px - pt_diff;
0561 float px_hi = r_track->px + pt_diff;
0562 float py_lo = r_track->py - pt_diff;
0563 float py_hi = r_track->py + pt_diff;
0564 float pz_diff = std::fabs(r_track->pz) * pz_search_scale;
0565 float pz_lo = r_track->pz - pz_diff;
0566 float pz_hi = r_track->pz + pz_diff;
0567
0568 true_sorted.append_list(pointer_list, px_lo, px_hi, py_lo, py_hi, pz_lo, pz_hi);
0569
0570 if (pointer_list.size() > 0)
0571 {
0572 float mom_reco = std::sqrt(pt * pt + (r_track->pz) * (r_track->pz));
0573 float best_ind = 0;
0574 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));
0575 float best_mom = mom_true;
0576 for (unsigned int i = 1; i < pointer_list.size(); ++i)
0577 {
0578 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));
0579 if (std::fabs(mom_reco - mom_true) < std::fabs(mom_reco - best_mom))
0580 {
0581 best_mom = mom_true;
0582 best_ind = i;
0583 }
0584 }
0585
0586 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};
0587 ntp_reco->Fill(ntp_data);
0588 }
0589 else
0590 {
0591 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};
0592 ntp_reco->Fill(ntp_data);
0593 }
0594
0595 r_track = reco_sorted.next();
0596 }
0597
0598 event_counter += 1;
0599 return Fun4AllReturnCodes::EVENT_OK;
0600 }
0601
0602 int MomentumEvaluator::End(PHCompositeNode* )
0603 {
0604 TFile outfile(file_name.c_str(), "recreate");
0605 outfile.cd();
0606 ntp_true->Write();
0607 ntp_reco->Write();
0608 outfile.Close();
0609
0610 return Fun4AllReturnCodes::EVENT_OK;
0611 }
