From be88ea16a82e420c0247362df35b4981322ec159 Mon Sep 17 00:00:00 2001
From: Stephen Nicholas Swatman <stephen.nicholas.swatman@cern.ch>
Date: Fri, 7 Nov 2025 16:11:58 +0100
Subject: [PATCH] =?UTF-8?q?Compute=20K=C3=A1lm=C3=A1n=20filter=20chi2=20fr?=
 =?UTF-8?q?om=20predicted=20state?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This commit migrates the computation of the chi2 value in our gain
matrix updater to use the predicted state rather than the filtered
state. This is mathematically equivalent but potentially more stable.
---
 .../kalman_filter/gain_matrix_updater.hpp     | 48 ++++++++++---------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp b/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp
index 413573d576..2ef177c174 100644
--- a/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp
+++ b/core/include/traccc/fitting/kalman_filter/gain_matrix_updater.hpp
@@ -66,7 +66,6 @@ struct gain_matrix_updater {
         // Some identity matrices
         // @TODO: Make constexpr work
         const auto I66 = matrix::identity<bound_matrix_type>();
-        const auto I_m = matrix::identity<matrix_type<D, D>>();
 
         // Measurement data on surface
         matrix_type<D, 1> meas_local;
@@ -113,6 +112,31 @@ struct gain_matrix_updater {
         TRACCC_DEBUG_HOST("Predicted residual: " << meas_local -
                                                         H * predicted_vec);
 
+        scalar chi2_val;
+
+        {
+            // Calculate the chi square
+            const matrix_type<D, D> R =
+                V + (H * predicted_cov * algebra::matrix::transpose(H));
+            // Residual between measurement and predicted vector
+            const matrix_type<D, 1> residual = meas_local - H * predicted_vec;
+            const matrix_type<1, 1> chi2_mat =
+                algebra::matrix::transposed_product<true, false>(
+                    residual, matrix::inverse(R)) *
+                residual;
+            chi2_val = getter::element(chi2_mat, 0, 0);
+
+            if (chi2_val < 0.f) {
+                TRACCC_ERROR_HOST_DEVICE("Chi2 negative");
+                return kalman_fitter_status::ERROR_UPDATER_CHI2_NEGATIVE;
+            }
+
+            if (!std::isfinite(chi2_val)) {
+                TRACCC_ERROR_HOST_DEVICE("Chi2 infinite");
+                return kalman_fitter_status::ERROR_UPDATER_CHI2_NOT_FINITE;
+            }
+        }
+
         const matrix_type<e_bound_size, D> projected_cov =
             algebra::matrix::transposed_product<false, true>(predicted_cov, H);
 
@@ -163,34 +187,12 @@ struct gain_matrix_updater {
             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;
-
-        // Calculate the chi square
-        const matrix_type<D, D> R = (I_m - H * K) * V;
-        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)};
-
         TRACCC_VERBOSE_HOST("Filtered residual: " << residual);
         TRACCC_DEBUG_HOST("R:\n" << R);
         TRACCC_DEBUG_HOST_DEVICE("det(R): %f", matrix::determinant(R));
         TRACCC_DEBUG_HOST("R_inv:\n" << matrix::inverse(R));
         TRACCC_VERBOSE_HOST_DEVICE("Chi2: %f", chi2_val);
 
-        if (chi2_val < 0.f) {
-            TRACCC_ERROR_HOST_DEVICE("Chi2 negative");
-            return kalman_fitter_status::ERROR_UPDATER_CHI2_NEGATIVE;
-        }
-
-        if (!std::isfinite(chi2_val)) {
-            TRACCC_ERROR_HOST_DEVICE("Chi2 infinite");
-            return kalman_fitter_status::ERROR_UPDATER_CHI2_NOT_FINITE;
-        }
-
         // Set the chi2 for this track and measurement
         trk_state.filtered_params().set_vector(filtered_vec);
         trk_state.filtered_params().set_covariance(filtered_cov);