### Filtering and smoothing for AR(1) with Gaussian and t-distributed
### observation noise

### Alternatively, use the package dlm in R


#Gaussian noise: time homogeneous case (s2v constant, y=constant)
n <- 30
y <- rep(1,n)
s2v <- rep(1,n)
s2eps <- 1
phi <- 0.8
al <- 1/s2eps
erg <- kalman.smooth(phi,s2eps,s2v,y)
plot(1:n,erg$mean[,1],type="l",ylim=c(0,1)) #prediction mean
lines(1:n,erg$mean[,2],col=2) #filter mean
lines(1:n,erg$mean[,3],col=4) #smoother mean
plot(1:n,erg$var[,1],type="l",ylim=c(0,erg$var[1,1])) #prediction variance
lines(1:n,erg$var[,2],col=2) #filter variance
lines(1:n,erg$var[,3],col=4) #smoother variance
# comparing with asymptotic values
Rf0 <- (sqrt(1 + ((1-phi^2)*al)^2 + 2*(1+phi^2)*al)- al*(1-phi^2) - 1)*
     s2eps*0.5/phi^2 #Limit of filter variance
Rf0-erg$var[n,2]
Rp0 <- phi^2*Rf0 + s2eps #Limit of prediction variance
Rp0-erg$var[n,1]
g <- Rp0/(1+Rp0) #Limit of the gain
m <- g/(1-phi+g*phi) #Limit of the filter mean
m-erg$mean[n,2]

#simulations with Gaussian noise
n <- 100
s2v <- rep(1,n)
s2eps <- 1
seps <- sqrt(s2eps)
phi <- 0.8
x <- rep(0,n)
for (i in (2:n)) x[i] = phi*x[i-1] + seps*rnorm(1)
y <- x+rnorm(n)
erg <- kalman.smooth(phi,s2eps,s2v,y)
plot(1:n,y,type="l",col=5) #observations
lines(1:n,x) #states
lines(1:n,erg$mean[,1],col=2) #prediction mean
lines(1:n,erg$mean[,1]+1.96*sqrt(erg$var[,1]),col=2,lty=2) #prediction quantile
lines(1:n,erg$mean[,1]-1.96*sqrt(erg$var[,1]),col=2,lty=2) #prediction quantile
lines(1:n,erg$mean[,2],col=3) #filter mean
lines(1:n,erg$mean[,2]+1.96*sqrt(erg$var[,2]),col=3,lty=2) #filter quantile
lines(1:n,erg$mean[,2]-1.96*sqrt(erg$var[,2]),col=3,lty=2) #filter quantile
lines(1:n,erg$mean[,3],col=4) #smoother mean
lines(1:n,erg$mean[,3]+1.96*sqrt(erg$var[,3]),col=4,lty=2) #smoother quantile
lines(1:n,erg$mean[,3]-1.96*sqrt(erg$var[,3]),col=4,lty=2) #smoother quantile

#Comparing the simulation smoother with the exact smoother
x.smooth <- kalman.sim(phi,s2eps,s2v,y)
for (i in (2:500)) {
  x.smooth <- rbind(x.smooth,kalman.sim(phi,s2eps,s2v,y))
}
x.mean <- apply(x.smooth,2,mean)
x.var <- apply(x.smooth,2,var)
plot(y,type="l")
lines(x,col=2)
lines(x.mean,col=4)
erg <- kalman.smooth(phi,s2eps,s2v,y)
lines(erg$mean[,3],col=3)
plot(x.mean-erg$mean[,3],type="l")
plot(x.var,type="l")
lines(erg$var[,3],col=2)

#simulations with t-distributed noise
s2eps <- 1
nu <- 1 #degrees of freedom for t-distribution
seps <- sqrt(s2eps)
phi <- 0.8
c <- 1
x <- rep(0,100)
for (i in (2:100)) x[i] = phi*x[i-1] + seps*rnorm(1)
y <- x+rt(100,df=nu)
# Sampling from the posterior 
erg <- mcmc.tnoise(phi,s2eps,nu,c,y,niter=1000)
#plotting the results for the states
plot(y,type="l",ylim=3*range(x),col=4)
for (i in (1:nrow(erg$states))) lines(erg$states[i,],col=2)
lines(x)
lines(y,col=4)
qt <- apply(erg$states,2,quantile, probs=c(0.1,0.5,0.9))
plot(y,type="l",col=5,ylim=3*range(x))
lines(x)
for (i in (1:3)) lines(qt[i,], col=i+1)
#Kalman smoother for comparison (choise of error variance s2v
# is somewhat arbitrary if the t-distribution has infinite variance)
lines(kalman.smooth(phi,s2eps,s2v=rep(c,100),y)$mean[,3],col=6)
#plotting the results for the scale variables
plot(abs(y-x),type="l")
for (i in (1:nrow(erg$scale))) lines(erg$scale[i,],col=2)
lines(abs(y-x))
abline(h=0)
abline(h=c)
#posterior of state value at a fixed time
i <- 83 #Choose a time point which might be an outlier
par(mfrow=c(2,3))
hist(erg$states[,i])
plot(erg$states[,i],type="l") #monitoring convergence
acf(erg$states[,i]) #monitoring dependence between replicates
hist(erg$scale[,i],nclass=100)
plot(erg$scale[,i],type="l")
acf(erg$scale[,i])

### Functions used #######################################


kalman.smooth <- function(phi,s2eps,s2v,y)
{
  ## Purpose: Kalman smoother for a Gaussian AR(1)
  ##          with Gaussian observation noise. 
  ## ----------------------------------------------------------------------
  ## Arguments: phi, s2eps = Parameters of the AR(1)
  ##            s2v = variance of the observation noise (time dependent)
  ##            y = observations
  ## ----------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 26 May 2005, 10:00
  n <- length(y)
  mf <- rep(0,n) #initializing filter means
  Rf <- rep(0,n) #initializing filter variances
  mp <- rep(0,n+1) #initializing prediction means
  Rp <- rep(s2eps/(1-phi^2),n+1) #initializing prediction variances
  ms <- rep(0,n) #initializing smoother
  Rs <- rep(0,n) #initializing prediction variances
  #forward filtering
  for (i in (1:n)){
    g <- Rp[i]/(s2v[i]+Rp[i]) #Kalman gain
    mf[i] <- mp[i] + g*(y[i]-mp[i]) #filter mean from prediction mean
    Rf[i] <- Rp[i]*(1-g) #filter variance from prediction variance
    Rp[i+1] <- phi^2*Rf[i] + s2eps #Prediction variance from filter variance
    mp[i+1] <- phi*mf[i] #Prediction mean from filter mean
  }
  #backward smoothing
  ms[n] <- mf[n]
  Rs[n] <- Rf[n]
  for (i in ((n-1):1)){
    g <- phi*Rf[i]/Rp[i+1]
    ms[i] <- mf[i]+g*(ms[i+1]-mp[i+1])
    Rs[i] <- Rf[i]-g^2*(Rp[i+1]-Rs[i+1])
  }
  list(mean=cbind(mp[1:n],mf,ms),var=cbind(Rp[1:n],Rf,Rs))
}

kalman.sim <- function(phi,s2eps,s2v,y)
{
  ## Purpose: Forward computation backward simulation for a Gaussian AR(1)
  ##          Gaussian observation noise. 
  ## ----------------------------------------------------------------------
  ## Arguments: phi, s2eps = Parameters of the AR(1)
  ##            s2v = variance of the observation noise (time dependent)
  ##            y = observations
  ## ----------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 26 May 2005, 10:00
  n <- length(y)
  mf <- rep(0,n) #initializing filter means
  Rf <- rep(0,n) #initializing filter variances
  mp <- rep(0,n+1) #initializing prediction means
  Rp <- rep(s2eps/(1-phi^2),n+1) #initializing prediction variances
  x <- rnorm(n) #initializing sample
  for (i in (1:n)){
    g <- Rp[i]/(s2v[i]+Rp[i]) #Kalman gain
    mf[i] <- mp[i] + g*(y[i]-mp[i]) #filter mean from prediction mean
    Rf[i] <- Rp[i]*(1-g) #filter variance from prediction variance
    Rp[i+1] <- phi^2*Rf[i] + s2eps #Prediction variance from filter variance
    mp[i+1] <- phi*mf[i] #Prediction mean from filter mean
  }
  x[n] <- mf[n] + sqrt(Rf[n])*x[n]
  for (i in ((n-1):1)){
    g <- phi*Rf[i]/Rp[i+1]
    x[i] <- mf[i]+g*(x[i+1]-mp[i+1]) + sqrt(Rf[i]*(1-g*phi))*x[i]
  }
  x
}

mcmc.tnoise <- function(phi,s2eps,nu,c,y,niter,thin=10)
{
  ## Purpose: blockwise mcmc for AR(1) with t-distributed observation noise
  ## ----------------------------------------------------------------------
  ## Arguments: phi,s2eps = Parameters of the AR(1)
  ##            nu, c = parameters of the noise distribution which is written
  ##                    as c*N(0,1)/sqrt{Gamma(nu/2,nu/2)}
  ##            niter=number of iterations of the Gibbs sampler
  ##            thin=thinning factor (only results for iteration equal a
  ##                 multiple of thin are returned)
  ## ----------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 26 May 2005, 10:33
  nu <- nu/2
  c <- c^2
  n <- length(y)
  x <- matrix(0,nrow=niter%/%thin,ncol=n)
  s <- matrix(1,nrow=niter%/%thin,ncol=n)
  s.act <- rep(1,n)
  for (i in (2:niter)){
    x.act <- kalman.sim(phi,s2eps,s2v=c/s.act,y)
    rate <- nu+0.5*(y-x.act)^2/c
    s.act <- rgamma(n,shape=nu+0.5,rate=rate)
    if ((i %% thin)==0) {
      x[i%/%thin,] <- x.act
      s[i%/%thin,] <- s.act
    }
  }
  list(states=x,scale=1/sqrt(s))
}