Gnn/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 "ActsPlugins/Gnn/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 using namespace Acts;
0020
0021 namespace {
0022 template <typename vertex_t, typename weight_t>
0023 auto weaklyConnectedComponents(vertex_t numNodes,
0024                                boost::beast::span<const vertex_t>& rowIndices,
0025                                boost::beast::span<const vertex_t>& colIndices,
0026                                boost::beast::span<const weight_t>& edgeWeights,
0027                                std::vector<vertex_t>& trackLabels) {
0028   using Graph =
0029       boost::adjacency_list<boost::vecS,         // edge list
0030                             boost::vecS,         // vertex list
0031                             boost::undirectedS,  // directedness
0032                             boost::no_property,  // property of vertices
0033                             weight_t             // property of edges
0034                             >;
0035
0036   Graph g(numNodes);
0037
0038   for (const auto [row, col, weight] :
0039        zip(rowIndices, colIndices, edgeWeights)) {
0040     boost::add_edge(row, col, weight, g);
0041   }
0042
0043   return boost::connected_components(g, &trackLabels[0]);
0044 }
0045 }  // namespace
0046
0047 namespace ActsPlugins {
0048
0049 std::vector<std::vector<int>> BoostTrackBuilding::operator()(
0050     PipelineTensors tensors, std::vector<int>& spacepointIDs,
0051     const ExecutionContext& execContext) {
0052   ACTS_DEBUG("Start track building");
0053
0054   using RTI = const Tensor<std::int64_t>&;
0055   const auto& edgeTensor = tensors.edgeIndex.device().isCpu()
0056                                ? static_cast<RTI>(tensors.edgeIndex)
0057                                : static_cast<RTI>(tensors.edgeIndex.clone(
0058                                      {Device::Cpu(), execContext.stream}));
0059
0060   assert(tensors.edgeScores.has_value());
0061   using RTF = const Tensor<float>&;
0062   const auto& scoreTensor = tensors.edgeScores->device().isCpu()
0063                                 ? static_cast<RTF>(*tensors.edgeScores)
0064                                 : static_cast<RTF>(tensors.edgeScores->clone(
0065                                       {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] : zip(trackLabels, spacepointIDs)) {
0106     trackCandidates[label].push_back(id);
0107   }
0108
0109   return trackCandidates;
0110 }
0111
0112 }  // namespace ActsPlugins