TrackFindingML/src/AmbiguityResolutionMLAlgorithm.cpp

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 #include "ActsExamples/TrackFindingML/AmbiguityResolutionMLAlgorithm.hpp"
0010
0011 #include "ActsExamples/Framework/ProcessCode.hpp"
0012
0013 #include <iterator>
0014 #include <map>
0015
0016 ActsExamples::AmbiguityResolutionMLAlgorithm::AmbiguityResolutionMLAlgorithm(
0017     ActsExamples::AmbiguityResolutionMLAlgorithm::Config cfg,
0018     Acts::Logging::Level lvl)
0019     : ActsExamples::AmbiguityResolutionML("AmbiguityResolutionMLAlgorithm",
0020                                           lvl),
0021       m_cfg(std::move(cfg)),
0022       m_duplicateClassifier(m_cfg.inputDuplicateNN.c_str()) {
0023   if (m_cfg.inputTracks.empty()) {
0024     throw std::invalid_argument("Missing trajectories input collection");
0025   }
0026   if (m_cfg.outputTracks.empty()) {
0027     throw std::invalid_argument("Missing trajectories output collection");
0028   }
0029   m_inputTracks.initialize(m_cfg.inputTracks);
0030   m_outputTracks.initialize(m_cfg.outputTracks);
0031 }
0032
0033 ActsExamples::ProcessCode ActsExamples::AmbiguityResolutionMLAlgorithm::execute(
0034     const AlgorithmContext& ctx) const {
0035   // Read input data
0036   const auto& tracks = m_inputTracks(ctx);
0037   // Associate measurement to their respective tracks
0038   std::multimap<int, std::pair<std::size_t, std::vector<std::size_t>>>
0039       trackMap = mapTrackHits(tracks, m_cfg.nMeasurementsMin);
0040   auto cluster = Acts::detail::clusterDuplicateTracks(trackMap);
0041   // Select the ID of the track we want to keep
0042   std::vector<std::size_t> goodTracks =
0043       m_duplicateClassifier.solveAmbiguity(cluster, tracks);
0044   // Prepare the output track collection from the IDs
0045   auto outputTracks = prepareOutputTrack(tracks, goodTracks);
0046   m_outputTracks(ctx, std::move(outputTracks));
0047
0048   return ActsExamples::ProcessCode::SUCCESS;
0049 }