Athena tool for DisplacedVertex graph-level ONNX inference. More...

#include <DVInferenceToolBase.h>

Inheritance diagram for MuonML::DVInferenceToolBase:

Collaboration diagram for MuonML::DVInferenceToolBase:

Classes
struct	InputTensorSpec

Public Member Functions
	~DVInferenceToolBase () override=default
StatusCode	initialize () override
StatusCode	runGraphInference (const EventContext &ctx, GraphRawData &graphData) const override
	IGraphInferenceTool entry point: build the DV event graph and run ONNX.
StatusCode	buildGraph (const EventContext &ctx, GraphRawData &graphData) const
	Build the DV ONNX input tensors: x, edge_index, edge_attr, n_muon_nodes.
StatusCode	runInference (GraphRawData &graphData) const
	Run the configured ONNX session on a graph already built by buildGraph.
StatusCode	inferEvent (const EventContext &ctx, DVInferenceResult &result) const
	Convenience event-classifier API used by DVInferenceAlg.

Protected Member Functions
StatusCode	setupModel ()
Ort::Session &	model () const
std::vector< std::string >	modelInputNames () const
std::vector< std::string >	modelOutputNames () const
StatusCode	runNamedInference (GraphRawData &graphData, const std::vector< InputTensorSpec > &inputSpecs, const std::vector< std::string > &outputNames) const
float	probabilityFromOutput (const Ort::Value &output, float &rawOutput) const

Protected Attributes
SG::ReadHandleKey< xAOD::MuonSegmentContainer >	m_segmentKey
SG::ReadHandleKeyArray< MuonR4::SpacePointContainer >	m_spacePointKeys
SG::ReadHandleKey< CaloTowerContainer >	m_towerKey
Gaudi::Property< float >	m_minTowerEnergyMeV
Gaudi::Property< float >	m_maxTowerSegmentDR
Gaudi::Property< float >	m_caloRMaxMm
Gaudi::Property< float >	m_caloZMaxMm
Gaudi::Property< int >	m_sectorModulo
Gaudi::Property< bool >	m_requireEdges
Gaudi::Property< bool >	m_useBucketSegmentSelection
Gaudi::Property< bool >	m_fallbackToAllSegments
Gaudi::Property< int >	m_maxEdges {this, "MaxEdges", -1, "Maximum number of directed segment-tower edges to create; negative means no cap"}
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_inputNMuonNodesName {this, "InputNMuonNodesName", "n_muon_nodes"}
Gaudi::Property< std::string >	m_outputName {this, "OutputName", "logits"}
Gaudi::Property< std::string >	m_singleOutputMode
Gaudi::Property< unsigned int >	m_debugDumpFirstNNodes {this, "DebugDumpFirstNNodes", 0}
Gaudi::Property< unsigned int >	m_debugDumpFirstNEdges {this, "DebugDumpFirstNEdges", 0}
Gaudi::Property< bool >	m_sanitizeNonFiniteInputs
Gaudi::Property< bool >	m_sanitizeNonFinitePredictions
bool	m_isCuda {false}
int	m_cudaDeviceId {0}

Static Protected Attributes
static constexpr std::size_t	kNodeFeatureCount = 7
static constexpr std::size_t	kEdgeFeatureCount = 5
static constexpr std::size_t	kInputTensorCount = 4
static constexpr std::array< std::string_view, kNodeFeatureCount >	kDefaultNodeFeatureNames
static constexpr std::array< std::string_view, kEdgeFeatureCount >	kDefaultEdgeFeatureNames

Private Attributes
ToolHandle< AthOnnx::IOnnxRuntimeSessionTool >	m_onnxSessionTool

Detailed Description

Athena tool for DisplacedVertex graph-level ONNX inference.

The current exported model embeds any training normalization and consumes raw tensors matching dv_converter_utils.py:

x [num_nodes, 7] edge_index [2, num_edges] edge_attr [num_edges, 5] n_muon_nodes [1] logits [1]

Nodes are ordered as muon-segment nodes first, followed by calorimeter tower nodes, because the model-side normalizer uses n_muon_nodes to split the raw x tensor into muon and calo node blocks.

Definition at line 56 of file DVInferenceToolBase.h.

Constructor & Destructor Documentation

◆ ~DVInferenceToolBase()

MuonML::DVInferenceToolBase::~DVInferenceToolBase ( )

overridedefault

Member Function Documentation

◆ buildGraph()

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

Build the DV ONNX input tensors: x, edge_index, edge_attr, n_muon_nodes.

Definition at line 259 of file DVInferenceToolBase.cxx.

                                                                           {
  graphData.graph.reset();
  graphData.featureLeaves.clear();
  graphData.srcEdges.clear();
  graphData.desEdges.clear();
  graphData.edgeIndexPacked.clear();
  graphData.spacePointsInBucket.clear();
  graphData.graph = std::make_unique<InferenceGraph>();
  graphData.graph->dataTensor.reserve(kInputTensorCount);
 
  std::vector<DVNodeAux> nodes;
 
  const xAOD::MuonSegmentContainer* segments{nullptr};
  ATH_CHECK(SG::get(segments, m_segmentKey, ctx));
 
  nodes.reserve(segments ? segments->size() : 0u);
 
  if (segments && m_useBucketSegmentSelection.value() && !m_spacePointKeys.empty()) {
    using SegmentsPerBucket_t =
        std::unordered_map<const MuonR4::SpacePointBucket*, SegmentList>;
 
    using SegmentsPerBucketSignature_t =
        std::unordered_map<std::string, SegmentList>;
 
    SegmentsPerBucket_t segmentsPerBucket{};
    SegmentsPerBucketSignature_t segmentsPerBucketSignature{};
    for (const xAOD::MuonSegment* seg : *segments) {
      const auto* detailed = MuonR4::detailedSegment(*seg);
      const MuonR4::SpacePointBucket* parentBucket = detailed->parent()->parentBucket();
      appendUniqueSegment(segmentsPerBucket[parentBucket], seg);
      const std::string parentSig = bucketSignatureKey(*parentBucket);
      if (!parentSig.empty()) appendUniqueSegment(segmentsPerBucketSignature[parentSig], seg);
    }
    std::size_t nSignatureMatchedBuckets = 0u;
    for (const SG::ReadHandleKey<MuonR4::SpacePointContainer>& spKey : m_spacePointKeys) {
      const MuonR4::SpacePointContainer* spContainer{nullptr};
      ATH_CHECK(SG::get(spContainer, spKey, ctx));
 
      for (const MuonR4::SpacePointBucket* bucket : *spContainer) {
        const auto it = segmentsPerBucket.find(bucket);
        const SegmentList* matchedSegments{nullptr};
        if (it != segmentsPerBucket.end() && !it->second.empty()) {
          matchedSegments = &it->second;
        } else {
          const std::string sig = bucketSignatureKey(*bucket);
          const auto sigIt = sig.empty() ? segmentsPerBucketSignature.end()
                                         : segmentsPerBucketSignature.find(sig);
          if (sigIt != segmentsPerBucketSignature.end() && !sigIt->second.empty()) {
            matchedSegments = &sigIt->second;
            ++nSignatureMatchedBuckets;
          }
        }
        if (!matchedSegments) continue;
 
        const int bucketSector = bucket->msSector() ? static_cast<int>(bucket->msSector()->sector()) : -1;
        const uint16_t bucketLayers = countLayersInBucket(*bucket);
        ATH_MSG_VERBOSE("DV bucket segment node source: key=" << spKey.key()
                        << " sector=" << bucketSector
                        << " layers=" << bucketLayers
                        << " segments=" << matchedSegments->size());
 
        for (const xAOD::MuonSegment* seg : *matchedSegments) {
          appendMuonSegmentNode(*seg, bucketSector, nodes);
        }
      }
    }
 
    ATH_MSG_DEBUG("DV graph built " << nodes.size()
                  << " muon nodes from BucketDumper-style SpacePointBucket-associated segments"
                  << " (signature-matched filtered buckets=" << nSignatureMatchedBuckets << ")");
  }
 
  if (segments && nodes.empty() &&
      (!m_useBucketSegmentSelection.value() || m_fallbackToAllSegments.value())) {
    if (m_useBucketSegmentSelection.value()) {
      ATH_MSG_WARNING("No bucket-associated segments were found for DV graph building; "
                      "falling back to all segments from " << m_segmentKey.key()
                      << ". This does not match the training converter exactly.");
    }
    for (const xAOD::MuonSegment* seg : *segments) {
      appendMuonSegmentNode(*seg, static_cast<int>(seg->sector()), nodes);
    }
  }
 
  if (segments && nodes.empty() && m_useBucketSegmentSelection.value() &&
      !m_fallbackToAllSegments.value()) {
    ATH_MSG_WARNING("No bucket-associated segments were found for DV graph building. "
                    "Not falling back to all segments because that does not match the training converter.");
  }
 
  const std::size_t nMuonNodes = nodes.size();
 
  if (!m_towerKey.empty() && nMuonNodes > 0u) {
    const CaloTowerContainer* towers{nullptr};
    ATH_CHECK(SG::get(towers, m_towerKey, ctx));
 
    nodes.reserve(nodes.size() + towers->size());
    for (const CaloTower* tower : *towers) {
      const float energyMeV = static_cast<float>(tower->energy());
      if (energyMeV < m_minTowerEnergyMeV) continue;
 
      const float eta = static_cast<float>(tower->eta());
      const float phi = static_cast<float>(tower->phi());
      float minDR = std::numeric_limits<float>::max();
      for (std::size_t i = 0; i < nMuonNodes; ++i) {
        minDR = std::min(
            minDR,
            static_cast<float>(xAOD::P4Helpers::deltaR(eta, phi, nodes[i].eta, nodes[i].phi)));
      }
 
      if (minDR >= m_maxTowerSegmentDR) continue;
      const Amg::Vector3D direction = Acts::makeDirectionFromPhiEta(
          static_cast<double>(phi), static_cast<double>(eta));
      const std::optional<Amg::Vector3D> posMm =
          firstIntersectionWithEnvelope(direction, m_caloRMaxMm.value(), m_caloZMaxMm.value());
      if (!posMm) continue;
 
      const Amg::Vector3D posM = (*posMm) / Gaudi::Units::m;
 
      DVNodeAux node{};
      node.kind = NodeKind::Calo;
      node.features[0] = static_cast<float>(posM.mag());
      node.features[1] = static_cast<float>(posM.theta());
      node.features[2] = static_cast<float>(posM.phi());
      node.features[3] = static_cast<float>(direction.theta());
      node.features[4] = static_cast<float>(direction.phi());
      node.features[5] = energyMeV;
      node.features[6] = static_cast<float>(tower->size());
      node.eta = eta;
      node.phi = phi;
      node.energyLike = energyMeV;
      node.direction = direction;
      node.sector = static_cast<int>(
          MuonR4::ExpandedSector{CxxUtils::wrapToPi(static_cast<double>(phi))}.msSector());
      nodes.push_back(node);
    }
  }
 
  const std::size_t nCaloNodes = nodes.size() - nMuonNodes;
  const std::size_t nNodes = nodes.size();
  graphData.spacePointsInBucket.push_back(static_cast<int64_t>(nMuonNodes));
  graphData.spacePointsInBucket.push_back(static_cast<int64_t>(nCaloNodes));
 
  if (nNodes == 0u) {
    ATH_MSG_WARNING("No muon segment or calo tower nodes found. Skipping DV inference.");
    return StatusCode::SUCCESS;
  }
 
  graphData.featureLeaves.reserve(nNodes * kNodeFeatureCount);
  for (const DVNodeAux& node : nodes) {
    graphData.featureLeaves.insert(graphData.featureLeaves.end(),
                                   node.features.begin(), node.features.end());
  }
 
  const int maxEdges = m_maxEdges.value();
  std::vector<float> edgeAttr;
  edgeAttr.reserve(2u * nMuonNodes * std::max<std::size_t>(nCaloNodes, 1u) * kEdgeFeatureCount);
 
  auto addEdge = [&graphData, &edgeAttr, maxEdges](std::size_t src,
                           std::size_t dst,
                           const DVNodeAux& a,
                           const DVNodeAux& b) -> bool {
    if (maxEdges >= 0 && static_cast<int>(graphData.srcEdges.size()) >= maxEdges) return false;
    const float dPhi = CxxUtils::deltaPhi(b.phi, a.phi);
    const float dEta = b.eta - a.eta;
    const float cosAng = std::clamp(static_cast<float>(a.direction.dot(b.direction)), -1.f, 1.f);
    std::array<float, kEdgeFeatureCount> attr{
        b.energyLike - a.energyLike,
        dPhi,
        dEta,
        cosAng,
        (a.sector == b.sector) ? 1.f : 0.f};
 
    graphData.srcEdges.push_back(static_cast<int64_t>(src));
    graphData.desEdges.push_back(static_cast<int64_t>(dst));
    edgeAttr.insert(edgeAttr.end(), attr.begin(), attr.end());
    return true;
  };
 
  bool edgeCapReached = false;
  for (std::size_t im = 0; im < nMuonNodes && !edgeCapReached; ++im) {
    for (std::size_t ic = nMuonNodes; ic < nNodes; ++ic) {
      if (xAOD::P4Helpers::deltaR(nodes[im].eta, nodes[im].phi,
                                    nodes[ic].eta, nodes[ic].phi) >= m_maxTowerSegmentDR.value()) continue;
      if (!addEdge(im, ic, nodes[im], nodes[ic])) {
        edgeCapReached = true;
        break;
      }
      if (!addEdge(ic, im, nodes[ic], nodes[im])) {
        edgeCapReached = true;
        break;
      }
    }
  }
 
  const std::size_t nEdges = graphData.srcEdges.size();
  if (m_requireEdges.value() && nEdges == 0u) {
    ATH_MSG_DEBUG("DV graph has no segment-tower edges and RequireEdges=True; skip inference.");
    graphData.graph.reset();
    return StatusCode::SUCCESS;
  }
 
  if (edgeAttr.size() != nEdges * kEdgeFeatureCount) {
    ATH_MSG_ERROR("DV edge attribute size mismatch: E=" << nEdges
                  << " edge_attr.size=" << edgeAttr.size());
    return StatusCode::FAILURE;
  }
 
  Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
  std::vector<int64_t> nodeShape{static_cast<int64_t>(nNodes), static_cast<int64_t>(kNodeFeatureCount)};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<float>(memInfo,
                                      graphData.featureLeaves.data(),
                                      graphData.featureLeaves.size(),
                                      nodeShape.data(),
                                      nodeShape.size()));
 
  graphData.edgeIndexPacked.clear();
  graphData.edgeIndexPacked.reserve(2u * nEdges);
  graphData.edgeIndexPacked.insert(graphData.edgeIndexPacked.end(), graphData.srcEdges.begin(), graphData.srcEdges.end());
  graphData.edgeIndexPacked.insert(graphData.edgeIndexPacked.end(), graphData.desEdges.begin(), graphData.desEdges.end());
 
  std::vector<int64_t> edgeIndexShape{2, static_cast<int64_t>(nEdges)};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<int64_t>(memInfo,
                                        graphData.edgeIndexPacked.data(),
                                        graphData.edgeIndexPacked.size(),
                                        edgeIndexShape.data(),
                                        edgeIndexShape.size()));
 
  Ort::AllocatorWithDefaultOptions allocator;
  std::vector<int64_t> edgeAttrShape{static_cast<int64_t>(nEdges), static_cast<int64_t>(kEdgeFeatureCount)};
  Ort::Value edgeAttrTensor = Ort::Value::CreateTensor(allocator,
                                                       edgeAttrShape.data(),
                                                       edgeAttrShape.size(),
                                                       ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);
  if (!edgeAttr.empty()) {
    float* edgeAttrData = edgeAttrTensor.GetTensorMutableData<float>();
    std::copy(edgeAttr.begin(), edgeAttr.end(), edgeAttrData);
  }
  graphData.graph->dataTensor.emplace_back(std::move(edgeAttrTensor));
 
  std::vector<int64_t> nMuonShape{1};
  graphData.graph->dataTensor.emplace_back(
      Ort::Value::CreateTensor<int64_t>(memInfo,
                                        graphData.spacePointsInBucket.data(),
                                        1,
                                        nMuonShape.data(),
                                        nMuonShape.size()));
 
  if (msgLvl(MSG::DEBUG)) {
    ATH_MSG_DEBUG("Built DV graph: N=" << nNodes << " (muon=" << nMuonNodes
                  << ", calo=" << nCaloNodes << "), E=" << nEdges
                  << ", n_muon_nodes=" << graphData.spacePointsInBucket[0]);
    const std::size_t dumpNodes = std::min<std::size_t>(m_debugDumpFirstNNodes.value(), nNodes);
    for (std::size_t i = 0; i < dumpNodes; ++i) {
      std::ostringstream row;
      row << "DVNode[" << i << "] kind=" << (nodes[i].kind == NodeKind::Muon ? "muon" : "calo") << ":";
      for (std::size_t f = 0; f < kNodeFeatureCount; ++f) {
        row << " f" << f << "=" << graphData.featureLeaves[i * kNodeFeatureCount + f];
      }
      ATH_MSG_DEBUG(row.str());
    }
    const std::size_t dumpEdges = std::min<std::size_t>(m_debugDumpFirstNEdges.value(), nEdges);
    for (std::size_t e = 0; e < dumpEdges; ++e) {
      ATH_MSG_DEBUG("DVEdge[" << e << "]: " << graphData.srcEdges[e]
                    << " -> " << graphData.desEdges[e]
                    << " edge_attr=["
                    << edgeAttr[e * kEdgeFeatureCount + 0] << ", "
                    << edgeAttr[e * kEdgeFeatureCount + 1] << ", "
                    << edgeAttr[e * kEdgeFeatureCount + 2] << ", "
                    << edgeAttr[e * kEdgeFeatureCount + 3] << ", "
                    << edgeAttr[e * kEdgeFeatureCount + 4] << "]");
    }
  }
 
  graphData.srcEdges.clear();
  graphData.desEdges.clear();
  return StatusCode::SUCCESS;
}

◆ inferEvent()

StatusCode DVInferenceToolBase::inferEvent	(	const EventContext &	ctx,
		DVInferenceResult &	result ) const

Convenience event-classifier API used by DVInferenceAlg.

Definition at line 224 of file DVInferenceToolBase.cxx.

                                                                             {
  result = DVInferenceResult{};
  GraphRawData graphData{};
  ATH_CHECK(buildGraph(ctx, graphData));
  if (!graphData.graph || graphData.graph->dataTensor.size() < kInputTensorCount) {
    ATH_MSG_WARNING("DV graph is empty; no event-classifier output will be produced.");
    return StatusCode::SUCCESS;
  }
 
  const auto xShape = graphData.graph->dataTensor[0].GetTensorTypeAndShapeInfo().GetShape();
  const auto edgeShape = graphData.graph->dataTensor[1].GetTensorTypeAndShapeInfo().GetShape();
  result.nNodes = !xShape.empty() && xShape[0] > 0 ? static_cast<std::size_t>(xShape[0]) : 0u;
  result.nEdges = edgeShape.size() > 1 && edgeShape[1] > 0 ? static_cast<std::size_t>(edgeShape[1]) : 0u;
  if (graphData.spacePointsInBucket.size() >= 2) {
    result.nMuonNodes = static_cast<std::size_t>(std::max<int64_t>(graphData.spacePointsInBucket[0], 0));
    result.nCaloNodes = static_cast<std::size_t>(std::max<int64_t>(graphData.spacePointsInBucket[1], 0));
  }
 
  ATH_CHECK(runInference(graphData));
  if (graphData.graph->dataTensor.size() <= kInputTensorCount) {
    ATH_MSG_ERROR("DV inference finished without an output tensor.");
    return StatusCode::FAILURE;
  }
 
  result.probability = probabilityFromOutput(graphData.graph->dataTensor.back(), result.rawOutput);
  result.valid = std::isfinite(result.probability);
  ATH_MSG_DEBUG("DV event classifier: N=" << result.nNodes
                << " (muon=" << result.nMuonNodes << ", calo=" << result.nCaloNodes
                << "), E=" << result.nEdges
                << ", raw=" << result.rawOutput
                << ", probability=" << result.probability);
  return StatusCode::SUCCESS;
}

◆ initialize()

StatusCode DVInferenceToolBase::initialize ( )

override

Definition at line 136 of file DVInferenceToolBase.cxx.

                                           {
  ATH_CHECK(setupModel());
  if (m_singleOutputMode.value() != "auto" &&
      m_singleOutputMode.value() != "logit" &&
      m_singleOutputMode.value() != "prob") {
    ATH_MSG_ERROR("SingleOutputMode must be one of auto, logit, prob; got " << m_singleOutputMode.value());
    return StatusCode::FAILURE;
  }
 
  if (m_useBucketSegmentSelection.value() && m_spacePointKeys.size() > 1u) {
    ATH_MSG_WARNING("DV SpacePointKeys has " << m_spacePointKeys.size()
                    << " entries. The MuonBucketDump training samples use the "
                    << "default SegmentKey array with segments attached only to "
                    << "the first SpacePointKeys entry. For parity with training, "
                    << "configure SpacePointKeys=['MuonSpacePoints'] unless the "
                    << "training dump was produced with matching segment keys for all entries.");
  }
 
  ATH_MSG_INFO("Initialized DVInferenceToolBase with SegmentKey=" << m_segmentKey.key()
               << ", SpacePointKeys=" << m_spacePointKeys.size()
               << ", " << m_useBucketSegmentSelection
               << ", TowerContainerKey="
               << (m_towerKey.empty() ? std::string("<disabled>") : m_towerKey.key())
               << ", " << m_minTowerEnergyMeV
               << ", " << m_maxTowerSegmentDR
               << ", " << m_caloRMaxMm
               << ", " << m_caloZMaxMm
               << ", " << m_fallbackToAllSegments
               << ", " << m_singleOutputMode);
  return StatusCode::SUCCESS;
}

◆ model()

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

protected

Definition at line 168 of file DVInferenceToolBase.cxx.

                                             {
  return m_onnxSessionTool->session();
}

◆ modelInputNames()

std::vector< std::string > DVInferenceToolBase::modelInputNames ( ) const

protected

Definition at line 172 of file DVInferenceToolBase.cxx.

                                                                {
  std::vector<std::string> names{};
  Ort::AllocatorWithDefaultOptions allocator;
  const std::size_t nInputs = model().GetInputCount();
  names.reserve(nInputs);
  for (std::size_t i = 0; i < nInputs; ++i) {
    auto name = model().GetInputNameAllocated(i, allocator);
    if (name) names.emplace_back(name.get());
  }
  return names;
}

◆ modelOutputNames()

std::vector< std::string > DVInferenceToolBase::modelOutputNames ( ) const

protected

Definition at line 184 of file DVInferenceToolBase.cxx.

                                                                 {
  std::vector<std::string> names{};
  Ort::AllocatorWithDefaultOptions allocator;
  const std::size_t nOutputs = model().GetOutputCount();
  names.reserve(nOutputs);
  for (std::size_t i = 0; i < nOutputs; ++i) {
    auto name = model().GetOutputNameAllocated(i, allocator);
    if (name) names.emplace_back(name.get());
  }
  return names;
}

◆ probabilityFromOutput()

float DVInferenceToolBase::probabilityFromOutput	(	const Ort::Value &	output,
		float &	rawOutput ) const

protected

Definition at line 728 of file DVInferenceToolBase.cxx.

                                                                                               {
  rawOutput = 0.f;
  const float* data = output.GetTensorData<float>();
  const auto shapeInfo = output.GetTensorTypeAndShapeInfo();
  const std::size_t nElem = shapeInfo.GetElementCount();
  if (nElem == 0 || data == nullptr) return std::numeric_limits<float>::quiet_NaN();
 
  if (nElem == 1) {
    rawOutput = data[0];
    if (m_singleOutputMode.value() == "prob") return rawOutput;
    if (m_singleOutputMode.value() == "logit" || m_singleOutputMode.value() == "auto") {
      return InferenceUtils::sigmoid(rawOutput);
    }
    return InferenceUtils::sigmoid(rawOutput);
  }
 
  if (nElem == 2) {
    rawOutput = data[1];
    const float z0 = data[0] - std::max(data[0], data[1]);
    const float z1 = data[1] - std::max(data[0], data[1]);
    const float e0 = std::exp(z0);
    const float e1 = std::exp(z1);
    return e1 / (e0 + e1);
  }
 
  ATH_MSG_WARNING("DV output tensor has " << nElem
                  << " elements; using element 0 with SingleOutputMode=" << m_singleOutputMode.value());
  rawOutput = data[0];
  if (m_singleOutputMode.value() == "prob") return rawOutput;
  if (m_singleOutputMode.value() == "logit" || m_singleOutputMode.value() == "auto") {
    return InferenceUtils::sigmoid(rawOutput);
  }
  return InferenceUtils::sigmoid(rawOutput);
}

◆ runGraphInference()

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

override

IGraphInferenceTool entry point: build the DV event graph and run ONNX.

Definition at line 214 of file DVInferenceToolBase.cxx.

                                                                                  {
  ATH_CHECK(buildGraph(ctx, graphData));
  if (!graphData.graph || graphData.graph->dataTensor.empty()) {
    ATH_MSG_DEBUG("DV graph has no input tensors; skip inference for this event.");
    return StatusCode::SUCCESS;
  }
  return runInference(graphData);
}

◆ runInference()

StatusCode DVInferenceToolBase::runInference ( GraphRawData & graphData ) const

Run the configured ONNX session on a graph already built by buildGraph.

Definition at line 655 of file DVInferenceToolBase.cxx.

                                                                          {
  const std::vector<std::string> availableInputs = modelInputNames();
  if (availableInputs.empty()) {
    ATH_MSG_ERROR("DV ONNX model has no inputs.");
    return StatusCode::FAILURE;
  }
 
  ATH_MSG_DEBUG("DV ONNX model inputs: " << joinNames(availableInputs));
 
  const std::string nodeName = m_inputNodeName.value();
  const std::string edgeIndexName = m_inputEdgeIndexName.value();
  const std::string edgeAttrName = m_inputEdgeAttrName.value();
  const std::string nMuonNodesName = m_inputNMuonNodesName.value();
 
  std::vector<InputTensorSpec> inputSpecs{};
  inputSpecs.reserve(kInputTensorCount);
 
  auto addIfPresent = [&availableInputs, &inputSpecs](const std::string& name, std::size_t tensorIndex) {
    if (hasName(availableInputs, name)) {
      inputSpecs.push_back(InputTensorSpec{name, tensorIndex});
      return true;
    }
    return false;
  };
 
  const bool hasNodeInput = addIfPresent(nodeName, 0u);
  const bool hasEdgeIndexInput = addIfPresent(edgeIndexName, 1u);
  const bool hasEdgeAttrInput = addIfPresent(edgeAttrName, 2u);
  const bool hasNMuonNodesInput = addIfPresent(nMuonNodesName, 3u);
 
  if (!hasNodeInput || !hasEdgeIndexInput) {
    ATH_MSG_ERROR("DV ONNX model is missing required inputs. Expected at least "
                  << nodeName << " and " << edgeIndexName
                  << "; model inputs are: " << joinNames(availableInputs));
    return StatusCode::FAILURE;
  }
 
  if (!hasEdgeAttrInput) {
    ATH_MSG_DEBUG("DV ONNX model has no input named " << edgeAttrName
                 << "; not binding edge_attr. This is expected for exports where "
                 << "the architecture does not consume edge attributes and ONNX pruned the input.");
  }
  if (!hasNMuonNodesInput) {
    ATH_MSG_DEBUG("DV ONNX model has no input named " << nMuonNodesName
                 << "; not binding n_muon_nodes. This is expected only if the exported "
                 << "model does not need model-side muon/calo normalization.");
  }
 
  for (const std::string& inputName : availableInputs) {
    if (inputName != nodeName && inputName != edgeIndexName &&
        inputName != edgeAttrName && inputName != nMuonNodesName) {
      ATH_MSG_ERROR("DV ONNX model has unsupported input " << inputName
                    << ". Configure the input-name properties or update the tool mapping.");
      return StatusCode::FAILURE;
    }
  }
 
  std::vector<std::string> outputNames{};
  const std::vector<std::string> availableOutputs = modelOutputNames();
  if (hasName(availableOutputs, m_outputName.value())) {
    outputNames.push_back(m_outputName.value());
  } else if (!availableOutputs.empty()) {
    ATH_MSG_WARNING("DV ONNX model has no output named " << m_outputName.value()
                    << "; using first model output " << availableOutputs.front() << ".");
    outputNames.push_back(availableOutputs.front());
  } else {
    ATH_MSG_ERROR("DV ONNX model has no outputs.");
    return StatusCode::FAILURE;
  }
 
  return runNamedInference(graphData, inputSpecs, outputNames);
}

◆ runNamedInference()

StatusCode DVInferenceToolBase::runNamedInference	(	GraphRawData &	graphData,
		const std::vector< InputTensorSpec > &	inputSpecs,
		const std::vector< std::string > &	outputNames ) const

protected

Definition at line 541 of file DVInferenceToolBase.cxx.

                                                   {
  if (!graphData.graph) {
    ATH_MSG_ERROR("Graph data is not built.");
    return StatusCode::FAILURE;
  }
  if (inputSpecs.empty()) {
    ATH_MSG_ERROR("No DV ONNX inputs were selected for inference.");
    return StatusCode::FAILURE;
  }
 
  for (const InputTensorSpec& spec : inputSpecs) {
    if (spec.tensorIndex >= graphData.graph->dataTensor.size()) {
      ATH_MSG_ERROR("Input " << spec.name << " requests tensor index " << spec.tensorIndex
                    << " but only " << graphData.graph->dataTensor.size()
                    << " tensors were prepared.");
      return StatusCode::FAILURE;
    }
  }
 
  std::vector<const char*> inputNamePtrs{};
  inputNamePtrs.reserve(inputSpecs.size());
  for (const InputTensorSpec& spec : inputSpecs) {
    inputNamePtrs.push_back(spec.name.c_str());
  }
 
  std::vector<const char*> outputNamePtrs{};
  outputNamePtrs.reserve(outputNames.size());
  for (const std::string& name : outputNames) {
    outputNamePtrs.push_back(name.c_str());
  }
 
  graphData.graph->dataTensor.reserve(graphData.graph->dataTensor.size() + outputNamePtrs.size()); 
 
  Ort::RunOptions runOptions;
  runOptions.SetRunLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);
 
  if (m_isCuda) {
    Ort::IoBinding binding(model());
    for (const InputTensorSpec& spec : inputSpecs) {
      binding.BindInput(spec.name.c_str(), graphData.graph->dataTensor[spec.tensorIndex]);
    }
    Ort::MemoryInfo cpuOut = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    for (const char* outName : outputNamePtrs) {
      binding.BindOutput(outName, cpuOut);
    }
 
    model().Run(runOptions, binding);
    binding.SynchronizeOutputs();
 
    std::vector<Ort::Value> outputs = binding.GetOutputValues();
    if (outputs.empty()) {
      ATH_MSG_ERROR("IoBinding inference returned empty output.");
      return StatusCode::FAILURE;
    }
 
    if (m_sanitizeNonFinitePredictions.value()) {
      float* outData = outputs[0].GetTensorMutableData<float>();
      const std::size_t outSize = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
      for (std::size_t i = 0; i < outSize; ++i) {
        if (!std::isfinite(outData[i])) {
          ATH_MSG_WARNING("Non-finite DV prediction detected at " << i << " -> set to -100.");
          outData[i] = -100.f;
        }
      }
    }
 
    for (auto& v : outputs) {
      graphData.graph->dataTensor.emplace_back(std::move(v));
    }
    return StatusCode::SUCCESS;
  }
 
  std::vector<Ort::Value> orderedInputs{};
  orderedInputs.reserve(inputSpecs.size());
  for (const InputTensorSpec& spec : inputSpecs) {
    orderedInputs.emplace_back(std::move(graphData.graph->dataTensor[spec.tensorIndex]));
  }
 
  std::vector<Ort::Value> outputs =
      model().Run(runOptions,
                  inputNamePtrs.data(),
                  orderedInputs.data(),
                  inputNamePtrs.size(),
                  outputNamePtrs.data(),
                  outputNamePtrs.size());
 
  if (outputs.empty()) {
    ATH_MSG_ERROR("Inference returned empty output.");
    return StatusCode::FAILURE;
  }
 
  ATH_MSG_DEBUG("DV ONNX raw output elementCount = "
                << outputs[0].GetTensorTypeAndShapeInfo().GetElementCount());
 
  if (m_sanitizeNonFinitePredictions.value()) {
    float* outData = outputs[0].GetTensorMutableData<float>();
    const std::size_t outSize = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();
    for (std::size_t i = 0; i < outSize; ++i) {
      if (!std::isfinite(outData[i])) {
        ATH_MSG_WARNING("Non-finite DV prediction detected at " << i << " -> set to -100.");
        outData[i] = -100.f;
      }
    }
  }
 
  for (auto& v : outputs) {
    graphData.graph->dataTensor.emplace_back(std::move(v));
  }
  return StatusCode::SUCCESS;
}

◆ setupModel()

StatusCode DVInferenceToolBase::setupModel ( )

protected

Definition at line 196 of file DVInferenceToolBase.cxx.

                                           {
  ATH_CHECK(m_onnxSessionTool.retrieve());
  ATH_CHECK(m_segmentKey.initialize());
  ATH_CHECK(m_spacePointKeys.initialize());
  ATH_CHECK(m_towerKey.initialize(SG::AllowEmpty));
 
  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;
}

Member Data Documentation

◆ kDefaultEdgeFeatureNames

std::array<std::string_view, kEdgeFeatureCount> MuonML::DVInferenceToolBase::kDefaultEdgeFeatureNames

staticconstexprprotected

Initial value:

= {

"d_energy_like", "d_phi", "d_eta", "cos_angle", "same_sector"}

Definition at line 82 of file DVInferenceToolBase.h.

82 {

83 "d_energy_like", "d_phi", "d_eta", "cos_angle", "same_sector"};

◆ kDefaultNodeFeatureNames

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

staticconstexprprotected

Initial value:

= {

"r_pos", "theta_pos", "phi_pos", "theta_dir", "phi_dir", "energy_like", "nCells_or_DoF"}

Definition at line 80 of file DVInferenceToolBase.h.

80 {

81 "r_pos", "theta_pos", "phi_pos", "theta_dir", "phi_dir", "energy_like", "nCells_or_DoF"};

◆ kEdgeFeatureCount

std::size_t MuonML::DVInferenceToolBase::kEdgeFeatureCount = 5

staticconstexprprotected

Definition at line 77 of file DVInferenceToolBase.h.

◆ kInputTensorCount

std::size_t MuonML::DVInferenceToolBase::kInputTensorCount = 4

staticconstexprprotected

Definition at line 78 of file DVInferenceToolBase.h.

◆ kNodeFeatureCount

std::size_t MuonML::DVInferenceToolBase::kNodeFeatureCount = 7

staticconstexprprotected

Definition at line 76 of file DVInferenceToolBase.h.

◆ m_caloRMaxMm

Gaudi::Property<float> MuonML::DVInferenceToolBase::m_caloRMaxMm

protected

Initial value:

{

this, "CaloRMaxMm", 4250.f, "Barrel radius used for the calo-envelope intersection in mm"}

Definition at line 114 of file DVInferenceToolBase.h.

114 {

115 this, "CaloRMaxMm", 4250.f, "Barrel radius used for the calo-envelope intersection in mm"};

◆ m_caloZMaxMm

Gaudi::Property<float> MuonML::DVInferenceToolBase::m_caloZMaxMm

protected

Initial value:

{

this, "CaloZMaxMm", 6500.f, "Endcap |z| used for the calo-envelope intersection in mm"}

Definition at line 116 of file DVInferenceToolBase.h.

116 {

117 this, "CaloZMaxMm", 6500.f, "Endcap |z| used for the calo-envelope intersection in mm"};

◆ m_cudaDeviceId

int MuonML::DVInferenceToolBase::m_cudaDeviceId {0}

protected

Definition at line 145 of file DVInferenceToolBase.h.

145{0};

◆ m_debugDumpFirstNEdges

Gaudi::Property<unsigned int> MuonML::DVInferenceToolBase::m_debugDumpFirstNEdges {this, "DebugDumpFirstNEdges", 0}

protected

Definition at line 138 of file DVInferenceToolBase.h.

138{this, "DebugDumpFirstNEdges", 0};

◆ m_debugDumpFirstNNodes

Gaudi::Property<unsigned int> MuonML::DVInferenceToolBase::m_debugDumpFirstNNodes {this, "DebugDumpFirstNNodes", 0}

protected

Definition at line 137 of file DVInferenceToolBase.h.

137{this, "DebugDumpFirstNNodes", 0};

◆ m_fallbackToAllSegments

Gaudi::Property<bool> MuonML::DVInferenceToolBase::m_fallbackToAllSegments

protected

Initial value:

{

this, "FallbackToAllSegments", false, "If bucket-segment matching fails, fall back to all SegmentKey segments."}

Definition at line 124 of file DVInferenceToolBase.h.

124 {

125 this, "FallbackToAllSegments", false, "If bucket-segment matching fails, fall back to all SegmentKey segments."};

◆ m_inputEdgeAttrName

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

protected

Definition at line 130 of file DVInferenceToolBase.h.

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

◆ m_inputEdgeIndexName

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

protected

Definition at line 129 of file DVInferenceToolBase.h.

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

◆ m_inputNMuonNodesName

Gaudi::Property<std::string> MuonML::DVInferenceToolBase::m_inputNMuonNodesName {this, "InputNMuonNodesName", "n_muon_nodes"}

protected

Definition at line 131 of file DVInferenceToolBase.h.

131{this, "InputNMuonNodesName", "n_muon_nodes"};

◆ m_inputNodeName

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

protected

Definition at line 128 of file DVInferenceToolBase.h.

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

◆ m_isCuda

bool MuonML::DVInferenceToolBase::m_isCuda {false}

protected

Definition at line 144 of file DVInferenceToolBase.h.

144{false};

◆ m_maxEdges

Gaudi::Property<int> MuonML::DVInferenceToolBase::m_maxEdges {this, "MaxEdges", -1, "Maximum number of directed segment-tower edges to create; negative means no cap"}

protected

Definition at line 126 of file DVInferenceToolBase.h.

126{this, "MaxEdges", -1, "Maximum number of directed segment-tower edges to create; negative means no cap"};

◆ m_maxTowerSegmentDR

Gaudi::Property<float> MuonML::DVInferenceToolBase::m_maxTowerSegmentDR

protected

Initial value:

{

this, "MaxTowerSegmentDR", 0.4f, "Maximum segment-calo deltaR used in the converter"}

Definition at line 112 of file DVInferenceToolBase.h.

112 {

113 this, "MaxTowerSegmentDR", 0.4f, "Maximum segment-calo deltaR used in the converter"};

◆ m_minTowerEnergyMeV

Gaudi::Property<float> MuonML::DVInferenceToolBase::m_minTowerEnergyMeV

protected

Initial value:

{

this, "MinTowerEnergyMeV", 1000.f, "Minimum calo tower energy used as a DV graph node"}

Definition at line 110 of file DVInferenceToolBase.h.

110 {

111 this, "MinTowerEnergyMeV", 1000.f, "Minimum calo tower energy used as a DV graph node"};

◆ m_onnxSessionTool

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

private

Initial value:

{

this, "ModelSession", "", "ONNX Runtime session tool for the DV classifier"}

Definition at line 148 of file DVInferenceToolBase.h.

148 {

149 this, "ModelSession", "", "ONNX Runtime session tool for the DV classifier"};

◆ m_outputName

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

protected

Definition at line 132 of file DVInferenceToolBase.h.

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

◆ m_requireEdges

Gaudi::Property<bool> MuonML::DVInferenceToolBase::m_requireEdges

protected

Initial value:

{

this, "RequireEdges", false, "Skip inference when the event graph has no segment-tower edges"}

Definition at line 120 of file DVInferenceToolBase.h.

120 {

121 this, "RequireEdges", false, "Skip inference when the event graph has no segment-tower edges"};

◆ m_sanitizeNonFiniteInputs

Gaudi::Property<bool> MuonML::DVInferenceToolBase::m_sanitizeNonFiniteInputs

protected

Initial value:

{

this, "SanitizeNonFiniteInputs", true, "Replace non-finite input features with zero before creating ONNX tensors"}

Definition at line 139 of file DVInferenceToolBase.h.

139 {

140 this, "SanitizeNonFiniteInputs", true, "Replace non-finite input features with zero before creating ONNX tensors"};

◆ m_sanitizeNonFinitePredictions

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

protected

Initial value:

{

this, "SanitizeNonFinitePredictions", false, "Replace non-finite ONNX outputs with -100 and log a warning"}

Definition at line 141 of file DVInferenceToolBase.h.

141 {

142 this, "SanitizeNonFinitePredictions", false, "Replace non-finite ONNX outputs with -100 and log a warning"};

◆ m_sectorModulo

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

protected

Initial value:

{

this, "SectorModulo", 16, "Number of sectors used by the calo phi->sector converter"}

Definition at line 118 of file DVInferenceToolBase.h.

118 {

119 this, "SectorModulo", 16, "Number of sectors used by the calo phi->sector converter"};

◆ m_segmentKey

SG::ReadHandleKey<xAOD::MuonSegmentContainer> MuonML::DVInferenceToolBase::m_segmentKey

protected

Initial value:

{

this, "SegmentKey", "MuonSegmentsFromR4", "Input R4 muon segment container"}

Definition at line 102 of file DVInferenceToolBase.h.

102 {

103 this, "SegmentKey", "MuonSegmentsFromR4", "Input R4 muon segment container"};

◆ m_singleOutputMode

Gaudi::Property<std::string> MuonML::DVInferenceToolBase::m_singleOutputMode

protected

Initial value:

{

this, "SingleOutputMode", "logit", "How to interpret a one-value output: auto, logit, or prob"}

Definition at line 134 of file DVInferenceToolBase.h.

134 {

135 this, "SingleOutputMode", "logit", "How to interpret a one-value output: auto, logit, or prob"};

◆ m_spacePointKeys

SG::ReadHandleKeyArray<MuonR4::SpacePointContainer> MuonML::DVInferenceToolBase::m_spacePointKeys

protected

Initial value:

{
      this, "SpacePointKeys", {"MuonSpacePoints"},
      "Default is MuonSpacePoints only, matching the training MuonBucketDump SegmentKey alignment."}

Definition at line 104 of file DVInferenceToolBase.h.

                                                                  {
      this, "SpacePointKeys", {"MuonSpacePoints"},
      "Default is MuonSpacePoints only, matching the training MuonBucketDump SegmentKey alignment."};

◆ m_towerKey

SG::ReadHandleKey<CaloTowerContainer> MuonML::DVInferenceToolBase::m_towerKey

protected

Initial value:

{

this, "TowerContainerKey", "CombinedTower", "Input calorimeter tower container"}

Definition at line 107 of file DVInferenceToolBase.h.

107 {

108 this, "TowerContainerKey", "CombinedTower", "Input calorimeter tower container"};

◆ m_useBucketSegmentSelection

Gaudi::Property<bool> MuonML::DVInferenceToolBase::m_useBucketSegmentSelection

protected

Initial value:

{

this, "UseBucketSegmentSelection", true, "Build muon nodes from segment-parent SpacePoint buckets"}

Definition at line 122 of file DVInferenceToolBase.h.

122 {

123 this, "UseBucketSegmentSelection", true, "Build muon nodes from segment-parent SpacePoint buckets"};

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

Classes

Public Member Functions

Protected Member Functions

Protected Attributes

Static Protected Attributes

Private Attributes

Detailed Description

Constructor & Destructor Documentation

◆ ~DVInferenceToolBase()

Member Function Documentation

◆ buildGraph()

◆ inferEvent()

◆ initialize()

◆ model()

◆ modelInputNames()

◆ modelOutputNames()

◆ probabilityFromOutput()

◆ runGraphInference()

◆ runInference()

◆ runNamedInference()

◆ setupModel()

Member Data Documentation

◆ kDefaultEdgeFeatureNames

◆ kDefaultNodeFeatureNames

◆ kEdgeFeatureCount

◆ kInputTensorCount

◆ kNodeFeatureCount

◆ m_caloRMaxMm

◆ m_caloZMaxMm

◆ m_cudaDeviceId

◆ m_debugDumpFirstNEdges

◆ m_debugDumpFirstNNodes

◆ m_fallbackToAllSegments

◆ m_inputEdgeAttrName

◆ m_inputEdgeIndexName

◆ m_inputNMuonNodesName

◆ m_inputNodeName

◆ m_isCuda

◆ m_maxEdges

◆ m_maxTowerSegmentDR

◆ m_minTowerEnergyMeV

◆ m_onnxSessionTool

◆ m_outputName

◆ m_requireEdges

◆ m_sanitizeNonFiniteInputs

◆ m_sanitizeNonFinitePredictions

◆ m_sectorModulo

◆ m_segmentKey

◆ m_singleOutputMode

◆ m_spacePointKeys

◆ m_towerKey

◆ m_useBucketSegmentSelection