TrackFindingML/src/AmbiguityResolutionML.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/AmbiguityResolutionML.hpp"
0010
0011 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0012 #include "ActsExamples/EventData/Measurement.hpp"
0013
0014 ActsExamples::AmbiguityResolutionML::AmbiguityResolutionML(
0015     std::string name, Acts::Logging::Level lvl)
0016     : ActsExamples::IAlgorithm(name, lvl) {}
0017
0018 std::multimap<int, std::pair<std::size_t, std::vector<std::size_t>>>
0019 ActsExamples::AmbiguityResolutionML::mapTrackHits(
0020     const ActsExamples::ConstTrackContainer& tracks,
0021     int nMeasurementsMin) const {
0022   std::multimap<int, std::pair<std::size_t, std::vector<std::size_t>>> trackMap;
0023   // Loop over all the trajectories in the events
0024   for (const auto& track : tracks) {
0025     std::vector<std::size_t> hits;
0026     int nbMeasurements = 0;
0027     // Store the hits id for the trajectory and compute the number of
0028     // measurement
0029     tracks.trackStateContainer().visitBackwards(
0030         track.tipIndex(), [&](const auto& state) {
0031           if (state.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
0032             std::size_t indexHit =
0033                 state.getUncalibratedSourceLink()
0034                     .template get<ActsExamples::IndexSourceLink>()
0035                     .index();
0036             hits.emplace_back(indexHit);
0037             ++nbMeasurements;
0038           }
0039         });
0040     if (nbMeasurements < nMeasurementsMin) {
0041       continue;
0042     }
0043     trackMap.emplace(nbMeasurements, std::make_pair(track.index(), hits));
0044   }
0045   return trackMap;
0046 }
0047
0048 ActsExamples::ConstTrackContainer
0049 ActsExamples::AmbiguityResolutionML::prepareOutputTrack(
0050     const ActsExamples::ConstTrackContainer& tracks,
0051     std::vector<std::size_t>& goodTracks) const {
0052   std::shared_ptr<Acts::ConstVectorMultiTrajectory> trackStateContainer =
0053       tracks.trackStateContainerHolder();
0054   auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
0055   trackContainer->reserve(goodTracks.size());
0056   // Temporary empty track state container: we don't change the original one,
0057   // but we need one for filtering
0058   auto tempTrackStateContainer =
0059       std::make_shared<Acts::VectorMultiTrajectory>();
0060
0061   TrackContainer solvedTracks{trackContainer, tempTrackStateContainer};
0062   solvedTracks.ensureDynamicColumns(tracks);
0063
0064   for (auto&& iTrack : goodTracks) {
0065     auto destProxy = solvedTracks.makeTrack();
0066     auto srcProxy = tracks.getTrack(iTrack);
0067     destProxy.copyFrom(srcProxy, false);
0068     destProxy.tipIndex() = srcProxy.tipIndex();
0069   }
0070
0071   ConstTrackContainer outputTracks{
0072       std::make_shared<Acts::ConstVectorTrackContainer>(
0073           std::move(*trackContainer)),
0074       trackStateContainer};
0075   return outputTracks;
0076 }