Fix bugs in the Kalman fitter

niermann999 · niermann999 · commit e05587d3ca72 · 2025-11-05T15:10:10.000+01:00
diff --git a/core/include/traccc/edm/track_parameters.hpp b/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
diff --git a/core/include/traccc/finding/details/combinatorial_kalman_filter.hpp b/core/include/traccc/finding/details/combinatorial_kalman_filter.hpp
@@ -536,11 +536,30 @@ combinatorial_kalman_filter(
                 assert(propagation._navigation.is_on_sensitive());
                 assert(!propagation._stepping.bound_params().is_invalid());
 
-                out_params.push_back(propagation._stepping.bound_params());
-                out_prop_data.emplace_back(
-                    s5.prev_surface_index,
-                    propagation._navigation.is_edge_candidate());
-                param_to_link[step].push_back(link_id);
+                const auto& out_param = propagation._stepping.bound_params();
+
+                bool valid_track{true};
+                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 (std::abs(out_param.qop()) == 0.f) {
+                    valid_track = false;
+                }
+
+                if (valid_track) {
+                    out_params.push_back(out_param);
+                    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
diff --git a/core/include/traccc/finding/propagation_data.hpp b/core/include/traccc/finding/propagation_data.hpp
@@ -23,4 +23,4 @@ struct propagation_data {
     bool is_edge{false};
 };
 
-}  // namespace traccc
+}  // namespace traccc
diff --git a/core/include/traccc/fitting/fitting_config.hpp b/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
diff --git a/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp b/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp
@@ -123,6 +123,19 @@ struct gain_matrix_updater {
             predicted_vec + K * (meas_local - H * predicted_vec);
         const matrix_type<6, 6> filtered_cov = (I66 - K * H) * predicted_cov;
 
+        // 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 (std::abs(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 +146,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();
+        const scalar chi2_val{getter::element(chi2, 0, 0)};
 
-        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;
-        }
-
-        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());
 
diff --git a/core/include/traccc/fitting/kalman_filter/kalman_actor.hpp b/core/include/traccc/fitting/kalman_filter/kalman_actor.hpp
@@ -170,6 +170,8 @@ 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,
diff --git a/core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp b/core/include/traccc/fitting/kalman_filter/kalman_fitter.hpp
@@ -280,16 +280,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());
@@ -327,6 +331,8 @@ class kalman_fitter {
 
         propagation._navigation.set_direction(
             detray::navigation::direction::e_backward);
+        propagation._navigation.safe_step_size =
+            10.f * traccc::unit<scalar>::um;
 
         // Synchronize the current barcode with the input track parameter
         while (propagation._navigation.get_target_barcode() !=
diff --git a/core/include/traccc/fitting/kalman_filter/two_filters_smoother.hpp b/core/include/traccc/fitting/kalman_filter/two_filters_smoother.hpp
@@ -107,8 +107,25 @@ struct two_filters_smoother {
              predicted_cov_inv * predicted_vec);
 
         trk_state.smoothed_params().set_vector(smoothed_vec);
+
+        // Return false if track is parallel to z-axis or phi is not finite
+        if (!std::isfinite(trk_state.smoothed_params().theta())) {
+            return kalman_fitter_status::ERROR_INVERSION;
+        }
+
+        if (!std::isfinite(trk_state.smoothed_params().phi())) {
+            return kalman_fitter_status::ERROR_INVERSION;
+        }
+
+        if (std::abs(trk_state.smoothed_params().qop()) == 0.f) {
+            return kalman_fitter_status::ERROR_QOP_ZERO;
+        }
+
         trk_state.smoothed_params().set_covariance(smoothed_cov);
 
+        // Wrap the phi and theta angles in their valid ranges
+        normalize_angles(trk_state.smoothed_params());
+
         matrix_type<D, e_bound_size> H =
             measurements.at(trk_state.measurement_index())
                 .subs.template projector<D>();
@@ -141,11 +158,13 @@ struct two_filters_smoother {
                 residual_smt, matrix::inverse(R_smt)) *
             residual_smt;
 
-        if (getter::element(chi2_smt, 0, 0) < 0.f) {
+        const scalar chi2_smt_value{getter::element(chi2_smt, 0, 0)};
+
+        if (chi2_smt_value < 0.f) {
             return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NEGATIVE;
         }
 
-        if (!std::isfinite(getter::element(chi2_smt, 0, 0))) {
+        if (!std::isfinite(chi2_smt_value)) {
             return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NOT_FINITE;
         }
 
@@ -180,26 +199,12 @@ struct two_filters_smoother {
             predicted_vec + K * (meas_local - H * predicted_vec);
         const matrix_type<6, 6> filtered_cov = (I66 - K * H) * predicted_cov;
 
-        // Residual between measurement and (projected) filtered vector
-        const matrix_type<D, 1> residual = meas_local - H * filtered_vec;
-
-        // Calculate backward chi2
-        const matrix_type<D, D> R = (I_m - H * K) * V;
-        // assert(matrix::determinant(R) != 0.f);
-        assert(std::isfinite(matrix::determinant(R)));
-        const matrix_type<1, 1> chi2 =
-            algebra::matrix::transposed_product<true, false>(
-                residual, matrix::inverse(R)) *
-            residual;
-
         // Update the bound track parameters
         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) {
-            return kalman_fitter_status::ERROR_THETA_ZERO;
+        if (!std::isfinite(bound_params.theta())) {
+            return kalman_fitter_status::ERROR_INVERSION;
         }
 
         if (!std::isfinite(bound_params.phi())) {
@@ -210,19 +215,40 @@ struct two_filters_smoother {
             return kalman_fitter_status::ERROR_QOP_ZERO;
         }
 
-        if (getter::element(chi2, 0, 0) < 0.f) {
-            return kalman_fitter_status::ERROR_UPDATER_CHI2_NEGATIVE;
+        bound_params.set_covariance(filtered_cov);
+
+        // Residual between measurement and (projected) filtered vector
+        const matrix_type<D, 1> residual = meas_local - H * filtered_vec;
+
+        // Calculate backward chi2
+        const matrix_type<D, D> R = (I_m - H * K) * V;
+        assert(matrix::determinant(R) != 0.f);
+        assert(std::isfinite(matrix::determinant(R)));
+        const matrix_type<1, 1> chi2 =
+            algebra::matrix::transposed_product<true, false>(
+                residual, matrix::inverse(R)) *
+            residual;
+
+        const scalar chi2_val{getter::element(chi2, 0, 0)};
+
+        if (chi2_val < 0.f) {
+            return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NEGATIVE;
         }
 
-        if (!std::isfinite(getter::element(chi2, 0, 0))) {
-            return kalman_fitter_status::ERROR_UPDATER_CHI2_NOT_FINITE;
+        if (!std::isfinite(chi2_val)) {
+            return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NOT_FINITE;
         }
 
         // Set backward chi2
-        trk_state.backward_chi2() = getter::element(chi2, 0, 0);
+        trk_state.backward_chi2() = chi2_val;
+
+        // Wrap the phi and theta angles in their valid ranges
+        normalize_angles(bound_params);
 
-        // Wrap the phi in the range of [-pi, pi]
-        wrap_phi(bound_params);
+        const scalar theta = bound_params.theta();
+        if (theta <= 0.f || theta >= 2.f * constant<traccc::scalar>::pi) {
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
 
         trk_state.set_smoothed();
 
diff --git a/core/include/traccc/fitting/status_codes.hpp b/core/include/traccc/fitting/status_codes.hpp
@@ -14,7 +14,7 @@ namespace traccc {
 enum class kalman_fitter_status : uint32_t {
     SUCCESS,
     ERROR_QOP_ZERO,
-    ERROR_THETA_ZERO,
+    ERROR_THETA_POLE,
     ERROR_INVERSION,
     ERROR_SMOOTHER_CHI2_NEGATIVE,
     ERROR_SMOOTHER_CHI2_NOT_FINITE,
@@ -36,8 +36,8 @@ struct fitter_debug_msg {
             case ERROR_QOP_ZERO: {
                 return msg + "Track qop is zero";
             }
-            case ERROR_THETA_ZERO: {
-                return msg + "Track theta is zero";
+            case ERROR_THETA_POLE: {
+                return msg + "Track theta hit pole";
             }
             case ERROR_INVERSION: {
                 return msg + "Failed matrix inversion";
@@ -49,10 +49,10 @@ struct fitter_debug_msg {
                 return msg + "Invalid chi2 in smoother";
             }
             case ERROR_UPDATER_CHI2_NEGATIVE: {
-                return msg + "Negative chi2 in gain matrix update";
+                return msg + "Negative chi2 in gain matrix updater";
             }
             case ERROR_UPDATER_CHI2_NOT_FINITE: {
-                return msg + "Invalid chi2 in gain matrix update";
+                return msg + "Invalid chi2 in gain matrix updater";
             }
             case ERROR_BARCODE_SEQUENCE_OVERFLOW: {
                 return msg + "Barcode sequence overflow in direct navigator";
diff --git a/core/include/traccc/utils/trigonometric_helpers.hpp b/core/include/traccc/utils/trigonometric_helpers.hpp
diff --git a/device/common/include/traccc/finding/device/impl/propagate_to_next_surface.ipp b/device/common/include/traccc/finding/device/impl/propagate_to_next_surface.ipp
