refactor!: Refactor CKF branch stopper to allow stop and keep tracks

Core/include/Acts/TrackFinding/CombinatorialKalmanFilter.hpp

@@ -73,14 +73,21 @@ struct CombinatorialKalmanFilterExtensions {
   using candidate_container_t =
       typename std::vector<typename traj_t::TrackStateProxy>;
 
+  enum class BranchStopperResult {
+    Continue,
+    StopAndDrop,
+    StopAndKeep,
+  };
+
   using Calibrator = typename KalmanFitterExtensions<traj_t>::Calibrator;
   using Updater = typename KalmanFitterExtensions<traj_t>::Updater;
   using MeasurementSelector =
       Delegate<Result<std::pair<typename candidate_container_t::iterator,
                                 typename candidate_container_t::iterator>>(
           candidate_container_t& trackStates, bool&, const Logger&)>;
-  using BranchStopper = Delegate<bool(const CombinatorialKalmanFilterTipState&,
-                                      typename traj_t::TrackStateProxy&)>;
+  using BranchStopper =
+      Delegate<BranchStopperResult(const CombinatorialKalmanFilterTipState&,
+                                   typename traj_t::TrackStateProxy&)>;
 
   /// The Calibrator is a dedicated calibration algorithm that allows to
   /// calibrate measurements using track information, this could be e.g. sagging
@@ -112,10 +119,10 @@ struct CombinatorialKalmanFilterExtensions {
 
   /// Default branch stopper which will never stop
   /// @return false
-  static bool voidBranchStopper(
+  static BranchStopperResult voidBranchStopper(
       const CombinatorialKalmanFilterTipState& /*tipState*/,
       typename traj_t::TrackStateProxy& /*trackState*/) {
-    return false;
+    return BranchStopperResult::Continue;
   }
 };
 
@@ -694,13 +701,23 @@ class CombinatorialKalmanFilter {
           auto nonSourcelinkState =
               result.fittedStates->getTrackState(currentTip);
 
+          using BranchStopperResult =
+              typename CombinatorialKalmanFilterExtensions<
+                  traj_t>::BranchStopperResult;
+          BranchStopperResult branchStopperResult =
+              m_extensions.branchStopper(tipState, nonSourcelinkState);
+
           // Check the branch
-          if (!m_extensions.branchStopper(tipState, nonSourcelinkState)) {
+          if (branchStopperResult == BranchStopperResult::Continue) {
             // Remember the active tip and its state
             result.activeTips.emplace_back(currentTip, tipState);
           } else {
             // No branch on this surface
             nBranchesOnSurface = 0;
+
+            if (branchStopperResult == BranchStopperResult::StopAndKeep) {
+              result.lastMeasurementIndices.emplace_back(currentTip);
+            }
           }
         }
 
@@ -904,13 +921,10 @@ class CombinatorialKalmanFilter {
           tipState.nMeasurements++;
         }
 
-        // Check if need to stop this branch
-        if (!m_extensions.branchStopper(tipState, trackState)) {
-          // Put tipstate back into active tips to continue with it
-          result.activeTips.emplace_back(currentTip, tipState);
-          // Record the number of branches on surface
-          nBranchesOnSurface++;
-        }
+        // Put tipstate back into active tips to continue with it
+        result.activeTips.emplace_back(currentTip, tipState);
+        // Record the number of branches on surface
+        nBranchesOnSurface++;
       }
       return Result<void>::success();
     }

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp

@@ -179,16 +179,18 @@ class BranchStopper {
   using Config =
       std::optional<std::variant<Acts::TrackSelector::Config,
                                  Acts::TrackSelector::EtaBinnedConfig>>;
+  using BranchStopperResult = Acts::CombinatorialKalmanFilterExtensions<
+      Acts::VectorMultiTrajectory>::BranchStopperResult;
 
   mutable std::atomic<std::size_t> m_nStoppedBranches{0};
 
   explicit BranchStopper(const Config& config) : m_config(config) {}
 
-  bool operator()(
+  BranchStopperResult operator()(
       const Acts::CombinatorialKalmanFilterTipState& tipState,
       Acts::VectorMultiTrajectory::TrackStateProxy& trackState) const {
     if (!m_config.has_value()) {
-      return false;
+      return BranchStopperResult::Continue;
     }
 
     const Acts::TrackSelector::Config* singleConfig = std::visit(
@@ -207,35 +209,21 @@ class BranchStopper {
 
     if (singleConfig == nullptr) {
       ++m_nStoppedBranches;
-      return true;
+      return BranchStopperResult::StopAndDrop;
     }
 
-    // Continue if the number of holes is below the maximum
-    if (tipState.nHoles <= singleConfig->maxHoles) {
-      return false;
-    }
-
-    // Continue if the number of outliers is below the maximum
-    if (tipState.nOutliers <= singleConfig->maxOutliers) {
-      return false;
-    }
+    bool enoughMeasurements =
+        tipState.nMeasurements >= singleConfig->minMeasurements;
+    bool tooManyHoles = tipState.nHoles > singleConfig->maxHoles;
+    bool tooManyOutliers = tipState.nOutliers > singleConfig->maxOutliers;
 
-    // If there are not enough measurements but more holes than allowed we stop
-    if (tipState.nMeasurements < singleConfig->minMeasurements) {
+    if (tooManyHoles || tooManyOutliers) {
       ++m_nStoppedBranches;
-      return true;
+      return enoughMeasurements ? BranchStopperResult::StopAndKeep
+                                : BranchStopperResult::StopAndDrop;
     }
 
-    // Getting another measurement guarantees that the holes are in the middle
-    // of the track
-    if (trackState.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
-      ++m_nStoppedBranches;
-      return true;
-    }
-
-    // We cannot be sure if the holes are just at the end of the track so we
-    // have to keep going
-    return false;
+    return BranchStopperResult::Continue;
   }
 
  private: