Maybe it is good to have a function to just update the Kalman filter. My sketch:
function update_kalman_filter(model::StateSpaceModel{Typ}, filter::FilterOutput, model_test::StateSpaceModel{Typ}; tol::Typ = Typ(1e-5)) where Typ
time_invariant = model.mode == "time-invariant"
# Predictive state and its covariance
a = cat(filter.a, Matrix{Typ}(undef, model_test.dim.n, model.dim.m), dims=1)
P = cat(filter.P, Array{Typ, 3}(undef, model.dim.m, model.dim.m, model_test.dim.n+1), dims=3)
att = cat(filter.att, Matrix{Typ}(undef, model_test.dim.n, model.dim.m), dims=1)
Ptt = cat(filter.Ptt, Array{Typ, 3}(undef, model.dim.m, model.dim.m, model_test.dim.n), dims=3)
# Innovation and its covariance
v = cat(filter.v, Matrix{Typ}(undef, model_test.dim.n, model.dim.p), dims=1)
F = cat(filter.F, Array{Typ, 3}(undef, model.dim.p, model.dim.p, model_test.dim.n) , dims=3)
# Kalman gain
K = Array{Typ, 3}(undef, model.dim.m, model.dim.p, model.dim.n + model_test.dim.n)
# Auxiliary structures
ZP = Array{Typ, 2}(undef, model.dim.p, model.dim.m)
P_Ztransp_invF = Array{Typ, 2}(undef, model.dim.m, model.dim.p)
RQR = Array{Typ, 2}(undef, model.dim.m, model.dim.m)
# Steady state initialization
steadystate = filter.steadystate
tsteady = filter.tsteady
# Initial state: big Kappa initialization
StateSpaceModels.fill_a1!(a)
StateSpaceModels.fill_P1!(P; bigkappa = Typ(1e6))
y = vcat(model.y, model_test.y)
d = vcat(model.d, model_test.d)
Z = cat(model.Z, model_test.Z,dims=3)
c = vcat(model.c, model_test.c)
StateSpaceModels.update_ZP!(ZP, Z, P, filter.tsteady) # ZP = Z[:, :, t] * P[:, :, t]
StateSpaceModels.update_F!(F, ZP, Z, model.H, filter.tsteady) # F_t = Z_t * P_t * Z_t + H
StateSpaceModels.update_P_Ztransp_Finv!(P_Ztransp_invF, ZP, F, filter.tsteady) # P_Ztransp_invF = ZP' * invF(F, t)
StateSpaceModels.update_K!(K, P_Ztransp_invF, model.T, filter.tsteady) # K_t = T * P_t * Z_t * F^-1_t
for t_ = 1:model_test.dim.n
t = model.dim.n+t_-1
if t in model_test.missing_observations
steadystate = false
v[t, :] = fill(NaN, model.dim.p)
F[:, :, t] = fill(NaN, (model.dim.p, model.dim.p))
StateSpaceModels.update_att!(att, a, t) # attt = a_t
StateSpaceModels.update_Ptt!(Ptt, P, t) # Pttt = P_t
StateSpaceModels.update_a!(a, att, model.T, c, t) # a_t+1 = T * attt + c_t
StateSpaceModels.update_P!(P, model.T, Ptt, RQR, t) # P_t+1 = T * Pttt * T' + RQR'
elseif steadystate
StateSpaceModels.update_v!(v, y, Z, d, a, t) # v_t = y_t - Z_t * a_t - d_t
StateSpaceModels.repeat_matrix_t_plus_1!(F, t-1) # F[:, :, t] = F[:, :, t-1]
StateSpaceModels.repeat_matrix_t_plus_1!(K, t-1) # K[:, :, t] = K[:, :, t-1]
StateSpaceModels.update_att!(att, a, P_Ztransp_invF, v, t) # attt = a_t + P_t * Z_t * F^-1_t * v_t
StateSpaceModels.repeat_matrix_t_plus_1!(Ptt, t-1) # Pttt = Pttt-1
StateSpaceModels.update_a!(a, att, model.T, c, t) # a_t+1 = T * attt + c_t
StateSpaceModels.repeat_matrix_t_plus_1!(P, t) # P__+1 = P_t
else
StateSpaceModels.update_v!(v, y, Z, d, a, t) # v_t = y_t - Z_t * a_t - d_t
StateSpaceModels.update_ZP!(ZP, Z, P, t) # ZP = Z[:, :, t] * P[:, :, t]
StateSpaceModels.update_F!(F, ZP, Z, model.H, t) # F_t = Z_t * P_t * Z_t + H
StateSpaceModels.update_P_Ztransp_Finv!(P_Ztransp_invF, ZP, F, t) # P_Ztransp_invF = ZP' * invF(F, t)
StateSpaceModels.update_K!(K, P_Ztransp_invF, model.T, t) # K_t = T * P_t * Z_t * F^-1_t
StateSpaceModels.update_att!(att, a, P_Ztransp_invF, v, t) # attt = a_t + P_t * Z_t * F^-1_t * v_t
StateSpaceModels.update_Ptt!(Ptt, P, P_Ztransp_invF, ZP, t) # Pttt = P_t - P_t * Z_t' * F^-1_t * Z_t * P_t
StateSpaceModels.update_a!(a, att, model.T, c, t) # a_t+1 = T * attt + c_t
StateSpaceModels.update_P!(P, model.T, Ptt, RQR, t) # P_t+1 = T * Pttt * T' + RQR'
if time_invariant && StateSpaceModels.check_steady_state(P, t, tol)
steadystate = true
tsteady = t
end
end
end
# Return the auxiliary filter structre
return KalmanFilter(a[model.dim.n+1:end,:], att[model.dim.n+1:end,:], v, P, Ptt, F, steadystate, tsteady, K)
end