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 "StoreGate/ReadDecorHandleKey.h"
#include "xAODMuon/MuonSegmentContainer.h"
 
#include <atomic>
#include <cstdint>
#include <mutex>
#include <string>
#include <unordered_set>
#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 finalize() 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;

    bool enableTruthDiagnostics() const override { return m_enableTruthDiagnostics; }
 
  private:
    StatusCode dumpDebugEvent(const EventContext& ctx,
                              const SegmentEdgeGraph& graph,
                              const std::vector<SegmentEdgeScore>& scores) const;

    void fillTruthDiagnostics(
        const xAOD::MuonSegmentContainer& segments,
        const std::unordered_set<const xAOD::MuonSegment*>& bucketRetained,
        SegmentEdgeGraph& graph) const;
 
    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,
        "Number of muon sectors used when applying wrap-around sector distance"};
    Gaudi::Property<unsigned int> m_maxSegmentsPerBucket{this, "MaxSegmentsPerBucket", 0,
        "Keep at most this many quality-ranked segments per (sector, chamber, eta) bucket before inference; 0 keeps all"};
    Gaudi::Property<unsigned int> m_maxEdgesPerNodeBeforeInference{this, "MaxEdgesPerNodeBeforeInference", 0,
        "Keep at most this many geometrical neighbour pairs per node before ONNX inference; 0 keeps all"};
    Gaudi::Property<unsigned int> m_maxEdgesPerTargetChamberBeforeInference{
        this, "MaxEdgesPerTargetChamberBeforeInference", 0,
        "Keep at most this many pre-ONNX neighbours from one target chamber per node; 0 keeps all"};
    Gaudi::Property<bool> m_dropSameChamberEdgesBeforeInference{this, "DropSameChamberEdgesBeforeInference", true,
        "Drop same-chamber segment pairs before ONNX inference"};
    Gaudi::Property<bool> m_dropIsolatedNodesBeforeInference{this, "DropIsolatedNodesBeforeInference", true,
        "Remove nodes without a retained pre-ONNX edge before creating ONNX tensors"};
    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"};
    Gaudi::Property<std::string> m_debugDumpFile{this, "DebugDumpFile", ""};
    Gaudi::Property<unsigned int> m_debugDumpMaxEvents{this, "DebugDumpMaxEvents", 0};
    Gaudi::Property<bool> m_enableTruthDiagnostics{
        this, "EnableTruthDiagnostics", false,
        "MC-only: fill SegmentEdgeGraph's truth diagnostics."};
    SG::ReadDecorHandleKey<xAOD::MuonSegmentContainer> m_truthLinkKey{
        this, "TruthLinkKey", "MuonSegmentsFromR4.truthParticleLink"};
    float m_cosMin{0.f};

    std::vector<std::string> m_nodeFeatureNames{};
    std::vector<SegmentNodeFeatureId> m_nodeFeatureIds{};
    mutable std::mutex m_debugDumpMutex;
    mutable std::atomic<unsigned int> m_debugDumpEvents{0};

    mutable std::atomic<std::size_t> m_sumInputSegments{0};
    mutable std::atomic<std::size_t> m_sumCandidatePairs{0};
    mutable std::atomic<std::size_t> m_sumRetainedPairs{0};
    mutable std::atomic<std::size_t> m_sumNodesBeforeIsolatedDrop{0};
    mutable std::atomic<std::size_t> m_sumNodesAfterIsolatedDrop{0};
  };
}
#endif