perf: Reduce memory consumption of track finding (#2206)

We use a lot of memory in busy events, e.g. for ttbar pu200 we use multiple gigabytes of track state storage.
This PR reduces this by rolling a track selection into the track finding algorithm. The CKF now operates on a temporary track container that is cleared between seeds and reused. Only tracks that pass the selection are copied into the output collection. 

In my testing this reduces the memory consumption by about 2x.

Before:
![main_PrMon_wtime_vs_vmem_pss_rss_swap](https://github.com/acts-project/acts/assets/1058585/38d424fc-6a36-4f01-a522-5bb89e756409)
After:
![feat_v3_PrMon_wtime_vs_vmem_pss_rss_swap](https://github.com/acts-project/acts/assets/1058585/36df26c7-701e-478c-b972-c1bd6ab2027c)


Blocked by:
- #2202
- #2201
- #2203
- #2204
- #2205
- #2267

Core/include/Acts/EventData/MultiTrajectory.hpp

@@ -21,6 +21,7 @@
 #include "Acts/Utilities/Concepts.hpp"
 #include "Acts/Utilities/HashedString.hpp"
 #include "Acts/Utilities/Helpers.hpp"
+#include "Acts/Utilities/ThrowAssert.hpp"
 #include "Acts/Utilities/TypeTraits.hpp"
 
 #include <bitset>
@@ -1436,6 +1437,9 @@ class MultiTrajectory {
   /// @note Is a noop if the track state already has an allocation
   ///       an the dimension is the same.
   void allocateCalibrated(IndexType istate, size_t measdim) {
+    throw_assert(measdim > 0 && measdim <= eBoundSize,
+                 "Invalid measurement dimension detected");
+
     self().allocateCalibrated_impl(istate, measdim);
   }
 

Core/include/Acts/EventData/TrackProxy.hpp

@@ -540,7 +540,7 @@ class TrackProxy {
 
     if (copyTrackStates) {
       // append track states (cheap), but they're in the wrong order
-      for (auto srcTrackState : other.trackStates()) {
+      for (const auto& srcTrackState : other.trackStates()) {
         auto destTrackState = appendTrackState(srcTrackState.getMask());
         if (srcTrackState.hasCalibrated()) {
           destTrackState.allocateCalibrated(srcTrackState.calibratedSize());

Core/include/Acts/EventData/TrackStateType.hpp

@@ -111,6 +111,7 @@ class ConstTrackStateType {
     for (size_t i = 0; i < TrackStateFlag::NumTrackStateFlags; i++) {
       trunc[i] = bs[i];
     }
+    // MeasurementParameterOutlierHoleMaterialSharedhit
     os << "MPOHMS=" << trunc;
     return os;
   }

Core/include/Acts/EventData/VectorMultiTrajectory.hpp

@@ -16,6 +16,8 @@
 #include "Acts/EventData/detail/DynamicColumn.hpp"
 #include "Acts/Utilities/Concepts.hpp"
 #include "Acts/Utilities/HashedString.hpp"
+#include "Acts/Utilities/Helpers.hpp"
+#include "Acts/Utilities/ThrowAssert.hpp"
 
 #include <any>
 #include <cassert>
@@ -156,7 +158,7 @@ class VectorMultiTrajectoryBase {
 
     IndexType iuncalibrated = kInvalid;
     IndexType icalibratedsourcelink = kInvalid;
-    IndexType measdim = 0;
+    IndexType measdim = kInvalid;
 
     TrackStatePropMask allocMask = TrackStatePropMask::None;
   };
@@ -448,6 +450,9 @@ class VectorMultiTrajectory final
   }
 
   void allocateCalibrated_impl(IndexType istate, size_t measdim) {
+    throw_assert(measdim > 0 && measdim <= eBoundSize,
+                 "Invalid measurement dimension detected");
+
     if (m_measOffset[istate] != kInvalid &&
         m_measCovOffset[istate] != kInvalid &&
         m_index[istate].measdim == measdim) {

Core/include/Acts/TrackFinding/CombinatorialKalmanFilter.hpp

@@ -41,6 +41,7 @@
 
 #include <functional>
 #include <memory>
+#include <type_traits>
 #include <unordered_map>
 
 namespace Acts {
@@ -466,6 +467,7 @@ class CombinatorialKalmanFilter {
       if (result.filtered) {
         // Return error if filtering finds no tracks
         if (result.lastTrackIndices.empty()) {
+          // @TODO: Tracks like this should not be in the final output!
           ACTS_WARNING("No tracks found");
           result.finished = true;
         } else {
@@ -832,6 +834,10 @@ class CombinatorialKalmanFilter {
           mask = PM::Calibrated;
         }
 
+        ACTS_VERBOSE(
+            "Create temp track state with mask: " << std::bitset<
+                sizeof(std::underlying_type_t<TrackStatePropMask>) * 8>(
+                static_cast<std::underlying_type_t<TrackStatePropMask>>(mask)));
         size_t tsi = result.stateBuffer->addTrackState(mask, prevTip);
         // CAREFUL! This trackstate has a previous index that is not in this
         // MultiTrajectory Visiting brackwards from this track state will
@@ -905,6 +911,13 @@ class CombinatorialKalmanFilter {
             result.fittedStates->getTrackState(
                 result.fittedStates->addTrackState(
                     mask, candidateTrackState.previous()));
+        ACTS_VERBOSE(
+            "Create SourceLink output track state #"
+            << trackState.index() << " with mask: "
+            << std::bitset<sizeof(std::underlying_type_t<TrackStatePropMask>) *
+                           8>{
+                   static_cast<std::underlying_type_t<TrackStatePropMask>>(
+                       mask)});
 
         if (it != begin) {
           // assign indices pointing to first track state
@@ -988,11 +1001,14 @@ class CombinatorialKalmanFilter {
                                  size_t prevTip) const {
       // Add a track state
       auto currentTip = result.fittedStates->addTrackState(stateMask, prevTip);
-      if (isSensitive) {
-        ACTS_VERBOSE("Creating Hole track state with tip = " << currentTip);
-      } else {
-        ACTS_VERBOSE("Creating Material track state with tip = " << currentTip);
-      }
+      ACTS_VERBOSE(
+          "Create "
+          << (isSensitive ? "Hole" : "Material") << " output track state #"
+          << currentTip << " with mask: "
+          << std::bitset<sizeof(std::underlying_type_t<TrackStatePropMask>) *
+                         8>{
+                 static_cast<std::underlying_type_t<TrackStatePropMask>>(
+                     stateMask)});
       // now get track state proxy back
       auto trackStateProxy = result.fittedStates->getTrackState(currentTip);
 

Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp

@@ -18,6 +18,7 @@
 #include "Acts/TrackFinding/CombinatorialKalmanFilter.hpp"
 #include "Acts/TrackFinding/MeasurementSelector.hpp"
 #include "Acts/TrackFinding/SourceLinkAccessorConcept.hpp"
+#include "Acts/TrackFinding/TrackSelector.hpp"
 #include "Acts/Utilities/Logger.hpp"
 #include "Acts/Utilities/Result.hpp"
 #include "ActsExamples/EventData/IndexSourceLink.hpp"
@@ -91,6 +92,8 @@ class TrackFindingAlgorithm final : public IAlgorithm {
     Acts::MeasurementSelector::Config measurementSelectorCfg;
     /// Compute shared hit information
     bool computeSharedHits = false;
+    /// Track selector config
+    std::optional<Acts::TrackSelector::Config> trackSelectorCfg = std::nullopt;
   };
 
   /// Constructor of the track finding algorithm
@@ -118,6 +121,7 @@ class TrackFindingAlgorithm final : public IAlgorithm {
 
  private:
   Config m_cfg;
+  std::optional<Acts::TrackSelector> m_trackSelector;
 
   ReadDataHandle<MeasurementContainer> m_inputMeasurements{this,
                                                            "InputMeasurements"};

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp

@@ -29,6 +29,7 @@
 #include <cmath>
 #include <functional>
 #include <memory>
+#include <optional>
 #include <ostream>
 #include <stdexcept>
 #include <system_error>
@@ -39,7 +40,14 @@
 ActsExamples::TrackFindingAlgorithm::TrackFindingAlgorithm(
     Config config, Acts::Logging::Level level)
     : ActsExamples::IAlgorithm("TrackFindingAlgorithm", level),
-      m_cfg(std::move(config)) {
+      m_cfg(std::move(config)),
+      m_trackSelector([this]() -> std::optional<Acts::TrackSelector> {
+        if (m_cfg.trackSelectorCfg.has_value()) {
+          return {m_cfg.trackSelectorCfg.value()};
+        } else {
+          return std::nullopt;
+        }
+      }()) {
   if (m_cfg.inputMeasurements.empty()) {
     throw std::invalid_argument("Missing measurements input collection");
   }
@@ -112,16 +120,24 @@ ActsExamples::ProcessCode ActsExamples::TrackFindingAlgorithm::execute(
   auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
   auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
 
+  auto trackContainerTemp = std::make_shared<Acts::VectorTrackContainer>();
+  auto trackStateContainerTemp =
+      std::make_shared<Acts::VectorMultiTrajectory>();
+
   TrackContainer tracks(trackContainer, trackStateContainer);
+  TrackContainer tracksTemp(trackContainerTemp, trackStateContainerTemp);
 
   tracks.addColumn<unsigned int>("trackGroup");
+  tracksTemp.addColumn<unsigned int>("trackGroup");
   Acts::TrackAccessor<unsigned int> seedNumber("trackGroup");
 
   unsigned int nSeed = 0;
 
   for (std::size_t iseed = 0; iseed < initialParameters.size(); ++iseed) {
+    trackContainerTemp->clear();
+    trackStateContainerTemp->clear();
     auto result =
-        (*m_cfg.findTracks)(initialParameters.at(iseed), options, tracks);
+        (*m_cfg.findTracks)(initialParameters.at(iseed), options, tracksTemp);
     m_nTotalSeeds++;
     nSeed++;
 
@@ -135,6 +151,11 @@ ActsExamples::ProcessCode ActsExamples::TrackFindingAlgorithm::execute(
     auto& tracksForSeed = result.value();
     for (auto& track : tracksForSeed) {
       seedNumber(track) = nSeed;
+      if (!m_trackSelector.has_value() ||
+          m_trackSelector->isValidTrack(track)) {
+        auto destProxy = tracks.getTrack(tracks.addTrack());
+        destProxy.copyFrom(track, true);  // make sure we copy track states!
+      }
     }
   }
 

Examples/Python/python/acts/examples/reconstruction.py

@@ -948,6 +948,27 @@ def addCKFTracks(
         measurementSelectorCfg=acts.MeasurementSelector.Config(
             [(acts.GeometryIdentifier(), ([], [15.0], [10]))]
         ),
+        trackSelectorCfg=acts.TrackSelector.Config(
+            **acts.examples.defaultKWArgs(
+                loc0Min=trackSelectorConfig.loc0[0],
+                loc0Max=trackSelectorConfig.loc0[1],
+                loc1Min=trackSelectorConfig.loc1[0],
+                loc1Max=trackSelectorConfig.loc1[1],
+                timeMin=trackSelectorConfig.time[0],
+                timeMax=trackSelectorConfig.time[1],
+                phiMin=trackSelectorConfig.phi[0],
+                phiMax=trackSelectorConfig.phi[1],
+                etaMin=trackSelectorConfig.eta[0],
+                etaMax=trackSelectorConfig.eta[1],
+                absEtaMin=trackSelectorConfig.absEta[0],
+                absEtaMax=trackSelectorConfig.absEta[1],
+                ptMin=trackSelectorConfig.pt[0],
+                ptMax=trackSelectorConfig.pt[1],
+                minMeasurements=trackSelectorConfig.nMeasurementsMin,
+            )
+        )
+        if trackSelectorConfig is not None
+        else None,
         inputMeasurements="measurements",
         inputSourceLinks="sourcelinks",
         inputInitialTrackParameters="estimatedparameters",
@@ -967,26 +988,6 @@ def addCKFTracks(
     s.addAlgorithm(trackConverter)
     s.addWhiteboardAlias("trajectories", trackConverter.config.outputTrajectories)
 
-    if trackSelectorConfig is not None:
-        trackSelector = addTrackSelection(
-            s,
-            trackSelectorConfig,
-            inputTracks=trackFinder.config.outputTracks,
-            outputTracks="selectedTracks",
-            logLevel=customLogLevel(),
-        )
-        s.addWhiteboardAlias("tracks", trackSelector.config.outputTracks)
-
-        selectedTrackConverter = acts.examples.TracksToTrajectories(
-            level=customLogLevel(),
-            inputTracks=trackSelector.config.outputTracks,
-            outputTrajectories="trajectories-from-selected-tracks",
-        )
-        s.addAlgorithm(selectedTrackConverter)
-        s.addWhiteboardAlias(
-            "trajectories", selectedTrackConverter.config.outputTrajectories
-        )
-
     addTrajectoryWriters(
         s,
         name="ckf",

Examples/Python/src/TrackFinding.cpp

@@ -295,6 +295,7 @@ void addTrackFinding(Context& ctx) {
     ACTS_PYTHON_MEMBER(outputTracks);
     ACTS_PYTHON_MEMBER(findTracks);
     ACTS_PYTHON_MEMBER(measurementSelectorCfg);
+    ACTS_PYTHON_MEMBER(trackSelectorCfg);
     ACTS_PYTHON_STRUCT_END();
   }