SegmentEdgeClassifierTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
#ifndef MUONINFERENCE_SEGMENTEDGECLASSIFIERTOOL_H
#define MUONINFERENCE_SEGMENTEDGECLASSIFIERTOOL_H
 
#include "BucketInferenceToolBase.h"
#include "MuonInferenceInterfaces/ISegmentEdgeClassifierTool.h"
#include "MuonMLEvent.h"
#include "Gaudi/Property.h"
 
#include <cstdint>
#include <string>
#include <vector>
 
namespace MuonML {

  struct BucketSegmentFeatures {
    int chamberIndex{0};  
    int layers{0};        
    int sector{0};        
    int nSegments{0};     
  };

  class SegmentEdgeClassifierTool final : public BucketInferenceToolBase,
                                          virtual public ISegmentEdgeClassifierTool {
  public:
    using BucketInferenceToolBase::BucketInferenceToolBase;

    StatusCode initialize() override;

    StatusCode runGraphInference(const EventContext& ctx,
                                 GraphRawData& graphData) const override;

    StatusCode buildGraph(const EventContext& ctx,
                          const xAOD::MuonSegmentContainer& segments,
                          SegmentEdgeGraph& graph) const override;

    StatusCode classifyEdges(const EventContext& ctx,
                             const SegmentEdgeGraph& graph,
                             std::vector<SegmentEdgeScore>& scores) const override;
 
  private:
    Gaudi::Property<float> m_maxDeltaThetaDeg{this, "MaxDeltaThetaDeg", 35.f};
    Gaudi::Property<int> m_maxDeltaSector{this, "MaxDeltaSector", 1};
    Gaudi::Property<int> m_sectorModulo{this, "SectorModulo", 16};
    Gaudi::Property<std::string> m_inputNodeName{this, "InputNodeName", "x"};
    Gaudi::Property<std::string> m_inputEdgeIndexName{this, "InputEdgeIndexName", "edge_index"};
    Gaudi::Property<std::string> m_inputEdgeAttrName{this, "InputEdgeAttrName", "edge_attr"};
    Gaudi::Property<std::string> m_outputName{this, "OutputName", "logits"};
    float m_cosMin{0.f};

    std::vector<std::string> m_nodeFeatureNames{};
    std::vector<SegmentNodeFeatureId> m_nodeFeatureIds{};
  };
}
#endif