Fix a number of assertions during track finding and reorder Kalman checks such that they perform an early reject before data is written to the track states (#1311)

Co-authored-by: Stephen Nicholas Swatman <stephen.nicholas.swatman@cern.ch>

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/logging.hpp"
 #include "traccc/utils/trigonometric_helpers.hpp"
 
 // detray include(s).
@@ -47,17 +48,45 @@ using bound_matrix = detray::bound_matrix<algebra_t>;
 using bound_track_parameters_collection_types =
     collection_types<bound_track_parameters<>>;
 
-// Wrap the phi of track parameters to [-pi,pi]
+// Wrap the phi of a track parameter vector to [-pi,pi]
 template <detray::concepts::algebra algebra_t>
-TRACCC_HOST_DEVICE constexpr void normalize_angles(
-    bound_track_parameters<algebra_t>& param) {
+TRACCC_HOST_DEVICE constexpr bool normalize_angles(
+    bound_vector<algebra_t>& vec) {
     traccc::scalar phi;
     traccc::scalar theta;
 
-    std::tie(phi, theta) = detail::wrap_phi_theta(param.phi(), param.theta());
+    const traccc::scalar in_theta{getter::element(vec, e_bound_theta, 0)};
+    if (math::fmod(in_theta, 2.f * constant<traccc::scalar>::pi) == 0.f ||
+        in_theta == 0.f) {
+        TRACCC_WARNING_HOST_DEVICE("Hit theta pole before normalization: %f",
+                                   in_theta);
+        return false;
+    }
+
+    std::tie(phi, theta) =
+        detail::wrap_phi_theta(getter::element(vec, e_bound_phi, 0), in_theta);
+
+    if (theta <= 0.f || theta >= 2.f * constant<traccc::scalar>::pi) {
+        TRACCC_WARNING_HOST_DEVICE("Hit theta pole after normalization: %f",
+                                   theta);
+        return false;
+    }
+
+    // Assertions of the detray bound track parameters
+    assert(math::fabs(phi) <= constant<traccc::scalar>::pi);
+    assert(theta <= constant<traccc::scalar>::pi);
 
-    param.set_phi(phi);
-    param.set_theta(theta);
+    getter::element(vec, e_bound_phi, 0) = phi;
+    getter::element(vec, e_bound_theta, 0) = theta;
+
+    return true;
+}
+
+// Wrap the phi of bound track parameters to [-pi,pi]
+template <detray::concepts::algebra algebra_t>
+TRACCC_HOST_DEVICE constexpr bool normalize_angles(
+    bound_parameters_vector<algebra_t>& param_vec) {
+    return normalize_angles<algebra_t>(param_vec.vector());
 }
 
 /// Covariance inflation used for track fitting

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

@@ -126,23 +126,10 @@ struct gain_matrix_updater {
         TRACCC_DEBUG_HOST("-> K:\n" << K);
 
         // Calculate the filtered track parameters
-        const matrix_type<6, 1> filtered_vec =
+        matrix_type<6, 1> filtered_vec =
             predicted_vec + K * (meas_local - H * predicted_vec);
-        const matrix_type<6, 6> i_minus_kh = I66 - K * H;
-        matrix_type<6, 6> filtered_cov =
-            i_minus_kh * predicted_cov * matrix::transpose(i_minus_kh) +
-            K * V * matrix::transpose(K);
 
         TRACCC_DEBUG_HOST("-> Filtered param:\n" << filtered_vec);
-        TRACCC_DEBUG_HOST("-> Filtered cov:\n" << filtered_cov);
-
-        // Check the covariance for consistency
-        // @TODO: Need to understand why negative variance happens
-        if (constexpr traccc::scalar min_var{-0.01f};
-            !details::regularize_covariance<algebra_t>(filtered_cov, min_var)) {
-            TRACCC_ERROR_HOST_DEVICE("Negative variance after filtering");
-            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))) {
@@ -163,6 +150,21 @@ struct gain_matrix_updater {
             return kalman_fitter_status::ERROR_QOP_ZERO;
         }
 
+        const matrix_type<6, 6> i_minus_kh = I66 - K * H;
+        matrix_type<6, 6> filtered_cov =
+            i_minus_kh * predicted_cov * matrix::transpose(i_minus_kh) +
+            K * V * matrix::transpose(K);
+
+        TRACCC_DEBUG_HOST("-> Filtered cov:\n" << filtered_cov);
+
+        // Check the covariance for consistency
+        // @TODO: Need to understand why negative variance happens
+        if (constexpr traccc::scalar min_var{-0.01f};
+            !details::regularize_covariance<algebra_t>(filtered_cov, min_var)) {
+            TRACCC_ERROR_HOST_DEVICE("Negative variance after filtering");
+            return kalman_fitter_status::ERROR_UPDATER_INVALID_COVARIANCE;
+        }
+
         // Residual between measurement and (projected) filtered vector
         const matrix_type<D, 1> residual = meas_local - H * filtered_vec;
 
@@ -192,29 +194,17 @@ struct gain_matrix_updater {
             return kalman_fitter_status::ERROR_UPDATER_CHI2_NOT_FINITE;
         }
 
+        // Wrap the phi and theta angles in their valid ranges
+        if (!normalize_angles<algebra_t>(filtered_vec)) {
+            TRACCC_ERROR_HOST_DEVICE("Hit theta pole in filtering!");
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
+
         // 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() = chi2_val;
 
-        if (math::fmod(trk_state.filtered_params().theta(),
-                       2.f * constant<traccc::scalar>::pi) == 0.f) {
-            TRACCC_ERROR_HOST_DEVICE(
-                "Hit theta pole after filtering : %f (unrecoverable error "
-                "pre-normalization)",
-                trk_state.filtered_params().theta());
-            return kalman_fitter_status::ERROR_THETA_POLE;
-        }
-
-        // 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) {
-            TRACCC_ERROR_HOST_DEVICE("Hit theta pole in filtering : %f", theta);
-            return kalman_fitter_status::ERROR_THETA_POLE;
-        }
-
         assert(!trk_state.filtered_params().is_invalid());
 
         return kalman_fitter_status::SUCCESS;

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

@@ -414,7 +414,7 @@ struct kalman_actor : detray::base_actor {
                               direction_e ==
                                   kalman_actor_direction::BIDIRECTIONAL) {
                     // Wrap the phi and theta angles in their valid ranges
-                    normalize_angles(bound_param);
+                    normalize_angles<algebra_t>(bound_param);
 
                     // Forward filter
                     TRACCC_DEBUG_HOST_DEVICE("Run filtering...");

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

@@ -101,35 +101,35 @@ struct two_filters_smoother {
 
         // Eq (3.38) of "Pattern Recognition, Tracking and Vertex
         // Reconstruction in Particle Detectors"
-        const matrix_type<e_bound_size, 1u> smoothed_vec =
+        matrix_type<e_bound_size, 1u> smoothed_vec =
             smoothed_cov *
             (filtered_cov_inv * trk_state.filtered_params().vector() +
              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())) {
+        if (!std::isfinite(getter::element(smoothed_vec, e_bound_theta, 0))) {
             TRACCC_ERROR_HOST_DEVICE(
                 "Theta is infinite after smoothing (Matrix inversion)");
             return kalman_fitter_status::ERROR_INVERSION;
         }
 
-        if (!std::isfinite(trk_state.smoothed_params().phi())) {
+        if (!std::isfinite(getter::element(smoothed_vec, e_bound_phi, 0))) {
             TRACCC_ERROR_HOST_DEVICE(
                 "Phi is infinite after smoothing (Matrix inversion)");
             return kalman_fitter_status::ERROR_INVERSION;
         }
 
-        if (math::fabs(trk_state.smoothed_params().qop()) == 0.f) {
+        if (math::fabs(getter::element(smoothed_vec, e_bound_qoverp, 0)) ==
+            0.f) {
             TRACCC_ERROR_HOST_DEVICE("q/p is zero after smoothing");
             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());
+        if (!normalize_angles<algebra_t>(smoothed_vec)) {
+            TRACCC_ERROR_HOST_DEVICE("Hit theta pole after smoothing!");
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
 
         const subspace<algebra_t, e_bound_size> subs(
             measurements.at(trk_state.measurement_index()).subspace());
@@ -192,8 +192,6 @@ struct two_filters_smoother {
             return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NOT_FINITE;
         }
 
-        trk_state.smoothed_chi2() = getter::element(chi2_smt, 0, 0);
-
         /*************************************
          *  Set backward filtered parameter
          *************************************/
@@ -222,45 +220,49 @@ struct two_filters_smoother {
         TRACCC_DEBUG_HOST("K:\n" << K);
 
         // Calculate the filtered track parameters
-        const matrix_type<6, 1> filtered_vec =
+        matrix_type<6, 1> filtered_vec =
             predicted_vec + K * (meas_local - H * predicted_vec);
-        const matrix_type<6, 6> i_minus_kh = I66 - K * H;
-        matrix_type<6, 6> filtered_cov =
-            i_minus_kh * predicted_cov * matrix::transpose(i_minus_kh) +
-            K * V * matrix::transpose(K);
-
-        // Check the covariance for consistency
-        // @TODO: Need to understand why negative variance happens
-        if (constexpr traccc::scalar min_var{-0.01f};
-            !details::regularize_covariance<algebra_t>(filtered_cov, min_var)) {
-            TRACCC_ERROR_HOST_DEVICE("Negative variance after filtering");
-            return kalman_fitter_status::ERROR_SMOOTHER_INVALID_COVARIANCE;
-        }
-
-        // Update the bound track parameters
-        bound_params.set_vector(filtered_vec);
 
         // Return false if track is parallel to z-axis or phi is not finite
-        if (!std::isfinite(bound_params.theta())) {
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_theta, 0))) {
             TRACCC_ERROR_HOST_DEVICE(
                 "Theta is infinite after filering in smoother (Matrix "
                 "inversion)");
             return kalman_fitter_status::ERROR_INVERSION;
         }
 
-        if (!std::isfinite(bound_params.phi())) {
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_phi, 0))) {
             TRACCC_ERROR_HOST_DEVICE(
                 "Phi is infinite after filering in smoother (Matrix "
                 "inversion)");
             return kalman_fitter_status::ERROR_INVERSION;
         }
 
-        if (math::fabs(bound_params.qop()) == 0.f) {
+        if (math::fabs(getter::element(filtered_vec, e_bound_qoverp, 0)) ==
+            0.f) {
             TRACCC_ERROR_HOST_DEVICE("q/p is zero after filering in smoother");
             return kalman_fitter_status::ERROR_QOP_ZERO;
         }
 
-        bound_params.set_covariance(filtered_cov);
+        // Wrap the phi and theta angles in their valid ranges
+        if (!normalize_angles<algebra_t>(filtered_vec)) {
+            TRACCC_ERROR_HOST_DEVICE(
+                "Hit theta pole after filtering in smoother!");
+            return kalman_fitter_status::ERROR_THETA_POLE;
+        }
+
+        const matrix_type<6, 6> i_minus_kh = I66 - K * H;
+        matrix_type<6, 6> filtered_cov =
+            i_minus_kh * predicted_cov * matrix::transpose(i_minus_kh) +
+            K * V * matrix::transpose(K);
+
+        // Check the covariance for consistency
+        // @TODO: Need to understand why negative variance happens
+        if (constexpr traccc::scalar min_var{-0.01f};
+            !details::regularize_covariance<algebra_t>(filtered_cov, min_var)) {
+            TRACCC_ERROR_HOST_DEVICE("Negative variance after filtering");
+            return kalman_fitter_status::ERROR_SMOOTHER_INVALID_COVARIANCE;
+        }
 
         // Residual between measurement and (projected) filtered vector
         const matrix_type<D, 1> residual = meas_local - H * filtered_vec;
@@ -291,21 +293,18 @@ struct two_filters_smoother {
             return kalman_fitter_status::ERROR_SMOOTHER_CHI2_NOT_FINITE;
         }
 
-        // Set backward chi2
+        // Update the smoothed track parameters
+        trk_state.smoothed_params().set_vector(smoothed_vec);
+        trk_state.smoothed_params().set_covariance(smoothed_cov);
+        trk_state.smoothed_chi2() = getter::element(chi2_smt, 0, 0);
         trk_state.backward_chi2() = chi2_val;
-
-        // Wrap the phi and theta angles in their valid ranges
-        normalize_angles(bound_params);
-
-        const scalar theta = bound_params.theta();
-        if (theta <= 0.f || theta >= 2.f * constant<traccc::scalar>::pi) {
-            TRACCC_ERROR_HOST_DEVICE("Hit theta pole after smoothing : %f",
-                                     theta);
-            return kalman_fitter_status::ERROR_THETA_POLE;
-        }
-
         trk_state.set_smoothed();
 
+        // Update the filtered track parameters
+        bound_params.set_vector(filtered_vec);
+        bound_params.set_covariance(filtered_cov);
+
+        assert(!trk_state.smoothed_params().is_invalid());
         assert(!bound_params.is_invalid());
 
         return kalman_fitter_status::SUCCESS;

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

@@ -220,9 +220,7 @@ TRACCC_HOST_DEVICE inline void build_tracks(
 #ifndef NDEBUG
     // Assert that we did not make any duplicate track states.
     for (const auto& i : track.constituent_links()) {
-        assert(i.type == edm::track_constituent_link::measurement);
         for (const auto& j : track.constituent_links()) {
-            assert(j.type == edm::track_constituent_link::measurement);
             if (i.index != j.index) {
                 // TODO: Re-enable me!
                 // assert(measurements.at(i->index).identifier() !=

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

@@ -379,7 +379,7 @@ TRACCC_HOST_DEVICE inline void find_tracks(
                          * $\chi^2$ value will be.
                          */
                         traccc::scalar highest =
-                            static_cast<traccc::scalar>(0.f);
+                            static_cast<traccc::scalar>(-1.f);
 
                         for (unsigned int i = 0;
                              i < cfg.max_num_branches_per_surface; ++i) {