TrackFindingML/src/AmbiguityResolutionMLAlgorithm.cpp

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 #include "ActsExamples/TrackFindingML/AmbiguityResolutionMLAlgorithm.hpp"
0010
0011 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0012 #include "ActsExamples/EventData/Measurement.hpp"
0013 #include "ActsExamples/Framework/ProcessCode.hpp"
0014
0015 #include <iterator>
0016 #include <map>
0017
0018 using namespace Acts;
0019 using namespace ActsPlugins;
0020
0021 namespace ActsExamples {
0022
0023 namespace {
0024
0025 static std::size_t sourceLinkHash(const SourceLink& a) {
0026   return static_cast<std::size_t>(a.get<IndexSourceLink>().index());
0027 }
0028
0029 static bool sourceLinkEquality(const SourceLink& a, const SourceLink& b) {
0030   return a.get<IndexSourceLink>().index() == b.get<IndexSourceLink>().index();
0031 }
0032
0033 }  // namespace
0034
0035 AmbiguityResolutionMLAlgorithm::AmbiguityResolutionMLAlgorithm(
0036     const Config& cfg, std::unique_ptr<const Acts::Logger> logger)
0037     : IAlgorithm("AmbiguityResolutionMLAlgorithm", std::move(logger)),
0038       m_cfg(cfg),
0039       m_ambiML(m_cfg.toAmbiguityResolutionMLConfig(), this->logger().clone()) {
0040   if (m_cfg.inputTracks.empty()) {
0041     throw std::invalid_argument("Missing trajectories input collection");
0042   }
0043   if (m_cfg.outputTracks.empty()) {
0044     throw std::invalid_argument("Missing trajectories output collection");
0045   }
0046   m_inputTracks.initialize(m_cfg.inputTracks);
0047   m_outputTracks.initialize(m_cfg.outputTracks);
0048 }
0049
0050 ProcessCode AmbiguityResolutionMLAlgorithm::execute(
0051     const AlgorithmContext& ctx) const {
0052   // Read input data
0053   const auto& tracks = m_inputTracks(ctx);
0054   // Associate measurement to their respective tracks to prepare the track
0055   // shared hits based clustering
0056   std::multimap<int, std::pair<std::size_t, std::vector<std::size_t>>>
0057       trackMap =
0058           m_ambiML.mapTrackHits(tracks, &sourceLinkHash, &sourceLinkEquality);
0059   // Cluster the tracks based on the shared hits
0060   auto cluster = Acts::detail::clusterDuplicateTracks(trackMap);
0061   // Select the ID of the track we want to keep
0062   std::vector<std::size_t> goodTracks =
0063       m_ambiML.solveAmbiguity(cluster, tracks);
0064   // Prepare the output track collection from the IDs
0065   TrackContainer solvedTracks{std::make_shared<VectorTrackContainer>(),
0066                               std::make_shared<VectorMultiTrajectory>()};
0067   solvedTracks.ensureDynamicColumns(tracks);
0068   for (auto iTrack : goodTracks) {
0069     auto destProxy = solvedTracks.makeTrack();
0070     auto srcProxy = tracks.getTrack(iTrack);
0071     destProxy.copyFromWithoutStates(srcProxy);
0072     destProxy.tipIndex() = srcProxy.tipIndex();
0073   }
0074
0075   ConstTrackContainer outputTracks{std::make_shared<ConstVectorTrackContainer>(
0076                                        std::move(solvedTracks.container())),
0077                                    tracks.trackStateContainerHolder()};
0078
0079   m_outputTracks(ctx, std::move(outputTracks));
0080
0081   return ProcessCode::SUCCESS;
0082 }
0083
0084 }  // namespace ActsExamples