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 {
 
@@ -121,22 +123,26 @@ struct gain_matrix_updater {
         const matrix_type<1, 1> chi2 = matrix_operator().transpose(residual) *
                                        matrix_operator().inverse(R) * residual;
 
-        // Make sure that the sign of qop does not change (This rarely happens
-        // when qop is set with a poor seed resolution)
-        assert(bound_params[e_bound_qoverp] *
-                   getter::element(filtered_vec, e_bound_qoverp, 0u) >
-               0.f);
-
         // 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);
+
         // Set the track state parameters
         trk_state.filtered().set_vector(filtered_vec);
         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

@@ -11,6 +11,7 @@
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/kalman_filter/gain_matrix_updater.hpp"
+#include "traccc/utils/particle.hpp"
 
 // detray include(s).
 #include "detray/propagator/base_actor.hpp"
@@ -100,13 +101,27 @@ struct kalman_actor : detray::actor {
             // This track state is not a hole
             trk_state.is_hole = false;
 
+            // Run Kalman Gain Updater
+            const auto sf = navigation.get_surface();
+
+            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;
 
-            // 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);
+            // Change the charge of hypothesized particles when the sign of qop
+            // is changed (This rarely happens when qop is set with a poor seed
+            // resolution)
+            propagation.set_particle(detail::correct_particle_hypothesis(
+                stepping._ptc, propagation._stepping._bound_params));
 
             // Update iterator
             actor_state.next();

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

@@ -188,13 +188,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);