addind diffuse smoother tests KO

src/KalmanFilterTools.jl

@@ -28,6 +28,7 @@ export DiffuseKalmanFilterWs
 export fast_kalman_likelihood, diffuse_kalman_likelihood
 export kalman_filter!, diffuse_kalman_filter!
 export KalmanSmootherWs, DiffuseKalmanSmootherWs, kalman_smoother!
+export diffuse_kalman_smoother!
 
 abstract type KalmanWs{T, U} end
 

src/kalman_base.jl

@@ -12,6 +12,15 @@ function get_alphah!(alphah::AbstractVector{T}, a::AbstractVector{T}, P::Abstrac
     mul!(alphah,P, r, 1.0, 1.0)
 end
 
+# alphah_t = a_t + Pstar_t*r0_{t-1} + Pinf_t*r1_{t-1}     (DK 5.24)
+function get_alphah!(valphah::AbstractVector{T}, va::AbstractVector{T},
+                     vPstar::AbstractMatrix{T}, vPinf::AbstractMatrix{T},
+                     r0::AbstractVector{T}, r1::AbstractVector{T}) where T <: AbstractFloat
+    copy!(valphah, va)
+    mul!(valphah, vPstar, r0, 1.0, 1.0)
+    mul!(valphah, vPinf, r1, 1.0, 1.0)
+end
+
 function get_cholF!(cholF::AbstractArray{T}, F::AbstractArray{T}) where T <: AbstractFloat
     cholF .= 0.5.*(F .+ transpose(F))
     info = LAPACK.potrf!('U', cholF)
@@ -42,6 +51,12 @@ function get_D!(D::AbstractMatrix{U}, iF::AbstractMatrix{U}, K::AbstractMatrix{U
     mul!(D, tmp, K, 1.0, 1.0)
 end
 
+# D_t = KDKinf_t'*N0_t*KDKinf_t    (DK p. 135)
+function get_D!(D::AbstractMatrix{T}, KDKinf::AbstractMatrix{T},  N0::AbstractMatrix{T}, Tmp::AbstractMatrix{T}) where T <: AbstractFloat
+    mul!(Tmp, Transpose(KDKinf), N0)
+    mul!(D, Tmp, KDKinf)
+end
+
 # epsilon_h = H*(iF_t*v_t - K(DK)_t'*r_t) (DK 4.69)
 function get_epsilonh!(epsilon::AbstractVector{U}, H::AbstractMatrix{U},
                        iFv::AbstractVector{U}, K::AbstractMatrix{U},
@@ -65,6 +80,14 @@ function get_epsilonh!(epsilon::AbstractVector{U}, H::AbstractMatrix{U},
     mul!(epsilon, H, tmp1)
 end
 
+# epsilon_t = -H_t*KDKinf'*r0_t         (DK p. 135)
+function get_epsilonh!(epsilon::AbstractVector{T}, H::AbstractMatrix{T},
+                       KDKinf::AbstractMatrix{T}, r0::AbstractVector{T},
+                       tmp::AbstractVector{T}) where T <: AbstractFloat
+    mul!(tmp, transpose(KDKinf), r0)
+    mul!(epsilon, H, tmp, -1.0, 0.0)
+end
+
 # etah = Q*R'*r_t
 function get_etah!(etah::AbstractVector{T}, Q::AbstractMatrix{T},
                    R::AbstractMatrix{T}, r::AbstractVector{T},
@@ -144,6 +167,15 @@ function get_iFZ!(iFZ::AbstractArray{T}, cholF::AbstractArray{T}, Z::AbstractArr
     LAPACK.potrs!('U', cholF, iFZ)
 end
 
+function get_iFZ!(iFZ::AbstractArray{T}, cholF::AbstractArray{T}, z::AbstractVector{U}) where {T <: AbstractFloat, U <: Integer}
+    n = length(z)
+    fill!(iFZ, 0.0)
+    @inbounds @simd for i = 1:n
+        iFZ[i, z[i]] = 1.0
+    end
+    LAPACK.potrs!('U', cholF, iFZ)
+end
+
 # K = iF*Z*P
 function get_K!(K::AbstractArray{T}, ZP::AbstractArray{T}, cholF::AbstractArray{T}) where T <: AbstractFloat
     copy!(K, ZP)
@@ -163,35 +195,52 @@ function get_Kstar!(Kstar::AbstractArray{T}, z::AbstractVector{U}, Pstar::Abstra
 end
 
 # L_t = T - K(DK)_t*Z (DK 4.29)
-function get_L!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, Z::AbstractArray{U}, L1::AbstractArray{U}) where U <: AbstractFloat
+function get_L!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, Z::AbstractArray{U}) where U <: AbstractFloat
+    copy!(L, T)
+    mul!(L, K, Z, -1.0, 1.0)
+end
+
+# L_t = T - K(DK)_t*z (DK 4.29)
+function get_L!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, z::AbstractArray{W}) where {U <: AbstractFloat, W <: Integer}
     copy!(L, T)
-    gemm!('N', 'N', -1.0, K, Z, 1.0, L)
+    for j = 1:length(z)
+        zj = z[j]
+        for i = 1:size(L, 1)
+            L[i, zj] += K[i, j]
+        end
+    end
 end
 
 # L = T(I - K'*Z)
-function get_L_alternative!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, Z::AbstractArray{U}, L1::AbstractArray{U}) where U <: AbstractFloat
-    fill!(L1, 0.0)
-    @inbounds @simd for i = 1:size(L1, 1)
-        L1[i,i] = 1.0
+function get_L_alternative!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, Z::AbstractArray{U}, Tmp::AbstractArray{U}) where U <: AbstractFloat
+    fill!(Tmp, 0.0)
+    @inbounds @simd for i = 1:size(Tmp, 1)
+        Tmp[i,i] = 1.0
     end
-    gemm!('T', 'N', -1.0, K, Z, 1.0, L1)
-    mul!(L, T, L1)
+    mul!(Tmp, transpose(K), Z, -1.0, 1.0)
+    mul!(L, T, Tmp)
 end
 
 # L = T(I - K'*z)
-function get_L!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, z::AbstractArray{W}, L1::AbstractArray{U}) where {U <: AbstractFloat, W <: Integer}
+function get_L!(L::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, z::AbstractArray{W}, Tmp::AbstractArray{U}) where {U <: AbstractFloat, W <: Integer}
     m, n = size(K)
-    fill!(L1, 0.0)
+    fill!(Tmp, 0.0)
     @inbounds for j = 1:m
         zj = z[j]
         @simd for k=1:n
-            L1[k, zj] = -K[j, k]
+            Tmp[k, zj] = -K[j, k]
         end
     end
     @inbounds @simd for i = 1:n
-        L1[i,i] += 1.0
+        Tmp[i,i] += 1.0
     end
-    mul!(L, T, L1)
+    mul!(L, T, Tmp)
+end
+
+# L1_t = - KDK*Z (DK 5.12)
+function get_L1!(L1::AbstractMatrix{T}, KDK::AbstractMatrix{T}, Zsmall::AbstractVector{U}) where {T <: AbstractFloat, U <: Integer}
+    vL1 = view(L1, :, Zsmall)
+    vL1 .= -KDK
 end
 
 function get_M!(y::AbstractArray{T}, x::AbstractArray{T}, work::AbstractArray{T}) where T <: AbstractFloat
@@ -328,6 +377,23 @@ function get_Valpha!(V::AbstractArray{T}, P::AbstractArray{T}, N::AbstractArray{
     gemm!('N', 'N', -1.0, Ptmp, P, 1.0, V)
 end
 
+# Valpha_t = Pstar_t - Pstar_t*N0_{t-1}*Pstar_t
+#            -(Pinf_t*N1_{t-1}*Pstar_t)' -Pinf_t*N1_{t-1}*Pstar_t -
+#            -Pinf_t*N2_{t-1}*Pinf_t                      (DK 5.30)
+function get_Valpha!(Valpha::AbstractMatrix{T},
+                     Pstar::AbstractMatrix{T},
+                     Pinf::AbstractMatrix{T}, N0::AbstractMatrix{T},
+                     N1::AbstractMatrix{T}, N2::AbstractMatrix{T},
+                     Tmp::AbstractMatrix{T}) where T <: AbstractFloat
+    copy!(Valpha, Pstar)
+    mul!(Tmp, Pstar, N0)
+    mul!(Tmp, Pinf, N1, 1.0, 1.0)
+    mul!(Valpha, Tmp, Pstar, -1.0, 1.0)
+    mul!(Tmp, Pstar, N1)
+    mul!(Tmp, Pinf, N2, 1.0, 1.0)
+    mul!(Valpha, Tmp, Pinf, -1.0, 1.0)
+end
+
 # Vepsilon_t = H - H*D_t*H
 function get_Vepsilon!(Vepsilon::AbstractArray{T}, H::AbstractArray{T}, D::AbstractArray{T}, tmp::AbstractArray{T}) where T <: AbstractFloat
     copy!(Vepsilon, H)
@@ -405,7 +471,7 @@ end
 # a = d + T*att
 function update_a!(a1::AbstractVector{U}, d::AbstractVector{U}, T::AbstractMatrix{U}, att::AbstractVector{U}) where U <: AbstractFloat
     copy!(a1, d)
-    mul!(a1, T, att)
+    mul!(a1, T, att, 1.0, 1.0)
 end
 
 # a = d + a + K'*v
@@ -442,15 +508,62 @@ end
 function update_N!(N1::AbstractArray{T}, Z::AbstractArray{T}, iFZ::AbstractArray{T}, L::AbstractArray{T}, N::AbstractArray{T}, Ptmp::AbstractArray{T}) where T <: AbstractFloat
     mul!(N1, transpose(Z), iFZ)
     mul!(Ptmp, transpose(L), N)
-    gemm!('N', 'N', 1.0, Ptmp, L, 1.0, N1)
+    mul!(N1, Ptmp, L, 1.0, 1.0)
 end
 
-function update_N!(N1::AbstractArray{T}, z::AbstractArray{U}, iFZ::AbstractArray{T}, L::AbstractArray{T}, N::AbstractArray{T}, Ptmp::AbstractArray{T}) where {T <: AbstractFloat, U <: Integer}
+function update_N!(N1::AbstractArray{T}, z::AbstractVector{U}, iFZ::AbstractArray{T}, L::AbstractArray{T}, N::AbstractArray{T}, Tmp::AbstractArray{T}) where {T <: AbstractFloat, U <: Integer}
     fill!(N1, 0.0)
     vN1 = view(N1, z, :)
     vN1 .= iFZ
-    mul!(Ptmp, transpose(L), N)
-    gemm!('N', 'N', 1.0, Ptmp, L, 1.0, N1)
+    mul!(Tmp, transpose(L), N)
+    mul!(N1, Tmp, L, 1.0, 1.0)
+end
+
+# N0_{t-1} = L0_t'N0_t*L0_t (DK 5.29)
+function update_N0!(N0, L0, N0_1, PTmp)
+    mul!(PTmp, transpose(L0), N0_1)
+    mul!(N0, PTmp, L0)
+end
+
+# F1 = inv(Finf)
+# N1_{t-1} = Z'*F1*Z + L0'*N1_t*L0 + L1'*N0_t*L0
+function update_N1!(N1::AbstractMatrix{T}, Z::AbstractMatrix{T},
+                    iFZ::AbstractMatrix{T}, L0::AbstractMatrix{T},
+                    N1_1::AbstractMatrix{T}, L1::AbstractMatrix{T},
+                    N0_1::AbstractMatrix{T}, PTmp::AbstractMatrix{T}) where T <: AbstractFloat
+    mul!(N1, transpose(Z), iFZ)
+    mul!(PTmp, transpose(L0), N1_1)
+    mul!(PTmp, transpose(L1), N0_1, 1.0, 1.0)
+    mul!(N1, PTmp, L0, 1.0, 1.0)
+end
+
+function update_N1!(N1::AbstractMatrix{T}, z::AbstractVector{U},
+                    iFZ::AbstractMatrix{T}, L0::AbstractMatrix{T},
+                    N1_1::AbstractMatrix{T}, L1::AbstractMatrix{T},
+                    N0_1::AbstractMatrix{T}, PTmp::AbstractMatrix{T}) where {T <: AbstractFloat, U <: Integer}
+    vN1 = view(N1, z, :)
+    vN1 .= iFZ
+    mul!(PTmp, transpose(L0), N1_1)
+    mul!(PTmp, transpose(L1), N0_1, 1.0, 1.0)
+    mul!(N1, PTmp, L0, 1.0, 1.0)
+end
+
+# F2 = -inv(Finf)*Fstar*inv(Finv)
+# N2_{t-1} = Z'*F2*Z + L0'*N2_t*L0 + L0'*N1_t*L1
+#            + L1'*N1_t*L0 + L1'*N0_t*L1
+function update_N2!(N2::AbstractMatrix{T}, iFZ::AbstractMatrix{T},
+                    Fstar::AbstractMatrix{T}, L0::AbstractMatrix{T},
+                    N2_1::AbstractMatrix{T}, N1_1::AbstractMatrix{T},
+                    L1::AbstractMatrix{T}, N0_1::AbstractMatrix{T},
+                    Tmp1::AbstractMatrix{T}, Tmp2::AbstractMatrix{T}) where T <: AbstractFloat
+    mul!(Tmp1, Fstar, iFZ) 
+    mul!(N2, transpose(iFZ), Tmp1, -1.0, 0.0)
+    mul!(Tmp2, transpose(L0), N2_1)
+    mul!(Tmp2, transpose(L1), N1_1, 1.0, 1.0)
+    mul!(N2, Tmp2, L0, 1.0, 1.0)
+    mul!(Tmp2, transpose(L0), N1_1)
+    mul!(Tmp2, transpose(L1), N0_1, 1.0, 1.0)
+    mul!(N2, Tmp2, L1, 1.0, 1.0)
 end
 
 # P = T*(P - K'*Z*P)*T'+ QQ
@@ -480,7 +593,7 @@ function update_P!(P::AbstractMatrix{U}, T::AbstractMatrix{U}, Ptt::AbstractMatr
     mul!(P, T, Ptmp)
 end
 
-# Pstar  = T*(Pstar-Pstar*Z'*Kinf-Pinf*Z'*Kstar)*T'+QQ;         %(5.14) DK(2012)
+# Pstar  = T*(Pstar-Pstar*Z'*Kinf-Pinf*Z'*Kstar)*T'+QQ         (5.14) DK(2012)
 function update_Pstar!(Pstar, T, ZPinf, ZPstar, Kinf, Kstar, QQ, PTmp)
     copy!(PTmp, Pstar)
     mul!(PTmp, transpose(ZPstar), Kinf, -1.0, 1.0)
@@ -491,7 +604,7 @@ function update_Pstar!(Pstar, T, ZPinf, ZPstar, Kinf, Kstar, QQ, PTmp)
     mul!(Pstar, PTmp, transpose(T), 1.0, 1.0)
 end
 
-# Pinf   = T*(Pinf-Pinf*Z'*Kinf)*T';                             %(5.14) DK(2012)
+# Pinf   = T*(Pinf-Pinf*Z'*Kinf)*T'                             (5.14) DK(2012)
 function update_Pinf!(Pinf, T, ZPinf, Kinf, PTmp)
     mul!(Pinf, transpose(ZPinf), Kinf, -1.0, 1.0) 
     mul!(PTmp, T, Pinf)
@@ -512,6 +625,27 @@ function update_r!(r1::AbstractVector{T}, z::AbstractVector{U}, iFv::AbstractVec
     gemm!('T', 'N', 1.0, L, r, 1.0, r1)
 end
 
+# rstar_{t-1} = Z_t'*iFinf_t*v_t + Linf_t'rstar_t + Lstar_t'*rinf_t      (DK 5.21)
+function update_r!(rstar::AbstractVector{T}, Z::AbstractMatrix{T},
+                   iFv::AbstractVector{T}, Linf::AbstractMatrix{T},
+                   rstar1::AbstractVector{T}, Lstar::AbstractMatrix{T},
+                   rinf1::AbstractVector{T}) where T <: AbstractFloat
+    mul!(rstar, transpose(Z), iFv)
+    mul!(rstar, transpose(Linf), rstar1, 1.0, 1.0)
+    mul!(rstar, transpose(Lstar), rinf1, 1.0, 1.0)
+end
+
+# rstar_{t-1} = z_t'*iFinf_t*v_t + Linf_t'rstar_t + Lstar_t'*rinf_t      (DK 5.21)
+function update_r!(rstar::AbstractVector{T}, z::AbstractVector{U},
+                   iFv::AbstractVector{T}, Linf::AbstractMatrix{T},
+                   rstar1::AbstractVector{T}, Lstar::AbstractMatrix{T},
+                   rinf1::AbstractVector{T}) where {T <: AbstractFloat, U <: Integer}
+    vrstar = view(rstar, z)
+    vrstar .= iFv
+    mul!(rstar, transpose(Linf), rstar1, 1.0, 1.0)
+    mul!(rstar, transpose(Lstar), rinf1, 1.0, 1.0)
+end
+
 # W = T(W - K'*iF*Z*W)
 function update_W!(W::AbstractArray{U}, ZW::AbstractArray{U}, cholF::AbstractArray{U}, T::AbstractArray{U}, K::AbstractArray{U}, iFZW::AbstractArray{U}, KtiFZW::AbstractArray{U}) where U <: AbstractFloat
     copy!(iFZW, ZW)

src/kalman_filter.jl

@@ -112,7 +112,6 @@ function kalman_filter!(Y::AbstractArray{U},
                        ws::KalmanWs,
                        data_pattern::Vector{Vector{V}}) where {U <: AbstractFloat, W <: Real, V <: Integer}
     changeC = ndims(c) > 1
-    changeZ = ndims(Z) > 2
     changeH = ndims(H) > 2
     changeD = ndims(d) > 1
     changeT = ndims(T) > 2
@@ -142,8 +141,7 @@ function kalman_filter!(Y::AbstractArray{U},
         ndata = length(pattern)
         vc = changeC ? view(c, :, t) : view(c, :)
         ws.csmall .= view(vc, pattern)
-        vZ = changeZ ? view(Z, :, :, t) : view(Z, :, :)
-        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, vZ, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny, t)
         vH = changeH ? view(H, :, :, t) : view(H, :, :)
         vT = changeT ? view(T, :, :, t) : view(T, :, :)
         vR = changeR ? view(R, :, :, t) : view(R, :, :)
@@ -232,7 +230,7 @@ struct DiffuseKalmanFilterWs{T, U} <: KalmanWs{T, U}
     iF::Matrix{T}
     iFv::Matrix{T}
     a1::Vector{T}
-    cholF::Matrix{T}
+    cholF::Array{T}
     cholH::Matrix{T}
     ZP::Matrix{T}
     Fstar::Matrix{T}
@@ -264,7 +262,7 @@ struct DiffuseKalmanFilterWs{T, U} <: KalmanWs{T, U}
         iF = Matrix{T}(undef, ny,ny )
         iFv = Matrix{T}(undef, ny, nobs)
         a1 = Vector{T}(undef, ns)
-        cholF = Matrix{T}(undef, ny, ny)
+        cholF = Array{T}(undef, ny, ny, nobs)
         cholH = Matrix{T}(undef, ny, ny)
         ZP = Matrix{T}(undef, ny, ns)
         Fstar = Matrix{T}(undef, ny, ny)
@@ -309,7 +307,7 @@ function diffuse_kalman_filter_init!(Y::AbstractArray{U},
                                      last::V,
                                      presample::V,
                                      tol::U,
-                                     ws::DiffuseKalmanFilterWs,
+                                     ws::KalmanWs,
                                      data_pattern::Vector{Vector{V}}) where {U <: AbstractFloat,
                                                                              V <: Integer,
                                                                              W <: Real}
@@ -338,7 +336,7 @@ function diffuse_kalman_filter_init!(Y::AbstractArray{U},
         ndata = length(pattern)
         vc = changeC ? view(c, :, t) : view(c, :)
         ws.csmall .= view(vc, pattern)
-        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny, t)
         vH = changeH ? view(H, :, :, t) : view(H, :, :)
         vT = changeT ? view(T, :, :, t) : view(T, :, :)
         vR = changeR ? view(R, :, :, t) : view(R, :, :)
@@ -358,8 +356,7 @@ function diffuse_kalman_filter_init!(Y::AbstractArray{U},
         vPstar1 = changePstar ? view(Pstar, :, :, t + 1) : view(Pstar, :, :)
         vv = view(ws.v, 1:ndata)
         vvH = view(vH, pattern, pattern)
-        vZP = view(ws.ZP, 1:ndata, :)
-        vcholF = view(ws.cholF, 1:ndata, 1:ndata, 1)
+        vcholF = view(ws.cholF, 1:ndata, 1:ndata, t)
         viFv = view(ws.iFv, 1:ndata)
         vFinf = view(ws.F, 1:ndata, 1:ndata)
         vFstar = view(ws.Fstar, 1:ndata, 1:ndata)
@@ -389,7 +386,7 @@ function diffuse_kalman_filter_init!(Y::AbstractArray{U},
         else
             ws.lik[t] = ndata*l2pi + log(det_from_cholesky(vcholF))
             # Kinf   = iFinf*Z*Pinf                                   %define Kinf'=T^{-1}*K_0 with M_{\infty}=Pinf*Z'
-            copy!(vKinf, vZP)
+            copy!(vKinf, vZPinf)
             LAPACK.potrs!('U', vcholF, vKinf)
             # Fstar  = Z*Pstar*Z' + H;                                        %(5.7) DK(2012)
             get_F!(vFstar, vZPstar, vZsmall, vPstar, vH)
@@ -433,7 +430,7 @@ function diffuse_kalman_filter!(Y::AbstractArray{U},
                                 last::V,
                                 presample::V,
                                 tol::U,
-                                ws::DiffuseKalmanFilterWs,
+                                ws::KalmanWs,
                                 data_pattern::Vector{Vector{V}}) where {U <: AbstractFloat,
                                                                         V <: Integer,
                                                                         W <: Real}
@@ -465,7 +462,7 @@ function diffuse_kalman_filter!(Y::AbstractArray{U},
                                 last::V,
                                 presample::V,
                                 tol::U,
-                                ws::DiffuseKalmanFilterWs) where {U <: AbstractFloat,
+                                ws::KalmanWs) where {U <: AbstractFloat,
                                                                         V <: Integer,
                                                                         W <: Real}
 

test/test_KalmanFilterTools.jl

@@ -289,7 +289,7 @@ end
 
 full_data_pattern = [collect(1:ny) for o = 1:nobs]
 
-@testset "Diffuse Kalman Filter" begin
+@testset "Diffuse Kalman Filter and Smoother" begin
     ws4 = DiffuseKalmanLikelihoodWs{Float64, Int64}(ny, ns, np, nobs)
     ws5 = DiffuseKalmanFilterWs{Float64, Int64}(ny, ns, np, nobs)
     
@@ -409,6 +409,31 @@ full_data_pattern = [collect(1:ny) for o = 1:nobs]
     Pstar = copy(Pstar_0)
     llk_5 = diffuse_kalman_likelihood(Y, z, H, T, R, Q, a, Pinf, Pstar, 1, nobs, 0, 1e-8, ws4, full_data_pattern)
     @test llk_5 ≈ llk_4 
+
+    ws6 = DiffuseKalmanSmootherWs{Float64, Int64}(ny, ns, np, nobs)
+
+    a = copy(a_0)
+    Pinf = zeros(ns, ns, nobs + 1)
+    Pinftt = zeros(ns, ns, nobs + 1)
+    Pstar = zeros(ns, ns, nobs + 1)
+    Pstartt = zeros(ns, ns, nobs + 1)
+    Pinf[:, :, 1] = Pinf_0
+    Pinftt[:, :, 1] = Pinf_0
+    Pstar[:, :, 1] = Pstar_0
+    Pstartt[:, :, 1] =  Pstar_0
+    alphah = zeros(ns, nobs)
+    epsilonh = zeros(ny, nobs)
+    etah = zeros(np, nobs)
+    Valphah = zeros(ns, ns, nobs)
+    Vepsilonh = zeros(ny, ny, nobs)
+    Vetah = zeros(np, np, nobs)
+    llk_6 = diffuse_kalman_smoother!(Y, c, z, H, d, T, R, Q, aa, att,
+                                     Pinf, Pinftt, Pstar, Pstartt,
+                                     alphah, epsilonh, etah, Valphah,
+                                     Vepsilonh, Vetah, 1, nobs, 0,
+                                     1e-8, ws6)
+    @test llk_6 ≈ llk_4 
+
 end
 
 @testset "start and last" begin

test/test_kalman_base.jl

@@ -22,6 +22,16 @@ r1 = randn(ns)
 KalmanFilterTools.get_alphah!(alphah, a, P, r1)
 @test alphah ≈ a + P*r1
 
+# alphah_t = a_t + Pstar_t*r0_{t-1} + Pinf_t*r1_{t-1}     (DK 5.24)
+alphah = randn(ns)
+a = randn(ns)
+Pstar = randn(ns, ns)
+Pinf = randn(ns, ns)
+r0 = randn(ns)
+r1 = randn(ns)
+KalmanFilterTools.get_alphah!(alphah, a, Pstar, Pinf, r0, r1)
+@test alphah ≈ a + Pstar*r0 + Pinf*r1
+
 F = randn(ny, ny)
 F = F'*F
 cholF = zeros(ny, ny)
@@ -42,6 +52,14 @@ tmp = randn(ny, ns)
 KalmanFilterTools.get_D!(D, cholF.U, K, T, N, KT, tmp)
 @test D ≈ inv(F) + K*T*N*T'*K'
 
+# D_t = KDKinf_t'*N0_t*KDKinf_t    (DK p. 135)
+D = randn(ny, ny)
+KDKinf = randn(ns, ny)
+N0 = randn(ns, ns)
+Tmp = randn(ny, ns)
+KalmanFilterTools.get_D!(D, KDKinf,  N0, Tmp)
+@test D ≈ KDKinf'*N0*KDKinf
+
 # epsilonh_t = H*(iF_t*v_t - K_t*T*r_t)
 epsilon = randn(ny)
 H = randn(ny, ny)
@@ -55,6 +73,14 @@ tmp2 = zeros(ns)
 KalmanFilterTools.get_epsilonh!(epsilon, H, iFv, K, T, r, tmp1, tmp2)
 @test epsilon ≈ H*(iFv - K*T*r)
 
+# epsilon_t = -H_t*KDKinf*r0_t         (DK p. 135)
+H = randn(ny, ny)
+KDKinf = randn(ns, ny)
+r0 = randn(ns)
+tmp = randn(ny)
+KalmanFilterTools.get_epsilonh!(epsilon, H, KDKinf, r0, tmp)
+@test epsilon ≈ - H*KDKinf'*r0
+
 # etah = Q*R'*r_t
 eta = randn(np)
 Q = randn(np, np)
@@ -91,6 +117,29 @@ KalmanFilterTools.get_F!(F, ZP, Z, Pinf)
 @test F ≈ Z*Pinf*Z'
 @test ZP ≈ Z*Pinf
 
+# iFZ = inv(F)*Z
+Z = randn(2, 5)
+F = randn(2, 2)
+F = F'*F
+iFZ = randn(2, 5)
+cholF = similar(F)
+KalmanFilterTools.get_cholF!(cholF, F)
+KalmanFilterTools.get_iFZ!(iFZ, cholF, Z)
+@test iFZ ≈ F\Z
+
+# iFZ = inv(F)*z
+z = [2, 4]
+F = randn(2, 2)
+F = F'*F
+iFZ = randn(2, 5)
+cholF = similar(F)
+KalmanFilterTools.get_cholF!(cholF, F)
+KalmanFilterTools.get_iFZ!(iFZ, cholF, z)
+Z = zeros(2, 5)
+Z[1, 2] = 1.0
+Z[2, 4] = 1.0
+@test iFZ ≈ F\Z
+
 # iFv = inv(F)*v
 F = randn(ny, ny)
 F = F'*F
@@ -158,7 +207,7 @@ Ptt = randn(ns, ns)
 P - randn(ns, ns)
 K = randn(ny, ns)
 ZP = randn(ny, ns)
-get_updated_Ptt!(Ptt, P, K, ZP)
+KalmanFilterTools.get_updated_Ptt!(Ptt, P, K, ZP)
 @test Ptt ≈ P - transpose(K)*ZP
 
 # Pstartt = Pstar-Pstar*Z'*Kinf-Pinf*Z'*Kstar                           %(5.14) DK(2012)
@@ -168,7 +217,9 @@ ZPstar = randn(ny, ns)
 Kinf = randn(ny, ns)
 ZPinf = randn(ny, ns)
 Kstar = randn(ny, ns)
-KalmanFilterTools.get_updated_Pstartt!(Pstartt, Pstar, ZPstar, Kinf, ZPinf, Kstar)
+Pinftt = randn(ns, ns)
+PTmp = randn(ns, ns)
+KalmanFilterTools.get_updated_Pstartt!(Pstartt, Pstar, ZPstar, Kinf, ZPinf, Kstar, Pinftt, PTmp)
 @test Pstartt ≈ Pstar - transpose(ZPstar)*Kinf - transpose(ZPinf)*Kstar
 
 # V_t = P_t - P_t*N_{t-1}*P_t
@@ -180,6 +231,27 @@ Ptmp = Matrix{Float64}(undef, ns, ns)
 KalmanFilterTools.get_Valpha!(V, P, N1, Ptmp)
 @test V ≈ P - P*N1*P
 
+# Valpha_t = Pstar_t - Pstar_t*N0_{t-1}*Pstar_t
+#            -(Pinf_t*N1_{t-1}*Pstar_t)'
+#            -Pinf_t*N1_{t-1}*Pstar_t
+#            -Pinf_t*N2_{t-1}*Pinf_t                       (DK 5.30)
+Valpha = randn(ns, ns)
+Pstar = randn(ns, ns)
+Pstar = Pstar'*Pstar
+Pinf = randn(ns, ns)
+Pinf = Pinf'*Pinf
+N0 = randn(ns, ns)
+N0 = N0'*N0
+N1 = randn(ns, ns)
+N1 = N1'*N1
+N2 = randn(ns, ns)
+N2 = N2'*N2
+Tmp = randn(ns, ns)
+KalmanFilterTools.get_Valpha!(Valpha, Pstar, Pinf,N0, N1, N2, Tmp)
+@test Valpha ≈ (Pstar - Pstar*N0*Pstar - (Pinf*N1*Pstar)'
+                - Pinf*N1*Pstar - Pinf*N2*Pinf)
+
+# QQ = R*Q*R'
 R = randn(ns, np)
 Q = randn(np, np)
 Q = Q'*Q
@@ -301,12 +373,62 @@ KalmanFilterTools.update_N!(N1, Z, iFZ, L, N, Ptmp)
 @test N1 ≈ Z'*iFZ + L'*N*L
 
 iFZ1 = iFZ[z, :]
+KalmanFilterTools.update_N!(N1, z, iFZ1, L, N, Ptmp)
 ZiFZ = zeros(ns, ns)
 ZiFZ[z, :] .= iFZ1
-KalmanFilterTools.update_N!(N1, z, iFZ1, L, N, Ptmp)
+Z = zeros(ny, ns)
+for i=1:length(z)
+    Z[i, z[i]] = 1.0
+end
+@test ZiFZ ≈ Z'iFZ1
 @test N1 ≈ ZiFZ + L'*N*L
 
 
+# N0_{t-1} = L0_t'N0_t*L0_t (DK 5.29)
+N0 = randn(ns, ns)
+L0 = randn(ns, ns)
+N0_1 = randn(ns, ns)
+PTmp = randn(ns, ns)
+KalmanFilterTools.update_N0!(N0, L0, N0_1, PTmp)
+@test N0 ≈ L0'*N0_1*L0
+
+# F1 = inv(Finf)
+# N1_{t-1} = Z'*F1*Z + L0'*N1_t*L0 + L1'*N0_t*L0
+N1 = randn(ns, ns)
+Z = randn(ny, ns)
+F = randn(ny, ny)
+F = 0.5*(F + F')
+iFZ = F\Z
+L0 = randn(ns, ns)
+N1_1 = randn(ns, ns)
+L1 = randn(ns, ns)
+N0_1 = randn(ns, ns)
+PTmp = randn(ns, ns)
+KalmanFilterTools.update_N1!(N1, Z, iFZ, L0, N1_1, L1, N0_1, PTmp)
+@test iFZ ≈ inv(F)*Z
+@test N1 ≈ Z'*inv(F)*Z + L0'*N1_1*L0 + L1'*N0_1*L0
+
+# F2 = -inv(Finf)*Fstar*inv(Finv)
+# N2_{t-1} = Z'*F2*Z + L0'*N2_t*L0 + L0'*N1_t*L1
+#            + L1'*N1_t*L0 + L1'*N0_t*L1
+N2 = randn(ns, ns)
+Z = randn(ny, ns)
+F = randn(ny, ny)
+F = 0.5*(F + F')
+iFZ = F\Z
+Fstar = randn(ny, ny)
+L0 = randn(ns, ns)
+N0_1 = randn(ns, ns)
+N1_1 = randn(ns, ns)
+N2_1 = randn(ns, ns)
+Tmp1 = randn(ny, ns)
+Tmp2 = randn(ns, ns)
+KalmanFilterTools.update_N2!(N2, iFZ, Fstar, L0, N2_1, N1_1,
+                             L1, N0_1, Tmp1, Tmp2)