Plugins/Onnx/AmbiguityTrackClassifier.hpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2023 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 http://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/Plugins/Onnx/OnnxRuntimeBase.hpp"
0014 #include "Acts/TrackFinding/detail/AmbiguityTrackClustering.hpp"
0015
0016 #include <map>
0017 #include <unordered_map>
0018 #include <vector>
0019
0020 #include <onnxruntime_cxx_api.h>
0021
0022 namespace Acts {
0023
0024 /// Onnx model implementation for track scoring and selection
0025 class AmbiguityTrackClassifier {
0026  public:
0027   /// Construct the ambiguity scoring algorithm.
0028   ///
0029   /// @param modelPath path to the model file
0030   AmbiguityTrackClassifier(const char* modelPath)
0031       : m_env(ORT_LOGGING_LEVEL_WARNING, "MLClassifier"),
0032         m_duplicateClassifier(m_env, modelPath) {}
0033
0034   /// Compute a score for each track to be used in the track selection
0035   ///
0036   /// @param clusters is a map of clusters, each cluster correspond to a vector of track ID
0037   /// @param tracks is the input track container
0038   /// @return a vector of vector of track score. Due to the architecture of the network each track only have a size 1 score vector.
0039   template <typename track_container_t, typename traj_t,
0040             template <typename> class holder_t>
0041   std::vector<std::vector<float>> inferScores(
0042       std::unordered_map<std::size_t, std::vector<std::size_t>>& clusters,
0043       const Acts::TrackContainer<track_container_t, traj_t, holder_t>& tracks)
0044       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   /// Select the track associated with each cluster
0106   ///
0107   /// @param clusters is a map of clusters, each cluster correspond to a vector of track ID
0108   /// @param tracks is the input track container
0109   /// @return a vector of trackID corresponding tho the good tracks
0110   template <typename track_container_t, typename traj_t,
0111             template <typename> class holder_t>
0112   std::vector<std::size_t> solveAmbiguity(
0113       std::unordered_map<std::size_t, std::vector<std::size_t>>& clusters,
0114       const Acts::TrackContainer<track_container_t, traj_t, holder_t>& tracks)
0115       const {
0116     std::vector<std::vector<float>> outputTensor =
0117         inferScores(clusters, tracks);
0118     std::vector<std::size_t> goodTracks =
0119         trackSelection(clusters, outputTensor);
0120
0121     return goodTracks;
0122   }
0123
0124  private:
0125   // ONNX environment
0126   Ort::Env m_env;
0127   // ONNX model for the duplicate neural network
0128   Acts::OnnxRuntimeBase m_duplicateClassifier;
0129 };
0130
0131 }  // namespace Acts