feat: Two-way CKF example (acts-project#3066)

Resurrect acts-project#2717 from @andiwand, updating to the latest `main`, and removing earlier test (`MyTrackFindingAlgorithm`). blocked by - acts-project#3074 - acts-project#3075 - acts-project#3076 - acts-project#3077 - acts-project#3078 - acts-project#3086 --------- Co-authored-by: andiwand <[email protected]> Co-authored-by: Tim Adye <[email protected]>
EleniXoch · May 6, 2024 · 305a098 · 305a098
1 parent f562f5e
commit 305a098
Show file tree

Hide file tree

Showing 25 changed files with 170 additions and 41 deletions.
diff --git a/CI/physmon/reference/performance_ambi_orthogonal.root b/CI/physmon/reference/performance_ambi_orthogonal.root
diff --git a/CI/physmon/reference/performance_ambi_seeded.root b/CI/physmon/reference/performance_ambi_seeded.root
diff --git a/CI/physmon/reference/performance_ambi_ttbar.root b/CI/physmon/reference/performance_ambi_ttbar.root
diff --git a/CI/physmon/reference/performance_amvf_gridseeder_seeded_hist.root b/CI/physmon/reference/performance_amvf_gridseeder_seeded_hist.root
diff --git a/CI/physmon/reference/performance_amvf_gridseeder_ttbar_hist.root b/CI/physmon/reference/performance_amvf_gridseeder_ttbar_hist.root
diff --git a/CI/physmon/reference/performance_amvf_orthogonal_hist.root b/CI/physmon/reference/performance_amvf_orthogonal_hist.root
diff --git a/CI/physmon/reference/performance_amvf_seeded_hist.root b/CI/physmon/reference/performance_amvf_seeded_hist.root
diff --git a/CI/physmon/reference/performance_amvf_truth_smeared_hist.root b/CI/physmon/reference/performance_amvf_truth_smeared_hist.root
diff --git a/CI/physmon/reference/performance_amvf_ttbar_hist.root b/CI/physmon/reference/performance_amvf_ttbar_hist.root
diff --git a/CI/physmon/reference/performance_ckf_orthogonal.root b/CI/physmon/reference/performance_ckf_orthogonal.root
diff --git a/CI/physmon/reference/performance_ckf_seeded.root b/CI/physmon/reference/performance_ckf_seeded.root
diff --git a/CI/physmon/reference/performance_ckf_truth_smeared.root b/CI/physmon/reference/performance_ckf_truth_smeared.root
diff --git a/CI/physmon/reference/performance_ckf_ttbar.root b/CI/physmon/reference/performance_ckf_ttbar.root
diff --git a/CI/physmon/reference/performance_ivf_orthogonal_hist.root b/CI/physmon/reference/performance_ivf_orthogonal_hist.root
diff --git a/CI/physmon/reference/performance_ivf_seeded_hist.root b/CI/physmon/reference/performance_ivf_seeded_hist.root
diff --git a/CI/physmon/reference/performance_ivf_truth_smeared_hist.root b/CI/physmon/reference/performance_ivf_truth_smeared_hist.root
diff --git a/CI/physmon/reference/tracksummary_ckf_orthogonal_hist.root b/CI/physmon/reference/tracksummary_ckf_orthogonal_hist.root
diff --git a/CI/physmon/reference/tracksummary_ckf_seeded_hist.root b/CI/physmon/reference/tracksummary_ckf_seeded_hist.root
diff --git a/CI/physmon/reference/tracksummary_ckf_truth_smeared_hist.root b/CI/physmon/reference/tracksummary_ckf_truth_smeared_hist.root
diff --git a/CI/physmon/reference/tracksummary_ckf_ttbar_hist.root b/CI/physmon/reference/tracksummary_ckf_ttbar_hist.root
diff --git a/Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp b/Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp
@@ -105,11 +105,14 @@ class TrackFindingAlgorithm final : public IAlgorithm {
     std::optional<std::variant<Acts::TrackSelector::Config,
                                Acts::TrackSelector::EtaBinnedConfig>>
         trackSelectorCfg = std::nullopt;
+
     /// Maximum number of propagation steps
     unsigned int maxSteps = 100000;
     /// Extrapolation strategy
     Acts::TrackExtrapolationStrategy extrapolationStrategy =
         Acts::TrackExtrapolationStrategy::firstOrLast;
+    /// Run finding in two directions
+    bool twoWay = true;
   };
 
   /// Constructor of the track finding algorithm

diff --git a/Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp b/Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp
@@ -16,6 +16,7 @@
 #include "Acts/EventData/TrackParameters.hpp"
 #include "Acts/EventData/VectorMultiTrajectory.hpp"
 #include "Acts/EventData/VectorTrackContainer.hpp"
+#include "Acts/Geometry/GeometryIdentifier.hpp"
 #include "Acts/Propagator/AbortList.hpp"
 #include "Acts/Propagator/EigenStepper.hpp"
 #include "Acts/Propagator/MaterialInteractor.hpp"
@@ -113,14 +114,23 @@ ActsExamples::ProcessCode ActsExamples::TrackFindingAlgorithm::execute(
       slAccessorDelegate;
   slAccessorDelegate.connect<&IndexSourceLinkAccessor::range>(&slAccessor);
 
-  Acts::PropagatorPlainOptions pOptions;
-  pOptions.maxSteps = m_cfg.maxSteps;
-  pOptions.direction = Acts::Direction::Forward;
+  Acts::PropagatorPlainOptions firstPropOptions;
+  firstPropOptions.maxSteps = m_cfg.maxSteps;
+  firstPropOptions.direction = Acts::Direction::Forward;
+
+  Acts::PropagatorPlainOptions secondPropOptions;
+  secondPropOptions.maxSteps = m_cfg.maxSteps;
+  secondPropOptions.direction = firstPropOptions.direction.invert();
 
   // Set the CombinatorialKalmanFilter options
-  ActsExamples::TrackFindingAlgorithm::TrackFinderOptions options(
+  ActsExamples::TrackFindingAlgorithm::TrackFinderOptions firstOptions(
+      ctx.geoContext, ctx.magFieldContext, ctx.calibContext, slAccessorDelegate,
+      extensions, firstPropOptions);
+
+  ActsExamples::TrackFindingAlgorithm::TrackFinderOptions secondOptions(
       ctx.geoContext, ctx.magFieldContext, ctx.calibContext, slAccessorDelegate,
-      extensions, pOptions);
+      extensions, secondPropOptions);
+  secondOptions.targetSurface = pSurface.get();
 
   Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator> extrapolator(
       Acts::EigenStepper<>(m_cfg.magneticField),
@@ -152,51 +162,164 @@ ActsExamples::ProcessCode ActsExamples::TrackFindingAlgorithm::execute(
 
   unsigned int nSeed = 0;
 
-  for (std::size_t iseed = 0; iseed < initialParameters.size(); ++iseed) {
+  for (std::size_t iSeed = 0; iSeed < initialParameters.size(); ++iSeed) {
     // Clear trackContainerTemp and trackStateContainerTemp
     tracksTemp.clear();
 
-    auto result =
-        (*m_cfg.findTracks)(initialParameters.at(iseed), options, tracksTemp);
+    const Acts::BoundTrackParameters& firstInitialParameters =
+        initialParameters.at(iSeed);
+
+    auto firstResult =
+        (*m_cfg.findTracks)(firstInitialParameters, firstOptions, tracksTemp);
     m_nTotalSeeds++;
     nSeed++;
 
-    if (!result.ok()) {
+    if (!firstResult.ok()) {
       m_nFailedSeeds++;
-      ACTS_WARNING("Track finding failed for seed " << iseed << " with error"
-                                                    << result.error());
+      ACTS_WARNING("Track finding failed for seed " << iSeed << " with error"
+                                                    << firstResult.error());
       continue;
     }
 
-    auto& tracksForSeed = result.value();
-    for (auto& track : tracksForSeed) {
-      auto smoothingResult = Acts::smoothTrack(ctx.geoContext, track, logger());
-      if (!smoothingResult.ok()) {
+    auto& firstTracksForSeed = firstResult.value();
+    for (auto& firstTrack : firstTracksForSeed) {
+      // TODO a copy of the track should not be necessary but is the safest way
+      //      with the current EDM
+      // TODO a lightweight copy without copying all the track state components
+      //      might be a solution
+      auto trackCandidate = tracksTemp.makeTrack();
+      trackCandidate.copyFrom(firstTrack, true);
+
+      auto firstSmoothingResult =
+          Acts::smoothTrack(ctx.geoContext, trackCandidate, logger());
+      if (!firstSmoothingResult.ok()) {
         m_nFailedSmoothing++;
-        ACTS_ERROR("Smoothing for seed "
-                   << iseed << " and track " << track.index()
-                   << " failed with error " << smoothingResult.error());
+        ACTS_ERROR("First smoothing for seed "
+                   << iSeed << " and track " << firstTrack.index()
+                   << " failed with error " << firstSmoothingResult.error());
         continue;
       }
 
-      auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
-          track, *pSurface, extrapolator, extrapolationOptions,
-          m_cfg.extrapolationStrategy, logger());
-      if (!extrapolationResult.ok()) {
-        m_nFailedExtrapolation++;
-        ACTS_ERROR("Extrapolation for seed "
-                   << iseed << " and track " << track.index()
-                   << " failed with error " << extrapolationResult.error());
-        continue;
-      }
+      // number of second tracks found
+      std::size_t nSecond = 0;
 
       // Set the seed number, this number decrease by 1 since the seed number
-      seedNumber(track) = nSeed - 1;
+      // has already been updated
+      seedNumber(trackCandidate) = nSeed - 1;
+
+      if (m_cfg.twoWay) {
+        std::optional<Acts::VectorMultiTrajectory::TrackStateProxy>
+            firstMeasurement;
+        for (auto trackState : trackCandidate.trackStatesReversed()) {
+          bool isMeasurement = trackState.typeFlags().test(
+              Acts::TrackStateFlag::MeasurementFlag);
+          bool isOutlier =
+              trackState.typeFlags().test(Acts::TrackStateFlag::OutlierFlag);
+          // We are excluding non measurement states and outlier here. Those can
+          // decrease resolution because only the smoothing corrected the very
+          // first prediction as filtering is not possible.
+          if (isMeasurement && !isOutlier) {
+            firstMeasurement = trackState;
+          }
+        }
+
+        if (firstMeasurement.has_value()) {
+          Acts::BoundTrackParameters secondInitialParameters =
+              trackCandidate.createParametersFromState(*firstMeasurement);
+
+          auto secondResult = (*m_cfg.findTracks)(secondInitialParameters,
+                                                  secondOptions, tracksTemp);
+
+          if (!secondResult.ok()) {
+            ACTS_WARNING("Second track finding failed for seed "
+                         << iSeed << " with error" << secondResult.error());
+          } else {
+            auto firstState =
+                *std::next(trackCandidate.trackStatesReversed().begin(),
+                           trackCandidate.nTrackStates() - 1);
+            assert(firstState.previous() == Acts::kTrackIndexInvalid);
+
+            auto& secondTracksForSeed = secondResult.value();
+            for (auto& secondTrack : secondTracksForSeed) {
+              if (secondTrack.nTrackStates() < 2) {
+                continue;
+              }
+
+              // TODO a copy of the track should not be necessary but is the
+              //      safest way with the current EDM
+              // TODO a lightweight copy without copying all the track state
+              //      components might be a solution
+              auto secondTrackCopy = tracksTemp.makeTrack();
+              secondTrackCopy.copyFrom(secondTrack, true);
+
+              // Note that this is only valid if there are no branches
+              // We disallow this by breaking this look after a second track was
+              // processed
+              secondTrackCopy.reverseTrackStates(true);
+
+              firstState.previous() =
+                  (*std::next(secondTrackCopy.trackStatesReversed().begin()))
+                      .index();
+
+              Acts::calculateTrackQuantities(trackCandidate);
+
+              // TODO This extrapolation should not be necessary
+              // TODO The CKF is targeting this surface and should communicate
+              //      the resulting parameters
+              // TODO Removing this requires changes in the core CKF
+              //      implementation
+              auto secondExtrapolationResult =
+                  Acts::extrapolateTrackToReferenceSurface(
+                      trackCandidate, *pSurface, extrapolator,
+                      extrapolationOptions, m_cfg.extrapolationStrategy,
+                      logger());
+              if (!secondExtrapolationResult.ok()) {
+                m_nFailedExtrapolation++;
+                ACTS_ERROR("Second extrapolation for seed "
+                           << iSeed << " and track " << secondTrack.index()
+                           << " failed with error "
+                           << secondExtrapolationResult.error());
+
+                continue;
+              }
+
+              if (!m_trackSelector.has_value() ||
+                  m_trackSelector->isValidTrack(trackCandidate)) {
+                auto destProxy = tracks.makeTrack();
+                destProxy.copyFrom(trackCandidate, true);
+              }
+
+              ++nSecond;
+            }
+
+            // restore `trackCandidate` to its original state in case we need it
+            // again
+            firstState.previous() = Acts::kTrackIndexInvalid;
+            Acts::calculateTrackQuantities(trackCandidate);
+          }
+        }
+      }
 
-      if (!m_trackSelector.has_value() ||
-          m_trackSelector->isValidTrack(track)) {
-        auto destProxy = tracks.makeTrack();
-        destProxy.copyFrom(track, true);  // make sure we copy track states!
+      // if no second track was found, we will use only the first track
+      if (nSecond == 0) {
+        auto firstExtrapolationResult =
+            Acts::extrapolateTrackToReferenceSurface(
+                trackCandidate, *pSurface, extrapolator, extrapolationOptions,
+                m_cfg.extrapolationStrategy, logger());
+        if (!firstExtrapolationResult.ok()) {
+          m_nFailedExtrapolation++;
+          ACTS_ERROR("Extrapolation for seed "
+                     << iSeed << " and track " << firstTrack.index()
+                     << " failed with error "
+                     << firstExtrapolationResult.error());
+          continue;
+        }
+
+        if (!m_trackSelector.has_value() ||
+            m_trackSelector->isValidTrack(trackCandidate)) {
+          auto destProxy = tracks.makeTrack();
+          destProxy.copyFrom(trackCandidate, true);
+        }
       }
     }
   }

diff --git a/Examples/Python/python/acts/examples/reconstruction.py b/Examples/Python/python/acts/examples/reconstruction.py
@@ -1164,6 +1164,7 @@ def addCKFTracks(
         Union[TrackSelectorConfig, List[TrackSelectorConfig]]
     ] = None,
     ckfConfig: CkfConfig = CkfConfig(),
+    twoWay: bool = True,
     outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
     writeTrajectories: bool = True,
@@ -1267,6 +1268,7 @@ def addCKFTracks(
             magneticField=field,
             trackSelectorCfg=trkSelCfg,
             maxSteps=ckfConfig.maxSteps,
+            twoWay=twoWay,
         ),
     )
     s.addAlgorithm(trackFinder)

diff --git a/Examples/Python/src/TrackFinding.cpp b/Examples/Python/src/TrackFinding.cpp
@@ -332,6 +332,7 @@ void addTrackFinding(Context& ctx) {
     ACTS_PYTHON_MEMBER(measurementSelectorCfg);
     ACTS_PYTHON_MEMBER(trackSelectorCfg);
     ACTS_PYTHON_MEMBER(maxSteps);
+    ACTS_PYTHON_MEMBER(twoWay);
     ACTS_PYTHON_STRUCT_END();
   }
 

diff --git a/Examples/Python/tests/root_file_hashes.txt b/Examples/Python/tests/root_file_hashes.txt
@@ -44,22 +44,22 @@ test_digitization_example_input[smeared]__particles.root: 7eec62018b6944fea565da
 test_digitization_example_input[smeared]__measurements.root: 0c168d371d0130c68d1ee44bd77eeeb3cf702a77c2afbf12bed8354b61a29262
 test_digitization_example_input[geometric]__particles.root: 7eec62018b6944fea565dad75aa41ef87d1f2737b2a814fbab189817ac8180fe
 test_digitization_example_input[geometric]__measurements.root: 0c6d88b4de3ee7365103b8f0d6be6b4db3d7b7f2a59d3db58a1e5f89fa8130b3
-test_ckf_tracks_example[generic-full_seeding]__trackstates_ckf.root: 75f429d8f1ee8198cf75bc3b685e2fee9354175893168279c14e67b6e6ab01c2
-test_ckf_tracks_example[generic-full_seeding]__tracksummary_ckf.root: 92a47dcc7caa490e960fa71db088d1a53fa016f5ccf2d3765de3376ca509b103
+test_ckf_tracks_example[generic-full_seeding]__trackstates_ckf.root: 495cbb5b2cf8a2d8a8efc4c3c45502049d142959c9aa9fec5482ca2faf7e43fc
+test_ckf_tracks_example[generic-full_seeding]__tracksummary_ckf.root: d4863282e7ebb07b223ba5ee6c608dbc8de26adecae52f3146f6f7144cf633f9
 test_ckf_tracks_example[generic-full_seeding]__performance_seeding_trees.root: 0e0676ffafdb27112fbda50d1cf627859fa745760f98073261dcf6db3f2f991e
 test_ckf_tracks_example[generic-truth_estimated]__trackstates_ckf.root: 2e772b7ab018667330b3633de81f898f5f65d01bdf9b20a243921454d5cbe843
 test_ckf_tracks_example[generic-truth_estimated]__tracksummary_ckf.root: e19a572abb138f5d6ce0716724fd10dcee62ac33a9749c27d62702355e86aa2d
 test_ckf_tracks_example[generic-truth_estimated]__performance_seeding.root: 1facb05c066221f6361b61f015cdf0918e94d9f3fce2269ec7b6a4dffeb2bc7e
-test_ckf_tracks_example[generic-truth_smeared]__trackstates_ckf.root: 5da4629bd647b5fff6b947fbd248e0ff97f6930ad727cc7b29c1252fe6c0d0e1
-test_ckf_tracks_example[generic-truth_smeared]__tracksummary_ckf.root: b58061b03e462d2b344fabcd45465e47b12c9c6ce4a39e5a23ae641c280e15b1
-test_ckf_tracks_example[odd-full_seeding]__trackstates_ckf.root: 8e9d52e4caf4b603b2e77d17ba4769b8839ab1222d95228424e8cfd9dab656ed
-test_ckf_tracks_example[odd-full_seeding]__tracksummary_ckf.root: 3f8e7b86c061b385580d9261cb1f32b463773a42c9ad4117a9edb5bef3fcdbc4
+test_ckf_tracks_example[generic-truth_smeared]__trackstates_ckf.root: f313dbe88b9d93be1faedd63b8b9ff7ef26ad6cc3d90f860e3a1f4654db58de1
+test_ckf_tracks_example[generic-truth_smeared]__tracksummary_ckf.root: d5d1a2c6151da064ff3020b1da29624abad4c58f2ee874b61fd410284103d8b2
+test_ckf_tracks_example[odd-full_seeding]__trackstates_ckf.root: 320dbc32a8b282f1f0af15c8bfe216b6fc266406f1a52adfd8cb82d9e40cb2bf
+test_ckf_tracks_example[odd-full_seeding]__tracksummary_ckf.root: a266500caf6666e020f53848e3a0f683fa6bf3566a893a4413c5bf09bee7b0a1
 test_ckf_tracks_example[odd-full_seeding]__performance_seeding_trees.root: 43c58577aafe07645e5660c4f43904efadf91d8cda45c5c04c248bbe0f59814f
 test_ckf_tracks_example[odd-truth_estimated]__trackstates_ckf.root: e3e6516c462c5cbfe49329173560b4b254e90c225d8cb7ed87135a3fc8e365b9
 test_ckf_tracks_example[odd-truth_estimated]__tracksummary_ckf.root: 1cc1f57a9b8c910b7c1da5c43e6c6098e06b4d51b6237445affdf89a53b7a6e6
 test_ckf_tracks_example[odd-truth_estimated]__performance_seeding.root: 1a36b7017e59f1c08602ef3c2cb0483c51df248f112e3780c66594110719c575
-test_ckf_tracks_example[odd-truth_smeared]__trackstates_ckf.root: b8026ab0be3feb71538c918656d424406d3ad16241e5421ddb7c536c703ce9fc
-test_ckf_tracks_example[odd-truth_smeared]__tracksummary_ckf.root: 6734320133c935f46d98476a4a6a4a3b6403a613c299088e07a3a01678da6bb6
+test_ckf_tracks_example[odd-truth_smeared]__trackstates_ckf.root: 15f742cd63df65dd553bc4a43aefe3a3fc376373159cf84af72fdef51a3171e7
+test_ckf_tracks_example[odd-truth_smeared]__tracksummary_ckf.root: 6232ef74addc49742aeca9833571c01b9b589c93405a1187505d1c3a0142b8ab
 test_vertex_fitting_reading[Truth-False-100]__performance_vertexing.root: 76ef6084d758dfdfc0151ddec2170e12d73394424e3dac4ffe46f0f339ec8293
 test_vertex_fitting_reading[Iterative-False-100]__performance_vertexing.root: 60372210c830a04f95ceb78c6c68a9b0de217746ff59e8e73053750c837b57eb
 test_vertex_fitting_reading[Iterative-True-100]__performance_vertexing.root: e34f217d524a5051dbb04a811d3407df3ebe2cc4bb7f54f6bda0847dbd7b52c3