ExaTrkX/src/BoostTrackBuilding.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 "Acts/Plugins/ExaTrkX/BoostTrackBuilding.hpp"
0010
0011 #include "Acts/Utilities/Zip.hpp"
0012
0013 #include <algorithm>
0014
0015 #include <boost/beast/core/span.hpp>
0016 #include <boost/graph/adjacency_list.hpp>
0017 #include <boost/graph/connected_components.hpp>
0018
0019 namespace {
0020 template <typename vertex_t, typename weight_t>
0021 auto weaklyConnectedComponents(vertex_t numNodes,
0022                                boost::beast::span<const vertex_t>& rowIndices,
0023                                boost::beast::span<const vertex_t>& colIndices,
0024                                boost::beast::span<const weight_t>& edgeWeights,
0025                                std::vector<vertex_t>& trackLabels) {
0026   using Graph =
0027       boost::adjacency_list<boost::vecS,         // edge list
0028                             boost::vecS,         // vertex list
0029                             boost::undirectedS,  // directedness
0030                             boost::no_property,  // property of vertices
0031                             weight_t             // property of edges
0032                             >;
0033
0034   Graph g(numNodes);
0035
0036   for (const auto [row, col, weight] :
0037        Acts::zip(rowIndices, colIndices, edgeWeights)) {
0038     boost::add_edge(row, col, weight, g);
0039   }
0040
0041   return boost::connected_components(g, &trackLabels[0]);
0042 }
0043 }  // namespace
0044
0045 namespace Acts {
0046
0047 std::vector<std::vector<int>> BoostTrackBuilding::operator()(
0048     PipelineTensors tensors, std::vector<int>& spacepointIDs,
0049     const ExecutionContext& execContext) {
0050   ACTS_DEBUG("Start track building");
0051
0052   using RTI = const Tensor<std::int64_t>&;
0053   const auto& edgeTensor =
0054       tensors.edgeIndex.device().isCpu()
0055           ? static_cast<RTI>(tensors.edgeIndex)
0056           : static_cast<RTI>(tensors.edgeIndex.clone(
0057                 {Acts::Device::Cpu(), execContext.stream}));
0058
0059   assert(tensors.scoreTensor.has_value());
0060   using RTF = const Tensor<float>&;
0061   const auto& scoreTensor =
0062       tensors.edgeScores->device().isCpu()
0063           ? static_cast<RTF>(*tensors.edgeScores)
0064           : static_cast<RTF>(tensors.edgeScores->clone(
0065                 {Acts::Device::Cpu(), execContext.stream}));
0066
0067   assert(edgeTensor.shape().at(0) == 2);
0068   assert(edgeTensor.shape().at(1) == scoreTensor.shape().at(0));
0069
0070   const auto numSpacepoints = spacepointIDs.size();
0071   const auto numEdges = edgeTensor.shape().at(1);
0072
0073   if (numEdges == 0) {
0074     ACTS_WARNING("No edges remained after edge classification");
0075     return {};
0076   }
0077
0078   using vertex_t = std::int64_t;
0079   using weight_t = float;
0080
0081   boost::beast::span<const vertex_t> rowIndices(edgeTensor.data(), numEdges);
0082   boost::beast::span<const vertex_t> colIndices(edgeTensor.data() + numEdges,
0083                                                 numEdges);
0084   boost::beast::span<const weight_t> edgeWeights(scoreTensor.data(), numEdges);
0085
0086   std::vector<vertex_t> trackLabels(numSpacepoints);
0087
0088   auto numberLabels = weaklyConnectedComponents<vertex_t, weight_t>(
0089       numSpacepoints, rowIndices, colIndices, edgeWeights, trackLabels);
0090
0091   ACTS_VERBOSE("Number of track labels: " << trackLabels.size());
0092   ACTS_VERBOSE("Number of unique track labels: " << [&]() {
0093     std::vector<vertex_t> sorted(trackLabels);
0094     std::ranges::sort(sorted);
0095     sorted.erase(std::unique(sorted.begin(), sorted.end()), sorted.end());
0096     return sorted.size();
0097   }());
0098
0099   if (trackLabels.size() == 0) {
0100     return {};
0101   }
0102
0103   std::vector<std::vector<int>> trackCandidates(numberLabels);
0104
0105   for (const auto [label, id] : Acts::zip(trackLabels, spacepointIDs)) {
0106     trackCandidates[label].push_back(id);
0107   }
0108
0109   return trackCandidates;
0110 }
0111
0112 }  // namespace Acts