refactor: optimize examples ambiguity resolution (acts-project#1938)

my first implementation was quite wasteful in terms of CPU resources because I just wanted to get it done to test downstream performance for the vertexing. I optimized the code in a way that it pre computes the shared hits and modifies these numbers when removing a track.

performance difference for the current `full_chain_odd.py` without Pythia8 and without Geant4

new

```
identifier	time_total_s	time_perevent_s
Algorithm:AmbiguityResolutionAlgorithm	0.5349	0.005349
```

vs old

```
identifier	time_total_s	time_perevent_s
Algorithm:AmbiguityResolutionAlgorithm	12.52	0.1252
```

Author: andiwand

CI/physmon/reference/performance_ambi_orthogonal.root

@@ -28,15 +28,15 @@ namespace ActsExamples {
 class AmbiguityResolutionAlgorithm final : public IAlgorithm {
  public:
   struct Config {
-    /// Input source links collection.
-    std::string inputSourceLinks;
     /// Input trajectories collection.
     std::string inputTrajectories;
     /// Output trajectories collection.
     std::string outputTrajectories;
 
     /// Maximum amount of shared hits per track.
     std::uint32_t maximumSharedHits = 1;
+    /// Maximum number of iterations
+    std::uint32_t maximumIterations = 1000;
 
     /// Minumum number of measurement to form a track.
     size_t nMeasurementsMin = 7;

CI/physmon/reference/performance_ambi_seeded.root

@@ -20,14 +20,14 @@
 #include <numeric>
 #include <stdexcept>
 
+#include <boost/container/flat_map.hpp>
+#include <boost/container/flat_set.hpp>
+
 ActsExamples::AmbiguityResolutionAlgorithm::AmbiguityResolutionAlgorithm(
     ActsExamples::AmbiguityResolutionAlgorithm::Config cfg,
     Acts::Logging::Level lvl)
     : ActsExamples::IAlgorithm("AmbiguityResolutionAlgorithm", lvl),
       m_cfg(std::move(cfg)) {
-  if (m_cfg.inputSourceLinks.empty()) {
-    throw std::invalid_argument("Missing source links input collection");
-  }
   if (m_cfg.inputTrajectories.empty()) {
     throw std::invalid_argument("Missing trajectories input collection");
   }
@@ -38,148 +38,145 @@ ActsExamples::AmbiguityResolutionAlgorithm::AmbiguityResolutionAlgorithm(
 
 namespace {
 
-// TODO this is somewhat duplicated in TrackFindingAlgorithm.hpp
-// TODO we should make a common implementation in the core at some point
-std::vector<std::size_t> computeSharedHits(
-    const ActsExamples::IndexSourceLinkContainer& sourceLinks,
+struct State {
+  std::size_t numberOfTracks{};
+
+  std::vector<std::pair<std::size_t, std::size_t>> trackTips;
+  std::vector<float> trackChi2;
+  std::vector<ActsExamples::TrackParameters> trackParameters;
+  std::vector<std::vector<std::size_t>> measurementsPerTrack;
+
+  boost::container::flat_map<std::size_t,
+                             boost::container::flat_set<std::size_t>>
+      tracksPerMeasurement;
+  std::vector<std::size_t> sharedMeasurementsPerTrack;
+
+  boost::container::flat_set<std::size_t> selectedTracks;
+};
+
+State computeInitialState(
     const ActsExamples::TrajectoriesContainer& trajectories,
-    const std::vector<uint32_t>& trackIndices,
-    const std::vector<std::pair<size_t, size_t>>& trackTips) {
-  std::vector<std::size_t> hitCountPerMeasurement(sourceLinks.size(), 0);
+    std::size_t nMeasurementsMin) {
+  State state;
 
-  for (auto indexTrack : trackIndices) {
-    const auto [indexTraj, tip] = trackTips[indexTrack];
-    const auto& traj = trajectories[indexTraj];
+  for (std::size_t iTrack = 0, iTraj = 0; iTraj < trajectories.size();
+       ++iTraj) {
+    const auto& traj = trajectories[iTraj];
+    for (auto tip : traj.tips()) {
+      if (!traj.hasTrackParameters(tip)) {
+        continue;
+      }
 
-    traj.multiTrajectory().visitBackwards(tip, [&](const auto& state) {
-      if (!state.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
-        return true;
+      auto trajState = Acts::MultiTrajectoryHelpers::trajectoryState(
+          traj.multiTrajectory(), tip);
+      if (trajState.nMeasurements < nMeasurementsMin) {
+        continue;
       }
 
-      const std::size_t indexHit =
-          state.getUncalibratedSourceLink()
-              .template get<ActsExamples::IndexSourceLink>()
-              .index();
+      std::vector<std::size_t> measurements;
+      traj.multiTrajectory().visitBackwards(tip, [&](const auto& hit) {
+        if (hit.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
+          std::size_t iMeasurement =
+              hit.getUncalibratedSourceLink()
+                  .template get<ActsExamples::IndexSourceLink>()
+                  .index();
+          measurements.push_back(iMeasurement);
+        }
+        return true;
+      });
 
-      ++hitCountPerMeasurement[indexHit];
+      ++state.numberOfTracks;
 
-      return true;
-    });
-  }
+      state.trackTips.emplace_back(iTraj, tip);
+      state.trackChi2.push_back(trajState.chi2Sum / trajState.NDF);
+      state.trackParameters.push_back(traj.trackParameters(tip));
+      state.measurementsPerTrack.push_back(std::move(measurements));
 
-  std::vector<std::size_t> sharedHitCountPerTrack(trackIndices.size(), 0);
+      state.selectedTracks.insert(iTrack);
 
-  for (std::size_t i = 0; i < trackIndices.size(); ++i) {
-    const auto indexTrack = trackIndices[i];
-    const auto [indexTraj, tip] = trackTips[indexTrack];
-    const auto& traj = trajectories[indexTraj];
+      ++iTrack;
+    }
+  }
 
-    traj.multiTrajectory().visitBackwards(tip, [&](const auto& state) {
-      if (!state.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
-        return true;
-      }
+  for (std::size_t iTrack = 0; iTrack < state.numberOfTracks; ++iTrack) {
+    for (auto iMeasurement : state.measurementsPerTrack[iTrack]) {
+      state.tracksPerMeasurement[iMeasurement].insert(iTrack);
+    }
+  }
 
-      const std::size_t indexHit =
-          state.getUncalibratedSourceLink()
-              .template get<ActsExamples::IndexSourceLink>()
-              .index();
+  state.sharedMeasurementsPerTrack =
+      std::vector<std::size_t>(state.trackTips.size(), 0);
 
-      if (hitCountPerMeasurement[indexHit] > 1) {
-        ++sharedHitCountPerTrack[i];
+  for (std::size_t iTrack = 0; iTrack < state.numberOfTracks; ++iTrack) {
+    for (auto iMeasurement : state.measurementsPerTrack[iTrack]) {
+      if (state.tracksPerMeasurement[iMeasurement].size() > 1) {
+        ++state.sharedMeasurementsPerTrack[iTrack];
       }
-
-      return true;
-    });
+    }
   }
 
-  return sharedHitCountPerTrack;
+  return state;
 }
 
-std::size_t computeTrackHits(
-    const Acts::ConstVectorMultiTrajectory& multiTrajectory,
-    const std::size_t tip) {
-  std::size_t result = 0;
+void removeTrack(State& state, std::size_t iTrack) {
+  for (auto iMeasurement : state.measurementsPerTrack[iTrack]) {
+    state.tracksPerMeasurement[iMeasurement].erase(iTrack);
 
-  multiTrajectory.visitBackwards(tip, [&](const auto&) { ++result; });
+    if (state.tracksPerMeasurement[iMeasurement].size() == 1) {
+      auto jTrack = *std::begin(state.tracksPerMeasurement[iMeasurement]);
+      --state.sharedMeasurementsPerTrack[jTrack];
+    }
+  }
 
-  return result;
+  state.selectedTracks.erase(iTrack);
 }
 
 }  // namespace
 
 ActsExamples::ProcessCode ActsExamples::AmbiguityResolutionAlgorithm::execute(
     const AlgorithmContext& ctx) const {
-  // Read input data
-  const auto& sourceLinks =
-      ctx.eventStore.get<IndexSourceLinkContainer>(m_cfg.inputSourceLinks);
   const auto& trajectories =
       ctx.eventStore.get<TrajectoriesContainer>(m_cfg.inputTrajectories);
 
-  TrackParametersContainer trackParameters;
-  std::vector<std::pair<size_t, size_t>> trackTips;
+  auto state = computeInitialState(trajectories, m_cfg.nMeasurementsMin);
 
-  for (std::size_t iTraj = 0; iTraj < trajectories.size(); ++iTraj) {
-    const auto& traj = trajectories[iTraj];
-    for (auto tip : traj.tips()) {
-      if (!traj.hasTrackParameters(tip)) {
-        continue;
-      }
-      auto trajState = Acts::MultiTrajectoryHelpers::trajectoryState(
-          traj.multiTrajectory(), tip);
-      if (trajState.nMeasurements < m_cfg.nMeasurementsMin) {
-        continue;
-      }
-      trackParameters.push_back(traj.trackParameters(tip));
-      trackTips.emplace_back(iTraj, tip);
-    }
-  }
+  auto sharedMeasurementsComperator = [&state](std::size_t a, std::size_t b) {
+    return state.sharedMeasurementsPerTrack[a] <
+           state.sharedMeasurementsPerTrack[b];
+  };
+  auto badTrackComperator = [&state](std::size_t a, std::size_t b) {
+    auto relativeSharedMeasurements = [&state](std::size_t i) {
+      return 1.0 * state.sharedMeasurementsPerTrack[i] /
+             state.measurementsPerTrack[i].size();
+    };
 
-  std::vector<uint32_t> hitCount(trackParameters.size(), 0);
-  for (std::size_t i = 0; i < trackParameters.size(); ++i) {
-    const auto [iTraj, tip] = trackTips[i];
-    const auto& traj = trajectories[iTraj];
-    hitCount[i] = computeTrackHits(traj.multiTrajectory(), tip);
-  }
-
-  std::vector<uint32_t> trackIndices(trackParameters.size());
-  std::iota(std::begin(trackIndices), std::end(trackIndices), 0);
-
-  while (true) {
-    const auto sharedHits =
-        computeSharedHits(sourceLinks, trajectories, trackIndices, trackTips);
-
-    if (sharedHits.empty() ||
-        *std::max_element(std::begin(sharedHits), std::end(sharedHits)) <
-            m_cfg.maximumSharedHits) {
-      break;
+    if (relativeSharedMeasurements(a) != relativeSharedMeasurements(b)) {
+      return relativeSharedMeasurements(a) < relativeSharedMeasurements(b);
     }
-
-    std::vector<float> relativeSharedHits(trackIndices.size(), 0);
-    for (std::size_t i = 0; i < trackIndices.size(); ++i) {
-      const auto indexTrack = trackIndices[i];
-      relativeSharedHits[i] = 1.0f * sharedHits[i] / hitCount[indexTrack];
+    return state.trackChi2[a] < state.trackChi2[b];
+  };
+
+  for (std::size_t i = 0; i < m_cfg.maximumIterations; ++i) {
+    auto maximumSharedMeasurements = *std::max_element(
+        state.selectedTracks.begin(), state.selectedTracks.end(),
+        sharedMeasurementsComperator);
+    ACTS_VERBOSE(
+        "maximum shared measurements "
+        << state.sharedMeasurementsPerTrack[maximumSharedMeasurements]);
+    if (state.sharedMeasurementsPerTrack[maximumSharedMeasurements] <
+        m_cfg.maximumSharedHits) {
+      break;
     }
 
-    const auto maxRelativeSharedHits = std::max_element(
-        std::begin(relativeSharedHits), std::end(relativeSharedHits));
-    const auto index =
-        std::distance(std::begin(relativeSharedHits), maxRelativeSharedHits);
-    trackIndices.erase(std::begin(trackIndices) + index);
-  }
-
-  if (trackIndices.size() == trackParameters.size()) {
-    const auto sharedHits =
-        computeSharedHits(sourceLinks, trajectories, trackIndices, trackTips);
-
-    std::vector<float> relativeSharedHits(trackIndices.size(), 0);
-    for (std::size_t i = 0; i < trackIndices.size(); ++i) {
-      const auto indexTrack = trackIndices[i];
-      relativeSharedHits[i] = 1.0f * sharedHits[i] / hitCount[indexTrack];
-    }
+    auto badTrack =
+        *std::max_element(state.selectedTracks.begin(),
+                          state.selectedTracks.end(), badTrackComperator);
+    ACTS_VERBOSE("remove track " << badTrack);
+    removeTrack(state, badTrack);
   }
 
-  ACTS_INFO("Resolved to " << trackIndices.size() << " tracks from "
-                           << trackParameters.size());
+  ACTS_INFO("Resolved to " << state.selectedTracks.size() << " tracks from "
+                           << state.trackTips.size());
 
   TrajectoriesContainer outputTrajectories;
   outputTrajectories.reserve(trajectories.size());
@@ -189,13 +186,13 @@ ActsExamples::ProcessCode ActsExamples::AmbiguityResolutionAlgorithm::execute(
     std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
     Trajectories::IndexedParameters parameters;
 
-    for (auto iTrack : trackIndices) {
-      if (trackTips[iTrack].first != iTraj) {
+    for (auto iTrack : state.selectedTracks) {
+      if (state.trackTips[iTrack].first != iTraj) {
         continue;
       }
-      const auto tip = trackTips[iTrack].second;
+      const auto tip = state.trackTips[iTrack].second;
       tips.push_back(tip);
-      parameters.emplace(tip, trackParameters[iTrack]);
+      parameters.emplace(tip, state.trackParameters[iTrack]);
     }
     if (!tips.empty()) {
       outputTrajectories.emplace_back(traj.multiTrajectory(), tips, parameters);

CI/physmon/reference/performance_vertexing_orthogonal_hist.root

@@ -1232,7 +1232,6 @@ def addAmbiguityResolution(
 
     alg = AmbiguityResolutionAlgorithm(
         level=customLogLevel(),
-        inputSourceLinks="sourcelinks",
         inputTrajectories="trajectories",
         outputTrajectories="filteredTrajectories",
         **acts.examples.defaultKWArgs(

CI/physmon/reference/performance_vertexing_seeded_hist.root

@@ -310,10 +310,10 @@ void addTrackFinding(Context& ctx) {
             .def(py::init(constructor));
   }
 
-  ACTS_PYTHON_DECLARE_ALGORITHM(
-      ActsExamples::AmbiguityResolutionAlgorithm, mex,
-      "AmbiguityResolutionAlgorithm", inputSourceLinks, inputTrajectories,
-      outputTrajectories, maximumSharedHits, nMeasurementsMin);
+  ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::AmbiguityResolutionAlgorithm, mex,
+                                "AmbiguityResolutionAlgorithm",
+                                inputTrajectories, outputTrajectories,
+                                maximumSharedHits, nMeasurementsMin);
 
   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::SeedsToPrototracks, mex,
                                 "SeedsToPrototracks", inputSeeds,