Run Kalman Smoother in CKF

core/include/traccc/edm/track_candidate.hpp

@@ -9,7 +9,7 @@
 
 // Project include(s).
 #include "traccc/edm/measurement.hpp"
-#include "traccc/edm/track_parameters.hpp"
+#include "traccc/edm/track_summary.hpp"
 
 // Detray include(s).
 #include "detray/geometry/barcode.hpp"
@@ -23,6 +23,6 @@ using track_candidate = measurement;
 using track_candidate_collection_types = collection_types<track_candidate>;
 /// Declare a track candidates container type
 using track_candidate_container_types =
-    container_types<bound_track_parameters, track_candidate>;
+    container_types<track_summary, track_candidate>;
 
 }  // namespace traccc

core/include/traccc/edm/track_state.hpp

@@ -12,6 +12,7 @@
 #include "traccc/edm/container.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate.hpp"
+#include "traccc/edm/track_summary.hpp"
 
 // detray include(s).
 #include "detray/navigation/navigator.hpp"

core/include/traccc/edm/track_summary.hpp

@@ -0,0 +1,34 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2024 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Project include(s).
+#include "traccc/edm/track_parameters.hpp"
+
+namespace traccc {
+
+struct track_summary {
+
+    /// (Mandatory) Seed track parameter
+    bound_track_parameters seed;
+
+    /// (Optional) Fitted track parameter at the both ends
+    thrust::pair<bound_track_parameters, bound_track_parameters>
+        fitted_params_at_tip;
+
+    /// (Optional) Number of degree of freedoms of the track
+    scalar ndf{0.f};
+
+    /// (Optional) Chi square from finding/fitting algorithm
+    scalar chi2{std::numeric_limits<float>::max()};
+
+    /// (Optional) The number of holes
+    unsigned int n_holes{0u};
+};
+
+}  // namespace traccc

core/include/traccc/finding/candidate_link.hpp

@@ -13,6 +13,9 @@
 // Thrust include(s).
 #include <thrust/pair.h>
 
+// System include(s).
+#include <limits>
+
 namespace traccc {
 
 // A link that contains the index of corresponding measurement and the index of
@@ -33,6 +36,9 @@ struct candidate_link {
 
     // How many times it skipped a surface
     unsigned int n_skipped;
+
+    // track state index
+    unsigned int track_state_idx{std::numeric_limits<unsigned int>::max()};
 };
 
 }  // namespace traccc

core/include/traccc/finding/details/find_tracks.hpp

@@ -14,6 +14,7 @@
 #include "traccc/finding/actors/interaction_register.hpp"
 #include "traccc/finding/candidate_link.hpp"
 #include "traccc/finding/finding_config.hpp"
+#include "traccc/fitting/kalman_filter/gain_matrix_smoother.hpp"
 #include "traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
 #include "traccc/sanity/contiguous_on.hpp"
 #include "traccc/utils/particle.hpp"
@@ -123,6 +124,8 @@ track_candidate_container_types::host find_tracks(
 
     std::vector<std::vector<candidate_link>> links;
     links.resize(config.max_track_candidates_per_track);
+    std::vector<std::vector<track_state<algebra_type>>> track_states;
+    track_states.resize(config.max_track_candidates_per_track);
 
     std::vector<std::vector<std::size_t>> param_to_link;
     param_to_link.resize(config.max_track_candidates_per_track);
@@ -247,10 +250,13 @@ track_candidate_container_types::host find_tracks(
                 if (res && trk_state.filtered_chi2() < config.chi2_max) {
                     n_branches++;
 
-                    links[step].push_back({{previous_step, in_param_id},
-                                           item_id,
-                                           orig_param_id,
-                                           skip_counter});
+                    track_states[step].push_back(trk_state);
+                    links[step].push_back(
+                        {{previous_step, in_param_id},
+                         item_id,
+                         orig_param_id,
+                         skip_counter,
+                         static_cast<unsigned int>(track_states.size() - 1u)});
                     updated_params.push_back(trk_state.filtered());
                 }
             }
@@ -388,7 +394,17 @@ track_candidate_container_types::host find_tracks(
         vecmem::vector<track_candidate> cands_per_track;
         cands_per_track.resize(n_cands);
 
-        // Reversely iterate to fill the track candidates
+        // Storage for states used in smoother
+        track_state<algebra_type> next_state(measurement{});
+        thrust::pair<bound_track_parameters, bound_track_parameters>
+            fitted_params_at_tip;
+
+        // Track summary variables
+        scalar ndf_sum = 0.f;
+        scalar chi2_sum = 0.f;
+
+        // Reversely iterate to fill the track candidates and run Kalman
+        // smoother
         for (auto it = cands_per_track.rbegin(); it != cands_per_track.rend();
              it++) {
 
@@ -407,14 +423,66 @@ track_candidate_container_types::host find_tracks(
             auto& cand = *it;
             cand = measurements.at(L.meas_idx);
 
+            unsigned int step =
+                (L.previous.first == std::numeric_limits<int>::max())
+                    ? 0u
+                    : L.previous.first + 1;
+            if (it == cands_per_track.rbegin()) {
+                // The smoothing algorithm requires the following:
+                // (1) the filtered track parameter of the current surface
+                // (2) the smoothed track parameter of the next surface
+                //
+                // Since the smoothed track parameter of the last surface can be
+                // considered to be the filtered one, we can reversly iterate
+                // the algorithm to obtain the smoothed parameter of other
+                // surfaces
+                auto& tip_state = track_states.at(step).at(L.track_state_idx);
+                tip_state.smoothed().set_parameter_vector(tip_state.filtered());
+                tip_state.smoothed().set_covariance(
+                    tip_state.filtered().covariance());
+                tip_state.smoothed_chi2() = tip_state.filtered_chi2();
+
+                // Set the back tip parameter
+                fitted_params_at_tip.second = tip_state.smoothed();
+
+                ndf_sum +=
+                    static_cast<scalar>(tip_state.get_measurement().meas_dim);
+                chi2_sum += tip_state.smoothed_chi2();
+
+                // Swap the state
+                next_state = tip_state;
+            } else {
+                // Current state
+                auto& cur_state = track_states.at(step).at(L.track_state_idx);
+
+                // Run kalman smoother
+                const detray::tracking_surface sf{det,
+                                                  cur_state.surface_link()};
+                sf.template visit_mask<gain_matrix_smoother<algebra_type>>(
+                    cur_state, next_state);
+
+                // Reset the front tip parameter
+                fitted_params_at_tip.first = cur_state.smoothed();
+
+                ndf_sum +=
+                    static_cast<scalar>(cur_state.get_measurement().meas_dim);
+                chi2_sum += cur_state.smoothed_chi2();
+
+                // Swap the state
+                next_state = cur_state;
+            }
+
             // Break the loop if the iterator is at the first candidate and
             // fill the seed
             if (it == cands_per_track.rend() - 1) {
 
                 auto cand_seed = seeds.at(L.previous.second);
-
+                
                 // Add seed and track candidates to the output container
-                output_candidates.push_back(cand_seed, cands_per_track);
+                output_candidates.push_back(
+                    track_summary{cand_seed, fitted_params_at_tip, ndf_sum,
+                                  chi2_sum, n_skipped},
+                    cands_per_track);
                 break;
             }