Back to home page

sPhenix code displayed by LXR
 
 

    

File indexing completed on 2025-12-17 09:21:49 

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

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!