Run Kalman Smoother in CKF

core/include/traccc/edm/fitting_result.hpp

@@ -0,0 +1,39 @@
+/** 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"
+
+// System include(s).
+#include <limits>
+
+namespace traccc {
+
+template <typename algebra_t>
+struct fitting_result {
+
+    using scalar_type = detray::dscalar<algebra_t>;
+
+    /// Fitted track parameter at the first track state
+    detray::bound_track_parameters<algebra_t> fitted_params_initial;
+
+    /// Fitted track parameter at the last track state
+    detray::bound_track_parameters<algebra_t> fitted_params_final;
+
+    /// Number of degree of freedoms of the track
+    scalar_type ndf{0.f};
+
+    /// Chi square from finding/fitting algorithm
+    scalar_type chi2{std::numeric_limits<scalar_type>::max()};
+
+    /// The number of holes
+    unsigned int n_holes{0u};
+};
+
+}  // namespace traccc

core/include/traccc/edm/track_candidate.hpp

@@ -8,21 +8,27 @@
 #pragma once
 
 // Project include(s).
+#include "traccc/edm/fitting_result.hpp"
 #include "traccc/edm/measurement.hpp"
-#include "traccc/edm/track_parameters.hpp"
-
-// Detray include(s).
-#include "detray/geometry/barcode.hpp"
 
 namespace traccc {
 
+struct track_summary {
+
+    /// (Mandatory) Seed track parameter
+    bound_track_parameters seed;
+
+    /// (Optional) Fitting result
+    fitting_result<traccc::default_algebra> fit_res{};
+};
+
 /// Track candidate is the measurement
 using track_candidate = measurement;
 
 /// Declare a track candidates collection types
 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

@@ -10,6 +10,7 @@
 // Project include(s).
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/container.hpp"
+#include "traccc/edm/fitting_result.hpp"
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate.hpp"
 
@@ -19,21 +20,6 @@
 
 namespace traccc {
 
-/// Fitting result per track
-template <typename algebra_t>
-struct fitting_result {
-    using scalar_type = detray::dscalar<algebra_t>;
-
-    /// Fitted track parameter
-    detray::bound_track_parameters<algebra_t> fit_params;
-
-    /// Number of degree of freedoms of fitted track
-    scalar_type ndf{0};
-
-    /// Chi square of fitted track
-    scalar_type chi2{0};
-};
-
 /// Fitting result per measurement
 template <typename algebra_t>
 struct track_state {

core/include/traccc/finding/actors/ckf_aborter.hpp

@@ -28,6 +28,8 @@ struct ckf_aborter : detray::actor {
 
         bool success = false;
         unsigned int count = 0;
+
+        scalar path_from_surface{0.f};
     };
 
     template <typename propagator_state_t>
@@ -38,14 +40,21 @@ struct ckf_aborter : detray::actor {
         auto &stepping = prop_state._stepping;
 
         abrt_state.count++;
+        abrt_state.path_from_surface += stepping.step_size();
 
         // Abort at the next sensitive surface
         if (navigation.is_on_sensitive() &&
-            stepping.path_from_surface() > abrt_state.min_step_length) {
+            abrt_state.path_from_surface > abrt_state.min_step_length) {
             prop_state._heartbeat &= navigation.abort();
             abrt_state.success = true;
         }
 
+        // Reset path from surface
+        if (navigation.is_on_sensitive() ||
+            navigation.encountered_sf_material()) {
+            abrt_state.path_from_surface = 0.f;
+        }
+
         if (abrt_state.count > abrt_state.max_count) {
             prop_state._heartbeat &= navigation.abort();
         }

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,14 @@ 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.at(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.at(step).size() - 1u)});
                     updated_params.push_back(trk_state.filtered());
                 }
             }
@@ -388,7 +395,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{});
+        bound_track_parameters fit_par_ini;
+        bound_track_parameters fit_par_fin;
+
+        // 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 +424,68 @@ 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
+                fit_par_fin = 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
+                // FIXME: Should be called only for the front tip
+                fit_par_ini = 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,
+                                  {fit_par_ini, fit_par_fin, ndf_sum, chi2_sum,
+                                   n_skipped}},
+                    cands_per_track);
                 break;
             }
 

core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp

@@ -122,8 +122,7 @@ class kalman_fitter {
     ///
     /// @param seed_params seed track parameter
     /// @param fitter_state the state of kalman fitter
-    template <typename seed_parameters_t>
-    TRACCC_HOST_DEVICE void fit(const seed_parameters_t& seed_params,
+    TRACCC_HOST_DEVICE void fit(const track_summary& trk_summary,
                                 state& fitter_state) {
 
         // Run the kalman filtering for a given number of iterations
@@ -133,7 +132,7 @@ class kalman_fitter {
             fitter_state.m_fit_actor_state.reset();
 
             if (i == 0) {
-                filter(seed_params, fitter_state);
+                filter(trk_summary.seed, fitter_state);
             }
             // From the second iteration, seed parameter is the smoothed track
             // parameter at the first surface
@@ -226,8 +225,10 @@ class kalman_fitter {
         auto& fit_res = fitter_state.m_fit_res;
         auto& track_states = fitter_state.m_fit_actor_state.m_track_states;
 
-        // Fit parameter = smoothed track parameter at the first surface
-        fit_res.fit_params = track_states[0].smoothed();
+        // Fit parameter = smoothed track parameter at the first and last
+        // surface
+        fit_res.fitted_params_initial = track_states.at(0).smoothed();
+        fit_res.fitted_params_final = track_states.back().smoothed();
 
         for (const auto& trk_state : track_states) {
 

device/common/include/traccc/finding/device/build_tracks.hpp

@@ -12,6 +12,7 @@
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate.hpp"
 #include "traccc/edm/track_parameters.hpp"
+#include "traccc/edm/track_state.hpp"
 #include "traccc/finding/candidate_link.hpp"
 
 namespace traccc::device {
@@ -33,6 +34,11 @@ struct build_tracks_payload {
      */
     vecmem::data::jagged_vector_view<const candidate_link> links_view;
 
+    /**
+     * @brief View object to the vector of track states
+     */
+    vecmem::data::jagged_vector_view<const track_state<default_algebra>> track_states_view;
+
     /**
      * @brief View object to the vector of tips
      */

device/common/include/traccc/finding/device/find_tracks.hpp

@@ -94,6 +94,11 @@ struct find_tracks_payload {
      */
     vecmem::data::vector_view<candidate_link> links_view;
 
+    /**
+     * @brief View object to the track state vector
+     */
+    vecmem::data::vector_view<track_state<default_algebra>> track_states_view;
+
     /**
      * @brief Pointer to the total of number of candidates; to be set to zero
      * before launching the kernel

device/common/include/traccc/finding/device/impl/build_tracks.ipp

@@ -30,6 +30,9 @@ TRACCC_DEVICE inline void build_tracks(std::size_t globalIndex,
     vecmem::jagged_device_vector<const candidate_link> links(
         payload.links_view);
 
+    vecmem::jagged_device_vector<const track_state<default_algebra>>
+        track_states(payload.track_states_view);
+
     vecmem::device_vector<const typename candidate_link::link_index_type> tips(
         payload.tips_view);
 
@@ -44,7 +47,7 @@ TRACCC_DEVICE inline void build_tracks(std::size_t globalIndex,
     }
 
     const auto tip = tips.at(globalIndex);
-    auto& seed = track_candidates[globalIndex].header;
+    auto& trk_summary = track_candidates[globalIndex].header;
     auto cands_per_track = track_candidates[globalIndex].items;
 
     // Get the link corresponding to tip
@@ -99,7 +102,7 @@ TRACCC_DEVICE inline void build_tracks(std::size_t globalIndex,
         // Break the loop if the iterator is at the first candidate and fill the
         // seed
         if (it == cands_per_track.rend() - 1) {
-            seed = seeds.at(L.previous.second);
+            trk_summary.seed = seeds.at(L.previous.second);
             break;
         }
 

device/common/include/traccc/finding/device/impl/find_tracks.ipp

@@ -61,6 +61,8 @@ TRACCC_DEVICE inline void find_tracks(
     vecmem::device_vector<unsigned int> out_params_liveness(
         payload.out_params_liveness_view);
     vecmem::device_vector<candidate_link> links(payload.links_view);
+    vecmem::device_vector<track_state<typename detector_t::algebra_type>>
+        track_states(payload.track_states_view);
     vecmem::device_atomic_ref<unsigned int> num_total_candidates(
         *payload.n_total_candidates);
     vecmem::device_vector<const detray::geometry::barcode> barcodes(
@@ -222,6 +224,7 @@ TRACCC_DEVICE inline void find_tracks(
                             prev_link.seed_idx,
                             prev_link.n_skipped};
                     }
+                    track_states.at(l_pos) = trk_state;
 
                     // Increase the number of candidates (or branches) per input
                     // parameter