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 #include <map>
0015
0016 #include <boost/beast/core/span.hpp>
0017 #include <boost/graph/adjacency_list.hpp>
0018 #include <boost/graph/connected_components.hpp>
0019 #include <torch/torch.h>
0020
0021 namespace {
0022 template <typename vertex_t, typename weight_t>
0023 auto weaklyConnectedComponents(vertex_t numNodes,
0024                                boost::beast::span<vertex_t>& rowIndices,
0025                                boost::beast::span<vertex_t>& colIndices,
0026                                boost::beast::span<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        Acts::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 Acts {
0048
0049 std::vector<std::vector<int>> BoostTrackBuilding::operator()(
0050     std::any /*nodes*/, std::any edges, std::any weights,
0051     std::vector<int>& spacepointIDs, const ExecutionContext& execContext) {
0052   ACTS_DEBUG("Start track building");
0053   const auto edgeTensor = std::any_cast<torch::Tensor>(edges).to(torch::kCPU);
0054   const auto edgeWeightTensor =
0055       std::any_cast<torch::Tensor>(weights).to(torch::kCPU);
0056
0057   assert(edgeTensor.size(0) == 2);
0058   assert(edgeTensor.size(1) == edgeWeightTensor.size(0));
0059
0060   const auto numSpacepoints = spacepointIDs.size();
0061   const auto numEdges = static_cast<std::size_t>(edgeWeightTensor.size(0));
0062
0063   if (numEdges == 0) {
0064     ACTS_WARNING("No edges remained after edge classification");
0065     return {};
0066   }
0067
0068   using vertex_t = std::int64_t;
0069   using weight_t = float;
0070
0071   boost::beast::span<vertex_t> rowIndices(edgeTensor.data_ptr<vertex_t>(),
0072                                           numEdges);
0073   boost::beast::span<vertex_t> colIndices(
0074       edgeTensor.data_ptr<vertex_t>() + numEdges, numEdges);
0075   boost::beast::span<weight_t> edgeWeights(edgeWeightTensor.data_ptr<float>(),
0076                                            numEdges);
0077
0078   std::vector<vertex_t> trackLabels(numSpacepoints);
0079
0080   auto numberLabels = weaklyConnectedComponents<vertex_t, weight_t>(
0081       numSpacepoints, rowIndices, colIndices, edgeWeights, trackLabels);
0082
0083   ACTS_VERBOSE("Number of track labels: " << trackLabels.size());
0084   ACTS_VERBOSE("Number of unique track labels: " << [&]() {
0085     std::vector<vertex_t> sorted(trackLabels);
0086     std::ranges::sort(sorted);
0087     sorted.erase(std::unique(sorted.begin(), sorted.end()), sorted.end());
0088     return sorted.size();
0089   }());
0090
0091   if (trackLabels.size() == 0) {
0092     return {};
0093   }
0094
0095   std::vector<std::vector<int>> trackCandidates(numberLabels);
0096
0097   for (const auto [label, id] : Acts::zip(trackLabels, spacepointIDs)) {
0098     trackCandidates[label].push_back(id);
0099   }
0100
0101   return trackCandidates;
0102 }
0103
0104 }  // namespace Acts