Plugins/Onnx/AmbiguityTrackClassifier.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/EventData/MultiTrajectoryHelpers.hpp"
0012 #include "Acts/EventData/TrackContainer.hpp"
0013 #include "Acts/EventData/TrackContainerFrontendConcept.hpp"
0014 #include "Acts/EventData/TrackProxyConcept.hpp"
0015 #include "Acts/Plugins/Onnx/OnnxRuntimeBase.hpp"
0016 #include "Acts/TrackFinding/detail/AmbiguityTrackClustering.hpp"
0017
0018 #include <map>
0019 #include <unordered_map>
0020 #include <vector>
0021
0022 #include <onnxruntime_cxx_api.h>
0023
0024 namespace Acts {
0025
0026 /// Onnx model implementation for track scoring and selection
0027 class AmbiguityTrackClassifier {
0028  public:
0029   /// Construct the ambiguity scoring algorithm.
0030   ///
0031   /// @param modelPath path to the model file
0032   AmbiguityTrackClassifier(const char* modelPath)
0033       : m_env(ORT_LOGGING_LEVEL_WARNING, "MLClassifier"),
0034         m_duplicateClassifier(m_env, modelPath) {}
0035
0036   /// Compute a score for each track to be used in the track selection
0037   ///
0038   /// @param clusters is a map of clusters, each cluster correspond to a vector of track ID
0039   /// @param tracks is the input track container
0040   /// @return a vector of vector of track score. Due to the architecture of the network each track only have a size 1 score vector.
0041   template <TrackContainerFrontend track_container_t>
0042   std::vector<std::vector<float>> inferScores(
0043       std::unordered_map<std::size_t, std::vector<std::size_t>>& clusters,
0044       const track_container_t& tracks) const {
0045     // Compute the number of entry (since it is smaller than the number of
0046     // track)
0047     int trackNb = 0;
0048     for (const auto& [_, val] : clusters) {
0049       trackNb += val.size();
0050     }
0051     // Input of the neural network
0052     Acts::NetworkBatchInput networkInput(trackNb, 8);
0053     std::size_t inputID = 0;
0054     // Get the input feature of the network for all the tracks
0055     for (const auto& [key, val] : clusters) {
0056       for (const auto& trackID : val) {
0057         auto track = tracks.getTrack(trackID);
0058         auto trajState = Acts::MultiTrajectoryHelpers::trajectoryState(
0059             tracks.trackStateContainer(), track.tipIndex());
0060         networkInput(inputID, 0) = trajState.nStates;
0061         networkInput(inputID, 1) = trajState.nMeasurements;
0062         networkInput(inputID, 2) = trajState.nOutliers;
0063         networkInput(inputID, 3) = trajState.nHoles;
0064         networkInput(inputID, 4) = trajState.NDF;
0065         networkInput(inputID, 5) = (trajState.chi2Sum * 1.0) /
0066                                    (trajState.NDF != 0 ? trajState.NDF : 1);
0067         networkInput(inputID, 6) = Acts::VectorHelpers::eta(track.momentum());
0068         networkInput(inputID, 7) = Acts::VectorHelpers::phi(track.momentum());
0069         inputID++;
0070       }
0071     }
0072     // Use the network to compute a score for all the tracks.
0073     std::vector<std::vector<float>> outputTensor =
0074         m_duplicateClassifier.runONNXInference(networkInput);
0075     return outputTensor;
0076   }
0077
0078   /// Select the track associated with each cluster based on the score vector
0079   ///
0080   /// @param clusters is a map of clusters, each cluster correspond to a vector of track ID
0081   /// @param outputTensor is the score vector obtained from inferScores.
0082   /// @return a vector of trackID corresponding tho the good tracks
0083   std::vector<std::size_t> trackSelection(
0084       std::unordered_map<std::size_t, std::vector<std::size_t>>& clusters,
0085       std::vector<std::vector<float>>& outputTensor) const {
0086     std::vector<std::size_t> goodTracks;
0087     std::size_t iOut = 0;
0088     // Loop over all the cluster and only keep the track with the highest score
0089     // in each cluster
0090     for (const auto& [key, val] : clusters) {
0091       std::size_t bestTrackID = 0;
0092       float bestTrackScore = 0;
0093       for (const auto& track : val) {
0094         if (outputTensor[iOut][0] > bestTrackScore) {
0095           bestTrackScore = outputTensor[iOut][0];
0096           bestTrackID = track;
0097         }
0098         iOut++;
0099       }
0100       goodTracks.push_back(bestTrackID);
0101     }
0102     return goodTracks;
0103   }
0104
0105  private:
0106   // ONNX environment
0107   Ort::Env m_env;
0108   // ONNX model for the duplicate neural network
0109   Acts::OnnxRuntimeBase m_duplicateClassifier;
0110 };
0111
0112 }  // namespace Acts