MuonML::SegmentEdgeClassifierTool Class Reference

Runs a segment-level GNN on reconstructed muon segments to classify segment-pair edges as "good" or "background". More...

#include <SegmentEdgeClassifierTool.h>

Inheritance diagram for MuonML::SegmentEdgeClassifierTool:

Collaboration diagram for MuonML::SegmentEdgeClassifierTool:

Public Member Functions
StatusCode	initialize () override
	Retrieve the ONNX model and resolve node feature ordering from metadata.
StatusCode	finalize () override
	Log a pre-ONNX candidate-edge pruning.
StatusCode	runGraphInference (const EventContext &ctx, GraphRawData &graphData) const override
	Not supported by this tool; returns FAILURE.
StatusCode	buildGraph (const EventContext &ctx, const xAOD::MuonSegmentContainer &segments, SegmentEdgeGraph &graph) const override
	Build a GNN graph from segments, computing node and edge features and storing the graph structure in graph.
StatusCode	classifyEdges (const EventContext &ctx, const SegmentEdgeGraph &graph, std::vector< SegmentEdgeScore > &scoresscores) const override
	Run ONNX inference on graph and populate scores with logit and probability for each edge; called after buildGraph().
bool	enableTruthDiagnostics () const override
	Whether this tool was configured to fill SegmentEdgeGraph's truth diagnostics.
StatusCode	buildGraph (const EventContext &ctx, GraphRawData &graphData) const
	GNN-style graph builder (features + edges). Kept for tools that want it.
StatusCode	runInference (GraphRawData &graphData) const
	Default ONNX run for GNN case: inputs {"features","edge_index"} -> outputs {"logits"}.
	DeclareInterfaceID (ISegmentEdgeClassifierTool, 1, 0)

Protected Member Functions
StatusCode	setupModel ()
Ort::Session &	model () const
StatusCode	buildFeaturesOnly (const EventContext &ctx, GraphRawData &graphData) const
	Build only features (N,6); attaches one tensor in graph.dataTensor[0].
StatusCode	buildTransformerInputs (const EventContext &ctx, GraphRawData &graphData) const
	Build Transformer inputs: features [1,S,6] and pad_mask [1,S] (False = valid), as tensors 0 and 1.
StatusCode	runNamedInference (GraphRawData &graphData, const std::vector< const char * > &inputNames, const std::vector< const char * > &outputNames) const
	Generic named inference, for tools with different I/O conventions.

Static Protected Member Functions
static std::string	trimFeatureToken (std::string s)
static std::vector< std::string >	parseFeatureNames (const std::string &raw)

Protected Attributes
SG::ReadHandleKey< MuonR4::SpacePointContainer >	m_readKey {this, "ReadSpacePoints", "MuonSpacePoints"}
ActsTrk::GeoContextReadKey_t	m_geoCtxKey {this, "AlignmentKey", "ActsAlignment", "cond handle key"}
Gaudi::Property< int >	m_minLayers {this, "MinLayersValid", 3}
Gaudi::Property< int >	m_maxChamberDelta {this, "MaxChamberDelta", 13}
Gaudi::Property< int >	m_maxSectorDelta {this, "MaxSectorDelta", 1}
Gaudi::Property< double >	m_maxDistXY {this, "MaxDistXY", 6800.0}
Gaudi::Property< double >	m_maxAbsDz {this, "MaxAbsDz", 15000.0}
Gaudi::Property< unsigned int >	m_debugDumpFirstNNodes {this, "DebugDumpFirstNNodes", 5}
Gaudi::Property< unsigned int >	m_debugDumpFirstNEdges {this, "DebugDumpFirstNEdges", 12}
Gaudi::Property< bool >	m_validateEdges {this, "ValidateEdges", true}
Gaudi::Property< bool >	m_sanitizeNonFinitePredictions
bool	m_isCuda {false}
int	m_cudaDeviceId {0}

Static Protected Attributes
static constexpr std::size_t	kBucketFeatureCount = 6
static constexpr std::size_t	kNodeFeatureCount = 10
static constexpr std::size_t	kEdgeFeatureCount = 7
static constexpr std::array< std::string_view, kNodeFeatureCount >	kDefaultNodeFeatureNames

Private Member Functions
StatusCode	dumpDebugEvent (const EventContext &ctx, const SegmentEdgeGraph &graph, const std::vector< SegmentEdgeScore > &scoresscores) const
void	fillTruthDiagnostics (const xAOD::MuonSegmentContainer &segments, const std::unordered_set< const xAOD::MuonSegment * > &bucketRetained, SegmentEdgeGraph &graph) const
	MC-only diagnostics: record, for every input segment and every pair of input segments sharing a truth particle, why it did or did not reach ONNX.

Private Attributes
Gaudi::Property< float >	m_maxDeltaThetaDeg {this, "MaxDeltaThetaDeg", 35.f}
Gaudi::Property< int >	m_maxDeltaSector {this, "MaxDeltaSector", 1}
Gaudi::Property< int >	m_sectorModulo
Gaudi::Property< unsigned int >	m_maxSegmentsPerBucket
Gaudi::Property< unsigned int >	m_maxEdgesPerNodeBeforeInference
Gaudi::Property< unsigned int >	m_maxEdgesPerTargetChamberBeforeInference
Gaudi::Property< bool >	m_dropSameChamberEdgesBeforeInference
Gaudi::Property< bool >	m_dropIsolatedNodesBeforeInference
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
SG::ReadDecorHandleKey< xAOD::MuonSegmentContainer >	m_truthLinkKey
float	m_cosMin {0.f}
std::vector< std::string >	m_nodeFeatureNames {}
	Node feature order expected by the model metadata (resolved at initialize).
std::vector< SegmentNodeFeatureId >	m_nodeFeatureIds {}
std::mutex	m_debugDumpMutex
std::atomic< unsigned int >	m_debugDumpEvents {0}
std::atomic< std::size_t >	m_sumInputSegments {0}
	Job-summed pre-ONNX pruning counters (see buildGraph()).
std::atomic< std::size_t >	m_sumCandidatePairs {0}
std::atomic< std::size_t >	m_sumRetainedPairs {0}
std::atomic< std::size_t >	m_sumNodesBeforeIsolatedDrop {0}
std::atomic< std::size_t >	m_sumNodesAfterIsolatedDrop {0}
ToolHandle< AthOnnx::IOnnxRuntimeSessionTool >	m_onnxSessionTool

Detailed Description

Runs a segment-level GNN on reconstructed muon segments to classify segment-pair edges as "good" or "background".

The tool reads a xAOD::MuonSegmentContainer and builds a graph where:

• Nodes are muon segments, each with 10 features:
  • Position and direction (6 floats)
  • Chamber index, layer count, sector, and segment multiplicity (4 floats)
• Edges connect all segment pairs within an angular threshold (cos(angle) >= cos(MaxDeltaThetaDeg)) and sector distance, with 7 features:
  • Spatial displacement (3 floats: dx, dy, dz)
  • Distance magnitude (1 float)
  • Angle (dot product, 1 float)
  • Chamber and sector match flags (2 flags)

The tool then runs an ONNX model (typically a GIN or GCN variant) to produce a logit or probability for each edge, enabling downstream algorithms to filter low-quality segment associations and improve reconstruction efficiency.

Key difference from GraphBucketFilterTool: operates at segment (edge) level rather than bucket (node) level, and the interface uses discrete graph structures (SegmentEdgeGraph) rather than tensors for input/output.

Note: runGraphInference() is not supported by this tool; use SegmentEdgeInferenceAlg and the ISegmentEdgeClassifierTool methods instead.

Definition at line 64 of file SegmentEdgeClassifierTool.h.

Member Function Documentation

buildFeaturesOnly()

StatusCode BucketInferenceToolBase::buildFeaturesOnly ( const EventContext & ctx, GraphRawData & graphData ) const

protectedinherited

Build only features (N,6); attaches one tensor in graph.dataTensor[0].

Definition at line 87 of file BucketInferenceToolBase.cxx.

                                                                                     {
  
  graphData.graph.reset();
  graphData.srcEdges.clear();
  graphData.desEdges.clear();
  graphData.edgeIndexPacked.clear();
  graphData.featureLeaves.clear();
  graphData.spacePointsInBucket.clear();
  graphData.graph = std::make_unique<InferenceGraph>();
  graphData.graph->dataTensor.reserve(1); // features input; outputs are reserved in runNamedInference()
 
  const MuonR4::SpacePointContainer* buckets{nullptr};
  ATH_CHECK(SG::get(buckets, m_readKey, ctx));
 
  const ActsTrk::GeometryContext* gctx = nullptr;
  ATH_CHECK(SG::get(gctx, m_geoCtxKey, ctx));
 
  std::vector<BucketGraphUtils::NodeAux> nodes;
  BucketGraphUtils::buildNodesAndFeatures(*buckets, *gctx, nodes,
                                          graphData.featureLeaves,
                                          graphData.spacePointsInBucket); // now int64_t-compatible
 
  const int64_t numNodes = static_cast<int64_t>(nodes.size());
  ATH_MSG_DEBUG("Total buckets: " << buckets->size()
                 << " -> nodes (size>0): " << numNodes
                 << " | features.size()=" << graphData.featureLeaves.size());
 
  if (numNodes == 0) {
    ATH_MSG_WARNING("No valid buckets found (all have size 0.0). Skipping inference.");
    return StatusCode::SUCCESS;
  }
 
  const int64_t nFeatPerNode = static_cast<int64_t>(kBucketFeatureCount);
  if (numNodes * nFeatPerNode != static_cast<int64_t>(graphData.featureLeaves.size())) {
    ATH_MSG_ERROR( "Feature size mismatch: expected " << (numNodes * nFeatPerNode)
                   << " got " << graphData.featureLeaves.size());
    return StatusCode::FAILURE;
  }
 
  Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  std::vector<int64_t> featShape{numNodes, nFeatPerNode};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      graphData.featureLeaves.data(),
                                      graphData.featureLeaves.size(),
                                      featShape.data(),
                                      featShape.size()));
  return StatusCode::SUCCESS;
}

buildGraph() [1/2]

StatusCode BucketInferenceToolBase::buildGraph ( const EventContext & ctx, GraphRawData & graphData ) const

inherited

GNN-style graph builder (features + edges). Kept for tools that want it.

Definition at line 200 of file BucketInferenceToolBase.cxx.

                                                                              {
  
  graphData.graph.reset();
  graphData.srcEdges.clear();
  graphData.desEdges.clear();
  graphData.featureLeaves.clear();
  graphData.spacePointsInBucket.clear();
  graphData.edgeIndexPacked.clear();
  graphData.graph = std::make_unique<InferenceGraph>();
  graphData.graph->dataTensor.reserve(2); // features and edge_index inputs; outputs are reserved in runNamedInference()
 
  const MuonR4::SpacePointContainer* buckets{nullptr};
  ATH_CHECK(SG::get(buckets, m_readKey, ctx));
 
  const ActsTrk::GeometryContext* gctx = nullptr;
  ATH_CHECK(SG::get(gctx, m_geoCtxKey, ctx));
 
  std::vector<BucketGraphUtils::NodeAux> nodes;
 
  BucketGraphUtils::buildNodesAndFeatures(*buckets, *gctx, nodes,
                                          graphData.featureLeaves,
                                          graphData.spacePointsInBucket);
 
  const int64_t numNodes = static_cast<int64_t>(nodes.size());
  ATH_MSG_DEBUG("Total buckets: " << buckets->size()
                << " -> nodes (size>0): " << numNodes
                << " | features.size()=" << graphData.featureLeaves.size());
 
  if (numNodes == 0) {
    ATH_MSG_WARNING("No valid buckets found (all have size 0.0). Skipping graph building.");
    return StatusCode::SUCCESS;
  }
 
  const int64_t nFeatPerNode = static_cast<int64_t>(kBucketFeatureCount);
  if (numNodes * nFeatPerNode != static_cast<int64_t>(graphData.featureLeaves.size())) {
    ATH_MSG_ERROR("Feature size mismatch: expected " << (numNodes * nFeatPerNode)
                  << " got " << graphData.featureLeaves.size());
    return StatusCode::FAILURE;
  }
 
  Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  std::vector<int64_t> featShape{numNodes, nFeatPerNode};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      graphData.featureLeaves.data(),
                                      graphData.featureLeaves.size(),
                                      featShape.data(),
                                      featShape.size()));
 
  BucketGraphUtils::buildSparseEdges(nodes,
                                     m_minLayers,
                                     m_maxChamberDelta,
                                     m_maxSectorDelta,
                                     m_maxDistXY,
                                     m_maxAbsDz,
                                     graphData.srcEdges, graphData.desEdges);
  if (m_validateEdges) {
    size_t bad = 0;
    size_t write = 0;
    for (size_t k = 0; k < graphData.srcEdges.size(); ++k) {
      const int64_t u = graphData.srcEdges[k];
      const int64_t v = graphData.desEdges[k];
      const bool okU = (u >= 0 && u < numNodes);
      const bool okV = (v >= 0 && v < numNodes);
      if (okU && okV) {
        graphData.srcEdges[write] = u;
        graphData.desEdges[write] = v;
        ++write;
      } else {
        ++bad;
        ATH_MSG_DEBUG( "Drop invalid edge " << k << ": (" << u << "->" << v
                        << "), valid node range [0," << (numNodes-1) << "]");
      }
    }
    if (bad) {
      ATH_MSG_WARNING( "Removed " << bad << " invalid edges out of "
                        << graphData.srcEdges.size());
      graphData.srcEdges.resize(write);
      graphData.desEdges.resize(write);
    }
  }
 
  const size_t E = graphData.srcEdges.size();
 
  if (msgLvl(MSG::DEBUG)) {
    // DEBUG: Count connections per node
    ATH_MSG_DEBUG("Edges built: " << E);
    const size_t dumpE = std::min<std::size_t>(m_debugDumpFirstNEdges.value(), E);
    for (size_t k = 0; k < dumpE; ++k) {
      ATH_MSG_DEBUG("EDGE[" << k << "]: "
                    << graphData.srcEdges[k] << " -> "
                    << graphData.desEdges[k]);
    }
 
    std::vector<int> nodeConnections(numNodes, 0);
    for (size_t k = 0; k < graphData.srcEdges.size(); ++k) {
      const int64_t u = graphData.srcEdges[k];
      const int64_t v = graphData.desEdges[k];
      if (u >= 0 && u < numNodes) nodeConnections[u]++;
      if (v >= 0 && v < numNodes) nodeConnections[v]++;
    }
 
    ATH_MSG_DEBUG("=== DEBUGGING: Node Connections (first 10 nodes) ===");
    const int64_t debugNodeCount = std::min(numNodes, static_cast<int64_t>(10));
    for (int64_t i = 0; i < debugNodeCount; ++i) {
      ATH_MSG_DEBUG("Node[" << i << "] connections: " << nodeConnections[i]);
    }
    ATH_MSG_DEBUG("=== END DEBUG NODE CONNECTIONS ===");
 
    ATH_MSG_DEBUG("=== DEBUGGING: Detailed Edge Connections (first 10 nodes) ===");
    for (int64_t nodeIdx = 0; nodeIdx < debugNodeCount; ++nodeIdx) {
      std::stringstream connections;
      connections << "Node[" << nodeIdx << "] connected to: ";
      bool foundAny = false;
 
      for (size_t k = 0; k < graphData.srcEdges.size(); ++k) {
        const int64_t u = graphData.srcEdges[k];
        const int64_t v = graphData.desEdges[k];
 
        if (u == nodeIdx) {
          if (foundAny) connections << ", ";
          connections << v;
          foundAny = true;
        } else if (v == nodeIdx) {
          if (foundAny) connections << ", ";
          connections << u;
          foundAny = true;
        }
      }
 
      if (!foundAny) connections << "none";
      ATH_MSG_DEBUG(connections.str());
    }
    ATH_MSG_DEBUG("=== END DEBUG DETAILED CONNECTIONS ===");
  }
 
  nodes = {};
 
  graphData.edgeIndexPacked.clear();
  const size_t Efinal = BucketGraphUtils::packEdgeIndex(graphData.srcEdges,
                                                       graphData.desEdges,
                                                       graphData.edgeIndexPacked);
 
  graphData.srcEdges.clear();
  graphData.desEdges.clear();
 
  std::vector<int64_t> edgeShape{2, static_cast<int64_t>(Efinal)};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<int64_t>(memInfo,
                                        graphData.edgeIndexPacked.data(),
                                        graphData.edgeIndexPacked.size(),
                                        edgeShape.data(),
                                        edgeShape.size()));
 
  ATH_MSG_DEBUG("Built sparse bucket graph: N=" << numNodes << ", E=" << Efinal);
  return StatusCode::SUCCESS;
}

buildGraph() [2/2]

StatusCode MuonML::SegmentEdgeClassifierTool::buildGraph ( const EventContext & ctx, const xAOD::MuonSegmentContainer & segments, SegmentEdgeGraph & graph ) const

overridevirtual

Build a GNN graph from segments, computing node and edge features and storing the graph structure in graph.

Implements MuonML::ISegmentEdgeClassifierTool.

Definition at line 234 of file SegmentEdgeClassifierTool.cxx.

                                   {
  graph = SegmentEdgeGraph{};
  graph.segments.reserve(segments.size());
  graph.nodeFeatures.reserve(segments.size() * kNodeFeatureCount);
  // Evaluated once per event; short-circuits without touching the message
  // service unless the property was explicitly enabled.
  const bool truthDiag = m_enableTruthDiagnostics.value() && msgLvl(MSG::DEBUG);
 
  /*
   * Keep the original bucket multiplicity in the node feature even when the
   * speed configuration retains only the best representatives of a bucket.
   * This preserves the model's occupancy input while removing duplicate node
   * and edge work before tensor construction.
   */
  std::map<SegmentBucketKey, std::vector<const xAOD::MuonSegment*>>
      segmentsByBucket;
  for (const xAOD::MuonSegment* segment : segments) {
    segmentsByBucket[segmentBucketKey(*segment)].push_back(segment);
  }
 
  const InferenceUtils::SegmentQualityOrder betterSegment{};
 
  std::unordered_set<const xAOD::MuonSegment*> retainedSegments;
  retainedSegments.reserve(segments.size());
  for (auto& [_, bucketSegments] : segmentsByBucket) {
    std::ranges::sort(bucketSegments, betterSegment);
    const std::size_t nKeep = m_maxSegmentsPerBucket.value() == 0
                                  ? bucketSegments.size()
                                  : std::min<std::size_t>(
                                        bucketSegments.size(),
                                        m_maxSegmentsPerBucket.value());
    retainedSegments.insert(bucketSegments.begin(),
                            bucketSegments.begin() + nKeep);
  }
 
  std::vector<Amg::Vector3D> pos;
  std::vector<Amg::Vector3D> dir;
  std::vector<BucketSegmentFeatures> bucket;
  pos.reserve(retainedSegments.size());
  dir.reserve(retainedSegments.size());
  bucket.reserve(retainedSegments.size());
 
  for (const xAOD::MuonSegment* segment : segments) {
    if (!retainedSegments.contains(segment)) continue;
 
    const Amg::Vector3D position = segment->position();
    const Amg::Vector3D direction = segment->direction();
    const SegmentBucketKey key = segmentBucketKey(*segment);
    const auto bucketIt = segmentsByBucket.find(key);
    const int multiplicity =
        bucketIt == segmentsByBucket.end()
            ? 1
            : static_cast<int>(bucketIt->second.size());
 
    const int chamberIndex = static_cast<int>(segment->chamberIndex());
    const int layers = layersInBucket(*MuonR4::detailedSegment(*segment)->parent()->parentBucket());
    const int sector = segment->sector();
 
    graph.segments.push_back(segment);
    pos.emplace_back(position / Gaudi::Units::m);
    dir.emplace_back(direction);
    bucket.emplace_back(BucketSegmentFeatures{
        chamberIndex, layers, sector, multiplicity});
    for (const SegmentNodeFeatureId featureId : m_nodeFeatureIds) {
      graph.nodeFeatures.push_back(
          nodeFeatureValue(featureId, pos.back(), dir.back(), bucket.back()));
    }
  }
  graph.nNodes = graph.segments.size();
 
  if (pos.size() != graph.nNodes || dir.size() != graph.nNodes ||
      bucket.size() != graph.nNodes) {
    ATH_MSG_ERROR("Inconsistent vector sizes during graph building: nodes="
                  << graph.nNodes << ", pos=" << pos.size()
                  << ", dir=" << dir.size() << ", bucket=" << bucket.size());
    return StatusCode::FAILURE;
  }
 
  if (graph.nNodes < 2) {
    graph.nEdges = 0;
    if (truthDiag) fillTruthDiagnostics(segments, retainedSegments, graph);
    return StatusCode::SUCCESS;
  }
 
  const auto wrapRegularSector = [&](int sector) {
    // MuonSegment::sector() is a regular sector number, not an ExpandedSector
    // coordinate. Wrap it to [0, modulo); <= 0 disables wrapping.
    if (m_sectorModulo.value() > 0) {
      sector %= m_sectorModulo.value();
      if (sector < 0) sector += m_sectorModulo.value();
    }
    return sector;
  };
 
  // The lookup key must use the same wrapping as the target sectors below:
  // ATLAS sectors are 1-based (1..16), so a raw key of 16 can never match a
  // wrapped target of 0, which silently dropped every edge into sector 16.
  // The per-pair sectorDistance check below enforces the true circular
  // distance on the raw sector numbers.
  std::unordered_map<int, std::vector<std::size_t>> nodesBySector;
  nodesBySector.reserve(graph.nNodes);
  for (std::size_t node = 0; node < graph.nNodes; ++node) {
    nodesBySector[wrapRegularSector(bucket[node].sector)].push_back(node);
  }
 
  std::unordered_map<int, std::vector<int>> targetSectorsBySourceSector;
  targetSectorsBySourceSector.reserve(nodesBySector.size());
  std::size_t sectorLocalEdgeUpperBound = 0;
  for (const auto& [sourceSector, sourceNodes] : nodesBySector) {
    std::vector<int> targetSectors;
    targetSectors.reserve(2 * m_maxDeltaSector.value() + 1);
    for (int delta = -m_maxDeltaSector.value();
         delta <= m_maxDeltaSector.value(); ++delta) {
      targetSectors.push_back(wrapRegularSector(sourceSector + delta));
    }
    for (const int targetSector : targetSectors) {
      const auto found = nodesBySector.find(targetSector);
      if (found == nodesBySector.end()) continue;
      sectorLocalEdgeUpperBound += sourceNodes.size() * found->second.size();
      if (targetSector == sourceSector) {
        sectorLocalEdgeUpperBound -= sourceNodes.size();
      }
    }
    targetSectorsBySourceSector.emplace(sourceSector,
                                        std::move(targetSectors));
  }
 
  /*
   * The model receives a directed graph, but the geometric candidate relation
   * is undirected.  Build each pair once, then emit both directions.  With a
   * non-zero input cap, each endpoint nominates its best candidates and the
   * union is made bidirectional before inference; this preserves the message
   * passing symmetry expected by the GNN.
   */
  struct UndirectedEdge {
    std::size_t first{0};
    std::size_t second{0};
    float dx{0.f};
    float dy{0.f};
    float dz{0.f};
    float distance{0.f};
    float cosAngle{0.f};
  };
  const auto betterEdge = [](const UndirectedEdge& first,
                             const UndirectedEdge& second) {
    const int cosOrder = InferenceUtils::compareFloatDescending(
        first.cosAngle, second.cosAngle);
    if (cosOrder != 0) {
      return cosOrder < 0;
    }
 
    const int distanceOrder =
        InferenceUtils::compareFloat(first.distance, second.distance);
    if (distanceOrder != 0) {
      return distanceOrder < 0;
    }
    if (first.first != second.first) return first.first < second.first;
    return first.second < second.second;
  };
  const auto edgeKey = [](const UndirectedEdge& edge) {
    return (static_cast<std::uint64_t>(edge.first) << 32) |
           static_cast<std::uint64_t>(edge.second);
  };
 
  const unsigned int maxEdgesPerNode =
      m_maxEdgesPerNodeBeforeInference.value();
  const unsigned int maxEdgesPerTargetChamber =
      m_maxEdgesPerTargetChamberBeforeInference.value();
  const bool usePreInferenceSelection =
      maxEdgesPerNode != 0 || maxEdgesPerTargetChamber != 0;
  std::vector<std::vector<UndirectedEdge>> bestEdgesByNode;
  if (usePreInferenceSelection) {
    bestEdgesByNode.resize(graph.nNodes);
    const unsigned int reservePerNode =
        maxEdgesPerNode != 0 ? maxEdgesPerNode : maxEdgesPerTargetChamber;
    for (std::vector<UndirectedEdge>& edges : bestEdgesByNode) {
      edges.reserve(reservePerNode);
    }
  } else {
    graph.edgeIndex.reserve(2 * sectorLocalEdgeUpperBound);
    graph.edgeFeatures.reserve(kEdgeFeatureCount * sectorLocalEdgeUpperBound);
  }
  const auto appendDirectedPair = [&](const UndirectedEdge& edge) {
    graph.edgeIndex.push_back(static_cast<int64_t>(edge.first));
    graph.edgeIndex.push_back(static_cast<int64_t>(edge.second));
    graph.edgeFeatures.insert(
        graph.edgeFeatures.end(),
        {edge.dx, edge.dy, edge.dz, edge.distance, edge.cosAngle,
         float(bucket[edge.first].chamberIndex ==
               bucket[edge.second].chamberIndex),
         float(bucket[edge.first].sector == bucket[edge.second].sector)});
 
    graph.edgeIndex.push_back(static_cast<int64_t>(edge.second));
    graph.edgeIndex.push_back(static_cast<int64_t>(edge.first));
    graph.edgeFeatures.insert(
        graph.edgeFeatures.end(),
        {-edge.dx, -edge.dy, -edge.dz, edge.distance, edge.cosAngle,
         float(bucket[edge.first].chamberIndex ==
               bucket[edge.second].chamberIndex),
         float(bucket[edge.first].sector == bucket[edge.second].sector)});
  };
 
  const auto retainForNode = [&](std::size_t node,
                                 const UndirectedEdge& candidate) {
    std::vector<UndirectedEdge>& retained = bestEdgesByNode[node];
    const std::size_t other = candidate.first == node ? candidate.second
                                                       : candidate.first;
    const int targetChamber = bucket[other].chamberIndex;
 
    if (maxEdgesPerTargetChamber != 0) {
      unsigned int sameChamberCount = 0;
      auto worstSameChamber = retained.end();
      for (auto it = retained.begin(); it != retained.end(); ++it) {
        const std::size_t retainedOther =
            it->first == node ? it->second : it->first;
        if (bucket[retainedOther].chamberIndex != targetChamber) continue;
        ++sameChamberCount;
        if (worstSameChamber == retained.end() ||
            betterEdge(*worstSameChamber, *it)) {
          worstSameChamber = it;
        }
      }
      if (sameChamberCount >= maxEdgesPerTargetChamber) {
        if (!betterEdge(candidate, *worstSameChamber)) return;
        *worstSameChamber = candidate;
      } else {
        retained.push_back(candidate);
      }
    } else {
      retained.push_back(candidate);
    }
 
    if (maxEdgesPerNode != 0 && retained.size() > maxEdgesPerNode) {
      auto worst = retained.begin();
      for (auto it = std::next(retained.begin()); it != retained.end(); ++it) {
        if (betterEdge(*worst, *it)) worst = it;
      }
      retained.erase(worst);
    }
  };
 
  std::size_t candidatePairs = 0;
  for (std::size_t first = 0; first < graph.nNodes; ++first) {
    const auto sectorsIt =
        targetSectorsBySourceSector.find(wrapRegularSector(bucket[first].sector));
    if (sectorsIt == targetSectorsBySourceSector.end()) continue;
    for (const int sector : sectorsIt->second) {
      const auto targetIt = nodesBySector.find(sector);
      if (targetIt == nodesBySector.end()) continue;
 
      for (const std::size_t second : targetIt->second) {
        // Every valid pair will be visited from the lower-index endpoint.
        if (second <= first) continue;
        if (sectorDistance(bucket[first].sector, bucket[second].sector,
                           m_sectorModulo.value()) >
            m_maxDeltaSector.value()) {
          continue;
        }
        if (m_dropSameChamberEdgesBeforeInference.value() &&
            bucket[first].chamberIndex == bucket[second].chamberIndex) {
          continue;
        }
        const float cosAngle = static_cast<float>(dir[first].dot(dir[second]));
        if (cosAngle < m_cosMin) continue;
 
        const Amg::Vector3D delta = pos[second] - pos[first];
        const UndirectedEdge candidate{
            first,
            second,
            static_cast<float>(delta.x()),
            static_cast<float>(delta.y()),
            static_cast<float>(delta.z()),
            static_cast<float>(delta.mag()),
            cosAngle};
        ++candidatePairs;
 
        if (!usePreInferenceSelection) {
          appendDirectedPair(candidate);
        } else {
          retainForNode(first, candidate);
          retainForNode(second, candidate);
        }
      }
    }
  }
  std::size_t retainedPairs = candidatePairs;
  if (usePreInferenceSelection) {
    const unsigned int selectedReservePerNode =
        maxEdgesPerNode != 0 ? maxEdgesPerNode : maxEdgesPerTargetChamber;
    std::unordered_set<std::uint64_t> selectedKeys;
    selectedKeys.reserve(graph.nNodes * selectedReservePerNode);
    std::vector<UndirectedEdge> selectedEdges;
    selectedEdges.reserve(graph.nNodes * selectedReservePerNode);
 
    for (const std::vector<UndirectedEdge>& nodeEdges : bestEdgesByNode) {
      for (const UndirectedEdge& edge : nodeEdges) {
        if (selectedKeys.insert(edgeKey(edge)).second) {
          selectedEdges.push_back(edge);
        }
      }
    }
    std::sort(selectedEdges.begin(), selectedEdges.end(),
              [](const UndirectedEdge& first,
                 const UndirectedEdge& second) {
                if (first.first != second.first) {
                  return first.first < second.first;
                }
                return first.second < second.second;
              });
 
    retainedPairs = selectedEdges.size();
    graph.edgeIndex.reserve(4 * retainedPairs);
    graph.edgeFeatures.reserve(2 * kEdgeFeatureCount * retainedPairs);
    for (const UndirectedEdge& edge : selectedEdges) {
      appendDirectedPair(edge);
    }
  }
  graph.nEdges = graph.edgeIndex.size() / 2;
  const std::size_t nodesBeforeIsolatedNodeDrop = graph.nNodes;
  if (m_dropIsolatedNodesBeforeInference.value() && graph.nEdges != 0) {
    std::vector<unsigned char> active(graph.nNodes, 0);
    for (const int64_t index : graph.edgeIndex) {
      active[static_cast<std::size_t>(index)] = 1;
    }
    const std::size_t activeNodes =
        std::count(active.begin(), active.end(), static_cast<unsigned char>(1));
    if (activeNodes != graph.nNodes) {
      std::vector<std::size_t> oldToNew(graph.nNodes, graph.nNodes);
      std::vector<const xAOD::MuonSegment*> compactedSegments;
      std::vector<float> compactedNodeFeatures;
      compactedSegments.reserve(activeNodes);
      compactedNodeFeatures.reserve(activeNodes * kNodeFeatureCount);
      for (std::size_t oldNode = 0; oldNode < graph.nNodes; ++oldNode) {
        if (!active[oldNode]) continue;
        oldToNew[oldNode] = compactedSegments.size();
        compactedSegments.push_back(graph.segments[oldNode]);
        const auto featureBegin = graph.nodeFeatures.begin() +
            oldNode * kNodeFeatureCount;
        compactedNodeFeatures.insert(compactedNodeFeatures.end(),
                                     featureBegin,
                                     featureBegin + kNodeFeatureCount);
      }
      for (int64_t& index : graph.edgeIndex) {
        index = static_cast<int64_t>(oldToNew[static_cast<std::size_t>(index)]);
      }
      graph.segments = std::move(compactedSegments);
      graph.nodeFeatures = std::move(compactedNodeFeatures);
      graph.nNodes = activeNodes;
    }
  }
  if (truthDiag) fillTruthDiagnostics(segments, retainedSegments, graph);
  ATH_MSG_DEBUG("buildGraph: input segments=" << segments.size()
                << ", kept nodes=" << graph.nNodes
                << ", nodes before isolated-node drop=" << nodesBeforeIsolatedNodeDrop
                << ", bucket cap=" << m_maxSegmentsPerBucket.value()
                << ", candidate pairs=" << candidatePairs
                << ", retained pairs=" << retainedPairs
                << ", built directed edges=" << graph.nEdges
                << ", pre-inference node cap=" << m_maxEdgesPerNodeBeforeInference.value()
                << ", per-target-chamber cap=" << maxEdgesPerTargetChamber
                << ", drop same chamber=" << m_dropSameChamberEdgesBeforeInference.value()
                << ", drop isolated nodes=" << m_dropIsolatedNodesBeforeInference.value()
                << ", sector-local reserve=" << sectorLocalEdgeUpperBound);
 
  // Job-summed diagnostics 
  if (msgLvl(MSG::DEBUG)) {
    m_sumInputSegments += segments.size();
    m_sumCandidatePairs += candidatePairs;
    m_sumRetainedPairs += retainedPairs;
    m_sumNodesBeforeIsolatedDrop += nodesBeforeIsolatedNodeDrop;
    m_sumNodesAfterIsolatedDrop += graph.nNodes;
  }
  return StatusCode::SUCCESS;
}

buildTransformerInputs()

StatusCode BucketInferenceToolBase::buildTransformerInputs ( const EventContext & ctx, GraphRawData & graphData ) const

protectedinherited

Build Transformer inputs: features [1,S,6] and pad_mask [1,S] (False = valid), as tensors 0 and 1.

Definition at line 138 of file BucketInferenceToolBase.cxx.

                                                                                          {
  // Start from (N,6)
  ATH_CHECK(buildFeaturesOnly(ctx, graphData));
 
  // Copy features flat buffer for lifetime management
  std::vector<float> featuresFlat = graphData.featureLeaves;
  const int64_t S = static_cast<int64_t>(featuresFlat.size() / kBucketFeatureCount);
 
  if (S == 0) {
    ATH_MSG_WARNING("No valid features for transformer input. Skipping inference.");
    return StatusCode::SUCCESS;
  }
 
  if (msgLvl(MSG::DEBUG)) {
    // DEBUG: Print transformer input features for first 10 nodes
    ATH_MSG_DEBUG("=== DEBUGGING: Transformer input features for first 10 nodes ===");
    const int64_t debugNodes = std::min(S, static_cast<int64_t>(10));
    for (int64_t nodeIdx = 0; nodeIdx < debugNodes; ++nodeIdx) {
      const int64_t baseIdx = nodeIdx * static_cast<int64_t>(kBucketFeatureCount);
      ATH_MSG_DEBUG("TransformerNode[" << nodeIdx << "]: "
                   << "x=" << featuresFlat[baseIdx + 0] << ", "
                   << "y=" << featuresFlat[baseIdx + 1] << ", "
                   << "z=" << featuresFlat[baseIdx + 2] << ", "
                   << "layers=" << featuresFlat[baseIdx + 3] << ", "
                   << "nSp=" << featuresFlat[baseIdx + 4] << ", "
                   << "bucketSize=" << featuresFlat[baseIdx + 5]);
    }
    ATH_MSG_DEBUG("=== END DEBUG TRANSFORMER FEATURES ===");
  }
 
  // Rebuild graph with exactly 2 inputs: features [1,S,6], pad_mask [1,S]
  graphData.graph.reset();
  graphData.graph = std::make_unique<InferenceGraph>();
  graphData.graph->dataTensor.reserve(2); // features and pad_mask inputs; outputs are reserved in runNamedInference()
 
  Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
 
  // features: [1,S,6] (backed by graphData.featureLeaves to keep alive)
  std::vector<int64_t> fShape{1, S, static_cast<int64_t>(kBucketFeatureCount)};
  graphData.featureLeaves.swap(featuresFlat);
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      graphData.featureLeaves.data(),
                                      graphData.featureLeaves.size(),
                                      fShape.data(),
                                      fShape.size()));
 
  // pad_mask: [1,S] (bool). Create ORT-owned tensor and fill with False (=valid).
  Ort::AllocatorWithDefaultOptions allocator;
  std::vector<int64_t> mShape{1, S};
  Ort::Value padVal = Ort::Value::CreateTensor(allocator,
                                               mShape.data(),
                                               mShape.size(),
                                               ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL);
  bool* maskPtr = padVal.GetTensorMutableData<bool>();
  for (int64_t i = 0; i < S; ++i) maskPtr[i] = false;
  graphData.graph->dataTensor.emplace_back(std::move(padVal));
 
  return StatusCode::SUCCESS;
}

classifyEdges()

StatusCode MuonML::SegmentEdgeClassifierTool::classifyEdges ( const EventContext & ctx, const SegmentEdgeGraph & graph, std::vector< SegmentEdgeScore > & scores ) const

overridevirtual

Run ONNX inference on graph and populate scores with logit and probability for each edge; called after buildGraph().

Implements MuonML::ISegmentEdgeClassifierTool.

Definition at line 710 of file SegmentEdgeClassifierTool.cxx.

                                                                                                {
  scores.clear();
  if (!graph.nNodes) return StatusCode::SUCCESS;
  if (!graph.nEdges) {
    ATH_CHECK(dumpDebugEvent(ctx, graph, scores));
    return StatusCode::SUCCESS;
  }
 
  if (graph.nodeFeatures.size() != graph.nNodes * kNodeFeatureCount) {
    ATH_MSG_ERROR("Unexpected node feature size " << graph.nodeFeatures.size()
                  << "; expected " << (graph.nNodes * kNodeFeatureCount));
    return StatusCode::FAILURE;
  }
  if (graph.edgeIndex.size() != 2 * graph.nEdges) {
    ATH_MSG_ERROR("Unexpected edge index size " << graph.edgeIndex.size()
                  << "; expected " << (2 * graph.nEdges));
    return StatusCode::FAILURE;
  }
  if (graph.edgeFeatures.size() != graph.nEdges * kEdgeFeatureCount) {
    ATH_MSG_ERROR("Unexpected edge feature size " << graph.edgeFeatures.size()
                  << "; expected " << (graph.nEdges * kEdgeFeatureCount));
    return StatusCode::FAILURE;
  }
 
  GraphRawData raw{};
  raw.graph = std::make_unique<InferenceGraph>();
  raw.edgeIndexPacked.resize(2 * graph.nEdges);
  for (std::size_t e = 0; e < graph.nEdges; ++e) {
    raw.edgeIndexPacked[e] = graph.edgeIndex[2 * e];
    raw.edgeIndexPacked[graph.nEdges + e] = graph.edgeIndex[2 * e + 1];
  }
 
  Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
 
  const std::vector<int64_t> nodeShape{static_cast<int64_t>(graph.nNodes), static_cast<int64_t>(kNodeFeatureCount)};
  // The graph outlives the synchronous ONNX call below.  Use its node
  // buffer directly instead of allocating and copying featureLeaves per event.
  ATLAS_THREAD_SAFE float* nodeFeaturesData =
      const_cast<float*>(graph.nodeFeatures.data());
  raw.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      nodeFeaturesData,
                                      graph.nodeFeatures.size(),
                                      nodeShape.data(),
                                      nodeShape.size()));
 
  const std::vector<int64_t> edgeIndexShape{2, static_cast<int64_t>(graph.nEdges)};
  raw.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<int64_t>(memInfo,
                                        raw.edgeIndexPacked.data(),
                                        raw.edgeIndexPacked.size(),
                                        edgeIndexShape.data(),
                                        edgeIndexShape.size()));
 
  // ONNX Runtime's CreateTensor API takes a non-const pointer, but it does not
  // mutate input buffers during inference. Avoid copying edge_attr every event.
  ATLAS_THREAD_SAFE float* edgeFeaturesData = const_cast<float*>(graph.edgeFeatures.data());
  const std::vector<int64_t> edgeAttrShape{static_cast<int64_t>(graph.nEdges), static_cast<int64_t>(kEdgeFeatureCount)};
  raw.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      edgeFeaturesData,
                                      graph.edgeFeatures.size(),
                                      edgeAttrShape.data(),
                                      edgeAttrShape.size()));
 
  const std::vector<const char*> inputNames{
      m_inputNodeName.value().c_str(),
      m_inputEdgeIndexName.value().c_str(),
      m_inputEdgeAttrName.value().c_str()};
  const std::vector<const char*> outputNames{m_outputName.value().c_str()};
  ATH_MSG_DEBUG("classifyEdges: ONNX inputs shapes x=[" << nodeShape[0] << "," << nodeShape[1]
                << "], edge_index=[" << edgeIndexShape[0] << "," << edgeIndexShape[1]
                << "], edge_attr=[" << edgeAttrShape[0] << "," << edgeAttrShape[1] << "]");
  ATH_CHECK(runNamedInference(raw, inputNames, outputNames));
 
  if (raw.graph->dataTensor.size() <= inputNames.size()) {
    ATH_MSG_ERROR("Missing ONNX output tensor for segment edge inference");
    return StatusCode::FAILURE;
  }
 
  const Ort::Value& outTensor = raw.graph->dataTensor[inputNames.size()];
  const auto outInfo = outTensor.GetTensorTypeAndShapeInfo();
  const std::vector<int64_t> outShape = outInfo.GetShape();
  const size_t outSize = outInfo.GetElementCount();
  if (!outShape.empty()) {
    ATH_MSG_DEBUG("classifyEdges: ONNX output rank=" << outShape.size()
                  << ", first dim=" << outShape.front()
                  << ", elements=" << outSize);
  } else {
    ATH_MSG_DEBUG("classifyEdges: ONNX scalar output, elements=" << outSize);
  }
  if (outSize < graph.nEdges) {
    ATH_MSG_ERROR("ONNX logits tensor has " << outSize << " entries for " << graph.nEdges << " edges");
    return StatusCode::FAILURE;
  }
 
  const float* logits = outTensor.GetTensorData<float>();
  scores.reserve(graph.nEdges);
  for (std::size_t e=0; e<graph.nEdges; ++e) {
    const float l = logits[e];
    scores.push_back({std::size_t(graph.edgeIndex[2 * e]),
                      std::size_t(graph.edgeIndex[2 * e + 1]),
                      l,
                      InferenceUtils::sigmoid(l)});
  }
 
  ATH_CHECK(dumpDebugEvent(ctx, graph, scores));
  return StatusCode::SUCCESS;
}

DeclareInterfaceID()

MuonML::ISegmentEdgeClassifierTool::DeclareInterfaceID ( ISegmentEdgeClassifierTool , 1 , 0  )

inherited

dumpDebugEvent()

StatusCode MuonML::SegmentEdgeClassifierTool::dumpDebugEvent ( const EventContext & ctx, const SegmentEdgeGraph & graph, const std::vector< SegmentEdgeScore > & scores ) const

private

Definition at line 822 of file SegmentEdgeClassifierTool.cxx.

                                                     {
  if (m_debugDumpFile.value().empty()) return StatusCode::SUCCESS;
 
  std::lock_guard<std::mutex> lock{m_debugDumpMutex};
  if (m_debugDumpMaxEvents.value() != 0 &&
      m_debugDumpEvents.load(std::memory_order_relaxed) >=
          m_debugDumpMaxEvents.value()) {
    return StatusCode::SUCCESS;
  }
 
  if (graph.nodeFeatures.size() != graph.nNodes * kNodeFeatureCount ||
      graph.edgeIndex.size() != graph.nEdges * 2 ||
      graph.edgeFeatures.size() != graph.nEdges * kEdgeFeatureCount ||
      scores.size() != graph.nEdges) {
    ATH_MSG_ERROR("Cannot write segment-edge debug dump: inconsistent graph/output sizes"
                  << " nodes=" << graph.nNodes
                  << " nodeFeatures=" << graph.nodeFeatures.size()
                  << " edges=" << graph.nEdges
                  << " edgeIndex=" << graph.edgeIndex.size()
                  << " edgeFeatures=" << graph.edgeFeatures.size()
                  << " scores=" << scores.size());
    return StatusCode::FAILURE;
  }
 
  nlohmann::json x = nlohmann::json::array();
  x.get_ref<nlohmann::json::array_t&>().reserve(graph.nodeFeatures.size());
  for (const float value : graph.nodeFeatures) {
    x.push_back(std::isfinite(value) ? nlohmann::json(value)
                                     : nlohmann::json(nullptr));
  }
 
  nlohmann::json edgeIndex = nlohmann::json::array();
  edgeIndex.get_ref<nlohmann::json::array_t&>().reserve(graph.nEdges * 2);
  // This is the actual ONNX [2,E] row-major buffer: all sources then all destinations.
  for (std::size_t edge = 0; edge < graph.nEdges; ++edge) {
    edgeIndex.push_back(graph.edgeIndex[2 * edge]);
  }
  for (std::size_t edge = 0; edge < graph.nEdges; ++edge) {
    edgeIndex.push_back(graph.edgeIndex[2 * edge + 1]);
  }
 
  nlohmann::json edgeAttr = nlohmann::json::array();
  edgeAttr.get_ref<nlohmann::json::array_t&>().reserve(graph.edgeFeatures.size());
  for (const float value : graph.edgeFeatures) {
    edgeAttr.push_back(std::isfinite(value) ? nlohmann::json(value)
                                            : nlohmann::json(nullptr));
  }
 
  nlohmann::json logits = nlohmann::json::array();
  nlohmann::json probabilities = nlohmann::json::array();
  nlohmann::json edgeSrc = nlohmann::json::array();
  nlohmann::json edgeDst = nlohmann::json::array();
  logits.get_ref<nlohmann::json::array_t&>().reserve(scores.size());
  probabilities.get_ref<nlohmann::json::array_t&>().reserve(scores.size());
  edgeSrc.get_ref<nlohmann::json::array_t&>().reserve(scores.size());
  edgeDst.get_ref<nlohmann::json::array_t&>().reserve(scores.size());
  for (const SegmentEdgeScore& score : scores) {
    edgeSrc.push_back(score.src);
    edgeDst.push_back(score.dst);
    logits.push_back(std::isfinite(score.logit) ? nlohmann::json(score.logit)
                                                : nlohmann::json(nullptr));
    probabilities.push_back(std::isfinite(score.probability)
                                ? nlohmann::json(score.probability)
                                : nlohmann::json(nullptr));
  }
 
  std::ofstream out{m_debugDumpFile.value(), std::ios::out | std::ios::app};
  if (!out) {
    ATH_MSG_ERROR("Could not append to segment-edge debug dump file: "
                  << m_debugDumpFile.value());
    return StatusCode::FAILURE;
  }
 
  const unsigned int dumpIndex =
      m_debugDumpEvents.fetch_add(1, std::memory_order_relaxed);
  nlohmann::ordered_json event;
  event["record_type"] = "event";
  event["format_version"] = 1;
  event["dump_index"] = dumpIndex;
  event["run_number"] = ctx.eventID().run_number();
  event["lumi_block"] = ctx.eventID().lumi_block();
  event["event_number"] = ctx.eventID().event_number();
  event["slot"] = ctx.slot();
  event["n_nodes"] = graph.nNodes;
  event["n_edges"] = graph.nEdges;
  event["x_shape"] = {graph.nNodes, kNodeFeatureCount};
  event["edge_index_shape"] = {2, graph.nEdges};
  event["edge_attr_shape"] = {graph.nEdges, kEdgeFeatureCount};
  event["logits_shape"] = {graph.nEdges};
  event["x"] = std::move(x);
  event["edge_index"] = std::move(edgeIndex);
  event["edge_attr"] = std::move(edgeAttr);
  event["edge_src"] = std::move(edgeSrc);
  event["edge_dst"] = std::move(edgeDst);
  event["logits"] = std::move(logits);
  event["probabilities"] = std::move(probabilities);
  out << event.dump() << '\n';
 
  ATH_MSG_DEBUG("Wrote segment-edge debug event " << dumpIndex
                << " to " << m_debugDumpFile.value());
 
  return StatusCode::SUCCESS;
}

enableTruthDiagnostics()

bool MuonML::SegmentEdgeClassifierTool::enableTruthDiagnostics ( ) const

inlineoverridevirtual

Whether this tool was configured to fill SegmentEdgeGraph's truth diagnostics.

Implements MuonML::ISegmentEdgeClassifierTool.

Definition at line 93 of file SegmentEdgeClassifierTool.h.

{ return m_enableTruthDiagnostics; }

fillTruthDiagnostics()

void MuonML::SegmentEdgeClassifierTool::fillTruthDiagnostics ( const xAOD::MuonSegmentContainer & segments, const std::unordered_set< const xAOD::MuonSegment * > & bucketRetained, SegmentEdgeGraph & graph ) const

private

MC-only diagnostics: record, for every input segment and every pair of input segments sharing a truth particle, why it did or did not reach ONNX.

Only called when EnableTruthDiagnostics and DEBUG output are set.

Definition at line 611 of file SegmentEdgeClassifierTool.cxx.

                                   {
  std::unordered_map<const xAOD::MuonSegment*, std::int32_t> nodeOf;
  nodeOf.reserve(graph.segments.size());
  for (std::size_t node = 0; node < graph.segments.size(); ++node) {
    nodeOf.emplace(graph.segments[node], static_cast<std::int32_t>(node));
  }
 
  graph.inputNodeIndex.assign(segments.size(), kDroppedAsIsolated);
  std::unordered_map<std::int32_t, std::vector<std::uint32_t>> byTruth;
  std::uint32_t inputIndex = 0;
  for (const xAOD::MuonSegment* segment : segments) {
    const auto found = nodeOf.find(segment);
    if (found != nodeOf.end()) {
      graph.inputNodeIndex[inputIndex] = found->second;
    } else if (!bucketRetained.contains(segment)) {
      graph.inputNodeIndex[inputIndex] = kDroppedByBucketCap;
    }
    if (const xAOD::TruthParticle* truthPart =
            MuonR4::getTruthMatchedParticle(*segment)) {
      byTruth[static_cast<std::int32_t>(truthPart->index())].push_back(inputIndex);
    }
    ++inputIndex;
  }
 
  const auto pairKey = [](std::uint64_t first, std::uint64_t second) {
    return first < second ? (first << 32) | second : (second << 32) | first;
  };
  std::unordered_set<std::uint64_t> scoredPairs;
  scoredPairs.reserve(graph.nEdges);
  for (std::size_t edge = 0; edge < graph.nEdges; ++edge) {
    scoredPairs.insert(
        pairKey(static_cast<std::uint64_t>(graph.edgeIndex[2 * edge]),
                static_cast<std::uint64_t>(graph.edgeIndex[2 * edge + 1])));
  }
 
  // Same order as the pair loop in buildGraph(): bucket cap (node level),
  // sector window, same-chamber drop, angle window; a pair that clears all of
  // them but was not scored can only have been evicted by the edge caps.
  for (const auto& entry : byTruth) {
    const std::vector<std::uint32_t>& members = entry.second;
    for (std::size_t x = 0; x < members.size(); ++x) {
      for (std::size_t y = x + 1; y < members.size(); ++y) {
        const std::uint32_t i = members[x];
        const std::uint32_t j = members[y];
        const xAOD::MuonSegment* first = segments[i];
        const xAOD::MuonSegment* second = segments[j];
        PairFate fate{i, j, PairGate::Scored, 0};
        const int sectorDelta = sectorDistance(first->sector(), second->sector(),
                                               m_sectorModulo.value());
        fate.sectorDelta = static_cast<std::uint8_t>(std::min(sectorDelta, 255));
        if (!bucketRetained.contains(first) || !bucketRetained.contains(second)) {
          fate.gate = PairGate::BucketCap;
        } else if (sectorDelta > m_maxDeltaSector.value()) {
          fate.gate = PairGate::SectorWindow;
        } else if (m_dropSameChamberEdgesBeforeInference.value() &&
                   first->chamberIndex() == second->chamberIndex()) {
          fate.gate = PairGate::SameChamber;
        } else if (static_cast<float>(first->direction().dot(second->direction())) <
                   m_cosMin) {
          fate.gate = PairGate::AngleWindow;
        } else {
          const std::int32_t firstNode = graph.inputNodeIndex[i];
          const std::int32_t secondNode = graph.inputNodeIndex[j];
          const bool scored =
              firstNode >= 0 && secondNode >= 0 &&
              scoredPairs.contains(pairKey(static_cast<std::uint64_t>(firstNode),
                                           static_cast<std::uint64_t>(secondNode)));
          if (!scored) fate.gate = PairGate::EdgeCaps;
        }
        graph.truthPairFates.push_back(fate);
      }
    }
  }
}

finalize()

StatusCode MuonML::SegmentEdgeClassifierTool::finalize ( )

override

Log a pre-ONNX candidate-edge pruning.

Definition at line 689 of file SegmentEdgeClassifierTool.cxx.

                                               {
  ATH_MSG_DEBUG(
      "SegmentEdgeClassifierTool pre-ONNX pruning summary (job-summed, "
      "independent of PairGateThreshold): "
      << "inputSegments=" << m_sumInputSegments
      << ", candidatePairs=" << m_sumCandidatePairs
      << " (geometric pairs before MaxEdgesPerNodeBeforeInference/"
         "MaxEdgesPerTargetChamberBeforeInference caps)"
      << ", retainedPairs=" << m_sumRetainedPairs
      << " (pairs actually sent to ONNX; MaxEdgesPerNodeBeforeInference="
      << m_maxEdgesPerNodeBeforeInference.value()
      << ", MaxEdgesPerTargetChamberBeforeInference="
      << m_maxEdgesPerTargetChamberBeforeInference.value() << ")"
      << ", nodesBeforeIsolatedDrop=" << m_sumNodesBeforeIsolatedDrop
      << ", nodesAfterIsolatedDrop=" << m_sumNodesAfterIsolatedDrop
      << " (DropIsolatedNodesBeforeInference="
      << m_dropIsolatedNodesBeforeInference.value()
      << "; nodes dropped here never reached ONNX or the pair-gate threshold)");
  return StatusCode::SUCCESS;
}

initialize()

StatusCode MuonML::SegmentEdgeClassifierTool::initialize ( )

override

Retrieve the ONNX model and resolve node feature ordering from metadata.

Definition at line 101 of file SegmentEdgeClassifierTool.cxx.

                                                 {
  if (m_sectorModulo.value() > 0 &&
      2ULL * static_cast<unsigned long long>(m_maxDeltaSector.value()) + 1ULL >
          static_cast<unsigned long long>(m_sectorModulo.value())) {
    ATH_MSG_ERROR("MaxDeltaSector=" << m_maxDeltaSector.value()
                  << " spans duplicate sectors for SectorModulo="
                  << m_sectorModulo.value());
    return StatusCode::FAILURE;
  }
  ATH_CHECK(setupModel());
 
  // Resolve node feature names from model metadata, matching the ONNX exporter.
  {
    Ort::AllocatorWithDefaultOptions allocator;
    Ort::ModelMetadata meta = model().GetModelMetadata();
    auto keys = meta.GetCustomMetadataMapKeysAllocated(allocator);
    std::vector<std::string> keyList;
    keyList.reserve(keys.size());
    for (const auto& k : keys) keyList.emplace_back(k.get());
 
    constexpr std::array<std::string_view, 4> candidates{
        "x_feature_names", "node_feature_names", "feature_names", "input_feature_names"};
    std::string usedKey;
    std::vector<std::string> names;
    for (std::string_view key : candidates) {
      const std::string keyStr{key};
      if (std::find(keyList.begin(), keyList.end(), keyStr) == keyList.end()) continue;
      names = parseFeatureNames(meta.LookupCustomMetadataMapAllocated(keyStr.c_str(), allocator).get());
      if (!names.empty()) {
        usedKey = keyStr;
        break;
      }
    }
 
    if (names.empty()) {
      m_nodeFeatureNames.assign(kDefaultNodeFeatureNames.begin(), kDefaultNodeFeatureNames.end());
      ATH_MSG_WARNING("Model metadata has no usable node feature name key"
                      " (tried x_feature_names/node_feature_names/feature_names/input_feature_names)."
                      " Falling back to default training order.");
    } else {
      if (names.size() != kNodeFeatureCount) {
        ATH_MSG_ERROR("Model metadata key '" << usedKey << "' has " << names.size()
                      << " features, expected " << kNodeFeatureCount);
        return StatusCode::FAILURE;
      }
      for (const std::string& n : names) {
        if (!nodeFeatureIdFromName(n).has_value()) {
          ATH_MSG_ERROR("Unsupported node feature name in model metadata ('" << usedKey
                        << "'): '" << n << "'."
                        " Add mapping in SegmentEdgeClassifierTool::nodeFeatureValue().");
          return StatusCode::FAILURE;
        }
      }
      m_nodeFeatureNames = std::move(names);
      ATH_MSG_DEBUG("Using node feature names from model metadata key '" << usedKey << "'.");
    }
 
    m_nodeFeatureIds.reserve(m_nodeFeatureNames.size());
    for (const std::string& n : m_nodeFeatureNames) {
      const auto id = nodeFeatureIdFromName(n);
      if (!id.has_value()) {
        ATH_MSG_ERROR("Internal feature-id resolution failed for node feature name '" << n << "'.");
        return StatusCode::FAILURE;
      }
      m_nodeFeatureIds.push_back(*id);
    }
 
    std::ostringstream order;
    order << "Node feature order:";
    for (std::size_t i = 0; i < m_nodeFeatureNames.size(); ++i) {
      order << " f" << i << "=" << m_nodeFeatureNames[i];
      if (i + 1 < m_nodeFeatureNames.size()) order << ",";
    }
    ATH_MSG_DEBUG(order.str());
  }
 
  if (m_nodeFeatureNames.size() != kNodeFeatureCount) {
    ATH_MSG_ERROR("Internal node feature setup has " << m_nodeFeatureNames.size()
                  << " entries, expected " << kNodeFeatureCount);
    return StatusCode::FAILURE;
  }
  if (m_nodeFeatureIds.size() != kNodeFeatureCount) {
    ATH_MSG_ERROR("Internal node feature id setup has " << m_nodeFeatureIds.size()
                  << " entries, expected " << kNodeFeatureCount);
    return StatusCode::FAILURE;
  }
 
  m_cosMin = std::cos(m_maxDeltaThetaDeg.value() * Gaudi::Units::deg);
 
  if (!m_debugDumpFile.value().empty()) {
    std::ofstream out{m_debugDumpFile.value(), std::ios::out | std::ios::trunc};
    if (!out) {
      ATH_MSG_ERROR("Could not create segment-edge debug dump file: "
                    << m_debugDumpFile.value());
      return StatusCode::FAILURE;
    }
 
    nlohmann::ordered_json metadata;
    metadata["record_type"] = "metadata";
    metadata["format_version"] = 1;
    metadata["tool"] = "SegmentEdgeClassifierTool";
    metadata["input_names"] = {m_inputNodeName.value(),
                               m_inputEdgeIndexName.value(),
                               m_inputEdgeAttrName.value()};
    metadata["output_name"] = m_outputName.value();
    metadata["x_feature_names"] = m_nodeFeatureNames;
    metadata["edge_attr_feature_names"] = {
        "deltaPositionX_m", "deltaPositionY_m", "deltaPositionZ_m",
        "distance_m", "cos_opening_angle", "same_chamber", "same_sector"};
    metadata["edge_index_layout"] = "row_major_2_by_E";
    metadata["edge_order"] = "directed src_to_dst; row 0 then row 1";
    metadata["max_delta_theta_deg"] = m_maxDeltaThetaDeg.value();
    metadata["max_delta_sector"] = m_maxDeltaSector.value();
    metadata["sector_modulo"] = m_sectorModulo.value();
    metadata["debug_dump_max_events"] = m_debugDumpMaxEvents.value();
    out << metadata.dump() << '\n';
 
    ATH_MSG_INFO("Writing segment-edge ONNX debug dump to "
                 << m_debugDumpFile.value()
                 << " (DebugDumpMaxEvents="
                 << m_debugDumpMaxEvents.value() << ")");
  }
 
  ATH_CHECK(m_truthLinkKey.initialize(m_enableTruthDiagnostics.value()));
 
  return StatusCode::SUCCESS;
}

model()

Ort::Session & BucketInferenceToolBase::model ( ) const

protectedinherited

Definition at line 65 of file BucketInferenceToolBase.cxx.

                                                 {
  return m_onnxSessionTool->session();
}

parseFeatureNames()

std::vector< std::string > BucketInferenceToolBase::parseFeatureNames ( const std::string & raw )

staticprotectedinherited

Definition at line 32 of file BucketInferenceToolBase.cxx.

                                                                                    {
  std::vector<std::string> out;
  const std::string s = trimFeatureToken(raw);
  if (s.empty()) return out;
 
  // Preferred exporter format: JSON list of strings.
  if (!s.empty() && s.front() == '[') {
    bool inQuote = false;
    std::string token;
    for (char c : s) {
      if (c == '"') {
        if (inQuote) {
          if (!token.empty()) out.push_back(token);
          token.clear();
        }
        inQuote = !inQuote;
        continue;
      }
      if (inQuote) token.push_back(c);
    }
    if (!out.empty()) return out;
  }
 
  // Backward-compatible format: comma-separated.
  std::istringstream ss(s);
  std::string tok;
  while (std::getline(ss, tok, ',')) {
    tok = trimFeatureToken(tok);
    if (!tok.empty()) out.push_back(tok);
  }
  return out;
}

runGraphInference()

StatusCode MuonML::SegmentEdgeClassifierTool::runGraphInference ( const EventContext & ctx, GraphRawData & graphData ) const

override

Not supported by this tool; returns FAILURE.

Use SegmentEdgeInferenceAlg + buildGraph() + classifyEdges() instead.

Definition at line 229 of file SegmentEdgeClassifierTool.cxx.

                                                                                                {
  ATH_MSG_ERROR("runGraphInference is not supported by SegmentEdgeClassifierTool. Use SegmentEdgeInferenceAlg + ISegmentEdgeClassifierTool methods.");
  return StatusCode::FAILURE;
}

runInference()

StatusCode BucketInferenceToolBase::runInference ( GraphRawData & graphData ) const

inherited

Default ONNX run for GNN case: inputs {"features","edge_index"} -> outputs {"logits"}.

Definition at line 532 of file BucketInferenceToolBase.cxx.

                                                                              {
  std::vector<const char*> inputNames  = {"features", "edge_index"};
  std::vector<const char*> outputNames = {m_outputName.value().c_str()};
  return runNamedInference(graphData, inputNames, outputNames);
}

runNamedInference()

StatusCode BucketInferenceToolBase::runNamedInference ( GraphRawData & graphData, const std::vector< const char * > & inputNames, const std::vector< const char * > & outputNames ) const

protectedinherited

Generic named inference, for tools with different I/O conventions.

Definition at line 359 of file BucketInferenceToolBase.cxx.

{
  if (!graphData.graph) {
    ATH_MSG_ERROR("Graph data is not built.");
    return StatusCode::FAILURE;
  }
  if (graphData.graph->dataTensor.empty()) {
    ATH_MSG_ERROR("No input tensors prepared for inference.");
    return StatusCode::FAILURE;
  }
 
  // Reserve the final size here from the actual I/O lists instead 
  // of hard-coding assumptions in the graph builders.
  graphData.graph->dataTensor.reserve(inputNames.size() + outputNames.size());
  if (graphData.graph->dataTensor.size() < inputNames.size()) {
    ATH_MSG_ERROR("Prepared " << graphData.graph->dataTensor.size()
                  << " tensors but inference expects " << inputNames.size() << " inputs.");
    return StatusCode::FAILURE;
  }
 
  if (msgLvl(MSG::DEBUG)) {
    // DEBUG: Print actual input tensor data for features tensor
 
    ATH_MSG_DEBUG("=== DEBUGGING: ONNX Input tensor data ===");
    if (!graphData.graph->dataTensor.empty()) {
      const auto& featureTensor = graphData.graph->dataTensor[0];
      auto featShape = featureTensor.GetTensorTypeAndShapeInfo().GetShape();
      ATH_MSG_DEBUG("Features tensor shape: [" << featShape[0] 
                   << (featShape.size()>1 ? ("," + std::to_string(featShape[1])) : "")
                   << (featShape.size()>2 ? ("," + std::to_string(featShape[2])) : "") << "]");
      
      float* featData = const_cast<Ort::Value&>(featureTensor).GetTensorMutableData<float>();
      const size_t totalElements = featureTensor.GetTensorTypeAndShapeInfo().GetElementCount();
      ATH_MSG_DEBUG("Features tensor total elements: " << totalElements);
      
      // Print up to 10 nodes; stride = nFeat from tensor shape
      const size_t nFeat = (featShape.size() > 1 && featShape[1] > 0) ? static_cast<size_t>(featShape[1]) : 1;
      const size_t nNodes = totalElements / nFeat;
      const size_t debugNodes = std::min(nNodes, static_cast<size_t>(10));
 
      // Try to read feature names from model custom metadata.
      // Prefer x_feature_names (current exporter), then fall back to legacy keys.
      std::vector<std::string> featNames;
      {
        Ort::AllocatorWithDefaultOptions allocator;
        Ort::ModelMetadata meta = model().GetModelMetadata();
        auto keys = meta.GetCustomMetadataMapKeysAllocated(allocator);
        std::vector<std::string> keyNames;
        keyNames.reserve(keys.size());
        for (const auto& k : keys) keyNames.emplace_back(k.get());
        const std::array<std::string, 4> candidates{
            "x_feature_names", "node_feature_names", "feature_names", "input_feature_names"};
        for (const std::string& key : candidates) {
          if (std::find(keyNames.begin(), keyNames.end(), key) != keyNames.end()) {
            std::string val = meta.LookupCustomMetadataMapAllocated(key.c_str(), allocator).get();
            featNames = parseFeatureNames(val);
            break;
          }
        }
        if (featNames.empty()) {
          ATH_MSG_DEBUG("No usable feature-name metadata key found in model; using generic fN labels.");
        }
      }
      auto featLabel = [&](size_t f) -> std::string {
        if (f < featNames.size()) return featNames[f];
        return "f" + std::to_string(f);
      };
 
      // Print legend
      {
        std::ostringstream legend;
        legend << "Node feature legend (" << nFeat << " features):";
        for (size_t f = 0; f < nFeat; ++f) {
          legend << " f" << f << "=" << featLabel(f);
          if (f + 1 < nFeat) legend << ",";
        }
        ATH_MSG_DEBUG(legend.str());
      }
 
      for (size_t n = 0; n < debugNodes; ++n) {
        std::ostringstream row;
        row << "ONNXNode[" << n << "]:";
        for (size_t f = 0; f < nFeat; ++f) {
          row << " f" << f << "=" << featData[n * nFeat + f];
          if (f + 1 < nFeat) row << ",";
        }
        ATH_MSG_DEBUG(row.str());
      }
    }
    ATH_MSG_DEBUG("=== END DEBUG ONNX INPUT ===");
  }
 
  Ort::RunOptions run_options;
  run_options.SetRunLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);
 
  if (m_isCuda) {
    // ---- CUDA path: use IoBinding so tensors stay on device ----
    Ort::IoBinding binding(model());
    for (std::size_t i = 0; i < inputNames.size(); ++i) {
      binding.BindInput(inputNames[i], graphData.graph->dataTensor[i]);
    }
    // Bind outputs to CPU so predictions are directly readable after sync.
    Ort::MemoryInfo cpuOut = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    for (const char* outName : outputNames) {
      binding.BindOutput(outName, cpuOut);
    }
 
    model().Run(run_options, binding);
    binding.SynchronizeOutputs();
 
    std::vector<Ort::Value> outputs = binding.GetOutputValues();
    if (outputs.empty()) {
      ATH_MSG_ERROR("IoBinding inference returned empty output.");
      return StatusCode::FAILURE;
    }
 
    float* outData = outputs[0].GetTensorMutableData<float>();
    const size_t outSize = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
    ATH_MSG_DEBUG("ONNX (IoBinding) raw output elementCount = " << outSize);
 
    if (m_sanitizeNonFinitePredictions.value()) {
      std::span<float> preds(outData, outData + outSize);
      for (size_t i = 0; i < outSize; ++i) {
        if (!std::isfinite(preds[i])) {
          ATH_MSG_WARNING("Non-finite prediction detected at " << i << " -> set to -100.");
          preds[i] = -100.0f;
        }
      }
    }
 
    for (auto& v : outputs) {
      graphData.graph->dataTensor.emplace_back(std::move(v));
    }
    return StatusCode::SUCCESS;
  }
 
  // ---- CPU path ----
  std::vector<Ort::Value> outputs =
      model().Run(run_options,
                  inputNames.data(),
                  graphData.graph->dataTensor.data(),
                  inputNames.size(),
                  outputNames.data(),
                  outputNames.size());
 
  if (outputs.empty()) {
    ATH_MSG_ERROR("Inference returned empty output.");
    return StatusCode::FAILURE;
  }
 
  float* outData = outputs[0].GetTensorMutableData<float>();
  const size_t outSize = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
  ATH_MSG_DEBUG("ONNX raw output elementCount = " << outSize);
 
  if (m_sanitizeNonFinitePredictions.value()) {
    std::span<float> preds(outData, outData + outSize);
    for (size_t i = 0; i < outSize; ++i) {
      if (!std::isfinite(preds[i])) {
        ATH_MSG_WARNING("Non-finite prediction detected at " << i << " -> set to -100.");
        preds[i] = -100.0f;
      }
    }
  }
 
  for (auto& v : outputs) {
    graphData.graph->dataTensor.emplace_back(std::move(v));
  }
  return StatusCode::SUCCESS;
}

setupModel()

StatusCode BucketInferenceToolBase::setupModel ( )

protectedinherited

Definition at line 69 of file BucketInferenceToolBase.cxx.

                                               {
  ATH_CHECK(m_onnxSessionTool.retrieve());
  ATH_CHECK(m_readKey.initialize());
  ATH_CHECK(m_geoCtxKey.initialize());
 
  const InferenceUtils::SessionBackend backend = InferenceUtils::sessionBackend(m_onnxSessionTool);
  m_isCuda = backend.isCuda;
  m_cudaDeviceId = backend.cudaDeviceId;
  if (m_isCuda) {
    ATH_MSG_INFO("ONNX session is running on CUDA device " << m_cudaDeviceId
                 << ". I/O binding will be used.");
  } else {
    ATH_MSG_INFO("ONNX session is running on CPU.");
  }
 
  return StatusCode::SUCCESS;
}

trimFeatureToken()

std::string BucketInferenceToolBase::trimFeatureToken ( std::string s )

staticprotectedinherited

Definition at line 25 of file BucketInferenceToolBase.cxx.

                                                               {
  auto notSpace = [](unsigned char c) { return !std::isspace(c); };
  s.erase(s.begin(), std::find_if(s.begin(), s.end(), notSpace));
  s.erase(std::find_if(s.rbegin(), s.rend(), notSpace).base(), s.end());
  return s;
}

Member Data Documentation

kBucketFeatureCount

std::size_t MuonML::BucketInferenceToolBase::kBucketFeatureCount = 6

staticconstexprprotectedinherited

Definition at line 53 of file BucketInferenceToolBase.h.

kDefaultNodeFeatureNames

std::array<std::string_view, kNodeFeatureCount> MuonML::BucketInferenceToolBase::kDefaultNodeFeatureNames

staticconstexprprotectedinherited

Initial value:

                                                 = {
"segmentPositionX_m", "segmentPositionY_m", "segmentPositionZ_m",
"segmentDirectionX", "segmentDirectionY", "segmentDirectionZ",
"bucket_chamberIndex", "bucket_layers", "bucket_sector", "bucket_segments"}

Definition at line 56 of file BucketInferenceToolBase.h.

                                                                                        {
      "segmentPositionX_m", "segmentPositionY_m", "segmentPositionZ_m",
      "segmentDirectionX", "segmentDirectionY", "segmentDirectionZ",
      "bucket_chamberIndex", "bucket_layers", "bucket_sector", "bucket_segments"};

kEdgeFeatureCount

std::size_t MuonML::BucketInferenceToolBase::kEdgeFeatureCount = 7

staticconstexprprotectedinherited

Definition at line 55 of file BucketInferenceToolBase.h.

kNodeFeatureCount

std::size_t MuonML::BucketInferenceToolBase::kNodeFeatureCount = 10

staticconstexprprotectedinherited

Definition at line 54 of file BucketInferenceToolBase.h.

m_cosMin

float MuonML::SegmentEdgeClassifierTool::m_cosMin {0.f}

private

Definition at line 134 of file SegmentEdgeClassifierTool.h.

{0.f};

m_cudaDeviceId

int MuonML::BucketInferenceToolBase::m_cudaDeviceId {0}

protectedinherited

Definition at line 102 of file BucketInferenceToolBase.h.

{0};

m_debugDumpEvents

std::atomic<unsigned int> MuonML::SegmentEdgeClassifierTool::m_debugDumpEvents {0}

mutableprivate

Definition at line 140 of file SegmentEdgeClassifierTool.h.

{0};

m_debugDumpFile

Gaudi::Property<std::string> MuonML::SegmentEdgeClassifierTool::m_debugDumpFile {this, "DebugDumpFile", ""}

private

Definition at line 127 of file SegmentEdgeClassifierTool.h.

{this, "DebugDumpFile", ""};

m_debugDumpFirstNEdges

Gaudi::Property<unsigned int> MuonML::BucketInferenceToolBase::m_debugDumpFirstNEdges {this, "DebugDumpFirstNEdges", 12}

protectedinherited

Definition at line 94 of file BucketInferenceToolBase.h.

{this, "DebugDumpFirstNEdges", 12};

m_debugDumpFirstNNodes

Gaudi::Property<unsigned int> MuonML::BucketInferenceToolBase::m_debugDumpFirstNNodes {this, "DebugDumpFirstNNodes", 5}

protectedinherited

Definition at line 93 of file BucketInferenceToolBase.h.

{this, "DebugDumpFirstNNodes", 5};

m_debugDumpMaxEvents

Gaudi::Property<unsigned int> MuonML::SegmentEdgeClassifierTool::m_debugDumpMaxEvents {this, "DebugDumpMaxEvents", 0}

private

Definition at line 128 of file SegmentEdgeClassifierTool.h.

{this, "DebugDumpMaxEvents", 0};

m_debugDumpMutex

std::mutex MuonML::SegmentEdgeClassifierTool::m_debugDumpMutex

mutableprivate

Definition at line 139 of file SegmentEdgeClassifierTool.h.

m_dropIsolatedNodesBeforeInference

Gaudi::Property<bool> MuonML::SegmentEdgeClassifierTool::m_dropIsolatedNodesBeforeInference

private

Initial value:

{this, "DropIsolatedNodesBeforeInference", true,
        "Remove nodes without a retained pre-ONNX edge before creating ONNX tensors"}

Definition at line 121 of file SegmentEdgeClassifierTool.h.

                                                          {this, "DropIsolatedNodesBeforeInference", true,
        "Remove nodes without a retained pre-ONNX edge before creating ONNX tensors"};

m_dropSameChamberEdgesBeforeInference

Gaudi::Property<bool> MuonML::SegmentEdgeClassifierTool::m_dropSameChamberEdgesBeforeInference

private

Initial value:

{this, "DropSameChamberEdgesBeforeInference", true,
        "Drop same-chamber segment pairs before ONNX inference"}

Definition at line 119 of file SegmentEdgeClassifierTool.h.

                                                             {this, "DropSameChamberEdgesBeforeInference", true,
        "Drop same-chamber segment pairs before ONNX inference"};

m_enableTruthDiagnostics

Gaudi::Property<bool> MuonML::SegmentEdgeClassifierTool::m_enableTruthDiagnostics

private

Initial value:

{
        this, "EnableTruthDiagnostics", false,
        "MC-only: fill SegmentEdgeGraph's truth diagnostics."}

Definition at line 129 of file SegmentEdgeClassifierTool.h.

                                                {
        this, "EnableTruthDiagnostics", false,
        "MC-only: fill SegmentEdgeGraph's truth diagnostics."};

m_geoCtxKey

ActsTrk::GeoContextReadKey_t MuonML::BucketInferenceToolBase::m_geoCtxKey {this, "AlignmentKey", "ActsAlignment", "cond handle key"}

protectedinherited

Definition at line 80 of file BucketInferenceToolBase.h.

{this, "AlignmentKey", "ActsAlignment", "cond handle key"};

m_inputEdgeAttrName

Gaudi::Property<std::string> MuonML::SegmentEdgeClassifierTool::m_inputEdgeAttrName {this, "InputEdgeAttrName", "edge_attr"}

private

Definition at line 125 of file SegmentEdgeClassifierTool.h.

{this, "InputEdgeAttrName", "edge_attr"};

m_inputEdgeIndexName

Gaudi::Property<std::string> MuonML::SegmentEdgeClassifierTool::m_inputEdgeIndexName {this, "InputEdgeIndexName", "edge_index"}

private

Definition at line 124 of file SegmentEdgeClassifierTool.h.

{this, "InputEdgeIndexName", "edge_index"};

m_inputNodeName

Gaudi::Property<std::string> MuonML::SegmentEdgeClassifierTool::m_inputNodeName {this, "InputNodeName", "x"}

private

Definition at line 123 of file SegmentEdgeClassifierTool.h.

{this, "InputNodeName", "x"};

m_isCuda

bool MuonML::BucketInferenceToolBase::m_isCuda {false}

protectedinherited

Definition at line 101 of file BucketInferenceToolBase.h.

{false};

m_maxAbsDz

Gaudi::Property<double> MuonML::BucketInferenceToolBase::m_maxAbsDz {this, "MaxAbsDz", 15000.0}

protectedinherited

Definition at line 90 of file BucketInferenceToolBase.h.

{this, "MaxAbsDz", 15000.0};

m_maxChamberDelta

Gaudi::Property<int> MuonML::BucketInferenceToolBase::m_maxChamberDelta {this, "MaxChamberDelta", 13}

protectedinherited

Definition at line 87 of file BucketInferenceToolBase.h.

{this, "MaxChamberDelta", 13};

m_maxDeltaSector

Gaudi::Property<int> MuonML::SegmentEdgeClassifierTool::m_maxDeltaSector {this, "MaxDeltaSector", 1}

private

Definition at line 109 of file SegmentEdgeClassifierTool.h.

{this, "MaxDeltaSector", 1};

m_maxDeltaThetaDeg

Gaudi::Property<float> MuonML::SegmentEdgeClassifierTool::m_maxDeltaThetaDeg {this, "MaxDeltaThetaDeg", 35.f}

private

Definition at line 108 of file SegmentEdgeClassifierTool.h.

{this, "MaxDeltaThetaDeg", 35.f};

m_maxDistXY

Gaudi::Property<double> MuonML::BucketInferenceToolBase::m_maxDistXY {this, "MaxDistXY", 6800.0}

protectedinherited

Definition at line 89 of file BucketInferenceToolBase.h.

{this, "MaxDistXY", 6800.0};

m_maxEdgesPerNodeBeforeInference

Gaudi::Property<unsigned int> MuonML::SegmentEdgeClassifierTool::m_maxEdgesPerNodeBeforeInference

private

Initial value:

{this, "MaxEdgesPerNodeBeforeInference", 0,
        "Keep at most this many geometrical neighbour pairs per node before ONNX inference; 0 keeps all"}

Definition at line 114 of file SegmentEdgeClassifierTool.h.

                                                                {this, "MaxEdgesPerNodeBeforeInference", 0,
        "Keep at most this many geometrical neighbour pairs per node before ONNX inference; 0 keeps all"};

m_maxEdgesPerTargetChamberBeforeInference

Gaudi::Property<unsigned int> MuonML::SegmentEdgeClassifierTool::m_maxEdgesPerTargetChamberBeforeInference

private

Initial value:

{
        this, "MaxEdgesPerTargetChamberBeforeInference", 0,
        "Keep at most this many pre-ONNX neighbours from one target chamber per node; 0 keeps all"}

Definition at line 116 of file SegmentEdgeClassifierTool.h.

                                                                         {
        this, "MaxEdgesPerTargetChamberBeforeInference", 0,
        "Keep at most this many pre-ONNX neighbours from one target chamber per node; 0 keeps all"};

m_maxSectorDelta

Gaudi::Property<int> MuonML::BucketInferenceToolBase::m_maxSectorDelta {this, "MaxSectorDelta", 1}

protectedinherited

Definition at line 88 of file BucketInferenceToolBase.h.

{this, "MaxSectorDelta", 1};

m_maxSegmentsPerBucket

Gaudi::Property<unsigned int> MuonML::SegmentEdgeClassifierTool::m_maxSegmentsPerBucket

private

Initial value:

{this, "MaxSegmentsPerBucket", 0,
        "Keep at most this many quality-ranked segments per (sector, chamber, eta) bucket before inference; 0 keeps all"}

Definition at line 112 of file SegmentEdgeClassifierTool.h.

                                                      {this, "MaxSegmentsPerBucket", 0,
        "Keep at most this many quality-ranked segments per (sector, chamber, eta) bucket before inference; 0 keeps all"};

m_minLayers

Gaudi::Property<int> MuonML::BucketInferenceToolBase::m_minLayers {this, "MinLayersValid", 3}

protectedinherited

Definition at line 86 of file BucketInferenceToolBase.h.

{this, "MinLayersValid", 3};

m_nodeFeatureIds

std::vector<SegmentNodeFeatureId> MuonML::SegmentEdgeClassifierTool::m_nodeFeatureIds {}

private

Definition at line 138 of file SegmentEdgeClassifierTool.h.

{};

m_nodeFeatureNames

std::vector<std::string> MuonML::SegmentEdgeClassifierTool::m_nodeFeatureNames {}

private

Node feature order expected by the model metadata (resolved at initialize).

Definition at line 137 of file SegmentEdgeClassifierTool.h.

{};

m_onnxSessionTool

ToolHandle<AthOnnx::IOnnxRuntimeSessionTool> MuonML::BucketInferenceToolBase::m_onnxSessionTool

privateinherited

Initial value:

{
      this, "ModelSession", ""}

Definition at line 105 of file BucketInferenceToolBase.h.

                                                              {
      this, "ModelSession", ""};

m_outputName

Gaudi::Property<std::string> MuonML::SegmentEdgeClassifierTool::m_outputName {this, "OutputName", "logits"}

private

Definition at line 126 of file SegmentEdgeClassifierTool.h.

{this, "OutputName", "logits"};

m_readKey

SG::ReadHandleKey<MuonR4::SpacePointContainer> MuonML::BucketInferenceToolBase::m_readKey {this, "ReadSpacePoints", "MuonSpacePoints"}

protectedinherited

Definition at line 79 of file BucketInferenceToolBase.h.

{this, "ReadSpacePoints", "MuonSpacePoints"};

m_sanitizeNonFinitePredictions

Gaudi::Property<bool> MuonML::BucketInferenceToolBase::m_sanitizeNonFinitePredictions

protectedinherited

Initial value:

{
      this, "SanitizeNonFinitePredictions", false,
      "When true, replace non-finite ONNX outputs with -100 and log a warning."}

Definition at line 96 of file BucketInferenceToolBase.h.

                                                            {
      this, "SanitizeNonFinitePredictions", false,
      "When true, replace non-finite ONNX outputs with -100 and log a warning."};

m_sectorModulo

Gaudi::Property<int> MuonML::SegmentEdgeClassifierTool::m_sectorModulo

private

Initial value:

{this, "SectorModulo", 16,
        "Number of muon sectors used when applying wrap-around sector distance"}

Definition at line 110 of file SegmentEdgeClassifierTool.h.

                                     {this, "SectorModulo", 16,
        "Number of muon sectors used when applying wrap-around sector distance"};

m_sumCandidatePairs

std::atomic<std::size_t> MuonML::SegmentEdgeClassifierTool::m_sumCandidatePairs {0}

mutableprivate

Definition at line 144 of file SegmentEdgeClassifierTool.h.

{0};

m_sumInputSegments

std::atomic<std::size_t> MuonML::SegmentEdgeClassifierTool::m_sumInputSegments {0}

mutableprivate

Job-summed pre-ONNX pruning counters (see buildGraph()).

Definition at line 143 of file SegmentEdgeClassifierTool.h.

{0};

m_sumNodesAfterIsolatedDrop

std::atomic<std::size_t> MuonML::SegmentEdgeClassifierTool::m_sumNodesAfterIsolatedDrop {0}

mutableprivate

Definition at line 147 of file SegmentEdgeClassifierTool.h.

{0};

m_sumNodesBeforeIsolatedDrop

std::atomic<std::size_t> MuonML::SegmentEdgeClassifierTool::m_sumNodesBeforeIsolatedDrop {0}

mutableprivate

Definition at line 146 of file SegmentEdgeClassifierTool.h.

{0};

m_sumRetainedPairs

std::atomic<std::size_t> MuonML::SegmentEdgeClassifierTool::m_sumRetainedPairs {0}

mutableprivate

Definition at line 145 of file SegmentEdgeClassifierTool.h.

{0};

m_truthLinkKey

SG::ReadDecorHandleKey<xAOD::MuonSegmentContainer> MuonML::SegmentEdgeClassifierTool::m_truthLinkKey

private

Initial value:

{
        this, "TruthLinkKey", "MuonSegmentsFromR4.truthParticleLink"}

Definition at line 132 of file SegmentEdgeClassifierTool.h.

                                                                 {
        this, "TruthLinkKey", "MuonSegmentsFromR4.truthParticleLink"};

m_validateEdges

Gaudi::Property<bool> MuonML::BucketInferenceToolBase::m_validateEdges {this, "ValidateEdges", true}

protectedinherited

Definition at line 95 of file BucketInferenceToolBase.h.

{this, "ValidateEdges", true};

---

The documentation for this class was generated from the following files:

• SegmentEdgeClassifierTool.h
• SegmentEdgeClassifierTool.cxx