From 86231fc3349465117b5b950ec9b0694cd97139fc Mon Sep 17 00:00:00 2001
From: Tim Hargreaves <38204689+THargreaves@users.noreply.github.com>
Date: Wed, 20 Nov 2024 13:04:22 +0000
Subject: [PATCH] Correct Cholesky factor to have positive diagonal (#17)

---
 src/srkf.jl  | 23 +++++++++++++++++++++++
 src/srukf.jl |  4 ++++
 2 files changed, 27 insertions(+)

diff --git a/src/srkf.jl b/src/srkf.jl
index cbf25fd..0072dd1 100644
--- a/src/srkf.jl
+++ b/src/srkf.jl
@@ -14,6 +14,25 @@ function calc_upper_triangular_of_qr!(A)
     R
 end
 
+"""
+    correct_cholesky_sign(R)
+
+    The upper triangle R of the QR decomposition doesn't necessarily have positive
+    elements on the diagonal elements. However, when used in the context of
+    a Cholesky decomposition the diagonal elements need to be positive. Hence, 
+    the function corrects the sign of the diagonal elements when necessary.
+"""
+correct_cholesky_sign(R) = sign.(diag(R)) .* R
+
+function correct_cholesky_sign!(R)
+    for i in axes(R,1)
+        if real(R[i,i]) < 0
+            R[i, :] = -R[i, :]
+        end
+    end
+    return R
+end
+
 # This implementation is based on
 # https://github.com/rlabbe/filterpy/blob/master/filterpy/kalman/square_root.py
 function calc_apriori_covariance(
@@ -23,6 +42,7 @@ function calc_apriori_covariance(
 )
     A = vcat(P.U * F', Q.U)
     R = calc_upper_triangular_of_qr(A)
+    R = correct_cholesky_sign(R)
     Cholesky(R, 'U', 0)
 end
 
@@ -60,6 +80,7 @@ function calc_cross_cov_innovation_posterior(P::Cholesky, H, R::Cholesky, consid
     num_y = size(R, 1)
     M = create_matrix_for_qr(P, H, R)
     RU = calc_upper_triangular_of_qr(M)
+    RU = correct_cholesky_sign(RU)
     PHᵀ = extract_cross_covariance(RU, num_y)
     S = extract_innovation_covariance(RU, num_y)
     P_post = extract_posterior_covariance(RU, num_y, P, consider)
@@ -103,6 +124,7 @@ function time_update!(tu::SRKFTUIntermediate, x, P::Cholesky, F::Union{Number, A
     tu.puft_vcat_q[1:size(F, 1),:] .= @~ P.U * F'
     tu.puft_vcat_q[size(F, 1) + 1:end,:] .= Q.U
     R = calc_upper_triangular_of_qr_inplace!(tu.p_apri, tu.puft_vcat_q, tu.qr_zeros, tu.qr_space)
+    correct_cholesky_sign!(R)
     P_apri = Cholesky(R, 'U', 0)
     KFTimeUpdate(x_apri, P_apri)
 end
@@ -149,6 +171,7 @@ function measurement_update!(
     M[dim_y + 1:end, 1:dim_y] .= @~ P.U * H'
     M[dim_y + 1:end, dim_y + 1:end] .= P.U
     RU = calc_upper_triangular_of_qr_inplace!(mu.R, M, mu.qr_zeros, mu.qr_space)
+    correct_cholesky_sign!(RU)
     PHᵀ = (@view(RU[1:dim_y, dim_y + 1:end]))'
     S = Cholesky(@view(RU[1:dim_y, 1:dim_y]), 'U', 0)
     K = calc_kalman_gain!(mu.kalman_gain, PHᵀ, S.L)
diff --git a/src/srukf.jl b/src/srukf.jl
index 3975068..76df91d 100644
--- a/src/srukf.jl
+++ b/src/srukf.jl
@@ -80,6 +80,7 @@ function cov(χ::TransformedSigmaPoints, noise::Cholesky)
     weight_0, weight_i = calc_cov_weights(χ.weight_params, (size(χ, 2) - 1) >> 1)
     A = vcat(sqrt(weight_i) * χ.xi', noise.uplo === 'U' ? noise.U : noise.L')
     R = calc_upper_triangular_of_qr!(A)
+    correct_cholesky_sign!(R)
     S = Cholesky(R, 'U', 0)
     if weight_0 < 0
         P = lowrankdowndate(S, sqrt(abs(weight_0)) * χ.x0)
@@ -94,6 +95,7 @@ function cov!(res, qr_A, qr_zeros, qr_space, x0_temp, χ::TransformedSigmaPoints
     qr_A[1:size(χ.xi, 2), :] .= sqrt(weight_i) .* χ.xi'
     qr_A[size(χ.xi, 2) + 1:end, :] = noise.uplo === 'U' ? noise.U : noise.L'
     R = calc_upper_triangular_of_qr_inplace!(res, qr_A, qr_zeros, qr_space)
+    correct_cholesky_sign!(R)
     S = Cholesky(R, 'U', 0)
     x0_temp .= sqrt(abs(weight_0)) .* χ.x0
     if weight_0 < 0
@@ -108,6 +110,7 @@ function cov(χ::TransformedSigmaPoints, noise::Augment{<:Cholesky})
     weight_0, weight_i = calc_cov_weights(χ.weight_params, (size(χ, 2) - 1) >> 1)
     A = sqrt(weight_i) * χ.xi'
     R = calc_upper_triangular_of_qr!(A)
+    correct_cholesky_sign!(R)
     S = Cholesky(R, 'U', 0)
     if weight_0 < 0
         P = lowrankdowndate(S, sqrt(abs(weight_0)) * χ.x0)
@@ -121,6 +124,7 @@ function cov!(res, qr_A, qr_zeros, qr_space, x0_temp, χ::TransformedSigmaPoints
     weight_0, weight_i = calc_cov_weights(χ.weight_params, (size(χ, 2) - 1) >> 1)
     qr_A .= sqrt(weight_i) .* χ.xi'
     R = calc_upper_triangular_of_qr_inplace!(res, qr_A, qr_zeros, qr_space)
+    correct_cholesky_sign!(R)
     S = Cholesky(R, 'U', 0)
     x0_temp .= sqrt(abs(weight_0)) .* χ.x0
     if weight_0 < 0