19 #include "CLHEP/Units/SystemOfUnits.h"
34 return StatusCode::FAILURE;
39 return StatusCode::FAILURE;
44 return StatusCode::FAILURE;
54 ATH_MSG_ERROR(
"Path to NN-based track extension ONNX file is empty! If you want to run this pipeline, you need to provide an input file.");
55 return StatusCode::FAILURE;
60 return StatusCode::SUCCESS;
65 const std::vector<std::shared_ptr<const FPGATrackSimTrack>> & tracks,
66 std::vector<std::shared_ptr<const FPGATrackSimRoad>> & roads) {
76 for (
int ireg = 0; ireg < rmap_2nd->
getNRegions(); ireg++) {
83 for (std::shared_ptr<const FPGATrackSimTrack>
track : tracks) {
85 if (
track->passedOR() == 0) {
88 const std::vector<FPGATrackSimHit> hitsOnTrack =
track->getFPGATrackSimHits();
92 for (
const auto &thit : hitsOnTrack) {
93 road.
addHit(std::make_shared<const FPGATrackSimHit>(thit));
98 for (
const std::shared_ptr<const FPGATrackSimHit>& hit:
hits) {
99 ATH_MSG_DEBUG(
"Hit " <<
" X: " << hit->getX() <<
" Y: " << hit->getY() <<
" Z: " << hit->getZ() <<
" R: " << hit->getR() <<
"phi = " << hit->getGPhi() <<
" hitType: " << hit->getHitType() <<
" getDetType: " << hit->getDetType());
102 std::vector<miniRoad> roadsToExtrapolate;
103 roadsToExtrapolate.push_back(road);
105 std::vector<miniRoad> completedRoads;
109 const int MAX_ROADS = 10000;
110 while(roadsToExtrapolate.size() > 0 &&
count < MAX_ROADS ) {
111 miniRoad currentRoad = *roadsToExtrapolate.begin();
114 roadsToExtrapolate.erase(roadsToExtrapolate.begin());
117 ATH_MSG_DEBUG(
"\033[1;31m-------------------------- extraploating road "<<
count <<
"------------------ \033[0m");
123 completedRoads.push_back(currentRoad);
127 std::vector<float> inputTensorValues;
132 std::vector<float> predhit;
142 if ((
m_useCartesian && (abs(predhit[0]) < 25 && abs(predhit[1]) < 25)) ||
145 completedRoads.push_back(currentRoad);
152 double radius = std::sqrt(predhit[0] * predhit[0] + predhit[1] * predhit[1]);
153 if ((
m_useCartesian && (abs(predhit[0]) > 1024 || abs(predhit[1]) > 1024 ||
radius > 1024 || abs(predhit[2]) > 3000)) ||
154 (!
m_useCartesian && (abs(predhit[0]) > 1024 || abs(predhit[2]) > 3000))) {
155 completedRoads.push_back(currentRoad);
160 ATH_MSG_DEBUG(
"Predicted hit at: " << predhit[0] <<
" " << predhit[1] <<
" " << predhit[2]);
164 bool foundhitForRoad =
false;
167 completedRoads.push_back(currentRoad);
171 unsigned lastLayerInRoad = 0;
172 std::shared_ptr<const FPGATrackSimHit> lastHit;
173 if(!
getLastLayer(currentRoad, lastLayerInRoad, lastHit) or !lastHit) {
177 unsigned layer = lastLayerInRoad+1;
178 bool lastHitWasReal = lastHit->
isReal();
179 float lastHitR = lastHit->
getR();
181 completedRoads.push_back(currentRoad);
184 unsigned int hitsInWindow = 0;
187 std::vector<std::vector<std::shared_ptr<const FPGATrackSimHit>>> listofHitsFound;
189 for (
const std::shared_ptr<const FPGATrackSimHit>& hit:
hits) {
192 if ((hit->getHitType() ==
HitType::spacepoint) && ((hit->getPhysLayer(
true)) %2 == 1))
continue;
194 ATH_MSG_DEBUG(
"In the hit loop hit at x: " << hit->getX() <<
" y " << hit->getY() <<
" z " << hit->getZ() <<
" phi " << hit->getGPhi());
197 if (
getFineID(*hit) == fineID && hit->isReal()) {
199 double hitz = hit->getZ();
200 double hitr = hit->getR();
201 double hitphi = hit->getGPhi();
202 double predr = (
m_useCartesian ? sqrt(predhit[0] * predhit[0] + predhit[1] * predhit[1]) : predhit[0]);
203 double predphi = predhit[1];
204 double predz = predhit[2];
208 int fineID_index = 0;
214 ATH_MSG_DEBUG(
"No windows for predicted fineID " << fineID <<
", using maximum in provided list instead!");
222 if (fineID_index == -1) {
230 if (fineID_index == -1) {
238 if (fineID_index == -1) {
255 double dr = abs(hitr - predr);
256 double dz = abs(hitz - predz);
257 double dphi = abs(hitphi - predphi);
258 while (dphi >
pi) dphi -=
pi;
262 std::vector<std::shared_ptr<const FPGATrackSimHit>> theseHits {hit};
263 hitsInWindow = hitsInWindow + 1;
273 std::shared_ptr<FPGATrackSimHit> guessedSecondHitPtr = std::make_shared<FPGATrackSimHit>();
274 guessedSecondHitPtr->
setX(0);
275 guessedSecondHitPtr->
setY(0);
276 guessedSecondHitPtr->
setZ(0);
282 theseHits.push_back(guessedSecondHitPtr);
286 listofHitsFound.push_back(theseHits);
292 std::sort(listofHitsFound.begin(), listofHitsFound.end(), [&predhit](
auto&
a,
auto&
b){
293 double predr = sqrt(predhit[0] * predhit[0] + predhit[1] * predhit[1]);
294 double predz = predhit[2];
297 double hitz = a[0]->getZ();
298 double hitr = a[0]->getR();
299 float distance_a = (hitr - predr)*(hitr - predr) + (hitz - predz)*(hitz - predz);
304 float distance_b = (hitr - predr)*(hitr - predr) + (hitz - predz)*(hitz - predz);
306 return distance_a < distance_b;
310 std::sort(listofHitsFound.begin(), listofHitsFound.end(), [&predhit](
auto&
a,
auto&
b){
312 double predr = predhit[0];
313 double predphi = predhit[1];
314 double predz = predhit[2];
316 double hitr = a[0]->getR();
317 double hitphi = a[0]->getGPhi();
318 double hitz = a[0]->getZ();
319 double dz = abs(hitz - predz);
320 double dr = abs(hitr - predr);
321 double dphi = abs(hitphi - predphi);
322 while (dphi > pi) dphi -= pi;
325 float distance_a = dz*dz/(getZScale()*getZScale()) + dphi*dphi/(getPhiScale()*getPhiScale()) + dr*dr/(getRScale()*getRScale());
329 hitphi = b[0]->getGPhi();
331 dz = abs(hitz - predz);
332 dr = abs(hitr - predr);
333 dphi = abs(hitphi - predphi);
334 while (dphi > pi) dphi -= pi;
337 float distance_b = dz*dz/(getZScale()*getZScale()) + dphi*dphi/(getPhiScale()*getPhiScale()) + dr*dr/(getRScale()*getRScale());
339 return distance_a < distance_b;
343 std::vector<std::vector<std::shared_ptr<const FPGATrackSimHit>>> cleanHitsToGrow;
348 cleanHitsToGrow.reserve(nHitsToChoose);
349 std::copy(listofHitsFound.begin(), listofHitsFound.begin() + nHitsToChoose, std::back_inserter(cleanHitsToGrow));
352 cleanHitsToGrow = std::move(listofHitsFound);
355 for (
auto& hitsFound: cleanHitsToGrow) {
359 if(!
addHitToRoad(newRoad, currentRoad, std::move(hitsFound))) {
363 roadsToExtrapolate.push_back(newRoad);
364 foundhitForRoad =
true;
371 if (!foundhitForRoad) {
378 std::vector<std::shared_ptr<const FPGATrackSimHit>> theseHits;
380 if (!
getFakeHit(currentRoad, predhit, fineID, theseHits)) {
390 if (!
addHitToRoad(newroad, currentRoad, std::move(theseHits))) {
394 roadsToExtrapolate.push_back(newroad);
399 for (
const auto &miniroad : completedRoads) {
412 std::vector<std::vector<std::shared_ptr<const FPGATrackSimHit>>> roadhits = miniroad.getVecHits();
414 if (roadhits.size() > nexpected) {
415 roadhits.resize(nexpected);
417 else if (roadhits.size() < nexpected) {
419 std::shared_ptr<FPGATrackSimHit> emptyHitPtr = std::make_shared<FPGATrackSimHit>();
420 emptyHitPtr->
setX(0);
421 emptyHitPtr->
setY(0);
422 emptyHitPtr->
setZ(0);
425 std::vector<std::shared_ptr<const FPGATrackSimHit>> hitVec;
426 hitVec.push_back(emptyHitPtr);
427 roadhits.push_back(hitVec);
429 wclayers |= (1 <<
layer);
433 road.
setHits(std::move(roadhits));
443 roads.emplace_back(std::make_shared<const FPGATrackSimRoad>(
r));
445 ATH_MSG_DEBUG(
"Found " << roads.size() <<
" new roads in second stage.");
447 return StatusCode::SUCCESS;
452 std::vector<std::shared_ptr<const FPGATrackSimHit>> hitsR;
454 std::vector<std::shared_ptr<const FPGATrackSimHit>>
hits = thisroad.
getHits();
457 for (
unsigned ihit = 0; ihit <
hits.size(); ihit++) {
459 hitsR.push_back(
hits[ihit]);
462 std::sort(hitsR.begin(), hitsR.end(), [](
auto&
a,
auto&
b){
463 double dist_a = std::hypot(a->getX(), a->getY(), a->getZ());
464 double dist_b = std::hypot(b->getX(), b->getY(), b->getZ());
465 return dist_a < dist_b;
469 for (
auto thit : hitsR)
475 std::vector<std::shared_ptr<const FPGATrackSimHit>> cleanHits;
476 bool skipHit =
false;
477 for (
auto thit : hitsR)
486 cleanHits.push_back(thit);
491 cleanHits.push_back(thit);
497 cleanHits.push_back(thit);
505 else if (thit->isStrip() && thit->isReal())
508 cleanHits.push_back(thit);
512 ATH_MSG_WARNING(
"No clue how to deal with this hit in the NN predicition ");
524 std::vector<std::shared_ptr<const FPGATrackSimHit>> hitsToEncode;
528 for (
auto thit : cleanHits)
530 if(
m_debugEvent)
ATH_MSG_DEBUG(thit->getX()<<
" "<<thit->getY()<<
" "<<thit->getZ()<<
" "<<thit->isStrip() <<
" and gphi = " << thit->getGPhi());
537 for (
auto thit : hitsToEncode)
540 inputTensorValues.push_back(thit->getX()/
getXScale());
541 inputTensorValues.push_back(thit->getY()/
getYScale());
542 inputTensorValues.push_back(thit->getZ()/
getZScale());
545 inputTensorValues.push_back(thit->getR()/
getRScale());
546 inputTensorValues.push_back(thit->getGPhi()/
getPhiScale());
547 inputTensorValues.push_back(thit->getZ()/
getZScale());
553 return StatusCode::SUCCESS;
561 fineID =
std::distance(NNVoloutput.begin(),std::max_element(NNVoloutput.begin(), NNVoloutput.end()));
563 inputTensorValues.insert(inputTensorValues.end(), NNVoloutput.begin(), NNVoloutput.end());
580 return StatusCode::SUCCESS;
587 return StatusCode::SUCCESS;
592 unsigned guessedLayer(0);
594 std::shared_ptr<FPGATrackSimHit> guessedHitPtr = std::make_shared<FPGATrackSimHit>();
596 guessedHitPtr->
setX(predhit[0]);
597 guessedHitPtr->
setY(predhit[1]);
598 guessedHitPtr->
setZ(predhit[2]);
601 double r = predhit[0];
602 double phi = predhit[1];
605 guessedHitPtr->
setZ(predhit[2]);
613 guessedLayer = currentRoad.
getNHits();
616 guessedHitPtr->
setLayer(guessedLayer);
629 hits.push_back(guessedHitPtr);
634 std::shared_ptr<FPGATrackSimHit> guessedSecondHitPtr = std::make_shared<FPGATrackSimHit>();
635 guessedSecondHitPtr->
setX(0);
636 guessedSecondHitPtr->
setY(0);
637 guessedSecondHitPtr->
setZ(0);
638 guessedSecondHitPtr->
setLayer( guessedLayer + 1 );
646 hits.push_back(guessedSecondHitPtr);
652 return StatusCode::SUCCESS;
659 for (
const std::shared_ptr<const FPGATrackSimHit>& hit: hitList)
664 return StatusCode::SUCCESS;
671 return StatusCode::FAILURE;
681 std::vector<std::shared_ptr<const FPGATrackSimHit>> hitsR;
682 for (
auto &hit : currentRoad.
getHits()) {
683 hitsR.push_back(hit);
689 std::sort(hitsR.begin(), hitsR.end(), [](
auto&
a,
auto&
b){
690 if(a->getR() == b->getR()) return a->getLayer() < b->getLayer();
691 return a->getR() < b->getR();
696 std::sort(hitsR.begin(), hitsR.end(), [](
auto&
a,
auto&
b){
697 return a->getR() > b->getR();
701 for (
unsigned long i = 0;
i < hitsR.size();
i++)
703 ATH_MSG_DEBUG(
"Hit i "<<
i<<
" X: "<<hitsR[
i]->getX()<<
" Y: "<<hitsR[
i]->getY()<<
" Z: "<<hitsR[
i]->getZ()<<
" R: "<<hitsR[
i]->getR()<<
" hitType: "<<hitsR[
i]->getHitType()<<
" getDetType: "<<hitsR[
i]->getDetType() <<
"phi = " << hitsR[
i]->getGPhi());
711 lastHitLayer = currentRoad.
getNHits()-1;
712 lastHit = currentRoad.
getHit(lastHitLayer);
713 return StatusCode::SUCCESS;