File ExaTrkXTrackFindingOnnx.hpp

namespace Acts

Note

This file is foreseen for the Geometry module to replace Extent

class ExaTrkXTrackFindingOnnx : public Acts::ExaTrkXTrackFindingBase
#include <Acts/Plugins/ExaTrkX/ExaTrkXTrackFindingOnnx.hpp>

Implementation of the Exa.TrkX track finding algorithm based on ONNX.

Uses cugraph as graph library.

Public Functions

ExaTrkXTrackFindingOnnx(const Config &config, std::unique_ptr<const Logger> _logger = Acts::getDefaultLogger("ETXTrkFnd", Logging::INFO))

Constructor of the track finding module.

Parameters
  • config – is the config struct to configure the module

  • _logger – A logger instance

~ExaTrkXTrackFindingOnnx()

Destructor.

inline const Config &config() const

Return the configuration object of the track finding module.

virtual std::optional<ExaTrkXTime> getTracks(std::vector<float> &inputValues, std::vector<int> &spacepointIDs, std::vector<std::vector<int>> &trackCandidates, const Logger &logger = getDummyLogger(), bool recordTiming = false) const override

Run the inference.

Note

The input values are not const, because the ONNX API takes only non-const pointers.

Parameters
  • inputValues – Spacepoint data as a flattened NxD array, where D is the dimensionality of a spacepoint (usually 3, but additional information like cell information can be provided).

  • spacepointIDs – The corresponding spacepoint IDs

  • trackCandidates – This vector is filled with the tracks as vectors of spacepoint IDs

  • logger – If provided, logging is enabled

  • recordTiming – If enabled, returns a ExaTrkXTime object with measured timings

Private Functions

void buildEdges(std::vector<float> &embedFeatures, std::vector<int64_t> &edgeList, int64_t numSpacepoints) const
inline const Logger &logger() const
void runSessionWithIoBinding(Ort::Session &sess, std::vector<const char*> &inputNames, std::vector<Ort::Value> &inputData, std::vector<const char*> &outputNames, std::vector<Ort::Value> &outputData) const

Private Members

Config m_cfg
std::unique_ptr<Ort::Session> m_embeddingSession
std::unique_ptr<Ort::Env> m_env
std::unique_ptr<Ort::Session> m_filterSession
std::unique_ptr<Ort::Session> m_gnnSession
std::unique_ptr<const Logger> m_logger
struct Config
#include <Acts/Plugins/ExaTrkX/ExaTrkXTrackFindingOnnx.hpp>

Configuration struct for the track finding.

Public Members

int embeddingDim = 8
float filterCut = 0.21
int knnVal = 500
std::string modelDir
float rVal = 1.6
int64_t spacepointFeatures = 3
namespace Ort