FPGATrackSimGNNEdgeClassifierTool.h

Go to the documentation of this file.

// Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
#ifndef FPGATRACKSIMGNNEDGECLASSIFIERTOOL_H
#define FPGATRACKSIMGNNEDGECLASSIFIERTOOL_H

 
#include "AthenaBaseComps/AthAlgTool.h"
 
#include "FPGATrackSimObjects/FPGATrackSimGNNHit.h"
#include "FPGATrackSimObjects/FPGATrackSimGNNEdge.h"
 
#include "AthOnnxInterfaces/IOnnxRuntimeInferenceTool.h"
#include <onnxruntime_cxx_api.h>
 
class FPGATrackSimGNNEdgeClassifierTool : public AthAlgTool
{
    public:

        // AthAlgTool
 
        FPGATrackSimGNNEdgeClassifierTool(const std::string&, const std::string&, const IInterface*);
 
        virtual StatusCode initialize() override;

        // Functions
 
        virtual StatusCode scoreEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
        int regionNum() const { return m_regionNum; }
 
    private:

        // Handles
 
        ToolHandle<AthOnnx::IOnnxRuntimeInferenceTool> m_GNNInferenceTool {this, "GNNInferenceTool", "AthOnnx::OnnxRuntimeInferenceTool"};
        Gaudi::Property<int> m_regionNum{this, "regionNum", -1, "Region number for this GNNEdgeClassifierTool"};
        Gaudi::Property<bool> m_doGNNPixelSeeding { this, "doGNNPixelSeeding", false, "Flag to configure for GNN Pixel Seeding" };

        // Helpers
        
        std::vector<float> getNodeFeatures(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits);
        std::vector<int64_t> getEdgeList(const std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
        std::vector<float> getEdgeFeatures(std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges, const std::vector<float> & gNodeFeatures);
        void computeEdgeFeatures(std::shared_ptr<FPGATrackSimGNNEdge>& edge, const int& hit1_index, const int& hit2_index, const std::vector<float> & gNodeFeatures);
 
        StringArrayProperty m_gnnFeatureNamesVec{
            this, "GNNFeatureNames",
            {"r", "phi", "z", "eta", "cluster_r_1", "cluster_phi_1", "cluster_z_1", "cluster_eta_1", "cluster_r_2", "cluster_phi_2", "cluster_z_2", "cluster_eta_2"},
            "Feature names for the GNN model"};
        FloatArrayProperty m_gnnFeatureScalesVec{
            this, "GNNFeatureScales",
            {1000.0, 3.14159265359, 1000.0, 1.0, 1000.0, 3.14159265359, 1000.0, 1.0, 1000.0, 3.14159265359, 1000.0, 1.0},
            "Feature scales for the GNN model"};
 
        StringArrayProperty m_gnnFeatureNamesVec_pixelOnly{
            this, "GNNFeatureNames",
            {"r", "phi", "z", "eta"},
            "Feature names for the GNN model"};
        FloatArrayProperty m_gnnFeatureScalesVec_pixelOnly{
            this, "GNNFeatureScales",
            {1000.0, 3.14159265359, 1000.0, 1.0},
            "Feature scales for the GNN model"};
};      
 
 
#endif // FPGATRACKSIMGNNEDGECLASSIFIERTOOL_H