ATLAS Offline Software
Loading...
Searching...
No Matches
InDet::GNNTrackFinderTritonTool Class Reference

InDet::GNNTrackFinderTritonTool is a tool that produces track candidates with graph neural networks-based pipeline using space points as inputs. More...

#include <GNNTrackFinderTritonTool.h>

Inheritance diagram for InDet::GNNTrackFinderTritonTool:
Collaboration diagram for InDet::GNNTrackFinderTritonTool:

Public Member Functions

virtual StatusCode initialize () override
virtual StatusCode getTracks (const std::vector< const Trk::SpacePoint * > &spacepoints, std::vector< std::vector< uint32_t > > &tracks) const override
 Get track candidates from a list of space points.
virtual MsgStream & dump (MsgStream &out) const override
virtual std::ostream & dump (std::ostream &out) const override

Private Member Functions

MsgStream & dumpevent (MsgStream &out) const

Private Attributes

ToolHandle< AthInfer::IAthInferenceToolm_gnnTrackingTritonTool
ToolHandle< ISpacepointFeatureToolm_spacepointFeatureTool
StringProperty m_featureNames
std::vector< std::string > m_featureNamesVec

Detailed Description

InDet::GNNTrackFinderTritonTool is a tool that produces track candidates with graph neural networks-based pipeline using space points as inputs.

Author
xiang.nosp@m.yang.nosp@m..ju@c.nosp@m.ern..nosp@m.ch

Definition at line 28 of file GNNTrackFinderTritonTool.h.

Member Function Documentation

◆ dump() [1/2]

MsgStream & InDet::GNNTrackFinderTritonTool::dump ( MsgStream & out) const
overridevirtual

Definition at line 75 of file GNNTrackFinderTritonTool.cxx.

76{
77 out<<std::endl;
78 return dumpevent(out);
79}
MsgStream & dumpevent(MsgStream &out) const

◆ dump() [2/2]

std::ostream & InDet::GNNTrackFinderTritonTool::dump ( std::ostream & out) const
overridevirtual

Definition at line 81 of file GNNTrackFinderTritonTool.cxx.

82{
83 return out;
84}

◆ dumpevent()

MsgStream & InDet::GNNTrackFinderTritonTool::dumpevent ( MsgStream & out) const
private

Definition at line 86 of file GNNTrackFinderTritonTool.cxx.

87{
88 out<<"|---------------------------------------------------------------------|"
89 <<std::endl;
90 out<<"| Number output tracks | "<<std::setw(12)
91 <<" |"<<std::endl;
92 out<<"|---------------------------------------------------------------------|"
93 <<std::endl;
94 return out;
95}

◆ getTracks()

StatusCode InDet::GNNTrackFinderTritonTool::getTracks ( const std::vector< const Trk::SpacePoint * > & spacepoints,
std::vector< std::vector< uint32_t > > & tracks ) const
overridevirtual

Get track candidates from a list of space points.

Parameters
spacepointsa list of spacepoints as inputs to the GNN-based track finder.
tracksa list of track candidates.
Returns

Definition at line 24 of file GNNTrackFinderTritonTool.cxx.

26 {
27 int64_t numSpacepoints = (int64_t)spacepoints.size();
28 std::vector<float> inputValues;
29 std::vector<uint32_t> spacepointIDs;
30
31 int64_t spacepointFeatures = m_featureNamesVec.size();
32 int sp_idx = 0;
33 for (const auto& sp : spacepoints) {
34 // depending on the trained embedding and GNN models, the input features
35 // may need to be updated.
36 auto featureMap = m_spacepointFeatureTool->getFeatures(sp);
37 for (int i = 0; i < spacepointFeatures; i++){
38 inputValues.push_back(featureMap[m_featureNamesVec[i]]);
39 }
40
41 spacepointIDs.push_back(sp_idx++);
42 }
43
44 AthInfer::InputDataMap inputData;
45 inputData["FEATURES"] = std::make_pair(
46 std::vector<int64_t>{numSpacepoints, spacepointFeatures}, std::move(inputValues));
47
48 AthInfer::OutputDataMap outputData;
49 outputData["LABELS"] = std::make_pair(std::vector<int64_t>{numSpacepoints, 1}, std::vector<int64_t>{});
50
51 ATH_CHECK(m_gnnTrackingTritonTool->inference(inputData, outputData));
52
53 auto& trackLabels = std::get<std::vector<int64_t>>(outputData["LABELS"].second);
54 if (trackLabels.size() == 0){
55 ATH_MSG_DEBUG("No tracks found in the event.");
56 return StatusCode::SUCCESS;
57 }
58
59 tracks.clear();
60 std::vector<uint32_t> this_track;
61 for (auto label : trackLabels) {
62 if (label == -1) {
63 if (this_track.size() > 0) {
64 tracks.push_back(this_track);
65 this_track.clear();
66 }
67 } else {
68 this_track.push_back(label);
69 }
70 }
71
72 return StatusCode::SUCCESS;
73}
#define ATH_CHECK
Evaluate an expression and check for errors.
#define ATH_MSG_DEBUG(x)
static Double_t sp
ToolHandle< AthInfer::IAthInferenceTool > m_gnnTrackingTritonTool
std::vector< std::string > m_featureNamesVec
ToolHandle< ISpacepointFeatureTool > m_spacepointFeatureTool
std::string label(const std::string &format, int i)
Definition label.h:19
std::map< std::string, InferenceData > OutputDataMap
std::map< std::string, InferenceData > InputDataMap

◆ initialize()

StatusCode InDet::GNNTrackFinderTritonTool::initialize ( )
overridevirtual

Definition at line 15 of file GNNTrackFinderTritonTool.cxx.

15 {
18
19 // tokenize the feature names by comma and push to the vector
21 return StatusCode::SUCCESS;
22}
std::vector< std::string > tokenize(const std::string &the_str, std::string_view delimiters)
Splits the string into smaller substrings.

Member Data Documentation

◆ m_featureNames

StringProperty InDet::GNNTrackFinderTritonTool::m_featureNames
private
Initial value:
{this, "FeatureNames",
"r,phi,z,cluster_x_1,cluster_y_1,cluster_z_1,cluster_x_2,cluster_y_2,cluster_z_2,count_1,charge_count_1,loc_eta_1,loc_phi_1,localDir0_1,localDir1_1,localDir2_1,lengthDir0_1,lengthDir1_1,lengthDir2_1,glob_eta_1,glob_phi_1,eta_angle_1,phi_angle_1,count_2,charge_count_2,loc_eta_2,loc_phi_2,localDir0_2,localDir1_2,localDir2_2,lengthDir0_2,lengthDir1_2,lengthDir2_2,glob_eta_2,glob_phi_2,eta_angle_2,phi_angle_2,eta,cluster_r_1,cluster_phi_1,cluster_eta_1,cluster_r_2,cluster_phi_2,cluster_eta_2",
"Feature names for the GNN pipeline"}

Definition at line 59 of file GNNTrackFinderTritonTool.h.

59 {this, "FeatureNames",
60 "r,phi,z,cluster_x_1,cluster_y_1,cluster_z_1,cluster_x_2,cluster_y_2,cluster_z_2,count_1,charge_count_1,loc_eta_1,loc_phi_1,localDir0_1,localDir1_1,localDir2_1,lengthDir0_1,lengthDir1_1,lengthDir2_1,glob_eta_1,glob_phi_1,eta_angle_1,phi_angle_1,count_2,charge_count_2,loc_eta_2,loc_phi_2,localDir0_2,localDir1_2,localDir2_2,lengthDir0_2,lengthDir1_2,lengthDir2_2,glob_eta_2,glob_phi_2,eta_angle_2,phi_angle_2,eta,cluster_r_1,cluster_phi_1,cluster_eta_1,cluster_r_2,cluster_phi_2,cluster_eta_2",
61 "Feature names for the GNN pipeline"};

◆ m_featureNamesVec

std::vector<std::string> InDet::GNNTrackFinderTritonTool::m_featureNamesVec
private

Definition at line 63 of file GNNTrackFinderTritonTool.h.

◆ m_gnnTrackingTritonTool

ToolHandle<AthInfer::IAthInferenceTool> InDet::GNNTrackFinderTritonTool::m_gnnTrackingTritonTool
private
Initial value:
{
this, "TritonTool", "AthInfer::TritonTool"}

Definition at line 54 of file GNNTrackFinderTritonTool.h.

54 {
55 this, "TritonTool", "AthInfer::TritonTool"};

◆ m_spacepointFeatureTool

ToolHandle<ISpacepointFeatureTool> InDet::GNNTrackFinderTritonTool::m_spacepointFeatureTool
private
Initial value:
{
this, "SpacepointFeatureTool", "InDet::SpacepointFeatureTool"}

Definition at line 56 of file GNNTrackFinderTritonTool.h.

56 {
57 this, "SpacepointFeatureTool", "InDet::SpacepointFeatureTool"};

The documentation for this class was generated from the following files: