feat: Skip the Kalman fitter on the first CKF step

core/CMakeLists.txt

@@ -92,6 +92,7 @@ traccc_add_library( traccc_core core TYPE SHARED
   "include/traccc/fitting/kalman_filter/kalman_actor.hpp"
   "include/traccc/fitting/kalman_filter/kalman_fitter.hpp"
   "include/traccc/fitting/kalman_filter/kalman_step_aborter.hpp"
+  "include/traccc/fitting/kalman_filter/measurement_selector.hpp"
   "include/traccc/fitting/kalman_filter/statistics_updater.hpp"
   "include/traccc/fitting/kalman_filter/two_filters_smoother.hpp"
   "include/traccc/fitting/details/kalman_fitting_types.hpp"

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

@@ -17,6 +17,7 @@
 #include "traccc/finding/finding_config.hpp"
 #include "traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
 #include "traccc/fitting/kalman_filter/is_line_visitor.hpp"
+#include "traccc/fitting/kalman_filter/measurement_selector.hpp"
 #include "traccc/fitting/status_codes.hpp"
 #include "traccc/sanity/contiguous_on.hpp"
 #include "traccc/utils/logging.hpp"
@@ -248,6 +249,9 @@
                                    bound_track_parameters<algebra_type>>>
                 best_links;
 
+            const bool is_line = detail::is_line(sf);
+            measurement_selector::config calib_cfg{};
+
             // Iterate over the measurements
             TRACCC_VERBOSE_HOST("No. measurements: " << (up - lo));
             for (unsigned int meas_id = lo; meas_id < up; meas_id++) {
@@ -256,24 +260,53 @@
                 // The measurement on surface to handle.
                 const auto meas = measurements.at(meas_id);
 
+                const scalar_type chi2 = measurement_selector::predicted_chi2(
+                    meas, in_param, calib_cfg, is_line);
+
+                // If the measurement is outside the chi2 cut, skip it
+                if (chi2 > config.chi2_max) {
+                    continue;
+                }
+
                 // Create a standalone track state object.
                 auto trk_state =
                     edm::make_track_state<algebra_type>(measurements, meas_id);
 
-                const bool is_line = detail::is_line(sf);
+                // Kalman filter status code
+                kalman_fitter_status res{kalman_fitter_status::ERROR_OTHER};
+
+                // Don't run the filter on the first measurement
+                if (step == 0 && !sf.has_material()) {
+                    res = kalman_fitter_status::SUCCESS;
 
-                // Run the Kalman update on a copy of the track parameters
-                const kalman_fitter_status res =
-                    gain_matrix_updater<algebra_type>{}(trk_state, measurements,
-                                                        in_param, is_line);
+                    // Copy the full track parameters
+                    // TODO: Apply calibration ?
+                    trk_state.filtered_params() = in_param;
+
+                    // Update measurement covariance
+                    const auto V =
+                        measurement_selector::calibrated_measurement_covariance<
+                            algebra_type, 2>(meas, calib_cfg);
+
+                    auto& filtered_cov =
+                        trk_state.filtered_params().covariance();
+                    getter::element(filtered_cov, e_bound_loc0, e_bound_loc0) =
+                        getter::element(V, 0, 0);
+                    getter::element(filtered_cov, e_bound_loc1, e_bound_loc1) =
+                        getter::element(V, 1, 1);
+                } else {
+                    // Run the Kalman update on the track state
+                    constexpr gain_matrix_updater<algebra_type>
+                        kalman_updater{};
+                    res = kalman_updater(trk_state, meas, in_param, is_line);
+                }
 
-                const traccc::scalar chi2 = trk_state.filtered_chi2();
+                trk_state.filtered_chi2() = chi2;
 
                 TRACCC_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)) {
+                if (res == kalman_fitter_status::SUCCESS) {
 
                     TRACCC_VERBOSE_HOST("Found measurement: " << meas_id);
 

core/include/traccc/fitting/fitting_config.hpp

@@ -29,6 +29,10 @@ struct fitting_config {
     float min_p = 100.f * traccc::unit<float>::MeV;
     float min_pT = 600.f * traccc::unit<float>::MeV;
 
+    // Maximal local radius a measurement is allowed to be away from the
+    // predicted track position in order to be considered as a condidate
+    float max_measurement_radius{5.f * traccc::unit<float>::mm};
+
     /// Particle hypothesis
     traccc::pdg_particle<traccc::scalar> ptc_hypothesis =
         traccc::muon<traccc::scalar>();

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

@@ -11,12 +11,10 @@
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/definitions/track_parametrization.hpp"
 #include "traccc/edm/measurement_collection.hpp"
-#include "traccc/edm/measurement_helpers.hpp"
-#include "traccc/edm/track_state_collection.hpp"
 #include "traccc/fitting/details/regularize_covariance.hpp"
+#include "traccc/fitting/kalman_filter/measurement_selector.hpp"
 #include "traccc/fitting/status_codes.hpp"
 #include "traccc/utils/logging.hpp"
-#include "traccc/utils/subspace.hpp"
 
 namespace traccc {
 
@@ -42,18 +40,16 @@ struct gain_matrix_updater {
     /// @param bound_params bound parameter
     ///
     /// @return true if the update succeeds
-    template <typename track_state_backend_t>
+    template <typename track_state_backend_t, typename measurement_backend_t>
     [[nodiscard]] TRACCC_HOST_DEVICE inline kalman_fitter_status operator()(
         typename edm::track_state<track_state_backend_t>& trk_state,
-        const edm::measurement_collection<default_algebra>::const_device&
-            measurements,
+        const edm::measurement<measurement_backend_t>& measurement,
         const bound_track_parameters<algebra_t>& bound_params,
         const bool is_line) const {
 
         static constexpr unsigned int D = 2;
 
-        [[maybe_unused]] const unsigned int dim{
-            measurements.at(trk_state.measurement_index()).dimensions()};
+        [[maybe_unused]] const unsigned int dim{measurement.dimensions()};
 
         TRACCC_VERBOSE_HOST_DEVICE("In gain-matrix-updater...");
         TRACCC_VERBOSE_HOST_DEVICE("Measurement dim: %d", dim);
@@ -70,8 +66,7 @@ struct gain_matrix_updater {
 
         // Measurement data on surface
         matrix_type<D, 1> meas_local;
-        edm::get_measurement_local<algebra_t>(
-            measurements.at(trk_state.measurement_index()), meas_local);
+        edm::get_measurement_local<algebra_t>(measurement, meas_local);
 
         assert((dim > 1) || (getter::element(meas_local, 1u, 0u) == 0.f));
 
@@ -83,8 +78,7 @@ struct gain_matrix_updater {
         // Predicted covaraince of bound track parameters
         const bound_matrix_type& predicted_cov = bound_params.covariance();
 
-        const subspace<algebra_t, e_bound_size> subs(
-            measurements.at(trk_state.measurement_index()).subspace());
+        const subspace<algebra_t, e_bound_size> subs(measurement.subspace());
         matrix_type<D, e_bound_size> H = subs.template projector<D>();
 
         // Flip the sign of projector matrix element in case the first element
@@ -101,8 +95,7 @@ struct gain_matrix_updater {
 
         // Spatial resolution (Measurement covariance)
         matrix_type<D, D> V;
-        edm::get_measurement_covariance<algebra_t>(
-            measurements.at(trk_state.measurement_index()), V);
+        edm::get_measurement_covariance<algebra_t>(measurement, V);
         // @TODO: Fix properly
         if (/*dim == 1*/ getter::element(meas_local, 1u, 0u) == 0.f) {
             getter::element(V, 1u, 1u) = 1000.f;

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

@@ -406,6 +406,8 @@ struct kalman_actor : detray::actor {
             const auto sf = navigation.current_surface();
             const bool is_line = detail::is_line(sf);
 
+            const auto measurement =
+                actor_state.m_measurements.at(trk_state.measurement_index());
             if (!actor_state.backward_mode) {
                 if constexpr (direction_e ==
                                   kalman_actor_direction::FORWARD_ONLY ||
@@ -417,8 +419,7 @@ struct kalman_actor : detray::actor {
                     // Forward filter
                     TRACCC_DEBUG_HOST_DEVICE("Run filtering...");
                     actor_state.fit_result = gain_matrix_updater<algebra_t>{}(
-                        trk_state, actor_state.m_measurements, bound_param,
-                        is_line);
+                        trk_state, measurement, bound_param, is_line);
 
                     // Update the propagation flow
                     bound_param = trk_state.filtered_params();
@@ -447,8 +448,7 @@ struct kalman_actor : detray::actor {
                     } else {
                         actor_state.fit_result =
                             two_filters_smoother<algebra_t>{}(
-                                trk_state, actor_state.m_measurements,
-                                bound_param, is_line);
+                                trk_state, measurement, bound_param, is_line);
                     }
                 } else {
                     assert(false);