###############################################################################
#  General linear model with AR(1) errors
#
################################################################################
# Author : Hedibert Freitas Lopes
#          Booth School of Business
#          University of Chicago
#          5807 South Woodlawn Avenue
#          Chicago, Illinois, 60637
#          Email : hlopes@ChicagoBooth.edu
#
###############################################################################
rm(list=ls())

# Simulated data
# --------------
set.seed(12434)
n      = 100
sigma2 = 1
p      = 2
beta   = c(2,4)
rho    = 0.9
sigma  = sqrt(sigma2)
x      = matrix(rnorm(p*n),n,p)
O      = matrix(0,n,n)
for (i in 1:n)
  O[i,] = rho^abs(((1-i):(n-i)))
O = O/(1-rho^2)  
y     = x%*%beta+sigma*t(chol(O))%*%rnorm(n)
error = y-x%*%beta

pdf(file="ar1-error1.pdf",width=10,height=8)
par(mfrow=c(2,2))
plot(x[,1],y,xlab=expression(x[1]),ylab=expression(y))
plot(x[,2],y,xlab=expression(x[2]),ylab=expression(y))
ts.plot(error,xlab="Time",ylab=expression(epsilon))
acf(error,xlab="Lag",ylab="ACF",main="")
dev.off()

# prior
# -----
b0        = rep(0,p)
V0        = diag(100,p)
nu0       = 5
s02       = 1.0
rho0      = 0.9
Vrho      = 100
iV0       = solve(V0)
iV0b0     = iV0%*%b0
nu0n      = nu0+n
nu0s02    = nu0*s02
iVrho     = 1/Vrho
iVrhorho0 = iVrho*rho0

# Initial values
# --------------
ols   = lm(y~x-1)
b     = ols$coef
s2    = mean(ols$res^2)
ols   = lm(ols$res[2:n]~ols$res[1:(n-1)]-1)
r     = as.numeric(ols$coef)
s2    = as.numeric(mean(ols$res^2))
pars = c(b,s2,r)
for (iter in 1:5){
  for (i in 1:n)
    O[i,] = r^abs(((1-i):(n-i)))
  O     = O/(1-r^2)
  iO    = solve(O)
  xiOx  = t(x)%*%iO%*%x
  xiOy  = t(x)%*%iO%*%y
  b     = solve(xiOx)%*%xiOy
  e     = y - x%*%b
  s2    = mean(e^2)*(1-r^2)
  r     = lm(e[2:n]~e[1:(n-1)]-1)$coef
  pars  = rbind(pars,c(b,s2,r))
}
b = pars[5,1:p]
s2 = pars[5,3]
r  = pars[5,4]

# MCMC scheme
# -----------
set.seed(23164)
burn  = 1000
M     = 2000
LAG   = 1
niter = burn+LAG*M
s2s   = rep(0,niter)
rs    = rep(0,niter)
bs    = matrix(0,niter,p)
for (iter in 1:niter){
  print(iter)
  e  = y-x%*%b
  s2 = 1/rgamma(1,(nu0n)/2,(nu0s02+t(e)%*%iO%*%e)/2)  
  V1 = solve(iV0+xiOx/s2)
  b1 = V1%*%(iV0b0+xiOy/s2)
  b  = b1+t(chol(V1))%*%rnorm(p)
  e  = y-x%*%b
  e1 = e[2:n]
  e2 = e[1:(n-1)]
  Vr = 1/(iVrho+sum(e2^2)/s2)
  r1 = Vr*(iVrhorho0+sum(e1*e2)/s2)
  A  = pnorm(-1,r1,sqrt(Vr))
  B  = pnorm(1,r1,sqrt(Vr))
  u  = runif(1)
  r  = r1+sqrt(Vr)*qnorm(u*B+(1-u)*A)
  for (i in 1:n)
    O[i,] = r^abs(((1-i):(n-i)))
  O        = O/(1-r^2)  
  iO       = solve(O)
  xOx      = t(x)%*%iO%*%x
  xiOy     = t(x)%*%iO%*%y
  s2s[iter]   = s2
  bs[iter,]   = t(b)
  rs[iter]    = r
}
s2s    = s2s[seq(burn+1,niter,by=LAG)]
bs     = bs[seq(burn+1,niter,by=LAG),1:p]
rs     = rs[seq(burn+1,niter,by=LAG)]

pdf(file="ar1-error2.pdf",width=10,height=8)
par(mfrow=c(p,3))
for (i in 1:p){
  ts.plot(bs[,i],ylab="",xlab="iteration",main=paste("beta",i,sep=""))
  abline(h=beta[i],col=2)
  acf(bs[,i],ylab="",main="")
  hist(bs[,i],prob=T,xlab="",ylab="",main="")
  abline(v=beta[i],col=2)
}
dev.off()

pdf(file="ar1-error3.pdf",width=10,height=8)
par(mfrow=c(2,3))
ts.plot(s2s,ylab="",xlab="iteration",main=expression(sigma^2))
abline(h=sigma2,col=2)
acf(s2s,ylab="",main="")
hist(s2s,prob=T,xlab="",ylab="",main="")
abline(v=sigma2,col=2)
ts.plot(rs,ylab="",xlab="iteration",main=expression(rho))
abline(h=rho,col=2)
acf(rs,ylab="",main="")
hist(rs,prob=T,xlab="",ylab="",main="")
abline(v=rho,col=2)
dev.off()