ref: Move Kalman smoother into CKF

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

@@ -130,6 +130,9 @@
 
     std::vector<std::pair<unsigned int, unsigned int>> tips;
 
+    std::vector<std::vector<typename detector_t::surface_type>> sf_sequences;
+    sf_sequences.resize(config.max_track_candidates_per_track);
+
     // Create propagator
     auto prop_cfg{config.propagation};
     prop_cfg.navigation.estimate_scattering_noise = false;
@@ -227,7 +230,7 @@
             if (sf.has_material()) {
                 const typename detector_t::geometry_context ctx{};
                 traccc::details::ckf_interactor_t::state interactor_state;
                 traccc::details::ckf_interactor_t{}.update(
                     ctx,
                     detail::correct_particle_hypothesis(config.ptc_hypothesis,
                                                         in_param),
@@ -373,10 +376,10 @@
                 last_meas_to_tracks_map[L.meas_idx].push_back(i);
             }
 
             for (const auto& it : last_meas_to_tracks_map) {
                 const auto& tracks = it.second;
 
                 for (std::size_t i = 0; i < tracks.size(); ++i) {
                     const auto& Lthisbase = links.at(step).at(tracks.at(i));
                     const scalar prob_this =
                         prob(Lthisbase.chi2_sum,
@@ -408,7 +411,7 @@
                             prob(Lthat.chi2_sum,
                                  static_cast<scalar>(Lthat.ndf_sum - 5));
 
                         while (true) {
                             while (Lthis.meas_idx >= n_meas &&
                                    Lthis.step != 0u) {
                                 const auto link_pos =
@@ -507,6 +510,10 @@
 
             typename detray::pathlimit_aborter<scalar_type>::state
                 aborter_state;
+            typename detray::surface_sequencer<
+                typename detector_t::surface_type>::state sequencer_state{
+                vecmem::device_vector<typename detector_t::surface_type>{
+                    vecmem::get_data(sf_sequences[link_id])}};
             typename detray::parameter_transporter<
                 typename detector_t::algebra_type>::state transporter_state;
             traccc::details::ckf_interactor_t::state interactor_state;
@@ -533,7 +540,7 @@
             TRACCC_DEBUG_HOST("Propagating... ");
             propagator.propagate(
                 propagation,
-                detray::tie(aborter_state, transporter_state,
+                detray::tie(aborter_state, sequencer_state, transporter_state,
                             interaction_register_state, interactor_state,
                             resetter_state, momentum_aborter_state,
                             ckf_aborter_state));

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

@@ -37,8 +37,10 @@ using ckf_interactor_t =
     detray::pointwise_material_interactor<traccc::default_algebra>;
 
 /// Actor chain used in the Combinatorial Kalman Filter (CKF)
+template <typename surface_t>
 using ckf_actor_chain_t =
     detray::actor_chain<detray::pathlimit_aborter<traccc::scalar>,
+                        detray::surface_sequencer<surface_t>,
                         detray::parameter_transporter<traccc::default_algebra>,
                         interaction_register<ckf_interactor_t>,
                         ckf_interactor_t,
@@ -53,6 +55,6 @@ template <typename detector_t, typename bfield_t>
 using ckf_propagator_t =
     detray::propagator<ckf_stepper_t<bfield_t>,
                        detray::caching_navigator<std::add_const_t<detector_t>>,
-                       ckf_actor_chain_t>;
+                       ckf_actor_chain_t<typename detector_t::surface_type>>;
 
 }  // namespace traccc::details

core/include/traccc/finding/finding_config.hpp

@@ -10,6 +10,7 @@
 // traccc include(s).
 #include "traccc/definitions/common.hpp"
 #include "traccc/definitions/primitives.hpp"
+#include "traccc/fitting/fitting_config.hpp"
 #include "traccc/utils/particle.hpp"
 
 // detray include(s).
@@ -46,6 +47,11 @@ struct finding_config {
     /// Enable the MBF smoother
     bool run_mbf_smoother = true;
 
+    /// Enable the Kalman smoother
+    bool run_kalman_smoother = false;
+    // Fitting config for the Kalman smoother
+    traccc::fitting_config kalman_smoother;
+
     /// Minimum step length that track should make to reach the next surface. It
     /// should be set higher than the overstep tolerance not to make it stay on
     /// the same surface

core/include/traccc/fitting/details/kalman_fitting.hpp

@@ -71,7 +71,8 @@ typename edm::track_container<algebra_t>::host kalman_fitting(
             seqs_buffer{
                 static_cast<vecmem::data::vector_buffer<
                     typename fitter_t::surface_type>::size_type>(
-                    std::max(fitted_track.constituent_links().size() *
+                    std::max(static_cast<unsigned int>(
+                                 fitted_track.constituent_links().size()) *
                                  fitter.config().surface_sequence_size_factor,
                              fitter.config().min_surface_sequence_capacity)),
                 mr, vecmem::data::buffer_type::resizable};

core/include/traccc/fitting/fitting_config.hpp

@@ -20,7 +20,7 @@ namespace traccc {
 /// Configuration struct for track fitting
 struct fitting_config {
 
-    std::size_t n_iterations = 1;
+    unsigned int n_iterations = 1u;
 
     /// Propagation configuration
     detray::propagation::config propagation{};
@@ -35,8 +35,8 @@ struct fitting_config {
 
     /// Smoothing with backward filter
     traccc::scalar covariance_inflation_factor = 1e3f;
-    std::size_t surface_sequence_size_factor = 5;
-    std::size_t min_surface_sequence_capacity = 100;
+    unsigned int surface_sequence_size_factor = 4u;
+    unsigned int min_surface_sequence_capacity = 100u;
 };
 
 }  // namespace traccc

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

@@ -341,12 +341,13 @@
     /// @param actor_state the actor state
     /// @param propagation the propagator state
     template <typename propagator_state_t>
     TRACCC_HOST_DEVICE void operator()(state& actor_state,
                                        propagator_state_t& propagation) const {
 
         auto& stepping = propagation._stepping;
         auto& navigation = propagation._navigation;
 
+        TRACCC_INFO_HOST_DEVICE("IN SMOOTHER");
         TRACCC_VERBOSE_HOST_DEVICE("In Kalman actor (status %d)...",
                                    actor_state.fit_result);
 
@@ -367,6 +368,7 @@
         // triggered only for sensitive surfaces
         if (navigation.is_on_sensitive()) {
 
+            TRACCC_INFO_HOST_DEVICE("On surface");
             TRACCC_DEBUG_HOST(
                 "-> on surface: " << navigation.current_surface());
 
@@ -379,9 +381,11 @@
                     actor_state.add_to_sequence(
                         std::as_const(navigation).current().sf_desc);
                 }
+                DETRAY_ERROR_HOST_DEVICE("Surface not matched!");
                 return;
             } else if (actor_state.fit_result !=
                        kalman_fitter_status::SUCCESS) {
+                DETRAY_ERROR_HOST_DEVICE("ABORTED");
                 // Surface matched but encountered error: Abort fit
                 navigation.abort(fitter_debug_msg{actor_state.fit_result});
                 propagation._heartbeat = false;
@@ -435,7 +439,7 @@
                               direction_e ==
                                   kalman_actor_direction::BIDIRECTIONAL) {
                     // Backward filter for smoothing
-                    TRACCC_DEBUG_HOST_DEVICE("Run smoothing...");
+                    TRACCC_INFO_HOST_DEVICE("Run smoothing...");
 
                     // Forward filter did not find this state: cannot smoothe
                     if (trk_state.filtered_params().is_invalid()) {
@@ -449,6 +453,12 @@
                             two_filters_smoother<algebra_t>{}(
                                 trk_state, actor_state.m_measurements,
                                 bound_param, is_line);
+                        if (!trk_state.is_smoothed()) {
+                            TRACCC_ERROR_HOST_DEVICE(
+                                "Track state NOT smoothed!");
+                        } else {
+                            TRACCC_INFO_HOST_DEVICE("Track state IS smoothed!");
+                        }
                     }
                 } else {
                     assert(false);
@@ -487,6 +497,7 @@
             if (actor_state.finished() && !actor_state.do_precise_hole_count) {
                 navigation.exit();
                 propagation._heartbeat = false;
+                TRACCC_INFO_HOST_DEVICE("Navigation exit");
                 return;
             }
 

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

@@ -295,13 +295,14 @@ class kalman_fitter {
     /// @param fitter_state the state of kalman fitter
     [[nodiscard]] TRACCC_HOST_DEVICE kalman_fitter_status
     smooth(state& fitter_state) const {
-        TRACCC_VERBOSE_HOST_DEVICE("Run smoothing...");
+        TRACCC_INFO_HOST_DEVICE("Run smoothing...");
 
         if (fitter_state.m_fit_actor_state.sequencer().overflow()) {
-            TRACCC_ERROR_HOST_DEVICE("Barcode sequence overlow");
+            TRACCC_ERROR_HOST_DEVICE("Surface sequence overlow");
             return kalman_fitter_status::ERROR_BARCODE_SEQUENCE_OVERFLOW;
         }
         if (fitter_state.m_fit_actor_state.sequencer().sequence().empty()) {
+            TRACCC_ERROR_HOST_DEVICE("Surface sequence empty");
             return kalman_fitter_status::ERROR_UPDATER_SKIPPED_STATE;
         }
 
@@ -314,6 +315,12 @@ class kalman_fitter {
         while (!fitter_state.m_fit_actor_state.finished() &&
                (!fitter_state.m_fit_actor_state.is_state() ||
                 fitter_state.m_fit_actor_state().is_hole())) {
+            if (!fitter_state.m_fit_actor_state.is_state()) {
+                TRACCC_ERROR_HOST_DEVICE("Not a track state");
+            }
+            if (fitter_state.m_fit_actor_state().is_hole()) {
+                TRACCC_ERROR_HOST_DEVICE("HOLE");
+            }
             fitter_state.m_fit_actor_state.next();
         }
 
@@ -367,6 +374,7 @@ class kalman_fitter {
         while (propagation._navigation.has_next_external() &&
                propagation._navigation.next_external().barcode() !=
                    last.smoothed_params().surface_link()) {
+            TRACCC_INFO_HOST_DEVICE("Barcode does not match!");
             TRACCC_DEBUG_HOST(
                 "Advancing to next external surface from: "
                 << propagation._navigation.next_external().barcode());
@@ -375,6 +383,7 @@ class kalman_fitter {
 
         // No valid states produced by forward pass
         if (propagation._navigation.finished()) {
+            TRACCC_ERROR_HOST_DEVICE("No matching barcodes!");
             return kalman_fitter_status::ERROR_UPDATER_SKIPPED_STATE;
         }
 
@@ -390,6 +399,7 @@ class kalman_fitter {
                            m_cfg.covariance_inflation_factor);
 
         // Run the smoothing
+        TRACCC_INFO_HOST_DEVICE("Propagating");
         propagator.propagate(propagation, fitter_state.backward_actor_state());
 
         // Reset the backward mode to false
@@ -398,10 +408,12 @@ class kalman_fitter {
         // Encountered error in the smoother?
         if (fitter_state.m_fit_actor_state.fit_result !=
             kalman_fitter_status::SUCCESS) {
+            TRACCC_ERROR_HOST_DEVICE("Fit failed!");
             return fitter_state.m_fit_actor_state.fit_result;
         }
         // Encountered error during propagation?
         if (!propagator.finished(propagation)) {
+            TRACCC_ERROR_HOST_DEVICE("-> Propagation failed!");
             return kalman_fitter_status::ERROR_PROPAGATION_FAILURE;
         }
 
@@ -519,6 +531,7 @@ class kalman_fitter {
                  fit_res.constituent_links()) {
                 assert(link.type == edm::track_constituent_link::track_state);
                 auto trk_state = track_states.at(link.index);
+
                 // Fitting fails if any of non-hole track states is not smoothed
                 if (!trk_state.is_hole() && !trk_state.is_smoothed()) {
                     TRACCC_ERROR_HOST_DEVICE("Not all smoothed");
@@ -529,7 +542,6 @@ class kalman_fitter {
             }
 
             // Fitting succeeds if any of non-hole track states is not smoothed
-            TRACCC_DEBUG_HOST_DEVICE("Fit status: SUCCESS");
             fit_res.fit_outcome() = track_fit_outcome::SUCCESS;
             return;
         }

device/alpaka/src/finding/combinatorial_kalman_filter.hpp

@@ -150,7 +150,7 @@ struct propagate_to_next_surface {
 struct build_tracks {
     template <typename TAcc>
     ALPAKA_FN_ACC void operator()(
-        TAcc const& acc, bool run_mbf,
+        TAcc const& acc, bool run_mbf, const bool run_kalman_smoother,
         const device::build_tracks_payload payload) const {
 
         device::global_index_t globalThreadIdx =

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

@@ -64,7 +64,7 @@ struct build_tracks_payload {
 /// @param[inout] payload      The function call payload
 ///
 TRACCC_HOST_DEVICE inline void build_tracks(
-    global_index_t globalIndex, bool run_mbf,
+    global_index_t globalIndex, bool run_mbf, const bool run_kalman_smoother,
     const build_tracks_payload& payload);
 
 }  // namespace traccc::device

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

@@ -20,7 +20,7 @@ namespace traccc::device {
 
 TRACCC_HOST_DEVICE inline void build_tracks(
     const global_index_t globalIndex, bool run_mbf,
-    const build_tracks_payload& payload) {
+    const bool run_kalman_smoother, const build_tracks_payload& payload) {
 
     const edm::measurement_collection<default_algebra>::const_device
         measurements(payload.tracks_view.measurements);
@@ -91,22 +91,33 @@ TRACCC_HOST_DEVICE inline void build_tracks(
 
         assert(L.meas_idx < n_meas);
 
-        if (run_mbf) {
-            const unsigned int track_state_index =
+        auto track_state_index{std::numeric_limits<unsigned int>::max()};
+        if (run_mbf || run_kalman_smoother) {
+            track_state_index =
                 track_states.push_back(edm::make_track_state<default_algebra>(
                     measurements, L.meas_idx));
             auto track_state = track_states.at(track_state_index);
 
+            // Get the filtered track parameters from the CKF
+            const auto& filtered_params = link_filtered_params.at(link_idx);
+            track_state.filtered_params() = filtered_params;
+            track_state.filtered_chi2() = L.chi2;
             track_state.set_hole(false);
+        }
+
+        // Run the MBF smoother
+        if (run_mbf) {
+            auto track_state = track_states.at(track_state_index);
+
             track_state.set_smoothed(true);
 
             // TODO: The fact that we store the chi2 three times is nonsense.
-            track_state.filtered_chi2() = L.chi2;
-            track_state.smoothed_chi2() = L.chi2;
-            track_state.backward_chi2() = L.chi2;
+            track_state.smoothed_chi2() =
+                std::numeric_limits<traccc::scalar>::max();
+            track_state.backward_chi2() =
+                std::numeric_limits<traccc::scalar>::max();
 
             const auto& predicted_params = link_predicted_params.at(link_idx);
-            const auto& filtered_params = link_filtered_params.at(link_idx);
 
             const auto& predicted_vec = predicted_params.vector();
             const auto& predicted_covariance = predicted_params.covariance();
@@ -174,10 +185,10 @@ TRACCC_HOST_DEVICE inline void build_tracks(
                 (predicted_covariance * big_lambda_tilde *
                  predicted_covariance);
 
-            track_state.filtered_params() = filtered_params;
-
             accumulated_jacobian = payload.jacobian_ptr[link_idx];
 
+            *it = {edm::track_constituent_link::track_state, track_state_index};
+        } else if (run_kalman_smoother) {
             *it = {edm::track_constituent_link::track_state, track_state_index};
         } else {
             *it = {edm::track_constituent_link::measurement, L.meas_idx};

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

@@ -77,28 +77,34 @@
             cfg.propagation.stepping.step_constraint);
 
     // Actor state
     // @TODO: simplify the syntax here
     // @NOTE: Post material interaction might be required here
     using actor_tuple_type =
         typename propagator_t::actor_chain_type::actor_tuple;
     // Pathlimit aborter
     typename detray::detail::tuple_element<0, actor_tuple_type>::type::state
         s0{};
-    typename detray::detail::tuple_element<1, actor_tuple_type>::type::state
-        s1{};
+    // Surface sequencer
+    typename detray::detail::tuple_element<1, actor_tuple_type>::type::state s1{
+        vecmem::device_vector<
+            typename propagator_t::detector_type::surface_type>(
+            *(payload.surfaces_view.ptr() + param_id))};
+    // Parameter transporter state
+    typename detray::detail::tuple_element<2, actor_tuple_type>::type::state
+        s2{};
     // CKF-interactor
-    typename detray::detail::tuple_element<3, actor_tuple_type>::type::state
-        s3{};
+    typename detray::detail::tuple_element<4, actor_tuple_type>::type::state
+        s4{};
     // Interaction register
-    typename detray::detail::tuple_element<2, actor_tuple_type>::type::state s2{
-        s3};
+    typename detray::detail::tuple_element<3, actor_tuple_type>::type::state s3{
+        s4};
     // Parameter resetter
-    typename detray::detail::tuple_element<4, actor_tuple_type>::type::state s4{
+    typename detray::detail::tuple_element<5, actor_tuple_type>::type::state s5{
         prop_cfg};
     // Momentum aborter
-    typename detray::detail::tuple_element<5, actor_tuple_type>::type::state s5;
-    // CKF aborter
     typename detray::detail::tuple_element<6, actor_tuple_type>::type::state s6;
+    // CKF aborter
+    typename detray::detail::tuple_element<7, actor_tuple_type>::type::state s7;
 
     /*
      * If we are running the MBF smoother, we need to accumulate the Jacobians
@@ -111,19 +117,23 @@
 
         payload.tmp_jacobian_ptr[param_id] = matrix::identity<
             bound_matrix<typename propagator_t::detector_type::algebra_type>>();
-        s1._full_jacobian_ptr = &payload.tmp_jacobian_ptr[param_id];
+        s2._full_jacobian_ptr = &payload.tmp_jacobian_ptr[param_id];
+    }
+
+    if (cfg.run_kalman_smoother) {
     }
 
-    s5.min_pT(static_cast<scalar_t>(cfg.min_pT));
-    s5.min_p(static_cast<scalar_t>(cfg.min_p));
-    s6.min_step_length = cfg.min_step_length_for_next_surface;
-    s6.max_count = cfg.max_step_counts_for_next_surface;
+    s6.min_pT(static_cast<scalar_t>(cfg.min_pT));
+    s6.min_p(static_cast<scalar_t>(cfg.min_p));
+    s7.min_step_length = cfg.min_step_length_for_next_surface;
+    s7.max_count = cfg.max_step_counts_for_next_surface;
 
     // Propagate to the next surface
-    propagator.propagate(propagation, detray::tie(s0, s1, s2, s3, s4, s5, s6));
+    propagator.propagate(propagation,
+                         detray::tie(s0, s1, s2, s3, s4, s5, s6, s7));
 
     // If a surface found, add the parameter for the next step
-    if (s6.success) {
+    if (s7.success) {
         assert(propagation._navigation.is_on_sensitive());
         assert(!propagation._stepping.bound_params().is_invalid());