# script used during presentation @ bern to illustrate a bivariate
# linear kalman filter.
rm(list=ls())
library(mvtnorm)
graphics.off()
x11(width=11.,height=7)
set.seed(14)


############################################################
# state-space parameters:
n <- 60                           # we use 60 time steps

obs <- c(1:43,seq(44,by=4,to=n))  # we observe at times obs

H <- diag(2)                      # measurement operator
G <- array(c(.7,-.1,-.1,.7),c(2,2))# system operator

R <- diag(2)                      # observation noise covariance
rm <- c(rep(.25,n/2),rep(1.5,n/2))# factor for observation noise 

Q <- array(c(1,.7,.7,1),c(2,2))   # state noise covariance
qm <- rep(1,n)                    # factor for observation noise 

x0 <- c(3,-1)                     # initialisation of first state
P <- diag(2)                      # initialisation of first forecast cov

############################################################
# program variables
dobs <- diff(c(obs))
nobs <- length(obs)               # number of observations
x <- array(NA,c(2,n))             # state
y <- array(NA,c(2,nobs))          # space
xf <- array(NA,c(2,nobs+1))       # filtered space

x[,1] <- x0                       # initialisation of first state
xf[,1] <- c(0,0)

############################################################
# generate state (unknown)

par(c(1,1))
plot(0,0,type='n',xlab='time',ylab='',ylim=c(-4,16),xlim=c(1,n),yaxt='n')
abline(h=c(7,12),col='gray')
abline(h=4)
for (i in 2:n){
  x[,i] <- G %*% x[,i-1]+t(rmvnorm(1,sigma=Q*qm[i]))
  
  points(c(i,i),x[,i],col=1,cex=.5,pch=20)
  lines(c(i-1,i),c(x[1,i-1],x[1,i]),lty=2)
  lines(c(i-1,i),c(x[2,i-1],x[2,i]),lty=2)
}
system('sleep 5')                       # small break

# we now have the true (unobserved) state. lets filter...
set.seed(14)





j <- 1
par(c(1,1))                                        # setup the plots 
plot(0,0,type='n',xlab='time',ylab='',ylim=c(-4,16),xlim=c(1,n),yaxt='n')
abline(h=c(7,12),col='gray')
abline(h=4)
lines(c(1:n),x[1,],lty=1,col='gray70')
lines(c(1:n),x[2,],lty=1,col='gray70')
for (i in 2:n){
  P <- G %*% P %*% t(G) +Q*qm[i]                   # propagate covariance 
  
  if (sum(i==obs)==1){                             # if we have an observation
    y[,j] <- x[,i]+t(rmvnorm(1,sigma=R*rm[i]))     # generate the observations
    points(c(i,i),y[,j],col=c(2,4),pch=16,cex=1.)  # plot them



    K <- P%*%t(H) %*% solve(H%*%P%*%t(H)+R*rm[i])  # kalman gain
    xf[,j+1] <- xf[,j]+ K %*% (y[,j]- H%*%xf[,j])  # filter
    P <- (diag(2)-K%*%H)%*%P                       # update covariance matrix
    
    points(c(i,i),xf[,j+1],col=1,cex=.5,pch=20)    # some plotting
    lines(c(i-dobs[j],i),c(xf[1,j],xf[1,j+1]),col=2)
    lines(c(i-dobs[j],i),c(xf[2,j],xf[2,j+1]),col=4)

    # for the state 1
    lines(c(i,i),+c(-P[1,1],+P[1,1])*4+7)          # magnitude of P^f
    lines(c(i,i)-.5,c(12,12+(y[,j]- H%*%xf[,j])[1]),col='gray70') # inovations
    lines(c(i,i),c(12,12+(xf[1,j+1]-x[1,i])))                     # errors
    
    j <- j+1
  }
  system('sleep .1')                    # small break

  lines(c(i-1,i),c(x[1,i-1],x[1,i]),lty=1,col='gray70') # more plotting
  lines(c(i-1,i),c(x[2,i-1],x[2,i]),lty=1,col='gray70')
}


# reinhard furrer, april 20th, 2004