#######################################################################################################
# Kalman filter and smoother function based on Koopman and Durbin's KFAS package
# y_t = Z_t * alpha_t + eps_t (observation equation)
# alpha_t+1 = T_t * alpha_t + R_t * eta_t(transition equation)
# Note the different time indexing for transition equation
# The y_t,y_{t-1} stacked vector with parameters vectors reconfigured to output the
# smoothed lag-1 covariances needed by the EM algorithm
#######################################################################################################
MARSSkfas = function( MLEobj, only.logLik=FALSE, return.lag.one=TRUE, return.kfas.model=FALSE ) {
modelObj = MLEobj$marss
control = MLEobj$control
diffuse=modelObj$diffuse
model.dims=attr(modelObj,"model.dims")
n=model.dims$data[1]; TT=model.dims$data[2]; m=model.dims$x[1]
par.1=parmat(MLEobj, t=1)
t.B=matrix(par.1$B,m,m,byrow=TRUE)
#create the YM matrix
YM=matrix(as.numeric(!is.na(modelObj$data)),n,TT)
#Make sure the missing vals in yt are NAs if there are any
yt=modelObj$data
yt[!YM]=as.numeric(NA)
#Stack the y so that we can get the lag-1 covariance smoother from the Kalman filter output
#stack.yt=rbind(yt[,1:TT,drop=FALSE],cbind(NA,yt[,1:(TT-1),drop=FALSE]))
#stack.yt=t(stack.yt)
#KFS needs time going down rows so yt can be properly converted to ts object
yt=t(yt)
#Build the Zt matrix which is n x (m+1) or n x (m+1) x T; (m+1) because A is in Z
if( model.dims$Z[3] == 1 & model.dims$A[3]==1 ){
#not time-varying
Zt=cbind(par.1$Z,par.1$A)
stack.Zt = matrix(0,n,2*(m+1))
stack.Zt[1:n,1:(m+1)]=Zt
}else{
Zt=array(0,dim=c(n,m+1,TT))
stack.Zt=array(0,dim=c(n,2*(m+1),TT))
pars=parmat(MLEobj,c("Z","A"),t=1:TT)
Zt[1:n,1:m,]=pars$Z
Zt[1:n,m+1,]=pars$A
stack.Zt[1:n,1:(m+1),]=Zt
}
#Build the Tt matrix which is (m+1) x (m+1) or (m+1) x (m+1) x T; (m+1) because x has extra row of 1s (for the A)
# Tt=cbind(rbind(parList$B,matrix(0,1,m)),matrix(c(parList$U,1),m+1,1))
if( model.dims$B[3] == 1 & model.dims$U[3]==1 ){
#not time-varying
Tt=cbind(rbind(par.1$B,matrix(0,1,m)),matrix(c(par.1$U,1),m+1,1))
stack.Tt=matrix(0,2*(m+1),2*(m+1))
stack.Tt[1:(m+1),1:(m+1)]=Tt
stack.Tt[(m+2):(2*m+2),1:(m+1)]=diag(1,m+1)
}else{
Tt=array(0,dim=c(m+1,m+1,TT))
stack.Tt=array(0,dim=c(2*(m+1),2*(m+1),TT))
pars=parmat(MLEobj,c("B","U"),t=1:TT)
#pars uses i+1 because in KFAS x(t)=B(t-1)x(t-1) versus MARSS where x(t)=B(t)x(t-1)
#KFAS sets prior on x(1) so recursions start at x(2); thus B(0) never appears in the KFAS recursions
Tt[1:m,1:m,1:(TT-1)]=pars$B[,,2:TT]
Tt[1:m,m+1,1:(TT-1)]=pars$U[,,2:TT]
Tt[m+1,m+1,]=1;
#in KFAS T[,,TT] is not used since x(T)=B(t-1)x(t-1) is the last computation
#Set to 0 since is.SSModel does like NA in any parameters
Tt[,,TT]=0
stack.Tt[(m+2):(2*m+2),1:(m+1),]=diag(1,m+1)
stack.Tt[1:(m+1),1:(m+1),]=Tt
stack.Tt[,,TT]=0 #see comment above re setting TT value
}
#Build the Ht (R) matrix which is n x n or n x n x T
if( model.dims$R[3] == 1 ){
#not time-varying
Ht=par.1$R
}else{
Ht=array(0,dim=c(n,n,TT))
for(i in 1:TT){
Ht[,,i]=parmat(MLEobj,"R",t=i)$R
}
}
#Build the Qt matrix which is m+1 x m+1 or m+1 x m+1 x T; m+1 since x includes extra row of 1 (for A)
if( model.dims$Q[3] == 1 ){
#not time-varying
Qt=matrix(0,m+1,m+1); Qt[1:m,1:m]=par.1$Q
stack.Qt=matrix(0,2*(m+1),2*(m+1))
stack.Qt[1:(m+1),1:(m+1)]=Qt
}else{
#See notes for Tt re differences in parameter indexing for the process equation
Qt=array(0,dim=c(m+1,m+1,TT))
stack.Qt=array(0,dim=c(2*(m+1),2*(m+1),TT))
for(i in 1:(TT-1)){
#i+1 since in KFAS my B(t) equals B(t-1)
Qt[,,i]=matrix(0,m+1,m+1); Qt[1:m,1:m,i]=parmat(MLEobj,"Q",t=i+1)$Q
stack.Qt[1:(m+1),1:(m+1),i]=Qt[,,i]
}
#see comments on setting of T re TT value; TT value never appears in the KFAS recursions
#but is.SSModel check does not like any NAs in the parameters
Qt[,,TT]=0; stack.Qt[,,TT]=0
}
#Build the a1, Rt and P1 matrices
#First compute x10 and V10 if tinitx=0
if(modelObj$tinitx==0){ # Compute needed x_1 | x_0
#B(1),U(1), and Q(1) correct here since equivalent to B(0), etc in KFAS terminology
x00=par.1$x0; V00=par.1$V0
x10 = par.1$B%*%x00 + par.1$U #Shumway and Stoffer treatment of initial states
V10 = par.1$B%*%V00%*%t.B + par.1$Q # eqn 6.20
}else{
x10=par.1$x0; x00=matrix(10,m,1) #dummy; not defined
V10=par.1$V0; V00=diag(0,m) #dummy; not defined
}
a1=rbind(x10,1); stack.a1=rbind(x10,1,x00,1)
Rt=diag(1,m+1); stack.Rt=diag(1,2*(m+1))
if( diffuse ) {
P1inf=matrix(0,m+1,m+1); P1inf[1:m,1:m]=V10;
stack.P1inf=matrix(0,2*(m+1),2*(m+1))
stack.P1inf[1:m,1:m]=V10; stack.P1inf[(m+2):(2*m+1),(m+2):(2*m+1)]=V00
P1=matrix(0,m+1,m+1)
stack.P1=matrix(0,2*(m+1),2*(m+1))
}else{
P1inf=matrix(0,m+1,m+1)
stack.P1inf=matrix(0,2*(m+1),2*(m+1))
P1=matrix(0,m+1,m+1); P1[1:m,1:m]=V10
#P1 is the var-cov matrix for the stacked x1,x0
#it is matrix(c(V10,B*V00,V00*t(B),V00),2,2)
#x1=B*x0+U+w; in the var-cov mat looks like
# E[x1*t(x1)] E[(Bx0+U+w1)*t(x0)] - E[x1]*E[x1] E[(Bx0+U+w1)]E[t(x0)]
# E[x0*t(Bx0+U+w1)] E[x0*t(x0)] E[(Bx0+U+w1)]E[t(x0)] - E[x0]E[t(x0)]
stack.P1=matrix(0,2*(m+1),2*(m+1))
stack.P1[1:m,1:m]=V10; stack.P1[(m+2):(2*m+1),(m+2):(2*m+1)]=V00
stack.P1[1:m,(m+2):(2*m+1)]=par.1$B%*%V00
stack.P1[(m+2):(2*m+1),1:m]=V00%*%t(par.1$B)
}
if(!return.lag.one){
kfas.model=SSModel(yt ~ -1+SSMcustom( Z=Zt, T=Tt, R=Rt, Q=Qt, a1=a1, P1=P1, P1inf=P1inf), H=Ht)
}else{
kfas.model=SSModel(yt ~ -1+SSMcustom( Z=stack.Zt, T=stack.Tt, R=stack.Rt, Q=stack.Qt, a1=stack.a1, P1=stack.P1, P1inf=stack.P1inf), H=Ht)
}
if(only.logLik){
return( list(logLik=logLik(kfas.model)) )
}
ks.out=KFS(kfas.model, simplify=FALSE) #simplify=FALSE so 1.0.0 returns LL
#because 1.0.0 has time down columns instead of across rows
if(packageVersion("KFAS")!="0.9.11"){
ks.out$a=t(ks.out$a)
ks.out$alphahat=t(ks.out$alphahat)
}
VtT = ks.out$V[1:m,1:m,,drop=FALSE]
Vtt1 = ks.out$P[1:m,1:m,1:TT,drop=FALSE]
if(!return.lag.one){ Vtt1T = NULL
}else{ Vtt1T = ks.out$V[1:m,(m+2):(2*m+1),,drop=FALSE] }
#zero out rows cols as needed when R diag = 0
#Check that if any R are 0 then model is solveable
#This works because I do not allow the location of 0s on the diagonal of R or Q to be time-varying
if(n==1){ diag.R=unname(par.1$R) }else{ diag.R = unname(par.1$R)[1 + 0:(n - 1)*(n + 1)] }
if( any(diag.R==0) ){
VtT[abs(VtT)<.Machine$double.eps]=0
Vtt1[abs(Vtt1)<.Machine$double.eps]=0
if(return.lag.one) Vtt1T[abs(Vtt1T)<.Machine$double.eps]=0
}
x10T = ks.out$alphahat[1:m,1,drop=FALSE]
V10T = matrix(VtT[,,1],m,m)
if(modelObj$tinitx==1){
x00=matrix(NA,m,1); V00=matrix(NA,m,m)
x0T=x10T; V0T=V10T
}else{ #modelObj$tinitx==0
Vtt1.1=sub3D(Vtt1,t=1)
Vinv=pcholinv(Vtt1.1)
if(m!=1) Vinv = symm(Vinv) #to enforce symmetry after chol2inv call
J0 = V00%*%t.B%*%Vinv # eqn 6.49 and 1s on diag when Q=0; Here it is t.B[1]
xtT.1=ks.out$alphahat[1:m,1,drop=FALSE]
x0T = x00 + J0%*%(ks.out$alphahat[1:m,1,drop=FALSE]-ks.out$a[1:m,1,drop=FALSE]); # eqn 6.47
V0T = V00 + J0%*%(VtT[,,1]-Vtt1[,,1])*t(J0) # eqn 6.48
}
if(!return.kfas.model) kfas.model=NULL
#not using ks.out$v (Innov) and ks.out$F (Sigma) since I think there might be a bug when R is not diagonal.
par.R=parmat(MLEobj,"R",t=1:model.dims$R[3])$R
rtn.list = list(
xtT = ks.out$alphahat[1:m,,drop=FALSE],
VtT = VtT,
Vtt1T = Vtt1T,
x0T = x0T,
V0T = V0T,
x10T = ks.out$alphahat[1:m,1,drop=FALSE],
V10T = V10T,
x00T = x00,
V00T = V00,
Vtt = "Use MARSSkfss to get Vtt",
Vtt1 = Vtt1,
J="Use MARSSkfss to get J",
J0="Use MARSSkfss to get J0",
Kt="Use MARSSkfss to get Kt",
xtt1 = ks.out$a[1:m,1:TT,drop=FALSE],
xtt= "Use MARSSkfss to get xtt",
Innov="Use MARSSkfss to get Innov",
Sigma="Use MARSSkfss to get Sigma",
kfas.model=kfas.model,
logLik=ks.out$logLik,
ok=TRUE,
errors = NULL
)
}
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.