Check if the result of Kalman update is OK

core/include/traccc/finding/finding_algorithm.ipp

@@ -186,14 +186,12 @@ finding_algorithm<stepper_t, navigator_t>::operator()(
 
                 // Run the Kalman update on a copy of the track parameters
                 bound_track_parameters bound_param(in_param);
-                sf.template visit_mask<gain_matrix_updater<algebra_type>>(
-                    trk_state, bound_param);
+                const bool res =
+                    sf.template visit_mask<gain_matrix_updater<algebra_type>>(
+                        trk_state, bound_param);
 
-                // Get the chi-square
-                const auto chi2 = trk_state.filtered_chi2();
-
-                // Found a good measurement
-                if (chi2 < m_cfg.chi2_max) {
+                // The chi2 from Kalman update should be less than chi2_max
+                if (res && trk_state.filtered_chi2() < m_cfg.chi2_max) {
                     n_branches++;
 
                     links[step].push_back({{previous_step, in_param_id},

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

@@ -36,7 +36,7 @@ struct gain_matrix_updater {
     ///
     /// @return true if the update succeeds
     template <typename mask_group_t, typename index_t>
-    TRACCC_HOST_DEVICE inline void operator()(
+    TRACCC_HOST_DEVICE inline bool operator()(
         const mask_group_t& /*mask_group*/, const index_t& /*index*/,
         track_state<algebra_t>& trk_state,
         bound_track_parameters& bound_params) const {
@@ -46,14 +46,16 @@ struct gain_matrix_updater {
         const auto D = trk_state.get_measurement().meas_dim;
         assert(D == 1u || D == 2u);
         if (D == 1u) {
-            update<1u, shape_type>(trk_state, bound_params);
+            return update<1u, shape_type>(trk_state, bound_params);
         } else if (D == 2u) {
-            update<2u, shape_type>(trk_state, bound_params);
+            return update<2u, shape_type>(trk_state, bound_params);
         }
+
+        return false;
     }
 
     template <size_type D, typename shape_t>
-    TRACCC_HOST_DEVICE inline void update(
+    TRACCC_HOST_DEVICE inline bool update(
         track_state<algebra_t>& trk_state,
         bound_track_parameters& bound_params) const {
 
@@ -124,6 +126,14 @@ struct gain_matrix_updater {
         // Set the stepper parameter
         bound_params.set_vector(filtered_vec);
         bound_params.set_covariance(filtered_cov);
+
+        // Return false if track is parallel to z-axis or phi is not finite
+        const scalar theta = bound_params.theta();
+        if (theta <= 0.f || theta >= constant<traccc::scalar>::pi ||
+            !std::isfinite(bound_params.phi())) {
+            return false;
+        }
+
         // Wrap the phi in the range of [-pi, pi]
         wrap_phi(bound_params);
 
@@ -132,7 +142,7 @@ struct gain_matrix_updater {
         trk_state.filtered().set_covariance(filtered_cov);
         trk_state.filtered_chi2() = matrix_operator().element(chi2, 0, 0);
 
-        return;
+        return true;
     }
 };
 

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

@@ -101,14 +101,21 @@ struct kalman_actor : detray::actor {
             // This track state is not a hole
             trk_state.is_hole = false;
 
-            // Set full jacobian
-            trk_state.jacobian() = stepping._full_jacobian;
-
             // Run Kalman Gain Updater
             const auto sf = navigation.get_surface();
 
-            sf.template visit_mask<gain_matrix_updater<algebra_t>>(
-                trk_state, propagation._stepping._bound_params);
+            const bool res =
+                sf.template visit_mask<gain_matrix_updater<algebra_t>>(
+                    trk_state, propagation._stepping._bound_params);
+
+            // Abort if the Kalman update fails
+            if (!res) {
+                propagation._heartbeat &= navigation.abort();
+                return;
+            }
+
+            // Set full jacobian
+            trk_state.jacobian() = stepping._full_jacobian;
 
             // Change the charge of hypothesized particles when the sign of qop
             // is changed (This rarely happens when qop is set with a poor seed

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

@@ -193,13 +193,12 @@ TRACCC_DEVICE inline void find_tracks(
             const detray::tracking_surface sf{det, in_par.surface_link()};
 
             // Run the Kalman update
-            sf.template visit_mask<
+            const bool res = sf.template visit_mask<
                 gain_matrix_updater<typename detector_t::algebra_type>>(
                 trk_state, in_par);
-            // Get the chi-square
-            const auto chi2 = trk_state.filtered_chi2();
 
-            if (chi2 < cfg.chi2_max) {
+            // The chi2 from Kalman update should be less than chi2_max
+            if (res && trk_state.filtered_chi2() < cfg.chi2_max) {
                 // Add measurement candidates to link
                 const unsigned int l_pos = num_total_candidates.fetch_add(1);