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