tidying things up

src/KalmanFilterTools.jl

@@ -88,7 +88,7 @@ struct KalmanLikelihoodWs{T, U} <: KalmanWs{T, U}
     end
 end
 
-KalmanLikelihoodWs(ny, ns, np, nobs) = KalmanLikelihoodWs{Float64, Integer}(ny, ns, np, nobs)
+KalmanLikelihoodWs(ny, ns, np, nobs) = KalmanLikelihoodWs{Float64, Int64}(ny, ns, np, nobs)
 
 # Z can be either a matrix or a selection vector
 function kalman_likelihood(Y::AbstractArray{U},
@@ -147,24 +147,6 @@ function kalman_likelihood(Y::AbstractArray{U},
     LIK
 end
 
-function get_vZsmall(ws::KalmanWs, Z::AbstractMatrix{T}, pattern::Vector{U}, n::U, ny::U) where {T <: AbstractFloat, U <: Integer}
-    vZsmall = view(ws.Zsmall, 1:n, :)
-    if n == ny
-        copyto!(vZsmall, Z)
-    else
-        vZsmall .= view(Z, pattern, :)
-    end
-end
-
-function get_vZsmall(ws::KalmanWs, Z::AbstractVector{U}, pattern::Vector{U}, n::U, ny::U) where {T <: AbstractFloat, U <: Integer}
-    vZsmall = view(ws.iZsmall, 1:n)
-    if n == ny
-        copyto!(vZsmall, Z)
-    else
-        vZsmall .= view(Z, pattern)
-    end
-end
-
 function kalman_likelihood(Y::AbstractArray{U},
                            Z::AbstractArray{W},
                            H::AbstractArray{U},
@@ -199,7 +181,7 @@ function kalman_likelihood(Y::AbstractArray{U},
         vcholH = view(ws.cholF, 1:ndata, 1:ndata)
         viFv = view(ws.iFv, 1:ndata)
         vK = view(ws.K, 1:ndata, :)
-        vZsmall = get_vZsmall(ws, Z, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny)
         
         # v  = Y[:,t] - Z*a
         get_v!(vv, Y, vZsmall, a, t, pattern)
@@ -441,6 +423,8 @@ struct FastKalmanLikelihoodWs{T, U} <: KalmanWs{T, U}
     end
 end
 
+FastKalmanLikelihoodWs(ny, ns, np, nobs) = FastKalmanLikelihoodWs{Float64, Int64}(ny, ns, np, nobs)
+
 """
 K doesn't represent the same matrix as above
 """
@@ -546,7 +530,7 @@ function fast_kalman_likelihood(Y::Matrix{U},
         vZWM = view(ws.ZWM, 1:ndata, 1:ndata)
         viFZWM = view(ws.iFZWM, 1:ndata, 1:ndata)
         vKtiFZW = view(ws.KtiFZW, :, 1:ndata)
-        vZsmall = get_vZsmall(ws, Z, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny)
         
         # v  = Y[:,t] - Z*a
         get_v!(vv, Y, vZsmall, a, t, pattern)
@@ -724,7 +708,7 @@ function diffuse_kalman_likelihood_init!(Y::Matrix{U},
         viFv = view(ws.iFv, 1:ndata)
         vK = view(ws.K, 1:ndata, :)
         vKstar = view(ws.Kstar, 1:ndata, :)
-        vZsmall = get_vZsmall(ws, Z, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, Z, pattern, ndata, ny)
         
         # v  = Y[:,t] - Z*a
         get_v!(vv, Y, vZsmall, a, t, pattern)
@@ -879,7 +863,7 @@ function kalman_filter!(Y::AbstractArray{U},
         vc = changeC ? view(c, :, t) : view(c, :)
         ws.csmall .= view(vc, pattern)
         vZ = changeZ ? view(Z, :, :, t) : view(Z, :, :)
-        get_vZsmall(ws, vZ, pattern, ndata, ny)
+        vZsmall = get_vZsmall(ws.Zsmall, ws.iZsmall, vZ, pattern, ndata, ny)
         vH = changeH ? view(H, :, :, t) : view(H, :, :)
         vT = changeT ? view(T, :, :, t) : view(T, :, :)
         vR = changeR ? view(R, :, :, t) : view(R, :, :)
@@ -901,11 +885,11 @@ function kalman_filter!(Y::AbstractArray{U},
         vK = view(ws.K, 1:ndata, :, 1)
         
         # v  = Y[:,t] - c - Z*a
-        get_v!(vv, Y, vc, ws.Zsmall, va, t, pattern)
+        get_v!(vv, Y, vc, vZsmall, va, t, pattern)
         iy += ny
         if !steady
             # F  = Z*P*Z' + H
-            get_F!(vF, vZP, ws.Zsmall, vP, vvH)
+            get_F!(vF, vZP, vZsmall, vP, vvH)
             info = get_cholF!(vcholF, vF)
             if info != 0
                 # F is near singular
@@ -1023,7 +1007,7 @@ struct KalmanSmootherWs{T, U} <: KalmanWs{T, U}
     end
 end
 
-KalmanSmootherWs(ny, ns, np, nobs) = KalmanSmootherWs{Float64, Integer}(ny, ns, np, nobs)
+KalmanSmootherWs(ny, ns, np, nobs) = KalmanSmootherWs{Float64, Int64}(ny, ns, np, nobs)
 
 function kalman_filter_2!(Y::AbstractArray{U},
                           c::AbstractArray{U},

src/kalman_base.jl

@@ -272,6 +272,24 @@ function get_Veta!(Veta, Q, R, N, RQ, tmp)
     mul!(Veta, RQ', tmp, -1.0, 1.0)
 end
 
+function get_vZsmall(Zsmall::AbstractMatrix{T}, iZsmall::AbstractVector{U}, Z::AbstractMatrix{T}, pattern::AbstractVector{Int64}, n::U, ny::U) where {T <: AbstractFloat, U <: Integer}
+    vZsmall = view(Zsmall, 1:n, :)
+    if n == ny
+        copyto!(vZsmall, Z)
+    else
+        vZsmall .= view(Z, pattern, :)
+    end
+end
+
+function get_vZsmall(Zsmall::AbstractMatrix{T}, iZsmall::AbstractVector{U}, Z::AbstractVector{Int64}, pattern::AbstractVector{Int64}, n::U, ny::U) where {T <: AbstractFloat, U <: Integer}
+    vZsmall = view(iZsmall, 1:n)
+    if n == ny
+        copyto!(vZsmall, Z)
+    else
+        vZsmall .= view(Z, pattern)
+    end
+end
+
 # a = T(a + K'*v)
 function update_a!(a::AbstractVector{U}, K::AbstractMatrix{U}, v::AbstractVector{U}, a1::Vector{U}, T::AbstractMatrix{U}) where U <: AbstractFloat
     copy!(a1, a)

src/univariate_step.jl

@@ -37,7 +37,7 @@ function univariate_step!(Y, t, Z, H, T, RQR, a, P, kalman_tol, ws)
         Zi = view(ws.Zstar, i, :)
         v = get_v!(ws.ystar, ws.Zstar, a, i)
         F = get_F(Zi, P, H[i,i], ws.PZi)
-        if F > kalman_tol
+        if abs(F) > kalman_tol
             a .+= (v/F) .* ws.PZi
             # P = P - PZi*PZi'/F 
             ger!(-1.0/F, ws.PZi, ws.PZi, P) 

test/runtests.jl

 include("test_kalman_base.jl")
+include("test_univariate_step.jl")
 include("test_KalmanFilterTools.jl")

test/test_kalman_base.jl

@@ -64,14 +64,13 @@ tmp = zeros(np)
 KalmanFilterTools.get_etah!(eta, Q, R, r, tmp)
 @test eta ≈ Q*R'*r
 
-Z = randn(ny, ns)
 H = randn(ny, ny)
 H = H'*H
-F = randn(ny, ny)
-ZP = randn(ny, ns)
-P = randn(ns, ns)
+F = zeros(ny, ny)
+ZP = zeros(ny, ns)
+Z = randn(ny, ns)
 KalmanFilterTools.get_F!(F, ZP, Z, P, H)
-@test ZP ≈ Z*P
+@test ZP == Z*P
 @test F ≈ Z*P*Z' + H
 
 # Z as selection matrix
@@ -207,6 +206,23 @@ tmp = zeros(ns, np)
 KalmanFilterTools.get_Veta!(Veta, Q, R, N, RQ, tmp)
 @test Veta ≈ Q - Q*R'*N*R*Q
 
+
+Zsmall = randn(ny, ns)
+iZsmall = rand(Int, ny)
+Z = randn(ny, ns)
+pattern = [1, 3]
+n = length(pattern)
+vZsmall = KalmanFilterTools.get_vZsmall(Zsmall, iZsmall, Z, pattern, n, ny)
+@test all(vZsmall .== Z[pattern, :])
+
+Zsmall = randn(ny, ns)
+iZsmall = rand(Int, ny)
+Z = [4, 3, 6]
+pattern = [1, 3]
+n = length(pattern)
+vZsmall = KalmanFilterTools.get_vZsmall(Zsmall, iZsmall, Z, pattern, n, ny)
+@test all(vZsmall .== Z[pattern])
+
 # a = T(a + K'*iFv)
 a_0 = copy(a)
 a1 = similar(a)
@@ -232,16 +248,19 @@ KalmanFilterTools.update_K!(K, ZWM, W)
 
 # M = M + M*W'*Z'iF*Z*W*M
 M_0 = copy(M)
+Z = randn(ny, ns)
 ZWM = similar(M)
 iFZWM = similar(M)
 KalmanFilterTools.update_M!(M, Z, W, cholF, ZW, ZWM, iFZWM)
 @test M ≈ M_0 + M_0*W'*Z'*inv(F)*Z*W*M_0
 
 # N_{t-1} = Z_t'iF_t*Z_t + L_t'N_t*L_t
+Z = randn(ny, ns)
 N = randn(ns, ns)
 N = N'*N
 N1 = similar(N)
 iFZ = randn(ny, ns)
+L = randn(ns, ns)
 Ptmp = similar(P)
 KalmanFilterTools.update_N!(N1, Z, iFZ, L, N, Ptmp)
 @test N1 ≈ Z'*iFZ + L'*N*L

test/test_univariate_step.jl

@@ -42,7 +42,7 @@ RQR = R*Q*R'
 a = randn(ns)
 P = randn(ns, ns)
 P = P'*P
-kalman_tol = eps()^(3/3)
+kalman_tol = eps()^(2/3)
 
 a0 = copy(a)
 P0 = copy(P)
@@ -64,6 +64,6 @@ lik1a = KalmanFilterTools.kalman_filter!(Y[:,1]', zeros(3), ZZ, H, zeros(ns), TT
 a0 = copy(a)
 P0 = copy(P)
 lik1 = KalmanFilterTools.kalman_likelihood(Y, Z, H, T, R, Q, a0, P0, 1, nobs, 0, ws)
-#@test a1 ≈ a0
-#@test P1 ≈ P0
+@test a1 ≈ a0
+@test P1 ≈ P0
 @test lik0 ≈ ws.lik[1]