SegmentEdgeInferenceAlg.cxx

Go to the documentation of this file.

/*
    Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
#include "SegmentEdgeInferenceAlg.h"
#include "InferenceUtils.h"
#include "AthContainers/ConstDataVector.h"
#include "StoreGate/ReadHandle.h"
#include "StoreGate/WriteHandle.h"
#include "xAODMuonViews/ContainerDecorator.h"
#include "MuonStationIndex/MuonStationIndex.h"
#include "MuonTruthHelpers/MuonSimHitHelpers.h"
#include "Acts/Utilities/Helpers.hpp"
#include <algorithm>
#include <bit>
#include <cmath>
#include <limits>
#include <memory>
#include <numeric>
#include <sstream>
#include <unordered_set>
#include <utility>
 
namespace MuonML {
 
namespace {
 
std::uint64_t
undirectedPairKey(std::size_t first, std::size_t second) {
  if (first > second) std::swap(first, second);
  const auto first32 = static_cast<std::uint32_t>(first);
  const auto second32 = static_cast<std::uint32_t>(second);
 
  return (static_cast<std::uint64_t>(first32) << 32) | static_cast<std::uint64_t>(second32);
}
 
std::pair<std::size_t, std::size_t>
unpackPairKey(std::uint64_t key) {
  const auto first = static_cast<std::uint32_t>(key >> 32);
  const auto second = static_cast<std::uint32_t>(key);
  return {first, second};
}
 
class DisjointSet {
public:
  explicit DisjointSet(std::size_t size) : m_parent(size), m_rank(size, 0) {
    std::iota(m_parent.begin(), m_parent.end(), 0);
  }
 
  std::size_t find(std::size_t node) {
    if (m_parent[node] != node) m_parent[node] = find(m_parent[node]);
    return m_parent[node];
  }
 
  void unite(std::size_t first, std::size_t second) {
    first = find(first);
    second = find(second);
    if (first == second) return;
    if (m_rank[first] < m_rank[second]) std::swap(first, second);
    m_parent[second] = first;
    if (m_rank[first] == m_rank[second]) ++m_rank[first];
  }
 
private:
  std::vector<std::size_t> m_parent;
  std::vector<unsigned char> m_rank;
};
 
// Explicit radial ranking (0 = closest to the IP), independent of the
// LayerIndex enum's underlying integer values. BarrelExtended is the
// barrel's BI-equivalent chamber, so it ranks alongside Inner.
int layerRadialRank(Muon::MuonStationIndex::LayerIndex layer) {
  using enum Muon::MuonStationIndex::LayerIndex;
  switch (layer) {
    case Inner:
    case BarrelExtended:
      return 0;
    case Middle:
      return 1;
    case Outer:
      return 2;
    case Extended:
      return 3;
    default:
      return 4;
  }
}
 
// Truth diagnostics only. Why a node that passed the threshold was not kept.
constexpr unsigned char kReasonNotSelected = 0;        // no selected pair (top-K / mutual gate)
constexpr unsigned char kReasonChamberDedup = 1;       // KeepBestSegmentPerChamber
constexpr unsigned char kReasonComponentRejected = 2;  // MinSegmentsPerComponent / no anchor
 
// Loss categories; the order matches the table printed by finalize().
enum LossCategory : std::uint8_t{
  kNoPartnerSingleton = 0,
  kNoPartnerSameChamber,
  kBucketSelf,
  kBucketPartners,
  kSectorWindow,
  kAngleWindow,
  kEdgeCaps,
  kBelowThreshold,
  kDedupChamber,
  kComponentRejected,
  kNotSelected,
  kNumLossCategories
};
 
enum class TruthRegion : std::uint8_t { kBarrel = 0, kTransition = 1, kEndcap = 2 };
 
TruthRegion regionOf(double absEta) {
  return absEta < 1.0 ? TruthRegion::kBarrel
                      : (absEta < 1.3 ? TruthRegion::kTransition : TruthRegion::kEndcap);
}
 
}  // namespace
 
StatusCode SegmentEdgeInferenceAlg::initialize() {
  ATH_CHECK(m_segmentKey.initialize());
  ATH_CHECK(m_pairGateDecorKey.initialize());
  ATH_CHECK(m_componentScoreDecorKey.initialize());
  ATH_CHECK(m_filteredSegmentKey.initialize(!m_filteredSegmentKey.empty()));
  ATH_CHECK(m_edgeClassifier.retrieve());
  m_truthDiagEnabled = m_edgeClassifier->enableTruthDiagnostics();
  return StatusCode::SUCCESS;
}
 
StatusCode SegmentEdgeInferenceAlg::execute(const EventContext& ctx) const {
  const xAOD::MuonSegmentContainer* segments{};
  ATH_CHECK(SG::get(segments, m_segmentKey, ctx));
  // Evaluated once per event; short-circuits without touching the message
  // service unless the property was explicitly enabled.
  const bool truthDiag = m_truthDiagEnabled && msgLvl(MSG::DEBUG);
  ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                << ": input segments in '" << m_segmentKey.key()
                << "' = " << segments->size());
 
  if (truthDiag) {
    std::size_t truthSegs = 0, bkgSegs = 0;
    for (const xAOD::MuonSegment* seg : *segments) {
      MuonR4::getTruthMatchedParticle(*seg) ? ++truthSegs : ++bkgSegs;
    }
    m_sumInputTruthSegments += truthSegs;
    m_sumInputBkgSegments += bkgSegs;
  }
 
  SegmentEdgeGraph graph{};
  std::vector<SegmentEdgeScore> scores{};
  ATH_CHECK(m_edgeClassifier->buildGraph(ctx, *segments, graph));
  ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                << ": built graph with nodes=" << graph.nNodes
                << ", edges=" << graph.nEdges);
 
  // Per-node truth-particle grouping key (-1 = unlabeled), used for the
  // segment-level funnel and for the edge-level true/background split.
  std::vector<int32_t> nodeTruthId;
  // Truth-particle id -> |eta|, filled alongside nodeTruthId.
  std::unordered_map<int32_t, double> truthIdToAbsEta;
  if (truthDiag) {
    nodeTruthId.assign(graph.nNodes, -1);
    std::size_t truthSegs = 0, bkgSegs = 0;
    for (std::size_t node = 0; node < graph.nNodes; ++node) {
      const xAOD::TruthParticle* truthPart =
          MuonR4::getTruthMatchedParticle(*graph.segments[node]);
      if (truthPart) {
        nodeTruthId[node] = static_cast<int32_t>(truthPart->index());
        truthIdToAbsEta.emplace(nodeTruthId[node], std::abs(truthPart->eta()));
        ++truthSegs;
      } else {
        ++bkgSegs;
      }
    }
    m_sumOnnxTruthSegments += truthSegs;
    m_sumOnnxBkgSegments += bkgSegs;
  }
 
  ATH_CHECK(m_edgeClassifier->classifyEdges(ctx, graph, scores));
  if (!scores.empty()) {
    float minProb = std::numeric_limits<float>::max();
    float maxProb = std::numeric_limits<float>::lowest();
    for (const SegmentEdgeScore& score : scores) {
      minProb = std::min(minProb, score.probability);
      maxProb = std::max(maxProb, score.probability);
    }
    ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                  << ": edge scores=" << scores.size()
                  << ", prob range=[" << minProb << ", " << maxProb << "]");
  } else {
    ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                  << ": no edge scores produced");
  }
 
  // Symmetrise directed model outputs once.  Everything downstream uses the
  // score of an undirected segment association.
  std::unordered_map<std::uint64_t, float> pairProbability;
  pairProbability.reserve(scores.size());
  for (const SegmentEdgeScore& score : scores) {
    if (score.src >= graph.nNodes || score.dst >= graph.nNodes ||
        score.src == score.dst) {
      continue;
    }
    const std::uint64_t key = undirectedPairKey(score.src, score.dst);
    auto [it, inserted] = pairProbability.emplace(key, score.probability);
    if (!inserted) it->second = std::max(it->second, score.probability);
  }
 
  // Edge-level recall/false-positive check: candidate pairs that reached ONNX 
  // true pairs reach PairGateThreshold, versus how many background pairs do. 
  if (truthDiag) {
    std::size_t trueTotal = 0, truePassed = 0, bkgTotal = 0, bkgPassed = 0;
    for (const auto& [key, probability] : pairProbability) {
      const auto [first, second] = unpackPairKey(key);
      if (first >= graph.nNodes || second >= graph.nNodes) continue;
      const bool isTrueEdge =
          nodeTruthId[first] >= 0 && nodeTruthId[first] == nodeTruthId[second];
      const bool passed = probability >= m_pairGateThreshold.value();
      if (isTrueEdge) {
        ++trueTotal;
        truePassed += passed;
      } else {
        ++bkgTotal;
        bkgPassed += passed;
      }
    }
    m_sumTrueEdgesTotal += trueTotal;
    m_sumTrueEdgesPassed += truePassed;
    m_sumBkgEdgesTotal += bkgTotal;
    m_sumBkgEdgesPassed += bkgPassed;
  }
 
  // Form a sparse, score-ranked topology. Mutual top-K associations retain
  // locally consistent paths while suppressing one-sided bridges.
  using WeightedEdge = std::pair<std::uint64_t, float>;
  const auto betterWeightedEdge = [](const WeightedEdge& first,
                                      const WeightedEdge& second) {
    const int probabilityOrder = InferenceUtils::compareFloatDescending(
        first.second, second.second);
    if (probabilityOrder != 0) {
      return probabilityOrder < 0;
    }
    return first.first < second.first;
  };
  std::vector<std::vector<WeightedEdge>> edgesByNode(graph.nNodes);
  std::size_t thresholdPairs = 0;
  for (const auto& [key, probability] : pairProbability) {
    if (probability < m_pairGateThreshold.value()) continue;
    const auto [first, second] = unpackPairKey(key);
    if (first >= graph.nNodes || second >= graph.nNodes) continue;
    edgesByNode[first].emplace_back(key, probability);
    edgesByNode[second].emplace_back(key, probability);
    ++thresholdPairs;
  }
 
  std::size_t thresholdedNodes = 0;
  for (const std::vector<WeightedEdge>& nodeEdges : edgesByNode) {
    thresholdedNodes += !nodeEdges.empty();
  }
 
  std::vector<Muon::MuonStationIndex::LayerIndex> nodeLayer(graph.nNodes);
  for (std::size_t node = 0; node < graph.nNodes; ++node) {
    nodeLayer[node] = Muon::MuonStationIndex::toLayerIndex(graph.segments[node]->chamberIndex());
  }
  const unsigned int maxEdgesPerLayer = m_maxEdgesPerNodePerLayer.value();
  std::unordered_map<std::uint64_t, unsigned char> nominations;
  nominations.reserve(thresholdPairs);
  for (std::size_t node = 0; node < graph.nNodes; ++node) {
    std::vector<WeightedEdge>& nodeEdges = edgesByNode[node];
    std::sort(nodeEdges.begin(), nodeEdges.end(), betterWeightedEdge);
    if (maxEdgesPerLayer != 0) {
      // Best edges towards each layer: a segment keeps its partner in every neighbouring
      // layer, and same-layer edges do not compete with cross-layer ones.
      std::array<unsigned int, Acts::toUnderlying(Muon::MuonStationIndex::LayerIndex::LayerIndexMax) + 1> taken{};
      std::erase_if(nodeEdges, [&](const WeightedEdge& edge) {
        const auto [first, second] = unpackPairKey(edge.first);
        const std::size_t other = first == node ? second : first;
        const auto layer = static_cast<std::size_t>(Acts::toUnderlying(nodeLayer[other]) + 1);
        return taken[layer]++ >= maxEdgesPerLayer;
      });
    } else if (m_maxEdgesPerNode.value() != 0 &&
               nodeEdges.size() > m_maxEdgesPerNode.value()) {
      nodeEdges.resize(m_maxEdgesPerNode.value());
    }
    for (const WeightedEdge& edge : nodeEdges) {
      ++nominations[edge.first];
    }
  }
 
  std::size_t mutualTopKPairs = 0;
  std::size_t oneSidedTopKPairs = 0;
  for (const auto& [_, count] : nominations) {
    if (count == 2) {
      ++mutualTopKPairs;
    } else {
      ++oneSidedTopKPairs;
    }
  }
 
  std::unordered_set<std::uint64_t> selectedPairKeys;
  selectedPairKeys.reserve(thresholdPairs);
  if (m_useDegreeCappedComponents.value()) {
    // Apply a score-ordered global degree cap when explicitly requested.
    std::vector<WeightedEdge> acceptedPairs;
    acceptedPairs.reserve(thresholdPairs);
    for (const auto& [key, probability] : pairProbability) {
      if (probability < m_pairGateThreshold.value()) continue;
      acceptedPairs.emplace_back(key, probability);
    }
    std::sort(acceptedPairs.begin(), acceptedPairs.end(), betterWeightedEdge);
 
    const unsigned int maxDegree = m_maxEdgesPerNode.value();
    std::vector<unsigned int> degree(graph.nNodes, 0);
    for (const WeightedEdge& edge : acceptedPairs) {
      const auto [first, second] = unpackPairKey(edge.first);
      if (maxDegree != 0 &&
          (degree[first] >= maxDegree || degree[second] >= maxDegree)) {
        continue;
      }
      selectedPairKeys.insert(edge.first);
      ++degree[first];
      ++degree[second];
    }
  } else {
    for (const auto& [key, count] : nominations) {
      if (m_requireMutualTopKEdges.value() && count != 2) continue;
      selectedPairKeys.insert(key);
    }
  }
 
  // A mutual top-K selection can leave an endpoint without an association.
  // Add its best thresholded edge, at most once per orphaned endpoint.
  std::size_t orphanRecoveryPairs = 0;
  if (!m_useDegreeCappedComponents.value() &&
      m_requireMutualTopKEdges.value() &&
      m_recoverOrphanNodes.value()) {
    std::vector<unsigned char> selectedNode(graph.nNodes, 0);
    for (const std::uint64_t key : selectedPairKeys) {
      const auto [first, second] = unpackPairKey(key);
      if (first < graph.nNodes) selectedNode[first] = 1;
      if (second < graph.nNodes) selectedNode[second] = 1;
    }
    for (std::size_t node = 0; node < graph.nNodes; ++node) {
      if (selectedNode[node] || edgesByNode[node].empty()) continue;
      const std::uint64_t key = edgesByNode[node].front().first;
      const auto [first, second] = unpackPairKey(key);
      if (first >= graph.nNodes || second >= graph.nNodes) continue;
      if (selectedPairKeys.insert(key).second) ++orphanRecoveryPairs;
      selectedNode[first] = 1;
      selectedNode[second] = 1;
    }
  }
 
  std::vector<std::uint64_t> selectedPairs{selectedPairKeys.begin(),
                                           selectedPairKeys.end()};
  std::sort(selectedPairs.begin(), selectedPairs.end());
 
  // Truth diagnostics only: true cross-layer edges above threshold that the selection keeps,
  // barrel / endcap (by the first endpoint's chamber)
  if (truthDiag) {
    for (const auto& [key, probability] : pairProbability) {
      if (probability < m_pairGateThreshold.value()) continue;
      const auto [first, second] = unpackPairKey(key);
      if (first >= graph.nNodes || second >= graph.nNodes) continue;
      if (nodeTruthId[first] < 0 || nodeTruthId[first] != nodeTruthId[second]) continue;
      if (nodeLayer[first] == nodeLayer[second]) continue;
      const std::size_t region =
          Muon::MuonStationIndex::isBarrel(graph.segments[first]->chamberIndex()) ? 0 : 1;
      ++m_sumTrueCrossLayerEdges[region];
      if (selectedPairKeys.contains(key)) ++m_sumTrueCrossLayerSelected[region];
    }
  }
 
  // Truth diagnostics only: best selected-edge score per node, so that the
  // strongest edge holding a component together is a max over its nodes.
  std::vector<float> nodeBestSelectedScore;
  if (truthDiag) {
    nodeBestSelectedScore.assign(graph.nNodes, -1.f);
    for (const std::uint64_t key : selectedPairs) {
      const auto found = pairProbability.find(key);
      if (found == pairProbability.end()) continue;
      const auto [first, second] = unpackPairKey(key);
      nodeBestSelectedScore[first] = std::max(nodeBestSelectedScore[first], found->second);
      nodeBestSelectedScore[second] = std::max(nodeBestSelectedScore[second], found->second);
    }
  }
 
  // Innermost-layer anchoring with an anchor cap: rank the inner candidates by
  // their strongest selected edge towards another layer, i.e. by how firmly
  // the model links them outwards. Same-chamber alternatives otherwise all
  // launch seeds built from the same outer segments.
  const auto layerOf = [&graph](std::size_t node) {
    return Muon::MuonStationIndex::toLayerIndex(graph.segments[node]->chamberIndex());
  };
  std::vector<float> nodeBestCrossLayerScore;
  std::vector<unsigned int> nodeCrossLayerEdges;  // truth diagnostics only
  {
    nodeBestCrossLayerScore.assign(graph.nNodes, -1.f);
    if (truthDiag) nodeCrossLayerEdges.assign(graph.nNodes, 0);
    for (const std::uint64_t key : selectedPairs) {
      const auto [first, second] = unpackPairKey(key);
      if (layerOf(first) == layerOf(second)) continue;
      const auto found = pairProbability.find(key);
      if (found == pairProbability.end()) continue;
      nodeBestCrossLayerScore[first] = std::max(nodeBestCrossLayerScore[first], found->second);
      nodeBestCrossLayerScore[second] = std::max(nodeBestCrossLayerScore[second], found->second);
      if (truthDiag) {
        ++nodeCrossLayerEdges[first];
        ++nodeCrossLayerEdges[second];
      }
    }
  }
 
 
  DisjointSet components{graph.nNodes};
  std::vector<bool> activeNode(graph.nNodes, false);
  for (const std::uint64_t key : selectedPairs) {
    const auto [first, second] = unpackPairKey(key);
    components.unite(first, second);
    activeNode[first] = true;
    activeNode[second] = true;
  }
 
  std::unordered_map<std::size_t, std::vector<std::size_t>> byRoot;
  byRoot.reserve(graph.nNodes);
  for (std::size_t node = 0; node < graph.nNodes; ++node) {
    if (activeNode[node]) byRoot[components.find(node)].push_back(node);
  }
 
  // Deterministic component IDs make debugging and validation reproducible.
  std::vector<std::vector<std::size_t>> componentNodes;
  componentNodes.reserve(byRoot.size());
  for (auto& [_, nodes] : byRoot) {
    std::sort(nodes.begin(), nodes.end());
    componentNodes.push_back(std::move(nodes));
  }
  std::ranges::sort(componentNodes,
            [](const auto& first, const auto& second) {
              return first.front() < second.front();
            });
 
  xAOD::ContainerDecorator<xAOD::MuonSegmentContainer, std::vector<unsigned>>
      decor{m_pairGateDecorKey, ctx};
  xAOD::ContainerDecorator<xAOD::MuonSegmentContainer, std::vector<float>>
      scoreDecor{m_componentScoreDecorKey, ctx};
 
  std::vector<std::uint8_t> keptNode(graph.nNodes, false);
  std::vector<std::uint8_t> nodeDropReason;
  if (truthDiag) {
     nodeDropReason.assign(graph.nNodes, kReasonNotSelected);
  }
 
  // Group-purity: which kept-component id(s) each truth muon's
  // segments landed in, and whether any of those components were mixed
  std::unordered_map<int32_t, std::unordered_set<unsigned>> muonComponents;
  std::unordered_map<int32_t, bool> muonEverMixed;
  // D8: weak clean components and, per truth muon, the number of strong multi-layer components
  std::vector<int32_t> weakCleanMuons;
  std::unordered_map<int32_t, unsigned int> strongComponentsPerMuon;
  const auto sizeBin = [](std::size_t size) -> std::size_t {
    if (size <= 4) return size - 2;  // 2, 3, 4 -> 0, 1, 2
    if (size < 10) return 3;         // 5-9
    if (size < 20) return 4;         // 10-19
    return 5;                        // 20+
  };
 
  std::size_t topologyNodes = 0;
  std::size_t retainedNodes = 0;
  std::size_t chamberSuppressedNodes = 0;
  std::size_t rejectedComponents = 0;
  std::size_t componentsKept = 0;
  std::size_t anchors = 0;
  std::size_t nodesRejectedByMinComponent = 0;
  unsigned nextComponentId = 1;
  const InferenceUtils::SegmentQualityOrder betterSegment{};
 
  const auto isBetterNode = [&](std::size_t candidate, std::size_t incumbent) {
    if (betterSegment(graph.segments[candidate], graph.segments[incumbent])) {
      return true;
    }
    if (betterSegment(graph.segments[incumbent], graph.segments[candidate])) {
      return false;
    }
    return candidate < incumbent;
  };
 
  for (const std::vector<std::size_t>& rawNodes : componentNodes) {
    topologyNodes += rawNodes.size();
    std::vector<std::size_t> retained = rawNodes;
 
    // The seeder resolves same-chamber alternatives while building a seed.
    // Retaining only the highest-ranked representative is therefore optional.
    if (m_keepBestSegmentPerChamber.value()) {
      std::unordered_map<int, std::size_t> bestByChamber;
      bestByChamber.reserve(rawNodes.size());
      for (const std::size_t node : rawNodes) {
        const int chamber =
            static_cast<int>(graph.segments[node]->chamberIndex());
        const auto found = bestByChamber.find(chamber);
        if (found == bestByChamber.end() || isBetterNode(node, found->second)) {
          bestByChamber[chamber] = node;
        }
      }
      retained.clear();
      retained.reserve(bestByChamber.size());
      for (const auto& [_, node] : bestByChamber) retained.push_back(node);
      std::sort(retained.begin(), retained.end());
      chamberSuppressedNodes += rawNodes.size() - retained.size();
    }
 
    // Provisional reasons; nodes that end up kept are identified by keptNode.
    if (truthDiag) {
      for (const std::size_t node : rawNodes) {
          nodeDropReason[node] = kReasonChamberDedup;
      }
      for (const std::size_t node : retained) {
          nodeDropReason[node] = kReasonComponentRejected;
      }
    }
 
    if (retained.size() < m_minSegmentsPerComponent.value()) {
      nodesRejectedByMinComponent += retained.size();
      ++rejectedComponents;
      continue;
    }
 
    // Only ranked component members launch seeds. The edge score therefore
    // reduces seed attempts directly rather than serving only as a label.
    std::vector<std::size_t> rankedNodes{retained};
    if (m_anchorInnermostLayer.value()) {
      // Anchor only on the inner segment(s)
      int bestRank = std::numeric_limits<int>::max();
      for (const std::size_t node : rankedNodes) {
        bestRank = std::min(
            bestRank,
            layerRadialRank(Muon::MuonStationIndex::toLayerIndex(
                graph.segments[node]->chamberIndex())));
      }
      std::erase_if(rankedNodes, [&](std::size_t node) {
        return layerRadialRank(Muon::MuonStationIndex::toLayerIndex(
                   graph.segments[node]->chamberIndex())) != bestRank;
      });
      // Rank the inner anchors (strongest cross-layer edge first, segment quality as
      // tie-break) and cap them.
      if (m_seedAnchorsPerComponent.value() != 0) {
        std::ranges::sort(rankedNodes, [&](std::size_t first, std::size_t second) {
          const int order = InferenceUtils::compareFloatDescending(
              nodeBestCrossLayerScore[first], nodeBestCrossLayerScore[second]);
          return order != 0 ? order < 0 : isBetterNode(first, second);
        });
        if (rankedNodes.size() > m_seedAnchorsPerComponent.value()) {
          rankedNodes.resize(m_seedAnchorsPerComponent.value());
        }
      }
    } else {
      std::ranges::sort(rankedNodes, isBetterNode);
      const std::size_t nAnchors = m_seedAnchorsPerComponent.value() == 0
          ? rankedNodes.size()
          : std::min<std::size_t>(m_seedAnchorsPerComponent.value(),
                                  rankedNodes.size());
      rankedNodes.resize(nAnchors);
    }
    if (rankedNodes.empty()) {
      ++rejectedComponents;
      continue;
    }
    // rankedNodes stays in rank order (best anchor first): the decoration
    // carries the rank so that the seeder can try anchors best-first.
 
    const unsigned componentId = nextComponentId++;
    float componentCrossScore = -1.f;
    for (const std::size_t node : rawNodes) {
      componentCrossScore = std::max(componentCrossScore, nodeBestCrossLayerScore[node]);
    }
    for (const std::size_t node : retained) {
      scoreDecor(*graph.segments[node]) = {componentCrossScore, nodeBestCrossLayerScore[node]};
      const auto anchorIt = std::ranges::find(rankedNodes, node);
      const bool isAnchor = anchorIt != rankedNodes.end();
      decor(*graph.segments[node]) = {
          componentId,
          isAnchor ? static_cast<unsigned int>(anchorIt - rankedNodes.begin()) + 1 : 0u};
      keptNode[node] = true;
    }
    retainedNodes += retained.size();
    anchors += rankedNodes.size();
    ++componentsKept;
 
    // Group purity: does this kept segments from more than one truth muon?
    if (truthDiag) {
      std::unordered_set<int32_t> componentTruthIds;
      for (const std::size_t node : retained) {
        if (nodeTruthId[node] >= 0) componentTruthIds.insert(nodeTruthId[node]);
      }
      const std::size_t outcome = componentTruthIds.empty() ? 2
                                  : componentTruthIds.size() == 1 ? 0 : 1;  // clean/mixed/bkgOnly
      ++m_groupPurity.componentsByOutcome[outcome];
      ++m_groupPurity.componentSizeHist[sizeBin(retained.size())];
 
      // Component descriptors that a component-level cut could use: the
      // strongest selected edge among the nodes that formed it (before the
      // chamber dedup), its retained size and its number of distinct layers.
      float strongest = -1.f;
      for (const std::size_t node : rawNodes) {
        strongest = std::max(strongest, nodeBestSelectedScore[node]);
      }
      const std::size_t scoreBin = strongest < 0.05f ? 0 : strongest < 0.1f ? 1
                                 : strongest < 0.3f ? 2 : strongest < 0.5f ? 3 : 4;
      std::uint32_t layerMask = 0;
      for (const std::size_t node : retained) {
        // +1 so that LayerUnknown (-1) maps onto bit 0
        layerMask |= 1u << (Acts::toUnderlying(Muon::MuonStationIndex::toLayerIndex(
                                graph.segments[node]->chamberIndex())) + 1);
      }
      const std::size_t layerBin =
          std::min<std::size_t>(std::popcount(layerMask), GroupPurityCounters::kLayerBins) - 1;
      ++m_groupPurity.scoreByOutcome[outcome * GroupPurityCounters::kScoreBins + scoreBin];
      ++m_groupPurity.sizeByOutcome[outcome * GroupPurityCounters::kSizeBins + sizeBin(retained.size())];
      ++m_groupPurity.layersByOutcome[outcome * GroupPurityCounters::kLayerBins + layerBin];
 
      // Components that can seed (>= 2 layers)
      if (std::popcount(layerMask) >= 2) {
        float crossScore = -1.f;
        unsigned int crossEdges = 0;
        for (const std::size_t node : rawNodes) {
          crossScore = std::max(crossScore, nodeBestCrossLayerScore[node]);
          crossEdges += nodeCrossLayerEdges[node];
        }
        crossEdges /= 2;
        const std::size_t crossBin = crossScore < 0.05f ? 0 : crossScore < 0.1f ? 1 : crossScore < 0.3f ? 2
                                   : crossScore < 0.5f ? 3 : crossScore < 0.8f ? 4 : 5;
        ++m_groupPurity.crossScoreByOutcome[outcome * GroupPurityCounters::kCrossScoreBins + crossBin];
        ++m_groupPurity.crossEdgesByOutcome[outcome * GroupPurityCounters::kCrossEdgeBins +
                                            std::min<std::size_t>(crossEdges, 3)];
        if (crossScore < 0.05f) {
          if (outcome == 0) weakCleanMuons.push_back(*componentTruthIds.begin());
        } else {
          for (const int32_t id : componentTruthIds) ++strongComponentsPerMuon[id];
        }
      }
      const bool mixed = componentTruthIds.size() >= 2;
      if (mixed) {
        m_groupPurity.mixedComponentSegments += retained.size();
        const std::uint64_t packed =
            (static_cast<std::uint64_t>(componentTruthIds.size()) << 32) | retained.size();
        std::uint64_t current = m_groupPurity.worstMixedPacked.load(std::memory_order_relaxed);
        while (packed > current && !m_groupPurity.worstMixedPacked.compare_exchange_weak(
                   current, packed, std::memory_order_relaxed)) {}
      }
      for (const int32_t id : componentTruthIds) {
        muonComponents[id].insert(componentId);
        if (mixed) muonEverMixed[id] = true;
      }
    }
  }
 
  if (truthDiag) {
    std::size_t truthSegs = 0, bkgSegs = 0;
    for (std::size_t node = 0; node < graph.nNodes; ++node) {
      if (!keptNode[node]) continue;
      nodeTruthId[node] >= 0 ? ++truthSegs : ++bkgSegs;
    }
    m_sumRetainedTruthSegments += truthSegs;
    m_sumRetainedBkgSegments += bkgSegs;
 
    std::vector<unsigned char> thresholded(graph.nNodes, 0);
    for (std::size_t node = 0; node < graph.nNodes; ++node) {
      thresholded[node] = !edgesByNode[node].empty();
    }
    classifyLostTruthSegments(*segments, graph, pairProbability, thresholded,
                              keptNode, nodeDropReason);
    classifyGroupPurity(muonComponents, muonEverMixed, truthIdToAbsEta);
    for (const int32_t id : weakCleanMuons) {
      ++(strongComponentsPerMuon[id] > 0 ? m_groupPurity.weakCleanWithOther : m_groupPurity.weakCleanAlone);
    }
  }
 
  if (!m_filteredSegmentKey.empty()) {
    auto connectedSegments =
        std::make_unique<ConstDataVector<xAOD::MuonSegmentContainer>>(
            SG::VIEW_ELEMENTS);
    connectedSegments->reserve(graph.nNodes);
    for (std::size_t node = 0; node < graph.nNodes; ++node) {
      if (!keptNode[node] || !graph.segments[node]) continue;
      connectedSegments->push_back(graph.segments[node]);
    }
 
    const std::size_t nConnectedSegments = connectedSegments->size();
    SG::WriteHandle<ConstDataVector<xAOD::MuonSegmentContainer>> connectedHandle{
        m_filteredSegmentKey, ctx};
    ATH_CHECK(connectedHandle.record(std::move(connectedSegments)));
    ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                  << ": wrote " << nConnectedSegments
                  << " ML-connected segment(s) to '"
                  << m_filteredSegmentKey.key() << "'");
  }
 
  ATH_MSG_DEBUG("Event " << ctx.eventID().event_number()
                << ": ML components graphNodes=" << graph.nNodes
                << ", thresholdedNodes=" << thresholdedNodes
                << ", thresholdPairs=" << thresholdPairs
                << ", mutualTopKPairs=" << mutualTopKPairs
                << ", oneSidedTopKPairs=" << oneSidedTopKPairs
                << ", selectedPairs=" << selectedPairKeys.size()
                << ", components=" << componentsKept
                << ", topologyNodes=" << topologyNodes
                << ", retainedNodes=" << retainedNodes
                << ", chamberSuppressedNodes=" << chamberSuppressedNodes
                << ", nodesRejectedByMinComponent=" << nodesRejectedByMinComponent
                << ", keepBestSegmentPerChamber=" << m_keepBestSegmentPerChamber.value()
                << ", seedAnchors=" << anchors
                << ", rejectedComponents=" << rejectedComponents
                << ", threshold=" << m_pairGateThreshold.value()
                << ", maxEdgesPerNode=" << m_maxEdgesPerNode.value()
                << ", orphanRecoveryPairs=" << orphanRecoveryPairs
                << ", mutualTopK=" << m_requireMutualTopKEdges.value()
                << ", degreeCapped=" << m_useDegreeCappedComponents.value());
 
  // Job-summed diagnostics
  if (msgLvl(MSG::DEBUG)) {
    m_sumGraphNodes += graph.nNodes;
    m_sumThresholdedNodes += thresholdedNodes;
    m_sumMutualTopKPairs += mutualTopKPairs;
    m_sumOneSidedTopKPairs += oneSidedTopKPairs;
    m_sumOrphanRecoveryPairs += orphanRecoveryPairs;
    m_sumTopologyNodes += topologyNodes;
    m_sumRetainedNodes += retainedNodes;
    m_sumChamberSuppressedNodes += chamberSuppressedNodes;
    m_sumNodesRejectedByMinComponent += nodesRejectedByMinComponent;
    m_sumRejectedComponents += rejectedComponents;
    m_sumComponentsKept += componentsKept;
  }
  return StatusCode::SUCCESS;
}
 
void SegmentEdgeInferenceAlg::classifyLostTruthSegments(
    const xAOD::MuonSegmentContainer& segments, const SegmentEdgeGraph& graph,
    const std::unordered_map<std::uint64_t, float>& pairProbability,
    const std::vector<unsigned char>& thresholded,
    const std::vector<std::uint8_t>& keptNode,
    const std::vector<unsigned char>& nodeDropReason) const {
  static_assert(kNumLossCategories == TruthLossCounters::kCategories);
  constexpr std::size_t kColumns = TruthLossCounters::kColumns;
  const std::size_t nInput = segments.size();
 
  if (graph.inputNodeIndex.size() != nInput) {
    if (!m_warnedMissingToolDiagnostics.exchange(true)) {
      ATH_MSG_WARNING("Truth diagnostics requested, but the edge classifier tool "
                      "did not provide its pre-ONNX fates; set "
                      "EnableTruthDiagnostics on the tool as well. Skipping the "
                      "lost-segment classification.");
    }
    return;
  }
 
  // Truth particle of every input segment and the segments grouped per particle.
  std::vector<const xAOD::TruthParticle*> truthPart(nInput, nullptr);
  std::vector<std::int32_t> truthId(nInput, -1);
  std::unordered_map<std::int32_t, std::vector<std::uint32_t>> byTruth;
  {
    std::uint32_t inputIndex = 0;
    for (const xAOD::MuonSegment* segment : segments) {
      if (const xAOD::TruthParticle* part = MuonR4::getTruthMatchedParticle(*segment)) {
        truthPart[inputIndex] = part;
        truthId[inputIndex] = static_cast<std::int32_t>(part->index());
        byTruth[truthId[inputIndex]].push_back(inputIndex);
      }
      ++inputIndex;
    }
  }
  const auto isRetained = [&](std::size_t input) {
    const std::int32_t node = graph.inputNodeIndex[input];
    return node >= 0 && keptNode[static_cast<std::size_t>(node)];
  };
 
  // Per segment: how far its best cross-chamber true pair got (1 = bucket cap,
  // 2 = sector window, 3 = angle window, 4 = edge caps, 5 = scored, 0 = only
  // same-chamber partners) and the score / index of its best scored pair.
  const auto rankOf = [](PairGate gate) -> unsigned char {
    switch (gate) {
      case PairGate::BucketCap: return 1;
      case PairGate::SectorWindow: return 2;
      case PairGate::AngleWindow: return 3;
      case PairGate::EdgeCaps: return 4;
      case PairGate::Scored: return 5;
      case PairGate::SameChamber: break;
    }
    return 0;
  };
  std::vector<unsigned char> bestRank(nInput, 0);
  std::vector<float> bestScore(nInput, -1.f);
  std::vector<std::int32_t> bestScoredPair(nInput, -1);
  for (std::size_t k = 0; k < graph.truthPairFates.size(); ++k) {
    const PairFate& fate = graph.truthPairFates[k];
    const unsigned char rank = rankOf(fate.gate);
    if (rank == 0) continue;
    const std::array<std::uint32_t, 2> ends{fate.first, fate.second};
    for (const std::uint32_t end : ends) bestRank[end] = std::max(bestRank[end], rank);
    if (fate.gate != PairGate::Scored) continue;
    const std::int32_t firstNode = graph.inputNodeIndex[fate.first];
    const std::int32_t secondNode = graph.inputNodeIndex[fate.second];
    if (firstNode < 0 || secondNode < 0) continue;
    const auto found = pairProbability.find(undirectedPairKey(
        static_cast<std::size_t>(firstNode), static_cast<std::size_t>(secondNode)));
    const float score = found != pairProbability.end() ? found->second : 0.f;
    for (const std::uint32_t end : ends) {
      if (score > bestScore[end]) {
        bestScore[end] = score;
        bestScoredPair[end] = static_cast<std::int32_t>(k);
      }
    }
  }
 
  // Seedability per truth muon: at least 2 distinct chambers in the input,
  // fewer than 2 among the retained segments.
  std::unordered_map<std::int32_t, bool> seedLostByTruth;
  std::array<std::size_t, 4> muonsSeedable{};
  std::array<std::size_t, 4> muonsSeedLost{};
  for (const auto& entry : byTruth) {
    std::unordered_set<int> inputChambers;
    std::unordered_set<int> retainedChambers;
    for (const std::uint32_t input : entry.second) {
      const int chamber = static_cast<int>(segments[input]->chamberIndex());
      inputChambers.insert(chamber);
      if (isRetained(input)) retainedChambers.insert(chamber);
    }
    const bool seedable = inputChambers.size() >= 2;
    const bool lost = seedable && retainedChambers.size() < 2;
    seedLostByTruth[entry.first] = lost;
    if (!seedable) continue;
    const std::size_t region = 1 + static_cast<std::size_t>(
        regionOf(std::abs(truthPart[entry.second.front()]->eta())));
    ++muonsSeedable[0];
    ++muonsSeedable[region];
    if (lost) {
      ++muonsSeedLost[0];
      ++muonsSeedLost[region];
    }
  }
 
  std::array<std::size_t, kNumLossCategories * kColumns> category{};
  std::array<std::size_t, 5> scoreBin{};
  std::array<std::size_t, 2> sectorDelta{};
  std::array<std::size_t, 4> layerPair{};
  std::array<std::size_t, 4> muonSegments{};
  std::array<std::size_t, 3> precisionHits{};
  for (std::size_t i = 0; i < nInput; ++i) {
    if (truthId[i] < 0 || isRetained(i)) continue;
    const std::vector<std::uint32_t>& members = byTruth[truthId[i]];
    const std::int32_t node = graph.inputNodeIndex[i];
 
    std::size_t cat{};
    if (members.size() == 1) {
      cat = kNoPartnerSingleton;
    } else if (bestRank[i] == 0) {
      cat = kNoPartnerSameChamber;
    } else if (node == kDroppedByBucketCap) {
      cat = kBucketSelf;
    } else if (node >= 0 && thresholded[static_cast<std::size_t>(node)]) {
      switch (nodeDropReason[static_cast<std::size_t>(node)]) {
        case kReasonChamberDedup: cat = kDedupChamber; break;
        case kReasonComponentRejected: cat = kComponentRejected; break;
        default: cat = kNotSelected; break;
      }
    } else {
      switch (bestRank[i]) {
        case 5: cat = kBelowThreshold; break;
        case 4: cat = kEdgeCaps; break;
        case 3: cat = kAngleWindow; break;
        case 2: cat = kSectorWindow; break;
        default: cat = kBucketPartners; break;
      }
    }
 
    const std::size_t region = 1 + static_cast<std::size_t>(
        regionOf(std::abs(truthPart[i]->eta())));
    ++category[cat * kColumns];
    ++category[cat * kColumns + region];
    if (seedLostByTruth[truthId[i]]) ++category[cat * kColumns + 4];
 
    if (cat != kBelowThreshold) continue;
    if (bestScoredPair[i] >= 0) {
      const PairFate& fate = graph.truthPairFates[static_cast<std::size_t>(bestScoredPair[i])];
      const float score = bestScore[i];
      ++scoreBin[score < 1e-3f ? 0 : score < 1e-2f ? 1 : score < 2e-2f ? 2
                                                         : score < 5e-2f ? 3 : 4];
      ++sectorDelta[fate.sectorDelta == 0 ? 0 : 1];
      const std::uint32_t other = fate.first == i ? fate.second : fate.first;
      int rankA = layerRadialRank(Muon::MuonStationIndex::toLayerIndex(
          segments[i]->chamberIndex()));
      int rankB = layerRadialRank(Muon::MuonStationIndex::toLayerIndex(
          segments[other]->chamberIndex()));
      if (rankA > rankB) std::swap(rankA, rankB);
      ++layerPair[(rankA == 0 && rankB == 1) ? 0
                  : (rankA == 1 && rankB == 2) ? 1
                  : (rankA == 0 && rankB == 2) ? 2 : 3];
    }
    ++muonSegments[std::min<std::size_t>(members.size(), 5) - 2];
    const unsigned int hits = segments[i]->nPrecisionHits();
    ++precisionHits[hits <= 4 ? 0 : (hits <= 6 ? 1 : 2)];
  }
 
  for (std::size_t c = 0; c < category.size(); ++c) m_truthLoss.category[c] += category[c];
  for (std::size_t c = 0; c < scoreBin.size(); ++c) m_truthLoss.scoreBin[c] += scoreBin[c];
  for (std::size_t c = 0; c < sectorDelta.size(); ++c) m_truthLoss.sectorDelta[c] += sectorDelta[c];
  for (std::size_t c = 0; c < layerPair.size(); ++c) m_truthLoss.layerPair[c] += layerPair[c];
  for (std::size_t c = 0; c < muonSegments.size(); ++c) m_truthLoss.muonSegments[c] += muonSegments[c];
  for (std::size_t c = 0; c < precisionHits.size(); ++c) m_truthLoss.precisionHits[c] += precisionHits[c];
  for (std::size_t c = 0; c < muonsSeedable.size(); ++c) {
    m_truthLoss.muonsSeedable[c] += muonsSeedable[c];
    m_truthLoss.muonsSeedLost[c] += muonsSeedLost[c];
  }
}
 
void SegmentEdgeInferenceAlg::classifyGroupPurity(
    const std::unordered_map<std::int32_t, std::unordered_set<unsigned>>& muonComponents,
    const std::unordered_map<std::int32_t, bool>& muonEverMixed,
    const std::unordered_map<std::int32_t, double>& truthIdToAbsEta) const {
  // Only muons with at least one retained segment in some kept component
  // appear in muonComponents; fully-lost muons are already covered by
  // classifyLostTruthSegments()/muonsSeedLost above, not double-counted here.
  std::array<std::size_t, GroupPurityCounters::kColumns * 3> outcome{};
  for (const auto& [id, componentIds] : muonComponents) {
    const auto mixedIt = muonEverMixed.find(id);
    const bool everMixed = mixedIt != muonEverMixed.end() && mixedIt->second;
    // 0 = clean (one component, alone), 1 = mixed (shared a component with
    // another muon), 2 = fragmented (split across components, never mixed).
    const std::size_t result = everMixed ? 1 : (componentIds.size() >= 2 ? 2 : 0);
    const auto etaIt = truthIdToAbsEta.find(id);
    const std::size_t region = 1 + (etaIt != truthIdToAbsEta.end()
        ? static_cast<std::size_t>(regionOf(etaIt->second)) : 0);
    ++outcome[0 * 3 + result];
    ++outcome[region * 3 + result];
  }
  for (std::size_t i = 0; i < outcome.size(); ++i) m_groupPurity.muonOutcome[i] += outcome[i];
}
 
StatusCode SegmentEdgeInferenceAlg::finalize() {
  // The job-summed counters are only filled, and the summaries only built, at DEBUG level.
  if (!msgLvl(MSG::DEBUG)) return StatusCode::SUCCESS;
  ATH_MSG_DEBUG(
      "SegmentEdgeInferenceAlg post-ONNX selection summary (job-summed): "
      << "graphNodes=" << m_sumGraphNodes
      << ", thresholdedNodes=" << m_sumThresholdedNodes
      << " (>=PairGateThreshold=" << m_pairGateThreshold.value() << ")"
      << ", mutualTopKPairs=" << m_sumMutualTopKPairs
      << ", oneSidedTopKPairs(rejected by RequireMutualTopKEdges="
      << m_requireMutualTopKEdges.value() << ")=" << m_sumOneSidedTopKPairs
      << " (MaxEdgesPerNode=" << m_maxEdgesPerNode.value()
      << ", MaxEdgesPerNodePerLayer=" << m_maxEdgesPerNodePerLayer.value() << " (0 = global top-K)"
      << "; a thresholded edge is dropped here purely by rank, "
         "independent of PairGateThreshold)"
      << ", orphanRecoveryPairs=" << m_sumOrphanRecoveryPairs
      << ", chamberSuppressedNodes=" << m_sumChamberSuppressedNodes
      << " (KeepBestSegmentPerChamber=" << m_keepBestSegmentPerChamber.value() << ")"
      << ", nodesRejectedByMinComponent=" << m_sumNodesRejectedByMinComponent
      << " (MinSegmentsPerComponent=" << m_minSegmentsPerComponent.value() << ")"
      << ", rejectedComponents=" << m_sumRejectedComponents
      << ", componentsKept=" << m_sumComponentsKept
      << ", topologyNodes=" << m_sumTopologyNodes
      << ", retainedNodes=" << m_sumRetainedNodes);
 
  if (m_truthDiagEnabled) {
    ATH_MSG_DEBUG("True cross-layer edges above PairGateThreshold kept by the selection, selected/total: barrel="
                  << m_sumTrueCrossLayerSelected[0] << "/" << m_sumTrueCrossLayerEdges[0]
                  << ", endcap=" << m_sumTrueCrossLayerSelected[1] << "/" << m_sumTrueCrossLayerEdges[1]);
    ATH_MSG_DEBUG(
        "SegmentEdgeInferenceAlg truth-vs-background funnel (job-summed; "
        "truth segment = getTruthMatchedParticle(seg) != nullptr, matching "
        "SegmentDumperAlg::m_segmentHasTruth, no isMuon() filter, no G4 "
        "pseudo-label fallback): "
        << "input: truth=" << m_sumInputTruthSegments
        << " bkg=" << m_sumInputBkgSegments
        << "; reachesONNX: truth=" << m_sumOnnxTruthSegments
        << " bkg=" << m_sumOnnxBkgSegments
        << "; retained: truth=" << m_sumRetainedTruthSegments
        << " bkg=" << m_sumRetainedBkgSegments
        << "; edge-level (same truth-particle index on both endpoints) at "
           "PairGateThreshold=" << m_pairGateThreshold.value()
        << ": trueEdges passed/total=" << m_sumTrueEdgesPassed << "/"
        << m_sumTrueEdgesTotal
        << ", bkgEdges passed/total=" << m_sumBkgEdgesPassed << "/"
        << m_sumBkgEdgesTotal);
 
    constexpr std::array<const char*, TruthLossCounters::kCategories> names{
        "A1 no possible partner: singleton muon segment",
        "A2 no possible partner: same-chamber partners only",
        "B0 never reached ONNX: segment dropped by MaxSegmentsPerBucket",
        "B1 never reached ONNX: partners dropped by MaxSegmentsPerBucket",
        "B2 never reached ONNX: outside sector window (MaxDeltaSector)",
        "B3 never reached ONNX: outside angle window (MaxDeltaThetaDeg)",
        "B4 never reached ONNX: pre-ONNX edge caps",
        "C  true pair(s) scored below PairGateThreshold",
        "D1 passed threshold, removed: KeepBestSegmentPerChamber",
        "D2 passed threshold, removed: component rejected (MinSegmentsPerComponent/anchors)",
        "D3 passed threshold, removed: not selected by top-K/mutual gate"};
    constexpr std::size_t kColumns = TruthLossCounters::kColumns;
    ATH_MSG_DEBUG("Lost truth segments by stage (job-summed; truth segment = "
                  "input segment with a truth particle, lost = not in the filtered "
                  "container). Columns: all | barrel | transition | endcap | "
                  "belonging to a muon that lost seedability");
    std::size_t lostTotal = 0;
    for (std::size_t c = 0; c < names.size(); ++c) {
      const auto value = [&](std::size_t col) { return m_truthLoss.category[c * kColumns + col].load(); };
      lostTotal += value(0);
      ATH_MSG_DEBUG("  " << names[c] << ": " << value(0) << " | " << value(1)
                    << " | " << value(2) << " | " << value(3) << " | " << value(4));
    }
    ATH_MSG_DEBUG("  check: sum of categories=" << lostTotal
                  << ", input truth - retained truth="
                  << (m_sumInputTruthSegments.load() - m_sumRetainedTruthSegments.load()));
    ATH_MSG_DEBUG("  truth muons with >=2 input chambers (all|barrel|transition|endcap): "
                  << m_truthLoss.muonsSeedable[0].load() << "|" << m_truthLoss.muonsSeedable[1].load()
                  << "|" << m_truthLoss.muonsSeedable[2].load() << "|" << m_truthLoss.muonsSeedable[3].load()
                  << "; of which <2 chambers retained: "
                  << m_truthLoss.muonsSeedLost[0].load() << "|" << m_truthLoss.muonsSeedLost[1].load()
                  << "|" << m_truthLoss.muonsSeedLost[2].load() << "|" << m_truthLoss.muonsSeedLost[3].load());
    ATH_MSG_DEBUG("  category C descriptors: best true-pair score <1e-3/1e-3..1e-2/1e-2..2e-2/2e-2..5e-2/>=5e-2: "
                  << m_truthLoss.scoreBin[0].load() << "/" << m_truthLoss.scoreBin[1].load() << "/"
                  << m_truthLoss.scoreBin[2].load() << "/" << m_truthLoss.scoreBin[3].load() << "/"
                  << m_truthLoss.scoreBin[4].load()
                  << "; pair sector delta same/adjacent: " << m_truthLoss.sectorDelta[0].load() << "/"
                  << m_truthLoss.sectorDelta[1].load()
                  << "; layer pair Inner-Middle/Middle-Outer/Inner-Outer/other: "
                  << m_truthLoss.layerPair[0].load() << "/" << m_truthLoss.layerPair[1].load() << "/"
                  << m_truthLoss.layerPair[2].load() << "/" << m_truthLoss.layerPair[3].load()
                  << "; muon input segments 2/3/4/5+: " << m_truthLoss.muonSegments[0].load() << "/"
                  << m_truthLoss.muonSegments[1].load() << "/" << m_truthLoss.muonSegments[2].load() << "/"
                  << m_truthLoss.muonSegments[3].load()
                  << "; nPrecisionHits <=4/5-6/>=7: " << m_truthLoss.precisionHits[0].load() << "/"
                  << m_truthLoss.precisionHits[1].load() << "/" << m_truthLoss.precisionHits[2].load());
 
    ATH_MSG_DEBUG(
        "SegmentEdgeInferenceAlg grouping purity (job-summed; components cover "
        "only kept components, i.e. those that survived MinSegmentsPerComponent "
        "and appear in the output): "
        << "clean=" << m_groupPurity.componentsByOutcome[0].load()
        << " mixed=" << m_groupPurity.componentsByOutcome[1].load()
        << " bkgOnly=" << m_groupPurity.componentsByOutcome[2].load()
        << "; segments inside mixed components=" << m_groupPurity.mixedComponentSegments.load()
        << "; component size histogram (2/3/4/5-9/10-19/20+): "
        << m_groupPurity.componentSizeHist[0].load() << "/" << m_groupPurity.componentSizeHist[1].load() << "/"
        << m_groupPurity.componentSizeHist[2].load() << "/" << m_groupPurity.componentSizeHist[3].load() << "/"
        << m_groupPurity.componentSizeHist[4].load() << "/" << m_groupPurity.componentSizeHist[5].load());
 
    // Kept components per outcome, binned by a candidate component-level cut variable.
    const auto printByOutcome = [&](const char* label, const auto& counters, std::size_t nBins) {
      constexpr std::array<const char*, 3> outcomes{"clean", "mixed", "bkgOnly"};
      std::stringstream sstr{};
      for (std::size_t outcome = 0; outcome < outcomes.size(); ++outcome) {
        sstr << " " << outcomes[outcome] << "=";
        for (std::size_t bin = 0; bin < nBins; ++bin) {
          sstr << (bin ? "/" : "") << counters[outcome * nBins + bin].load();
        }
      }
      ATH_MSG_DEBUG("  kept components by " << label << ":" << sstr.str());
    };
    printByOutcome("strongest selected-edge score (<0.05/0.05-0.1/0.1-0.3/0.3-0.5/>=0.5)",
                   m_groupPurity.scoreByOutcome, GroupPurityCounters::kScoreBins);
    printByOutcome("retained segments (2/3/4/5-9/10-19/20+)",
                   m_groupPurity.sizeByOutcome, GroupPurityCounters::kSizeBins);
    printByOutcome("distinct layers (1/2/3/4+)",
                   m_groupPurity.layersByOutcome, GroupPurityCounters::kLayerBins);
    printByOutcome("strongest selected cross-layer edge, components with >=2 layers "
                   "(<0.05/0.05-0.1/0.1-0.3/0.3-0.5/0.5-0.8/>=0.8)",
                   m_groupPurity.crossScoreByOutcome, GroupPurityCounters::kCrossScoreBins);
    printByOutcome("number of selected cross-layer edges, components with >=2 layers (0/1/2/3+)",
                   m_groupPurity.crossEdgesByOutcome, GroupPurityCounters::kCrossEdgeBins);
    ATH_MSG_DEBUG("  D8 clean components with >=2 layers and strongest cross-layer edge < 0.05: "
                  << "their muon has another component with >=2 layers and a stronger edge="
                  << m_groupPurity.weakCleanWithOther.load()
                  << ", it is the muon's only such component=" << m_groupPurity.weakCleanAlone.load());
 
    const std::uint64_t worstMixed = m_groupPurity.worstMixedPacked.load();
    ATH_MSG_DEBUG("  worst mixed component this job: "
                  << (worstMixed >> 32) << " distinct truth muons in "
                  << (worstMixed & 0xffffffffull) << " segments"
                  << (worstMixed == 0 ? " (no mixed components found)" : ""));
 
    ATH_MSG_DEBUG(
        "  truth-muon grouping outcome (all|barrel|transition|endcap; muons with "
        ">=1 retained segment only, fully-lost muons excluded -- see muonsSeedLost above): "
        << "clean=" << m_groupPurity.muonOutcome[0 * 3 + 0].load() << "|"
        << m_groupPurity.muonOutcome[1 * 3 + 0].load() << "|"
        << m_groupPurity.muonOutcome[2 * 3 + 0].load() << "|"
        << m_groupPurity.muonOutcome[3 * 3 + 0].load()
        << "; mixed=" << m_groupPurity.muonOutcome[0 * 3 + 1].load() << "|"
        << m_groupPurity.muonOutcome[1 * 3 + 1].load() << "|"
        << m_groupPurity.muonOutcome[2 * 3 + 1].load() << "|"
        << m_groupPurity.muonOutcome[3 * 3 + 1].load()
        << "; fragmented=" << m_groupPurity.muonOutcome[0 * 3 + 2].load() << "|"
        << m_groupPurity.muonOutcome[1 * 3 + 2].load() << "|"
        << m_groupPurity.muonOutcome[2 * 3 + 2].load() << "|"
        << m_groupPurity.muonOutcome[3 * 3 + 2].load());
  }
  return StatusCode::SUCCESS;
}
 
} // namespace MuonML