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);
 
    private:
 
        // Handles
 
        ToolHandle<AthOnnx::IOnnxRuntimeInferenceTool> m_GNNInferenceTool {this, "GNNInferenceTool", "AthOnnx::OnnxRuntimeInferenceTool"};
 
        // 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(const std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
 
        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"};
 
};      
 
 
#endif // FPGATRACKSIMGNNEDGECLASSIFIERTOOL_H