Add additional error checking in Kálmán filter

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

@@ -12,6 +12,7 @@
 #include "traccc/edm/track_parameters.hpp"
 #include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/status_codes.hpp"
+#include "traccc/utils/debug.hpp"
 
 // Detray inlcude(s)
 #include <detray/geometry/shapes/line.hpp>
@@ -105,23 +106,28 @@ struct gain_matrix_smoother {
             cur_filtered_cov +
             A * (next_smoothed_cov - next_predicted_cov) * matrix::transpose(A);
 
+        if (!std::isfinite(getter::element(smt_vec, e_bound_phi, 0)))
+            [[unlikely]] { return kalman_fitter_status::ERROR_INVERSION; }
+        else
+            [[likely]] {
+                // Assert that the entire matrix is finite, i.e. that checking
+                // only the phi value didn't let any other non-finite values
+                // through.
+                assert(matrix_is_finite(smt_vec));
+                assert(matrix_is_finite(smt_cov));
+            }
+
         cur_state.smoothed().set_vector(smt_vec);
         cur_state.smoothed().set_covariance(smt_cov);
 
         // Return false if track is parallel to z-axis or phi is not finite
         const scalar theta = cur_state.smoothed().theta();
 
-        if (theta <= 0.f || theta >= constant<traccc::scalar>::pi) {
-            return kalman_fitter_status::ERROR_THETA_ZERO;
-        }
-
-        if (!std::isfinite(cur_state.smoothed().phi())) {
-            return kalman_fitter_status::ERROR_INVERSION;
-        }
+        if (theta <= 0.f || theta >= constant<traccc::scalar>::pi)
+            [[unlikely]] { return kalman_fitter_status::ERROR_THETA_ZERO; }
 
-        if (std::abs(cur_state.smoothed().qop()) == 0.f) {
-            return kalman_fitter_status::ERROR_QOP_ZERO;
-        }
+        if (std::abs(cur_state.smoothed().qop()) == 0.f)
+            [[unlikely]] { return kalman_fitter_status::ERROR_QOP_ZERO; }
 
         // Wrap the phi in the range of [-pi, pi]
         wrap_phi(cur_state.smoothed());

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

@@ -12,6 +12,7 @@
 #include "traccc/definitions/track_parametrization.hpp"
 #include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/status_codes.hpp"
+#include "traccc/utils/debug.hpp"
 
 // Detray inlcude(s)
 #include <detray/geometry/shapes/line.hpp>
@@ -130,17 +131,25 @@ struct gain_matrix_updater {
         // 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 (theta <= 0.f || theta >= constant<traccc::scalar>::pi)
+            [[unlikely]] { return kalman_fitter_status::ERROR_THETA_ZERO; }
 
-        if (!std::isfinite(bound_params.phi())) {
-            return kalman_fitter_status::ERROR_INVERSION;
-        }
+        if (!std::isfinite(bound_params.phi()))
+            [[unlikely]] { return kalman_fitter_status::ERROR_INVERSION; }
 
-        if (std::abs(bound_params.qop()) == 0.f) {
-            return kalman_fitter_status::ERROR_QOP_ZERO;
-        }
+        if (std::abs(bound_params.qop()) == 0.f)
+            [[unlikely]] { return kalman_fitter_status::ERROR_QOP_ZERO; }
+
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_phi, 0)))
+            [[unlikely]] { return kalman_fitter_status::ERROR_INVERSION; }
+        else
+            [[likely]] {
+                // Assert that the entire matrix is finite, i.e. that checking
+                // only the phi value didn't let any other non-finite values
+                // through.
+                assert(matrix_is_finite(filtered_vec));
+                assert(matrix_is_finite(filtered_cov));
+            }
 
         // Set the track state parameters
         trk_state.filtered().set_vector(filtered_vec);

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

@@ -158,9 +158,8 @@ class kalman_fitter {
 
             if (kalman_fitter_status res =
                     filter(seed_params_cpy, fitter_state);
-                res != kalman_fitter_status::SUCCESS) {
-                return res;
-            }
+                res != kalman_fitter_status::SUCCESS)
+                [[unlikely]] { return res; }
         }
 
         return kalman_fitter_status::SUCCESS;
@@ -206,9 +205,11 @@ class kalman_fitter {
 
         // Run smoothing
         if (kalman_fitter_status res = smooth(fitter_state);
-            res != kalman_fitter_status::SUCCESS) {
-            return res;
-        }
+            res != kalman_fitter_status::SUCCESS)
+            [[unlikely]] { return res; }
+
+        if (fitter_state.m_fit_res.fit_params.theta() == 0.f)
+            [[unlikely]] { return kalman_fitter_status::ERROR_THETA_ZERO; }
 
         // Update track fitting qualities
         update_statistics(fitter_state);
@@ -265,11 +266,12 @@ class kalman_fitter {
             fitter_state.m_fit_actor_state.backward_mode = true;
 
             const auto& dir = propagation._stepping().dir();
-            if (dir[0] == 0.f && dir[1] == 0.f) {
-                // Particle is exactly parallel to the beampipe, which we
-                // cannot represent.
-                return kalman_fitter_status::ERROR_THETA_ZERO;
-            }
+            if (dir[0] == 0.f && dir[1] == 0.f)
+                [[unlikely]] {
+                    // Particle is exactly parallel to the beampipe, which we
+                    // cannot represent.
+                    return kalman_fitter_status::ERROR_THETA_ZERO;
+                }
 
             propagator.propagate(propagation,
                                  fitter_state.backward_actor_state());
@@ -289,9 +291,8 @@ class kalman_fitter {
                 if (kalman_fitter_status res =
                         sf.template visit_mask<
                             gain_matrix_smoother<algebra_type>>(*it, *(it - 1));
-                    res != kalman_fitter_status::SUCCESS) {
-                    return res;
-                }
+                    res != kalman_fitter_status::SUCCESS)
+                    [[unlikely]] { return res; }
             }
         }
 

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

@@ -12,6 +12,7 @@
 #include "traccc/definitions/track_parametrization.hpp"
 #include "traccc/edm/track_state.hpp"
 #include "traccc/fitting/status_codes.hpp"
+#include "traccc/utils/debug.hpp"
 
 namespace traccc {
 
@@ -151,13 +152,22 @@ struct two_filters_smoother {
 
         // 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;
-        }
+        if (theta <= 0.f || theta >= constant<traccc::scalar>::pi)
+            [[unlikely]] { return kalman_fitter_status::ERROR_THETA_ZERO; }
+
+        if (!std::isfinite(bound_params.phi()))
+            [[unlikely]] { return kalman_fitter_status::ERROR_INVERSION; }
+
+        if (!std::isfinite(getter::element(filtered_vec, e_bound_phi, 0)))
+            [[unlikely]] { return kalman_fitter_status::ERROR_INVERSION; }
+        else
+            [[likely]] {
+                // Assert that the entire matrix is finite, i.e. that checking
+                // only the phi value didn't let any other non-finite values
+                // through.
+                assert(matrix_is_finite(filtered_vec));
+                assert(matrix_is_finite(filtered_cov));
+            }
 
         // Update the bound track parameters
         bound_params.set_vector(filtered_vec);

core/include/traccc/utils/debug.hpp

@@ -0,0 +1,29 @@
+/** TRACCC library, part of the ACTS project (R&D line)
+ *
+ * (c) 2021 CERN for the benefit of the ACTS project
+ *
+ * Mozilla Public License Version 2.0
+ */
+
+#pragma once
+
+#include <algebra/concepts.hpp>
+#include <algebra/type_traits.hpp>
+
+#include "traccc/definitions/qualifiers.hpp"
+
+namespace traccc {
+template <::algebra::concepts::matrix T>
+TRACCC_HOST_DEVICE bool matrix_is_finite(const T& mat) {
+    for (std::size_t i = 0; i < ::algebra::traits::columns<T>; ++i) {
+        for (std::size_t j = 0; j < ::algebra::traits::rows<T>; ++j) {
+            if (!std::isfinite(getter::element(mat, i, j))) {
+                return false;
+            }
+        }
+    }
+    return true;
+}
+}  // namespace traccc
+//
+//

device/common/include/traccc/fitting/device/impl/fit.ipp

@@ -60,6 +60,10 @@ TRACCC_HOST_DEVICE inline void fit(
     // TODO: Process fit failures more elegantly
     assert(fit_status == kalman_fitter_status::SUCCESS);
 
+    if (fit_status == kalman_fitter_status::SUCCESS) {
+        assert(fitter_state.m_fit_res.fit_params.theta() != 0.f);
+    }
+
     // Get the final fitting information
     track_states.at(param_id).header = fitter_state.m_fit_res;
 }