Fix overlaps handling in the CKF

Author: niermann999

core/CMakeLists.txt

@@ -80,6 +80,7 @@ traccc_add_library( traccc_core core TYPE SHARED
   # Finding algorithmic code
   "include/traccc/finding/candidate_link.hpp"
   "include/traccc/finding/finding_config.hpp"
+  "include/traccc/finding/propagation_data.hpp"
   "include/traccc/finding/actors/ckf_aborter.hpp"
   "include/traccc/finding/actors/interaction_register.hpp"
   "include/traccc/finding/details/combinatorial_kalman_filter_types.hpp"

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

@@ -10,8 +10,10 @@
 // Project include(s)
 #include "traccc/definitions/primitives.hpp"
 #include "traccc/definitions/qualifiers.hpp"
+#include "traccc/utils/logging.hpp"
 
 // detray include(s)
+#include <detray/definitions/indexing.hpp>
 #include <detray/propagator/base_actor.hpp>
 
 // System include(s)
@@ -24,36 +26,45 @@ struct ckf_aborter : detray::actor {
     struct state {
         // minimal step length to prevent from staying on the same surface
         scalar min_step_length = 0.5f;
+
         /// Maximum step counts that track can make to reach the next surface
-        unsigned int max_count = 100;
+        unsigned int max_count = 100u;
+        unsigned int count = 0u;
 
-        bool success = false;
-        unsigned int count = 0;
+        /// Previous sensitive surface index
+        detray::dindex prev_surface_index{detray::dindex_invalid};
+        /// Current sensitive surface index
+        detray::dindex surface_index{detray::dindex_invalid};
 
-        scalar path_from_surface{0.f};
+        /// Whether a surface was found
+        bool success = false;
     };
 
     template <typename propagator_state_t>
     TRACCC_HOST_DEVICE void operator()(state &abrt_state,
                                        propagator_state_t &prop_state) const {
 
         auto &navigation = prop_state._navigation;
-        const auto &stepping = prop_state._stepping;
 
         abrt_state.count++;
-        abrt_state.path_from_surface += stepping.step_size();
 
-        // Stop at the next sensitive surface
-        if (navigation.is_on_sensitive() &&
-            abrt_state.path_from_surface > abrt_state.min_step_length) {
-            prop_state._heartbeat &= navigation.pause();
-            abrt_state.success = navigation.is_alive();
+        // Is this a valid sensitive surface to run the CKF on?
+        const detray::dindex sf_idx{navigation.barcode().index()};
+        if (!navigation.is_on_sensitive() ||
+            (abrt_state.surface_index == sf_idx) ||
+            (abrt_state.prev_surface_index == sf_idx)) {
+            return;
         }
 
-        // Reset path from surface
-        if (navigation.is_on_sensitive()) {
-            abrt_state.path_from_surface = 0.f;
-        }
+        // Update the visited sensitive surfaces and pause the propagation
+        abrt_state.prev_surface_index = abrt_state.surface_index;
+        abrt_state.surface_index = sf_idx;
+
+        prop_state._heartbeat &= navigation.pause();
+        abrt_state.success = true;
+
+        assert(abrt_state.surface_index <
+               navigation.detector().surfaces().size());
 
         if (abrt_state.count > abrt_state.max_count) {
             prop_state._heartbeat &= navigation.abort(

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

@@ -7,8 +7,6 @@
 #pragma once
 
 // Project include(s).
-#include <detray/utils/log.hpp>
-
 #include "traccc/edm/measurement.hpp"
 #include "traccc/edm/track_candidate_collection.hpp"
 #include "traccc/edm/track_state_helpers.hpp"
@@ -17,6 +15,7 @@
 #include "traccc/finding/candidate_link.hpp"
 #include "traccc/finding/details/combinatorial_kalman_filter_types.hpp"
 #include "traccc/finding/finding_config.hpp"
+#include "traccc/finding/propagation_data.hpp"
 #include "traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
 #include "traccc/fitting/kalman_filter/is_line_visitor.hpp"
 #include "traccc/fitting/status_codes.hpp"
@@ -81,6 +80,10 @@ combinatorial_kalman_filter(
     // Create a logger.
     auto logger = [&log]() -> const Logger& { return log; };
 
+    // Minimum chi2 cut to use for measurements on surfaces that are edge hits
+    // (don't accumulate pileup measurements)
+    const float chi2_min{math::max(10.f, 0.1f * config.chi2_max)};
+
     /*****************************************************************
      * Measurement Operations
      *****************************************************************/
@@ -144,20 +147,18 @@ combinatorial_kalman_filter(
 
     std::vector<bound_track_parameters<algebra_type>> out_params;
 
-    std::vector<std::uint8_t> in_is_edge(seeds.size(), false);
-    std::vector<std::uint8_t> out_is_edge;
+    std::vector<propagation_data> in_prop_data(seeds.size());
+    std::vector<propagation_data> out_prop_data;
 
     for (unsigned int step = 0u; step < config.max_track_candidates_per_track;
          step++) {
 
-        DETRAY_DEBUG_HOST("PRPGAGATION STEP: " << step);
         TRACCC_VERBOSE("Starting step "
                        << step + 1 << " / "
                        << config.max_track_candidates_per_track);
 
         // Iterate over input parameters
         const std::size_t n_in_params = in_params.size();
-        DETRAY_DEBUG_HOST("# PARAMS: " << n_in_params);
 
         // Terminate if there is no parameter to proceed
         if (n_in_params == 0) {
@@ -166,23 +167,22 @@ combinatorial_kalman_filter(
 
         // Rough estimation on out parameters size
         out_params.reserve(n_in_params);
-        out_is_edge.reserve(n_in_params);
+        out_prop_data.reserve(n_in_params);
 
         // Previous step ID
         std::fill(n_trks_per_seed.begin(), n_trks_per_seed.end(), 0u);
 
         // Parameters updated by Kalman fitter
         std::vector<bound_track_parameters<algebra_type>> updated_params;
+        std::vector<propagation_data> updated_prop_data;
 
         for (unsigned int in_param_id = 0; in_param_id < n_in_params;
              in_param_id++) {
 
             bound_track_parameters<algebra_type>& in_param =
                 in_params[in_param_id];
 
-            DETRAY_DEBUG_HOST("PARAMS: " << in_param);
-
-            const bool is_edge{in_is_edge[in_param_id] > 0u};
+            const propagation_data& prop_data = in_prop_data[in_param_id];
 
             assert(!in_param.is_invalid());
 
@@ -259,9 +259,7 @@ combinatorial_kalman_filter(
                 best_links;
 
             // Iterate over the measurements
-            DETRAY_DEBUG_HOST("# MEAS: " << (up - lo));
             for (unsigned int meas_id = lo; meas_id < up; meas_id++) {
-                DETRAY_DEBUG_HOST("TESTING MEAS: " << meas_id);
 
                 // The measurement on surface to handle.
                 const measurement& meas = measurements.at(meas_id);
@@ -279,13 +277,10 @@ combinatorial_kalman_filter(
 
                 const traccc::scalar chi2 = trk_state.filtered_chi2();
 
-                DETRAY_DEBUG_HOST("KF STATUS: " << fitter_debug_msg{res}());
-
                 // The chi2 from Kalman update should be less than chi2_max
                 if (res == kalman_fitter_status::SUCCESS &&
-                    chi2 < config.chi2_max) {
-
-                    DETRAY_DEBUG_HOST("FOUND MEAS: " << meas_id);
+                    ((!prop_data.is_edge && chi2 <= config.chi2_max) ||
+                     (chi2 <= chi2_min))) {
 
                     best_links.push_back(
                         {{.step = step,
@@ -322,6 +317,7 @@ combinatorial_kalman_filter(
 
                 // Add the updated parameter to the updated parameters
                 updated_params.push_back(filtered_params);
+                updated_prop_data.push_back(prop_data);
                 TRACCC_VERBOSE("updated_params["
                                << updated_params.size() - 1
                                << "] = " << updated_params.back());
@@ -340,15 +336,15 @@ combinatorial_kalman_filter(
                      .meas_idx = std::numeric_limits<unsigned int>::max(),
                      .seed_idx = orig_param_id,
                      .n_cand = n_cand,
-                     .n_skipped = is_edge ? skip_counter : skip_counter + 1,
+                     .n_skipped =
+                         prop_data.is_edge ? skip_counter : skip_counter + 1,
                      .n_consecutive_skipped = consecutive_skipped + 1,
                      .chi2 = std::numeric_limits<traccc::scalar>::max(),
                      .chi2_sum = prev_chi2_sum,
                      .ndf_sum = prev_ndf_sum});
 
-                DETRAY_DEBUG_HOST("HOLE: " << std::boolalpha << !is_edge);
-
                 updated_params.push_back(in_param);
+                updated_prop_data.push_back(prop_data);
                 TRACCC_VERBOSE("updated_params["
                                << updated_params.size() - 1
                                << "] = " << updated_params.back());
@@ -468,7 +464,6 @@ combinatorial_kalman_filter(
                         }
 
                         if (this_is_dominated) {
-                            DETRAY_DEBUG_HOST("Track is DEAD!");
                             param_liveness.at(tracks.at(i)) = 0u;
                             break;
                         }
@@ -497,12 +492,14 @@ combinatorial_kalman_filter(
             if (links.at(step).at(link_id).n_skipped >
                 config.max_num_skipping_per_cand) {
 
-                DETRAY_DEBUG_HOST("MAKE TIP: MAX NO. HOLES");
                 tips.push_back({step, link_id});
                 continue;
             }
 
+            assert(updated_params.size() == updated_prop_data.size());
             const auto& param = updated_params[link_id];
+            const propagation_data& prop_data = updated_prop_data[link_id];
+
             // Create propagator state
             typename traccc::details::ckf_propagator_t<
                 detector_t, bfield_t>::state propagation(param, field, det);
@@ -524,50 +521,46 @@ combinatorial_kalman_filter(
             // Update the actor config
             s4.min_pT(static_cast<scalar_type>(config.min_pT));
             s4.min_p(static_cast<scalar_type>(config.min_p));
+            s5.prev_surface_index = prop_data.prev_surface;
+            s5.surface_index = param.surface_link().index();
             s5.min_step_length = config.min_step_length_for_next_surface;
             s5.max_count = config.max_step_counts_for_next_surface;
 
-            DETRAY_DEBUG_HOST("PROPAGATING... ");
             // Propagate to the next surface
             propagator.propagate(propagation,
                                  detray::tie(s0, s1, s2, s3, s4, s5));
 
-            DETRAY_DEBUG_HOST(
-                "FINISHED PROPAGATING. PATH:  " << s5.path_from_surface);
-
             // If a surface found, add the parameter for the next
             // step
             if (s5.success) {
                 assert(propagation._navigation.is_on_sensitive());
                 assert(!propagation._stepping.bound_params().is_invalid());
 
                 out_params.push_back(propagation._stepping.bound_params());
-                out_is_edge.push_back(
+                out_prop_data.emplace_back(
+                    s5.prev_surface_index,
                     propagation._navigation.is_edge_candidate());
                 param_to_link[step].push_back(link_id);
             }
             // Unless the track found a surface, it is considered a
             // tip
             else if (!s5.success &&
                      (step >= (config.min_track_candidates_per_track - 1u))) {
-
-                DETRAY_DEBUG_HOST("MAKE TIP: NO SENSITIVE");
                 tips.push_back({step, link_id});
             }
 
             // If no more CKF step is expected, current candidate is
             // kept as a tip
             if (s5.success &&
                 (step == (config.max_track_candidates_per_track - 1u))) {
-                DETRAY_DEBUG_HOST("MAKE TIP: MAX CANDIDATES");
                 tips.push_back({step, link_id});
             }
         }
 
         in_params = std::move(out_params);
-        in_is_edge = std::move(out_is_edge);
+        in_prop_data = std::move(out_prop_data);
         out_params.clear();
-        out_is_edge.clear();
+        out_prop_data.clear();
     }
 
     /**********************

core/include/traccc/finding/propagation_data.hpp

@@ -0,0 +1,26 @@
+
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2025 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+// Detray include(s).
+#include <detray/definitions/indexing.hpp>
+
+namespace traccc {
+
+/// Data from the propagation loop that has to be kept between CKF steps
+struct propagation_data {
+    /// The surface that was visited before the current one (overlaps
+    /// resolution)
+    detray::dindex prev_surface{detray::dindex_invalid};
+
+    /// Is the current surface hit in the extended tolerance band?
+    bool is_edge{false};
+};
+
+}  // namespace traccc

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

@@ -212,9 +212,9 @@ struct kalman_actor : detray::actor {
             actor_state.next();
 
             // Flag renavigation of the current candidate
-            // if (math::fabs(navigation()) > 1.f * unit<float>::um) {
-            navigation.set_high_trust();
-            //}
+            if (math::fabs(navigation()) > 1.f * unit<float>::um) {
+                navigation.set_high_trust();
+            }
         }
     }
 };

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

@@ -185,7 +185,7 @@ struct two_filters_smoother {
 
         // Calculate backward chi2
         const matrix_type<D, D> R = (I_m - H * K) * V;
-        assert(matrix::determinant(R) != 0.f);
+        // assert(matrix::determinant(R) != 0.f);
         assert(std::isfinite(matrix::determinant(R)));
         const matrix_type<1, 1> chi2 =
             algebra::matrix::transposed_product<true, false>(

device/alpaka/src/finding/combinatorial_kalman_filter.hpp

@@ -237,9 +237,9 @@ combinatorial_kalman_filter(
     copy.memset(param_liveness_buffer, 1)->wait();
 
     // On the first measurement, no parameter can be a hole
-    vecmem::data::vector_buffer<std::uint8_t> is_edges_buffer(n_seeds, mr.main);
-    copy.setup(is_edges_buffer)->wait();
-    copy.memset(is_edges_buffer, 0)->wait();
+    vecmem::data::vector_buffer<propagation_data> prop_data_buffer(n_seeds,
+                                                                   mr.main);
+    copy.setup(prop_data_buffer)->wait();
 
     // Number of tracks per seed
     vecmem::data::vector_buffer<unsigned int> n_tracks_per_seed_buffer(n_seeds,
@@ -306,9 +306,9 @@ combinatorial_kalman_filter(
         copy.setup(updated_liveness_buffer)->wait();
 
         // Updated edges buffer after branching
-        vecmem::data::vector_buffer<std::uint8_t> is_edges_updated_buffer(
+        vecmem::data::vector_buffer<propagation_data> prop_data_updated_buffer(
             n_max_candidates, mr.main);
-        copy.setup(is_edges_updated_buffer)->wait();
+        copy.setup(prop_data_updated_buffer)->wait();
 
         // Reset the number of tracks per seed
         copy.memset(n_tracks_per_seed_buffer, 0)->wait();
@@ -353,7 +353,7 @@ combinatorial_kalman_filter(
                 .in_params_view = in_params_buffer,
                 .in_params_liveness_view = param_liveness_buffer,
                 .n_in_params = n_in_params,
-                .is_edges_view = is_edges_buffer,
+                .prop_data_view = prop_data_buffer,
                 .measurement_ranges_view = meas_ranges_buffer,
                 .links_view = links_buffer,
                 .prev_links_idx =
@@ -517,7 +517,7 @@ combinatorial_kalman_filter(
                     .params_view = in_params_buffer,
                     .params_liveness_view = param_liveness_buffer,
                     .param_ids_view = param_ids_buffer,
-                    .is_edges_view = is_edges_updated_buffer,
+                    .prop_data_view = prop_data_updated_buffer,
                     .links_view = links_buffer,
                     .prev_links_idx = step_to_link_idx_map[step],
                     .step = step,
@@ -549,7 +549,7 @@ combinatorial_kalman_filter(
                     config, device_payload.ptr());
                 ::alpaka::wait(queue);
 
-                std::swap(is_edges_buffer, is_edges_updated_buffer);
+                std::swap(prop_data_buffer, prop_data_updated_buffer);
             }
         }