################################################################################
# 
# SECOND ORDER DYNAMIC LINEAR MODEL
#
################################################################################
# Author : Hedibert Freitas Lopes
#          The University of Chicago Booth School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email: hlopes@ChicagoBooth.edu
#
###############################################################################
# Simulating 2nd order DLM
# ------------------------
set.seed(6543442)
n  = 100
W1 = 0.01
W2 = 0.001
V  = 0.5
true = c(V,W1,W2)
x0 = c(0.0,0.0)
iW = diag(c(1/W1,1/W2))
w  = cbind(rnorm(n,0,sqrt(W1)),rnorm(n,0,sqrt(W2)))
v  = rnorm(n,0,sqrt(V))
y = rep(0,n)
x = matrix(0,n,2)
x[1,2] = x0[2]          + w[1,2]
x[1,1] = x0[1] + x0[2]  + w[1,1]
y[1]   = x[1,1]         + v[1]
for (t in 2:n){
  x[t,2] =            x[t-1,2] + w[t,2]
  x[t,1] = x[t-1,1] + x[t-1,2] + w[t,1]
  y[t]   = x[t,1]            + v[t]
}

par(mfrow=c(1,1))
plot(y,xlab="time",ylab="",pch=16)
title(paste("data, x1 and x2 \n V=",V," - W1=",W1," - W2=",W2,sep=""))
lines(x[,1],col=2,lwd=2)
lines(x[,2],col=4,lwd=2)

# Forward and backward distributions for the 2nd order DLM
# --------------------------------------------------------
F  = matrix(c(1,0),2,1)
G  = matrix(c(1,0,1,1),2,2)
Ws = diag(c(W1,W2))
a1 = matrix(0,2,1)
R1 = diag(0.1,2)
a  = matrix(0,n,2)
R  = array(0,c(n,2,2))
m  = matrix(0,n,2)
C  = array(0,c(n,2,2))
# time t=1
a[1,] = a1[,1]
R[1,,] = R1
f = t(F)%*%a[1,]
Q = t(F)%*%R[1,,]%*%F+V
A = R[1,,]%*%F/Q[1,1]
m[1,]  = a[1,]+A%*%(y[1]-f)
C[1,,] = R[1,,]-A%*%t(A)*Q[1,1]
# forward filtering
for (t in 2:n){
  a[t,]  = G%*%m[t-1,]
  R[t,,] = G%*%C[t-1,,]%*%t(G) + Ws
  f      = t(F)%*%a[t,]
  Q      = t(F)%*%R[t,,]%*%F+V
  A      = R[t,,]%*%F/Q[1,1]
  m[t,]  = a[t,]+A%*%(y[t]-f)
  C[t,,] = R[t,,]-A%*%t(A)*Q[1,1]
}
mm = matrix(0,n,2)
CC = array(0,c(n,2,2))
mm[n,] = m[n,]
CC[n,,] = C[n,,]
for (t in (n-1):1){
  B = C[t,,]%*%t(G)%*%solve(R[t+1,,])
  mm[t,] = m[t,] + B%*%(mm[t+1,]-a[t+1,])
  CC[t,,] = C[t,,] - B%*%(R[t+1,,]-CC[t+1,,])%*%t(B)
}

par(mfrow=c(2,3))
L1 = sqrt(min(C[,1,1],CC[,1,1]))
U1 = sqrt(max(C[,1,1],CC[,1,1]))
L2 = sqrt(min(C[,2,2],CC[,2,2]))
U2 = sqrt(max(C[,2,2],CC[,2,2]))

ts.plot(sqrt(C[,1,1]),xlab="Time",ylab="",lwd=2,ylim=c(L1,U1))
title("Standard deviation of X1 \n forward")
ts.plot(sqrt(C[,2,2]),xlab="Time",ylab="",lwd=2,ylim=c(L2,U2))
title("Standard deviation of X2 \n forward")
ts.plot(C[,1,2]/sqrt(C[,1,1]*C[,2,2]),xlab="Time",ylab="",lwd=2)
title("Correlation between X1 and X2 \n forward")
ts.plot(sqrt(CC[,1,1]),xlab="Time",ylab="",lwd=2,ylim=c(L1,U1))
title("Standard deviation of X1 \n backward")
ts.plot(sqrt(CC[,2,2]),xlab="Time",ylab="",lwd=2,ylim=c(L2,U2))
title("Standard deviation of X2 \n backward")
ts.plot(CC[,1,2]/sqrt(CC[,1,1]*CC[,2,2]),xlab="Time",ylab="",lwd=2)
title("Correlation between X1 and X2 \n backward")

         
L  = cbind(m[,1]-2*sqrt(C[,1,1]),m[,2]-2*sqrt(C[,2,2]))
U  = cbind(m[,1]+2*sqrt(C[,1,1]),m[,2]+2*sqrt(C[,2,2]))
LL = cbind(mm[,1]-2*sqrt(CC[,1,1]),mm[,2]-2*sqrt(CC[,2,2]))
UU = cbind(mm[,1]+2*sqrt(CC[,1,1]),mm[,2]+2*sqrt(CC[,2,2]))
L1 = c(min(L[,1],LL[,1]),min(L[,2],LL[,2]))
U1 = c(max(U[,1],UU[,1]),max(U[,2],UU[,2]))
par(mfrow=c(2,2))
for (i in 1:2){
  plot(m[,i],xlab="Time",ylab="",main="",col=2,lwd=2,ylim=c(L1[i],U1[i]),type="l")
  lines(L[,i],col=2,lwd=2)
  lines(U[,i],col=2,lwd=2)
  title(paste("p(x",i,"t|y^t)",sep=""))
  lines(x[,i],col=4)
}
for (i in 1:2){
  plot(mm[,i],xlab="Time",ylab="",main="",col=2,lwd=2,ylim=c(L1[i],U1[i]),type="l")
  lines(LL[,i],col=2,lwd=2)
  lines(UU[,i],col=2,lwd=2)
  title(paste("p(x",i,"t|y^n)",sep=""))
  lines(x[,i],col=4)
}

# Drawing from the backward distributions
# ---------------------------------------
M        = 100
CCC      = array(0,c(n,2,2))
mmm      = array(0,c(n,M,2))
xb       = array(0,c(n,M,2))
mmm[n,,] = matrix(m[n,],M,2,byrow=T)
CCC[n,,] = C[n,,]
xb[n,,]  = mmm[n,,] + matrix(rnorm(M*2),M,2)%*%chol(CCC[n,,])
for (t in (n-1):1){
  iC = solve(C[t,,])
  CCC[t,,] = solve(t(G)%*%iW%*%G + iC)
  mmm[t,,] = matrix(CCC[t,,]%*%iC%*%m[t,],M,2,byrow=T)+
            t(CCC[t,,]%*%t(G)%*%iW%*%t(xb[t+1,,]))
  xb[t,,] = mmm[t,,] + matrix(rnorm(M*2),M,2)%*%chol(CCC[t,,])
}

mxb = cbind(apply(xb[,,1],1,mean),apply(xb[,,2],1,mean))
lxb = cbind(apply(xb[,,1],1,quantile,0.02275013),apply(xb[,,2],1,quantile,0.02275013))
uxb = cbind(apply(xb[,,1],1,quantile,0.9772499),apply(xb[,,2],1,quantile,0.9772499))
par(mfrow=c(1,2))
for (i in 1:2){
  plot(mm[,i],xlab="Time",ylab="",main="",col=2,lwd=2,ylim=c(L1[i],U1[i]),type="l")
  lines(LL[,i],col=2,lwd=2)
  lines(UU[,i],col=2,lwd=2)
  lines(mxb[,i],col=4,lwd=2)
  lines(lxb[,i],col=4,lwd=2)
  lines(uxb[,i],col=4,lwd=2)
  title(paste("p(x",i,"t|y^n)",sep=""))
}

# FFBS scheme: sampling (x^n|V,W1,W2) and (V,W1,W2|x^n)
# -----------------------------------------------------
ffbs = function(y,F,G,V,W1,W2,a1,R1){
  n  = length(y)
  Ws = diag(c(W1,W2))
  iW = diag(c(1/W1,1/W2))
  a  = matrix(0,n,2)
  R  = array(0,c(n,2,2))
  m  = matrix(0,n,2)
  C  = array(0,c(n,2,2))
  a[1,]  = a1[,1]
  R[1,,] = R1
  f      = t(F)%*%a[1,]
  Q      = t(F)%*%R[1,,]%*%F+V
  A      = R[1,,]%*%F/Q[1,1]
  m[1,]  = a[1,]+A%*%(y[1]-f)
  C[1,,] = R[1,,]-A%*%t(A)*Q[1,1]
  for (t in 2:n){
    a[t,]  = G%*%m[t-1,]
    R[t,,] = G%*%C[t-1,,]%*%t(G) + Ws
    f      = t(F)%*%a[t,]
    Q      = t(F)%*%R[t,,]%*%F+V
    A      = R[t,,]%*%F/Q[1,1]
    m[t,]  = a[t,]+A%*%(y[t]-f)
    C[t,,] = R[t,,]-A%*%t(A)*Q[1,1]
  }
  xb       = matrix(0,n,2)
  xb[n,]  = m[n,] + t(chol(C[n,,]))%*%rnorm(2)
  for (t in (n-1):1){
    iC  = solve(C[t,,])
    CCC = solve(t(G)%*%iW%*%G+iC)
    mmm = CCC%*%(t(G)%*%iW%*%xb[t+1,]+iC%*%m[t,])
    xb[t,] = mmm + t(chol(CCC))%*%rnorm(2)
  }
  return(xb)
}

# MCMC scheme
set.seed(12566)
V      = 1
W1     = 0.1
W2     = 0.01
burnin  = 1000
step   = 20
M      = 1000
niter  = burnin+step*M
pars   = matrix(0,niter,3)
xbs    = array(0,c(niter,n,2))
time   = system.time({
for (iter in 1:niter){
  xb = ffbs(y,F,G,V,W1,W2,a1,R1)
  V  = 1/rgamma(1,n/2,sum((y-xb[,1])^2)/2)
  W1 = 1/rgamma(1,(n-1)/2,sum((xb[2:n,1]-xb[1:(n-1),1]-xb[1:(n-1),2])^2)/2)
  W2 = 1/rgamma(1,(n-1)/2,sum((xb[2:n,2]-xb[1:(n-1),2])^2)/2)
  xbs[iter,,] = xb
  pars[iter,] = c(V,W1,W2)
}
})

ind = seq(burnin+1,niter,by=step)

names = c("V","W1","W2")
par(mfrow=c(3,3))
for (i in 1:3){
  ts.plot(pars[ind,i],xlab="iterations",ylab="",main=names[i])
  acf(pars[ind,i],main="")
  hist(pars[ind,i],main="",xlab="")
}

par(mfrow=c(1,3))
hist(pars[ind,1],xlab="",main="V",breaks=seq(min(pars[ind,1]),max(pars[ind,1]),length=20))
abline(v=true[1],col=2,lwd=2)
hist(pars[ind,2],xlab="",main="W1",breaks=seq(min(pars[ind,2]),max(pars[ind,2]),length=20))
abline(v=true[2],col=2,lwd=2)
hist(pars[ind,3],xlab="",main="W2",breaks=seq(min(pars[ind,3]),max(pars[ind,3]),length=20))
abline(v=true[3],col=2,lwd=2)

mxbs = cbind(apply(xbs[ind,,1],2,mean),apply(xbs[ind,,2],2,mean))
lxbs = cbind(apply(xbs[ind,,1],2,quantile,0.02275013),apply(xbs[ind,,2],2,quantile,0.02275013))
uxbs = cbind(apply(xbs[ind,,1],2,quantile,0.9772499),apply(xbs[ind,,2],2,quantile,0.9772499))
par(mfrow=c(1,2))
for (i in 1:2){
  plot(mm[,i],xlab="Time",ylab="",main="",ylim=c(L1[i],U1[i]),type="l")
  lines(LL[,i])
  lines(UU[,i])
  lines(mxbs[,i],col=2)
  lines(lxbs[,i],col=2)
  lines(uxbs[,i],col=2)
  title(paste("p(x",i,"t|y^n)",sep=""))
}
legend(0,0.6,legend=c("p(x^n|y^n,theta)","p(x^n|y^n)"),col=1:2,bty="n",lty=c(1,1))