BucketInferenceToolBase.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN
  for the benefit of the ATLAS collaboration
*/
#ifndef MUONINFERENCETOOLS_BUCKETINFERENCETOOLBASE_H
#define MUONINFERENCETOOLS_BUCKETINFERENCETOOLBASE_H
 
#include "MuonInferenceInterfaces/IGraphInferenceTool.h"
#include "MuonInferenceInterfaces/GraphData.h"
#include "AthOnnxInterfaces/IOnnxRuntimeSessionTool.h"
#include "MuonSpacePoint/SpacePointContainer.h"
 
#include "AthenaBaseComps/AthAlgTool.h"
#include "StoreGate/ReadHandleKey.h"
#include "StoreGate/ReadCondHandleKey.h"
 
#include <onnxruntime_cxx_api.h>
#include <vector>
 
class ActsGeometryContext;
 
namespace MuonML {

class BucketInferenceToolBase : public extends<AthAlgTool, IGraphInferenceTool> {
public:
  using base_class::base_class;
  ~BucketInferenceToolBase() override = default;

  StatusCode buildGraph(const EventContext& ctx, GraphRawData& graphData) const;

  StatusCode runInference(GraphRawData& graphData) const;
 
protected:
  StatusCode setupModel();
  Ort::Session& model() const;

  StatusCode buildFeaturesOnly(const EventContext& ctx, GraphRawData& graphData) const;

  StatusCode buildTransformerInputs(const EventContext& ctx, GraphRawData& graphData) const;

  StatusCode runNamedInference(GraphRawData& graphData,
                               const std::vector<const char*>& inputNames,
                               const std::vector<const char*>& outputNames) const;
 
  SG::ReadHandleKey<MuonR4::SpacePointContainer> m_readKey{this, "ReadSpacePoints", "MuonSpacePoints"};
  SG::ReadHandleKey<ActsTrk::GeometryContext> m_geoCtxKey{this, "AlignmentKey", "ActsAlignment", "cond handle key"};
 
  // Sparse-graph parameters (GNN)
  Gaudi::Property<int>    m_minLayers{this, "MinLayersValid", 3};
  Gaudi::Property<int>    m_maxChamberDelta{this, "MaxChamberDelta", 13};
  Gaudi::Property<int>    m_maxSectorDelta{this, "MaxSectorDelta", 1};
  Gaudi::Property<double> m_maxDistXY{this, "MaxDistXY", 6800.0};
  Gaudi::Property<double> m_maxAbsDz{this, "MaxAbsDz", 15000.0};
 
  // Debug/validation knobs
  Gaudi::Property<unsigned int> m_debugDumpFirstNNodes{this, "DebugDumpFirstNNodes", 5};
  Gaudi::Property<unsigned int> m_debugDumpFirstNEdges{this, "DebugDumpFirstNEdges", 12};
  Gaudi::Property<bool>         m_validateEdges{this, "ValidateEdges", true};
 
private:
  ToolHandle<AthOnnx::IOnnxRuntimeSessionTool> m_onnxSessionTool{
      this, "ModelSession", ""};
};
 
} // namespace MuonML
 
#endif