Merge branch 'main' into fix-kf-reverse-navigation

Plugins/ExaTrkX/src/BoostTrackBuilding.cpp

@@ -8,34 +8,38 @@
 
 #include "Acts/Plugins/ExaTrkX/BoostTrackBuilding.hpp"
 
+#include "Acts/Utilities/Zip.hpp"
+
 #include <map>
 
+#include <boost/beast/core/span.hpp>
 #include <boost/graph/adjacency_list.hpp>
 #include <boost/graph/connected_components.hpp>
 #include <torch/torch.h>
 
 namespace {
-template <typename vertex_t, typename edge_t, typename weight_t>
-void weaklyConnectedComponents(vertex_t numNodes,
-                               std::vector<vertex_t>& rowIndices,
-                               std::vector<vertex_t>& colIndices,
-                               std::vector<weight_t>& edgeWeights,
+template <typename vertex_t, typename weight_t>
+auto weaklyConnectedComponents(vertex_t numNodes,
+                               boost::beast::span<vertex_t>& rowIndices,
+                               boost::beast::span<vertex_t>& colIndices,
+                               boost::beast::span<weight_t>& edgeWeights,
                                std::vector<vertex_t>& trackLabels) {
-  typedef boost::adjacency_list<boost::vecS,         // edge list
-                                boost::vecS,         // vertex list
-                                boost::undirectedS,  // directedness
-                                boost::no_property,  // property of vertices
-                                float                // property of edges
-                                >
-      Graph;
+  using Graph =
+      boost::adjacency_list<boost::vecS,         // edge list
+                            boost::vecS,         // vertex list
+                            boost::undirectedS,  // directedness
+                            boost::no_property,  // property of vertices
+                            weight_t             // property of edges
+                            >;
 
   Graph g(numNodes);
-  for (size_t idx = 0; idx < rowIndices.size(); ++idx) {
-    boost::add_edge(rowIndices[idx], colIndices[idx], edgeWeights[idx], g);
+
+  for (const auto [row, col, weight] :
+       Acts::zip(rowIndices, colIndices, edgeWeights)) {
+    boost::add_edge(row, col, weight, g);
   }
 
-  [[maybe_unused]] size_t num_components =
-      boost::connected_components(g, &trackLabels[0]);
+  return boost::connected_components(g, &trackLabels[0]);
 }
 }  // namespace
 
@@ -60,27 +64,23 @@ std::vector<std::vector<int>> BoostTrackBuilding::operator()(
     return {};
   }
 
-  using vertex_t = int32_t;
-  std::vector<vertex_t> rowIndices(numEdges);
-  std::copy(edgeTensor.data_ptr<int64_t>(),
-            edgeTensor.data_ptr<int64_t>() + numEdges, rowIndices.begin());
+  using vertex_t = int64_t;
+  using weight_t = float;
 
-  std::vector<vertex_t> colIndices(numEdges);
-  std::copy(edgeTensor.data_ptr<int64_t>() + numEdges,
-            edgeTensor.data_ptr<int64_t>() + numEdges + numEdges,
-            colIndices.begin());
-
-  std::vector<float> edgeWeights(numEdges);
-  std::copy(edgeWeightTensor.data_ptr<float>(),
-            edgeWeightTensor.data_ptr<float>() + numEdges, edgeWeights.begin());
+  boost::beast::span<vertex_t> rowIndices(edgeTensor.data_ptr<vertex_t>(),
+                                          numEdges);
+  boost::beast::span<vertex_t> colIndices(
+      edgeTensor.data_ptr<vertex_t>() + numEdges, numEdges);
+  boost::beast::span<weight_t> edgeWeights(edgeWeightTensor.data_ptr<float>(),
+                                           numEdges);
 
   std::vector<vertex_t> trackLabels(numSpacepoints);
 
-  weaklyConnectedComponents<int32_t, int32_t, float>(
+  auto numberLabels = weaklyConnectedComponents<vertex_t, weight_t>(
       numSpacepoints, rowIndices, colIndices, edgeWeights, trackLabels);
 
   ACTS_VERBOSE("Number of track labels: " << trackLabels.size());
-  ACTS_VERBOSE("NUmber of unique track labels: " << [&]() {
+  ACTS_VERBOSE("Number of unique track labels: " << [&]() {
     std::vector<vertex_t> sorted(trackLabels);
     std::sort(sorted.begin(), sorted.end());
     sorted.erase(std::unique(sorted.begin(), sorted.end()), sorted.end());
@@ -91,28 +91,10 @@ std::vector<std::vector<int>> BoostTrackBuilding::operator()(
     return {};
   }
 
-  std::vector<std::vector<int>> trackCandidates;
-
-  int existTrkIdx = 0;
-  // map labeling from MCC to customized track id.
-  std::map<int32_t, int32_t> trackLableToIds;
-
-  for (auto idx = 0ul; idx < numSpacepoints; ++idx) {
-    int32_t trackLabel = trackLabels[idx];
-    int spacepointID = spacepointIDs[idx];
-
-    int trkId;
-    if (trackLableToIds.find(trackLabel) != trackLableToIds.end()) {
-      trkId = trackLableToIds[trackLabel];
-      trackCandidates[trkId].push_back(spacepointID);
-    } else {
-      // a new track, assign the track id
-      // and create a vector
-      trkId = existTrkIdx;
-      trackCandidates.push_back(std::vector<int>{spacepointID});
-      trackLableToIds[trackLabel] = trkId;
-      existTrkIdx++;
-    }
+  std::vector<std::vector<int>> trackCandidates(numberLabels);
+
+  for (const auto [label, id] : Acts::zip(trackLabels, spacepointIDs)) {
+    trackCandidates[label].push_back(id);
   }
 
   return trackCandidates;

Tests/UnitTests/Plugins/ExaTrkX/CMakeLists.txt

@@ -2,3 +2,4 @@ set(unittest_extra_libraries ActsPluginExaTrkX ${TORCH_LIBRARIES})
 
 add_unittest(ExaTrkXTensorConversion ExaTrkXTensorConversionTests.cpp)
 add_unittest(ExaTrkXEdgeBuilding ExaTrkXEdgeBuildingTest.cpp)
+add_unittest(ExaTrkXBoostTrackBuilding ExaTrkXBoostTrackBuildingTests.cpp)

Tests/UnitTests/Plugins/ExaTrkX/ExaTrkXBoostTrackBuildingTests.cpp

@@ -0,0 +1,62 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2022 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include <boost/test/unit_test.hpp>
+
+#include "Acts/Plugins/ExaTrkX/BoostTrackBuilding.hpp"
+#include "Acts/Plugins/ExaTrkX/detail/TensorVectorConversion.hpp"
+
+#include <algorithm>
+
+BOOST_AUTO_TEST_CASE(test_track_building) {
+  // Make some spacepoint IDs
+  // The spacepoint ids are [100, 101, 102, ...]
+  // They should not be zero based to check if the thing also works if the
+  // spacepoint IDs do not match the node IDs used for the edges
+  std::vector<int> spacepointIds(16);
+  std::iota(spacepointIds.begin(), spacepointIds.end(), 100);
+
+  // Build 4 tracks with 4 hits
+  std::vector<std::vector<int>> refTracks;
+  for (auto t = 0ul; t < 4; ++t) {
+    refTracks.emplace_back(spacepointIds.begin() + 4 * t,
+                           spacepointIds.begin() + 4 * (t + 1));
+  }
+
+  // Make edges
+  std::vector<int64_t> edges;
+  for (const auto &track : refTracks) {
+    for (auto it = track.begin(); it != track.end() - 1; ++it) {
+      // edges must be 0 based, so subtract 100 again
+      edges.push_back(*it - 100);
+      edges.push_back(*std::next(it) - 100);
+    }
+  }
+
+  auto edgeTensor =
+      Acts::detail::vectorToTensor2D(edges, 2).t().contiguous().clone();
+  auto dummyWeights = torch::ones(edges.size() / 2, torch::kFloat32);
+
+  // Run Track building
+  auto logger = Acts::getDefaultLogger("TestLogger", Acts::Logging::ERROR);
+  Acts::BoostTrackBuilding trackBuilder(std::move(logger));
+
+  auto testTracks = trackBuilder({}, edgeTensor, dummyWeights, spacepointIds);
+
+  // Sort tracks, so we can find them
+  std::for_each(testTracks.begin(), testTracks.end(),
+                [](auto &t) { std::sort(t.begin(), t.end()); });
+  std::for_each(refTracks.begin(), refTracks.end(),
+                [](auto &t) { std::sort(t.begin(), t.end()); });
+
+  // Check what we have here
+  for (const auto &refTrack : refTracks) {
+    auto found = std::find(testTracks.begin(), testTracks.end(), refTrack);
+    BOOST_CHECK(found != testTracks.end());
+  }
+}