From f1d839e7b9ab00c924b2d852632453d1e2f58d80 Mon Sep 17 00:00:00 2001
From: Dmitry Kalinkin <dmitry.kalinkin@gmail.com>
Date: Sun, 1 Dec 2024 10:32:14 -0500
Subject: [PATCH] CKFTracking: remove Acts tracks that failed to extrapolate,
 not just the trajectories (#1677)

### Briefly, what does this PR introduce?

This is re-implements #1663 in such a way that downstream algorithms,
such as the AmbiguitySolver, that operate on Acts EDM do not see the
failed tracks.
Resolves: #1672

### What kind of change does this PR introduce?
- [x] Bug fix (issue #1672)
- [ ] New feature (issue #__)
- [ ] Documentation update
- [ ] Other: __

### Please check if this PR fulfills the following:
- [ ] Tests for the changes have been added
- [ ] Documentation has been added / updated
- [ ] Changes have been communicated to collaborators

### Does this PR introduce breaking changes? What changes might users
need to make to their code?
No

### Does this PR change default behavior?
No
---
 src/algorithms/tracking/CKFTracking.cc | 28 +++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/algorithms/tracking/CKFTracking.cc b/src/algorithms/tracking/CKFTracking.cc
index 0f809fff0..ba07f9356 100644
--- a/src/algorithms/tracking/CKFTracking.cc
+++ b/src/algorithms/tracking/CKFTracking.cc
@@ -7,6 +7,10 @@
 #include <Acts/Definitions/TrackParametrization.hpp>
 #include <Acts/Definitions/Units.hpp>
 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
+#if Acts_VERSION_MAJOR >= 32
+#include <Acts/EventData/TrackStateProxy.hpp>
+#endif
+#include <Acts/EventData/Types.hpp>
 #if Acts_VERSION_MAJOR < 36
 #include <Acts/EventData/Measurement.hpp>
 #endif
@@ -56,6 +60,7 @@
 #include <functional>
 #include <list>
 #include <optional>
+#include <ranges>
 #include <utility>
 
 #include "ActsGeometryProvider.h"
@@ -260,6 +265,7 @@ namespace eicrecon {
 #else
         Acts::TrackAccessor<unsigned int> seedNumber("seed");
 #endif
+        std::vector<Acts::TrackIndexType> failed_tracks;
 
         // Loop over seeds
         for (std::size_t iseed = 0; iseed < acts_init_trk_params.size(); ++iseed) {
@@ -291,6 +297,7 @@ namespace eicrecon {
                     ACTS_ERROR("Extrapolation for seed "
                         << iseed << " and track " << track.index()
                         << " failed with error " << extrapolationResult.error());
+                    failed_tracks.push_back(track.index());
                     continue;
                 }
 #endif
@@ -299,6 +306,18 @@ namespace eicrecon {
             }
         }
 
+        for (Acts::TrackIndexType track_index : std::ranges::reverse_view(failed_tracks)) {
+          // NOTE This does not remove track states corresponding to the
+          // removed tracks. Doing so would require implementing some garbage
+          // collection. We'll just assume no algorithm will access them
+          // directly.
+          acts_tracks.removeTrack(track_index);
+#if Acts_VERSION_MAJOR < 36
+          // Workaround an upstream bug in Acts::VectorTrackContainer::removeTrack_impl()
+          // https://github.com/acts-project/acts/commit/94cf81f3f1109210b963977e0904516b949b1154
+          trackContainer->m_particleHypothesis.erase(trackContainer->m_particleHypothesis.begin() + track_index);
+#endif
+        }
 
         // Move track states and track container to const containers
         // NOTE Using the non-const containers leads to references to
@@ -336,15 +355,6 @@ namespace eicrecon {
 
         std::optional<unsigned int> lastSeed;
         for (const auto& track : constTracks) {
-#if Acts_VERSION_MAJOR >= 34
-          // Some B0 tracks fail to extrapolate to the perigee surface. The
-          // Acts::extrapolateTrackToReferenceSurface will not set
-          // referenceSurface in that case, which is what we check here.
-          if (not track.hasReferenceSurface()) {
-            m_log->warn("Skipping a track not on perigee surface");
-            continue;
-          }
-#endif
           if (!lastSeed) {
             lastSeed = constSeedNumber(track);
           }