Fix bugs in the Kalman fitter

Author: niermann999

core/include/traccc/edm/track_parameters.hpp

@@ -14,6 +14,7 @@
 #include "traccc/definitions/qualifiers.hpp"
 #include "traccc/definitions/track_parametrization.hpp"
 #include "traccc/edm/container.hpp"
+#include "traccc/utils/trigonometric_helpers.hpp"
 
 // detray include(s).
 #include <detray/tracks/tracks.hpp>
@@ -48,19 +49,15 @@ using bound_track_parameters_collection_types =
 
 // Wrap the phi of track parameters to [-pi,pi]
 template <detray::concepts::algebra algebra_t>
-TRACCC_HOST_DEVICE inline void wrap_phi(
+TRACCC_HOST_DEVICE constexpr void normalize_angles(
     bound_track_parameters<algebra_t>& param) {
+    traccc::scalar phi;
+    traccc::scalar theta;
+
+    std::tie(phi, theta) = detail::wrap_phi_theta(param.phi(), param.theta());
 
-    traccc::scalar phi = param.phi();
-    static constexpr traccc::scalar TWOPI =
-        2.f * traccc::constant<traccc::scalar>::pi;
-    phi = math::fmod(phi, TWOPI);
-    if (phi > traccc::constant<traccc::scalar>::pi) {
-        phi -= TWOPI;
-    } else if (phi < -traccc::constant<traccc::scalar>::pi) {
-        phi += TWOPI;
-    }
     param.set_phi(phi);
+    param.set_theta(theta);
 }
 
 /// Covariance inflation used for track fitting

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

@@ -143,7 +143,7 @@ combinatorial_kalman_filter(
 
     for (unsigned int step = 0u; step < config.max_track_candidates_per_track;
          step++) {
-            
+
         TRACCC_VERBOSE("Starting step "
                        << step + 1 << " / "
                        << config.max_track_candidates_per_track);
@@ -242,7 +242,6 @@ combinatorial_kalman_filter(
 
             // Iterate over the measurements
             for (unsigned int meas_id = lo; meas_id < up; meas_id++) {
-
                 // The measurement on surface to handle.
                 const measurement& meas = measurements.at(meas_id);
 
@@ -502,17 +501,36 @@ combinatorial_kalman_filter(
 
             // If a surface found, add the parameter for the next
             // step
-            if (s5.success) {
+            bool valid_track{s5.success};
+            if (valid_track) {
                 assert(propagation._navigation.is_on_sensitive());
                 assert(!propagation._stepping.bound_params().is_invalid());
 
-                out_params.push_back(propagation._stepping.bound_params());
-                param_to_link[step].push_back(link_id);
+                const auto& out_param = propagation._stepping.bound_params();
+
+                const scalar theta = out_param.theta();
+                if (theta <= 0.f ||
+                    theta >= 2.f * constant<traccc::scalar>::pi) {
+                    valid_track = false;
+                }
+
+                if (!std::isfinite(out_param.phi())) {
+                    valid_track = false;
+                }
+
+                if (math::fabs(out_param.qop()) == 0.f) {
+                    valid_track = false;
+                }
+
+                if (valid_track) {
+                    out_params.push_back(out_param);
+                    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))) {
+            if (!valid_track &&
+                (step >= (config.min_track_candidates_per_track - 1u))) {
                 tips.push_back({step, link_id});
             }
 

core/include/traccc/fitting/fitting_config.hpp

@@ -34,7 +34,7 @@ struct fitting_config {
     std::size_t barcode_sequence_size_factor = 5;
     std::size_t min_barcode_sequence_capacity = 100;
     traccc::scalar backward_filter_mask_tolerance =
-        5.f * traccc::unit<scalar>::mm;
+        10.f * traccc::unit<scalar>::m;
 };
 
 }  // namespace traccc

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

@@ -123,6 +123,30 @@ struct gain_matrix_updater {
             predicted_vec + K * (meas_local - H * predicted_vec);
         const matrix_type<6, 6> filtered_cov = (I66 - K * H) * predicted_cov;
 
+        // Check the covariance for consistency
+        if (getter::element(filtered_cov, 0, 0) <= 0.f ||
+            getter::element(filtered_cov, 1, 1) <= 0.f ||
+            getter::element(filtered_cov, 2, 2) <= 0.f ||
+            getter::element(filtered_cov, 3, 3) <= 0.f ||
+            getter::element(filtered_cov, 4, 4) <= 0.f ||
+            getter::element(filtered_cov, 5, 5) <= 0.f) {
+            return kalman_fitter_status::ERROR_UPDATER_INVALID_COVARIANCE;
+        }
+
+        // Return false if track is parallel to z-axis or phi is not finite
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_theta, 0))) {
+            return kalman_fitter_status::ERROR_INVERSION;
+        }
+
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_phi, 0))) {
+            return kalman_fitter_status::ERROR_INVERSION;
+        }
+
+        if (math::fabs(getter::element(filtered_vec, e_bound_qoverp, 0)) ==
+            0.f) {
+            return kalman_fitter_status::ERROR_QOP_ZERO;
+        }
+
         // Residual between measurement and (projected) filtered vector
         const matrix_type<D, 1> residual = meas_local - H * filtered_vec;
 
@@ -133,36 +157,28 @@ struct gain_matrix_updater {
                 residual, matrix::inverse(R)) *
             residual;
 
-        // 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) {
-            return kalman_fitter_status::ERROR_THETA_ZERO;
-        }
-
-        if (!std::isfinite(bound_params.phi())) {
-            return kalman_fitter_status::ERROR_INVERSION;
-        }
+        const scalar chi2_val{getter::element(chi2, 0, 0)};
 
-        if (std::abs(bound_params.qop()) == 0.f) {
-            return kalman_fitter_status::ERROR_QOP_ZERO;
-        }
-
-        if (getter::element(chi2, 0, 0) < 0.f) {
+        if (chi2_val < 0.f) {
             return kalman_fitter_status::ERROR_UPDATER_CHI2_NEGATIVE;
         }
 
-        if (!std::isfinite(getter::element(chi2, 0, 0))) {
+        if (!std::isfinite(chi2_val)) {
             return kalman_fitter_status::ERROR_UPDATER_CHI2_NOT_FINITE;
         }
 
-        // Set the track state parameters
+        // Set the chi2 for this track and measurement
         trk_state.filtered_params().set_vector(filtered_vec);
         trk_state.filtered_params().set_covariance(filtered_cov);
-        trk_state.filtered_chi2() = getter::element(chi2, 0, 0);
+        trk_state.filtered_chi2() = chi2_val;
 
-        // Wrap the phi in the range of [-pi, pi]
-        wrap_phi(trk_state.filtered_params());
+        // Wrap the phi and theta angles in their valid ranges
+        normalize_angles(trk_state.filtered_params());
+
+        const scalar theta = trk_state.filtered_params().theta();
+        if (theta <= 0.f || theta >= 2.f * constant<traccc::scalar>::pi) {
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
 
         assert(!trk_state.filtered_params().is_invalid());
 

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

@@ -173,6 +173,9 @@ struct kalman_actor : detray::actor {
                                   kalman_actor_direction::FORWARD_ONLY ||
                               direction_e ==
                                   kalman_actor_direction::BIDIRECTIONAL) {
+                    // Wrap the phi and theta angles in their valid ranges
+                    normalize_angles(propagation._stepping.bound_params());
+
                     // Forward filter
                     res = gain_matrix_updater<algebra_t>{}(
                         trk_state, actor_state.m_measurements,

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

@@ -55,6 +55,7 @@ class kalman_fitter {
 
     // Actor types
     using aborter = detray::pathlimit_aborter<scalar_type>;
+    using momentum_aborter = detray::momentum_aborter<scalar_type>;
     using transporter = detray::parameter_transporter<algebra_type>;
     using interactor = detray::pointwise_material_interactor<algebra_type>;
     using forward_fit_actor =
@@ -71,11 +72,13 @@ class kalman_fitter {
 
     using forward_actor_chain_type =
         detray::actor_chain<aborter, transporter, interactor, forward_fit_actor,
-                            resetter, barcode_sequencer, kalman_step_aborter>;
+                            resetter, momentum_aborter, barcode_sequencer,
+                            kalman_step_aborter>;
 
     using backward_actor_chain_type =
         detray::actor_chain<aborter, transporter, backward_fit_actor,
-                            interactor, resetter, kalman_step_aborter>;
+                            interactor, resetter, momentum_aborter,
+                            kalman_step_aborter>;
 
     // Navigator type for backward propagator
     using direct_navigator_type = detray::direct_navigator<detector_type>;
@@ -125,7 +128,8 @@ class kalman_fitter {
         typename forward_actor_chain_type::state_ref_tuple operator()() {
             return detray::tie(m_aborter_state, m_interactor_state,
                                m_fit_actor_state, m_parameter_resetter,
-                               m_sequencer_state, m_step_aborter_state);
+                               m_montum_aborter_state, m_sequencer_state,
+                               m_step_aborter_state);
         }
 
         /// @return the actor chain state
@@ -134,7 +138,7 @@ class kalman_fitter {
         backward_actor_state() {
             return detray::tie(m_aborter_state, m_fit_actor_state,
                                m_interactor_state, m_parameter_resetter,
-                               m_step_aborter_state);
+                               m_montum_aborter_state, m_step_aborter_state);
         }
 
         /// Individual actor states
@@ -144,6 +148,7 @@ class kalman_fitter {
         typename barcode_sequencer::state m_sequencer_state;
         kalman_step_aborter::state m_step_aborter_state{};
         typename resetter::state m_parameter_resetter{};
+        typename momentum_aborter::state m_montum_aborter_state{};
 
         /// Fitting result per track
         typename edm::track_collection<algebra_type>::device::proxy_type
@@ -281,16 +286,20 @@ class kalman_fitter {
         }
 
         auto last = fitter_state.m_fit_actor_state();
-
         const scalar theta = last.filtered_params().theta();
-        if (theta <= 0.f || theta >= constant<traccc::scalar>::pi) {
-            return kalman_fitter_status::ERROR_THETA_ZERO;
+
+        if (!std::isfinite(last.filtered_params().theta())) {
+            return kalman_fitter_status::ERROR_INVERSION;
         }
 
         if (!std::isfinite(last.filtered_params().phi())) {
             return kalman_fitter_status::ERROR_INVERSION;
         }
 
+        if (theta <= 0.f || theta >= 2.f * constant<traccc::scalar>::pi) {
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
+
         last.smoothed_params().set_parameter_vector(last.filtered_params());
         last.smoothed_params().set_covariance(
             last.filtered_params().covariance());
@@ -328,6 +337,8 @@ class kalman_fitter {
 
         propagation._navigation.set_direction(
             detray::navigation::direction::e_backward);
+        propagation._navigation.safe_step_size =
+            0.1f * traccc::unit<scalar>::mm;
 
         // Synchronize the current barcode with the input track parameter
         while (propagation._navigation.get_target_barcode() !=