FPGATrackSimGNNGraphConstructionTool.h

Go to the documentation of this file.

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

 
#include "AthenaBaseComps/AthAlgTool.h"
 
#include "FPGATrackSimMaps/IFPGATrackSimMappingSvc.h"
#include "FPGATrackSimObjects/FPGATrackSimGNNEdge.h"
#include "FPGATrackSimObjects/FPGATrackSimGNNHit.h"
 
#include "AthOnnxInterfaces/IOnnxRuntimeInferenceTool.h"
#include "AthOnnxUtils/OnnxUtils.h"
#include <onnxruntime_cxx_api.h>
 
#include <vector>
#include <string>
#include <memory>
#include <array>
#include <unordered_map>
 
class FPGATrackSimGNNGraphConstructionTool : public AthAlgTool
{
    public:

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

        // Functions
 
        virtual StatusCode getEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits,
                                    std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
 
    private:
        
        // Handles
 
        ToolHandle<AthOnnx::IOnnxRuntimeInferenceTool> m_MLInferenceTool {this, "MLInferenceTool", "AthOnnx::OnnxRuntimeInferenceTool"};
        ServiceHandle<IFPGATrackSimMappingSvc> m_FPGATrackSimMapping{this, "FPGATrackSimMappingSvc", ""};

        // Properties
 
        Gaudi::Property<std::string> m_graphTool { this, "graphTool", "", "Tool for graph construction" };
        Gaudi::Property<std::string> m_moduleMapType { this, "moduleMapType", "", "Type for Module Map for graph construction" };
        Gaudi::Property<std::string> m_moduleMapFunc { this, "moduleMapFunc", "", "Function for Module Map for graph construction" };
        Gaudi::Property<float> m_moduleMapTol { this, "moduleMapTol", 0.0, "Tolerance value for Module Map cut calculations" };
        Gaudi::Property<float> m_moduleMapRMSThresholdFactor { this, "moduleMapRMSThresholdFactor", 0.0, "RMS Threshold value for Module Map cut calculations" };
        Gaudi::Property<float> m_metricLearningR { this, "metricLearningR", 0.0, "Clustering radius for Metric Learning"};
        Gaudi::Property<int> m_metricLearningMaxN { this, "metricLearningMaxN", 1, "Max number of neighbours for Metric Learning"};

        // Convenience
 
        // Module Map Information
        std::string m_moduleMapPath;
 
        struct ModuleMapConfig {
            unsigned mid1 = 0;
            unsigned mid2 = 0;
            unsigned mid3 = 0;        // Only exists for triplet map
            unsigned occurence = 0;   // Only exists for triplet map
 
            struct FeatureCuts {
                float min = 0.0f;
                float max = 0.0f;
                float sum = 0.0f;
                float sumSq = 0.0f;
                float mean = 0.0f;
                float rms = 0.0f;
            };
 
            struct DoubletCuts {
                FeatureCuts z0;
                FeatureCuts dphi;
                FeatureCuts deta;
                FeatureCuts phiSlope;
            };
 
            struct TripletCuts {
                FeatureCuts diff_dydx;
                FeatureCuts diff_dzdr;
            };
 
            std::array<DoubletCuts, 2> doubletCuts = {};
            TripletCuts tripletCuts = {};
        };
 
        struct TripletKey {
            unsigned mid1;
            unsigned mid2;
            unsigned mid3;
 
            bool operator==(const TripletKey& other) const {
                return mid1 == other.mid1 && mid2 == other.mid2 && mid3 == other.mid3;
            }
        };
 
        struct TripletKeyHash {
            std::size_t operator()(const TripletKey& k) const {
                return std::hash<unsigned>()(k.mid1) ^
                    (std::hash<unsigned>()(k.mid2) << 1) ^
                    (std::hash<unsigned>()(k.mid3) << 2);
            }
        };
        
        std::vector<ModuleMapConfig> m_cfgs;
        std::unordered_map<TripletKey, const ModuleMapConfig*, TripletKeyHash> m_tripletMap;

        // Helpers
 
        void loadDoubletModuleMap();
        void loadTripletModuleMap();
        void doModuleMap(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits,
                               std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
        void getTripletEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, 
                               std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
        void getDoubletEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits,
                             std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges,
                             int cutIndex);
        bool applyDoubletCuts(const std::shared_ptr<FPGATrackSimGNNHit> & hit1, 
                              const std::shared_ptr<FPGATrackSimGNNHit> & hit2, 
                              const ModuleMapConfig::DoubletCuts& cuts);
        bool applyTripletCuts(const std::shared_ptr<FPGATrackSimGNNHit> & hit1, 
                              const std::shared_ptr<FPGATrackSimGNNHit> & hit2, 
                              const std::shared_ptr<FPGATrackSimGNNHit> & hit3, 
                              const ModuleMapConfig::TripletCuts& cuts);
        bool doMask(float val, const ModuleMapConfig::FeatureCuts& cuts);
        bool doMinMaxMask(float val, const ModuleMapConfig::FeatureCuts& cuts);
        bool doMeanRMSMask(float val, const ModuleMapConfig::FeatureCuts& cuts);
        float featureSign(float feature);
        void doMetricLearning(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges);
        std::vector<float> getNodeFeatures(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits);
        std::vector<float> embed(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits);
        void doClustering(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges, std::vector<float> & gEmbedded);
        // Metric Learning Properties
        StringArrayProperty m_MLFeatureNamesVec{
            this, "MLFeatureNames",
            {"r", "phi", "z"},
            "Feature names for the Metric Learning model"};
        FloatArrayProperty m_MLFeatureScalesVec{
            this, "MLFeatureScales",
            {1000.0, 3.14159265359, 1000.0},
            "Feature scales for the Metric Learning model"};
};
 
#endif // FPGATRACKSIMGNNGRAPHCONSTRUCTIONTOOL_H