AmbiguityTrackClassifier.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file AmbiguityTrackClassifier.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2023 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once

#include "Acts/EventData/MultiTrajectoryHelpers.hpp"
#include "Acts/EventData/TrackContainer.hpp"
#include "Acts/Plugins/Onnx/OnnxRuntimeBase.hpp"
#include "Acts/TrackFinding/detail/AmbiguityTrackClustering.hpp"

#include <map>
#include <unordered_map>
#include <vector>

#include <onnxruntime_cxx_api.h>

namespace Acts {

class AmbiguityTrackClassifier {
 public:
  AmbiguityTrackClassifier(const char* modelPath)
      : m_env(ORT_LOGGING_LEVEL_WARNING, "MLClassifier"),
        m_duplicateClassifier(m_env, modelPath) {}

  template <typename track_container_t, typename traj_t,
            template <typename> class holder_t>
  std::vector<std::vector<float>> inferScores(
      std::unordered_map<int, std::vector<int>>& clusters,
      const Acts::TrackContainer<track_container_t, traj_t, holder_t>& tracks)
      const {
    // Compute the number of entry (since it is smaller than the number of
    // track)
    int trackNb = 0;
    for (const auto& [_, val] : clusters) {
      trackNb += val.size();
    }
    // Input of the neural network
    Acts::NetworkBatchInput networkInput(trackNb, 8);
    int inputID = 0;
    // Get the input feature of the network for all the tracks
    for (const auto& [key, val] : clusters) {
      for (const auto& trackID : val) {
        auto track = tracks.getTrack(trackID);
        auto trajState = Acts::MultiTrajectoryHelpers::trajectoryState(
            tracks.trackStateContainer(), track.tipIndex());
        networkInput(inputID, 0) = trajState.nStates;
        networkInput(inputID, 1) = trajState.nMeasurements;
        networkInput(inputID, 2) = trajState.nOutliers;
        networkInput(inputID, 3) = trajState.nHoles;
        networkInput(inputID, 4) = trajState.NDF;
        networkInput(inputID, 5) = (trajState.chi2Sum * 1.0) /
                                   (trajState.NDF != 0 ? trajState.NDF : 1);
        networkInput(inputID, 6) = Acts::VectorHelpers::eta(track.momentum());
        networkInput(inputID, 7) = Acts::VectorHelpers::phi(track.momentum());
        inputID++;
      }
    }
    // Use the network to compute a score for all the tracks.
    std::vector<std::vector<float>> outputTensor =
        m_duplicateClassifier.runONNXInference(networkInput);
    return outputTensor;
  }

  std::vector<int> trackSelection(
      std::unordered_map<int, std::vector<int>>& clusters,
      std::vector<std::vector<float>>& outputTensor) const {
    std::vector<int> goodTracks;
    int iOut = 0;
    // Loop over all the cluster and only keep the track with the highest score
    // in each cluster
    for (const auto& [key, val] : clusters) {
      int bestTrackID = 0;
      float bestTrackScore = 0;
      for (const auto& track : val) {
        if (outputTensor[iOut][0] > bestTrackScore) {
          bestTrackScore = outputTensor[iOut][0];
          bestTrackID = track;
        }
        iOut++;
      }
      goodTracks.push_back(bestTrackID);
    }
    return goodTracks;
  }

  template <typename track_container_t, typename traj_t,
            template <typename> class holder_t>
  std::vector<int> solveAmbuguity(
      std::unordered_map<int, std::vector<int>>& clusters,
      const Acts::TrackContainer<track_container_t, traj_t, holder_t>& tracks)
      const {
    std::vector<std::vector<float>> outputTensor =
        inferScores(clusters, tracks);
    std::vector<int> goodTracks = trackSelection(clusters, outputTensor);

    return goodTracks;
  }

 private:
  // ONNX environment
  Ort::Env m_env;
  // ONNX model for the duplicate neural network
  Acts::OnnxRuntimeBase m_duplicateClassifier;
};

}  // namespace Acts