############################################################################### 
#
# Solution to midterm take home exame 
# Bayesian Econometrics
# Spring 2013
# Booth School of Business
# University of Chicago
#
############################################################################### 
#
# Hedibert Freitas Lopes
# Booth School of Business 
# University of Chicago
#
############################################################################### 
rm(list=ls())
logdt = function(x,m,s2,df){
  dt((x-m)/sqrt(s2),df,log=TRUE)-0.5*log(s2)
}

logdmt = function(x,m,iS,df){
  p = length(x)
  lgamma((df+p)/2)-lgamma(df/2)-0.5*p*log(pi*df)+
  0.5*abs(determinant(iS,log=TRUE)$modulus)-
  0.5*((df+p)/2)*log(1+t(x-m)%*%iS%*%(x-m)/df)	
}

logdmn = function(x,m,iS,p){
  -0.5*p*log(2*pi)-0.5*abs(determinant(iS,log=TRUE)$modulus)-
   0.5*t(x-m)%*%iS%*%(x-m)
}

summary=function(draws){
  round(cbind(apply(draws,2,mean),sqrt(apply(draws,2,var)),
        t(apply(draws,2,quantile,c(0.025,0.5,0.975)))),1)
}

# Data and sufficient statistics
# ------------------------------
data = read.table("school-spending.txt",header=TRUE)
y   = data[,2]
x   = data[,3]/10000
n   = length(y)
X   = cbind(1,x)
XtX = t(X)%*%X
Xty = t(X)%*%y
In  = diag(1,n)
x1  = 8000/10000

# Other definitions
names    = c("beta0","beta1","sig2")

# ------------------------------------
# Model 1
# ------------------------------------

# Prior hyperparameters
b0    = matrix(c(-70,600),2,1)
B0    = diag(2.667,2)
eta0  = 4.0048
s02   = 1877
iB0   = solve(B0)
iB0b0 = iB0%*%b0

# Posterior hyperparameters
B1    = solve(iB0+XtX)
b1    = B1%*%(iB0b0+Xty)
eta1  = eta0 + n
s12   = (eta0*s02 + t(y-X%*%b1)%*%y + t(b0-b1)%*%iB0b0)/eta1

# Posterior summaries
m = b1[1,1]         
s = sqrt(eta1/(eta1-2)*s12*B1[1,1])
summary1 = c(m,s,m+s*qt(c(0.025,0.5,0.975),eta1))
m = b1[2,1]         
s = sqrt(eta1/(eta1-2)*s12*B1[2,2])
summary1 = rbind(summary1,c(m,s,m+s*qt(c(0.025,0.5,0.975),eta1)))
p1 = eta1/2
p2 = eta1*s12/2
E  = p2/(p1-1)
SD = E/sqrt(p1-2)
summary1 = rbind(summary1,c(E,SD,1/qgamma(c(0.975,0.5,0.025),p1,p2)))
summary1 = round(summary1,1)

# Prior predictive
m  = X%*%b0
iS = solve(s02*(In+X%*%B0%*%t(X)))
logpriorpred1 = logdmt(y,m,iS,eta0)[1,1]

# Posterior predictive
xx = matrix(c(1,x1),2,1)
m  = t(xx)%*%b1
s2 = s12*(1+t(xx)%*%B1%*%xx)
yy = seq(0,800,length=1000)
logpostpred1 = rep(0,1000)
for (i in 1:1000) 
  logpostpred1[i] = logdt(yy[i],m,s2,eta1)

# DIC
M1   = 10000
sig2 = 1/rgamma(M1,eta1/2,eta1*s12/2)
beta = cbind(rnorm(M1,b1[1,1],sqrt(sig2)),rnorm(M1,b1[2,1],sqrt(sig2)))%*%chol(B1)
DIC1 = -4*(-0.5*n*log(2*pi)-0.5*(mean(log(sig2))+mean(beta[,1]/sig2))-
        0.5*sum(y)/s12-0.5*sum(x)*mean(beta[,2]/sig2))+
        2*sum(dnorm(y,X%*%b1,sqrt((eta1*s12/(eta1-2))),log=TRUE))  

# ------------------------------------
# Model 2
# ------------------------------------
set.seed(123253)

# Prior hyperparameters
b0      = matrix(c(-70,600),2,1)
B0      = diag(10000,2)
eta0    = 4.0048
s02     = 1877
iB0     = solve(B0)
iB0b0   = iB0%*%b0
eta0n   = (eta0+n)/2
eta0s02 = eta0*s02

# Gibbs sampler
ols    = lm(y~x)
beta   = matrix(ols$coef,2,1)
burnin = 1000
thin   = 10
M      = 1000
niter  = burnin+thin*M
draws2 = matrix(0,niter,3)
b1s    = matrix(0,niter,2)
B1s    = array(0,c(niter,2,2))
for (iter in 1:niter){
  sig2  = 1/rgamma(1,eta0n,(eta0s02+sum((y-X%*%beta)^2))/2)
  B1    = solve(iB0+XtX/sig2)
  b1    = B1%*%(iB0b0+Xty/sig2)
  beta  = b1+t(chol(B1))%*%rnorm(2)
  draws2[iter,] = c(beta[1],beta[2],sig2)
  b1s[iter,]  = b1
  B1s[iter,,] = B1
}
ind      = seq(burnin+1,niter,by=thin)
draws2   = draws2[ind,]
b1s      = b1s[ind,]
B1s      =  B1s[ind,,]

# Posterior summaries
summary2 = summary(draws2)

# Prior predictive
M1 = 10000
m  = X%*%b0
sig2 = 1/rgamma(M1,eta0/2,eta0*s02/2)
logpriorpred2 = rep(0,M1)
for (i in 1:M1){
  iS = (1/sig2[i])*solve(In+X%*%(sig2[i]*iB0+XtX)%*%t(X))
  logpriorpred2[i] = logdmn(y,m,iS,n)[1,1]
}
logpriorpred2 = log(mean(exp(logpriorpred2)))

# Posterior predictive
sig2 = draws2[,3]
xx   = matrix(c(1,x1),2,1)
yy   = seq(0,800,length=1000)
logpostpred2 = matrix(0,1000,M)
for (j in 1:M){
  m  = t(xx)%*%b1s[j,]
  s2 = sig2[j]+t(xx)%*%B1s[j,,]%*%xx
  logpostpred2[,j] = dnorm(yy,m,sqrt(s2),log=TRUE)
}
logpostpred2=log(apply(exp(logpostpred2),1,mean))

# DIC
beta = draws2[,1:2]
sig2 = draws2[,3]
mb   = matrix(apply(beta,2,mean),2,1)
ms2  = mean(sig2)
DIC2 = -4*(-0.5*n*log(2*pi)-0.5*(mean(log(sig2))+mean(beta[,1]/sig2))-
        0.5*sum(y)*mean(1/sig2)-0.5*sum(x)*mean(beta[,2]/sig2))+
        2*sum(dnorm(y,X%*%mb,sqrt(ms2),log=TRUE))

# ------------------------------------
# Model 3
# ------------------------------------
set.seed(56758)
v     = 4.46

# Prior hyperparameters
b0    = matrix(c(-70,600),2,1)
B0    = diag(2.667,2)
eta0  = 4.0048
s02   = 1877
iB0   = solve(B0)
iB0b0 = iB0%*%b0
eta0n   = (eta0+n)/2
eta0s02 = eta0*s02

# Gibbs sampler
ols    = lm(y~x)
sig2   = mean(ols$res^2)
beta   = matrix(ols$coef,2,1)
burnin = 1000
thin   = 10
M      = 1000
niter  = burnin+thin*M
draws3 = matrix(0,niter,3)
eta1s  = rep(0,niter)
s12s   = rep(0,niter)
b1s    = matrix(0,niter,2)
B1s    = array(0,c(niter,2,2))
for (iter in 1:niter){
  lambda = 1/rgamma(n,(v+1)/2,(v+(y-X%*%beta)^2/sig2)/2)
  iL     = diag(1/lambda)
  B1    = solve(iB0+t(X)%*%iL%*%X)
  b1    = B1%*%(iB0b0+t(X)%*%iL%*%y)
  eta1  = eta0 + n
  s12   = (eta0s02 + t(y-X%*%b1)%*%iL%*%y + t(b0-b1)%*%iB0b0)/eta1
  sig2  = 1/rgamma(1,eta1/2,eta1*s12/2)
  beta  = b1+t(chol(B1))%*%rnorm(2,0,sqrt(sig2))
  eta1s[iter] = eta1
  s12s[iter]  = s12
  b1s[iter,]  = b1
  B1s[iter,,] = B1
  draws3[iter,] = c(beta[1],beta[2],sig2)
}
ind      = seq(burnin+1,niter,by=thin)
draws3   = draws3[ind,]
eta1s    = eta1s[ind]
s12s     = s12s[ind]
b1s      = b1s[ind,]
B1s      = B1s[ind,,]

# Posterior summaries
summary3 = summary(draws3)

# Prior predictive
M1 = 10000
lambdas = matrix(1/rgamma(n*M,v/2,v/2),n,M1)
m  = X%*%b0
logpriorpred3 = rep(0,M1)
for (i in 1:M1){
  iL = diag(1/lambdas[,i])
  iS = (1/s02)*(iL-iL%*%X%*%solve(iB0+t(X)%*%iL%*%X)%*%t(X)%*%iL)
  logpriorpred3[i] = logdmt(y,m,iS,eta0)[1,1]
}
logpriorpred3 = log(mean(exp(logpriorpred3)))

# Posterior predictive
xx = matrix(c(1,x1),2,1)
lambda = 1/rgamma(M,v/2,v/2)
yy = seq(0,800,length=1000)
logpostpred3 = matrix(0,1000,M)
for (j in 1:M){
  m  = t(xx)%*%b1s[j,]
  s2 = s12s[j]*(lambda[j]+t(xx)%*%B1s[j,,]%*%xx)
  logpostpred3[,j] = logdt(yy,m,s2,eta1s[j])
}
logpostpred3=log(apply(exp(logpostpred3),1,mean))

# DIC
beta = draws3[,1:2]
sig2 = draws3[,3]
mb   = matrix(apply(beta,2,mean),2,1)
ms2  = mean(sig2)
DIC3 = 0.0
for (i in 1:M)
  DIC3 = DIC3 + sum(logdt(y,X%*%beta[i,],sig2[i],v))
DIC3 = -4*DIC3/M + 2*sum(logdt(y,X%*%mb,ms2,v))
DIC3

# --------------------------------
# RESULTS
# --------------------------------
summary1
summary2
summary3

BF31 = exp(logpriorpred3-logpriorpred1)
BF32 = exp(logpriorpred3-logpriorpred2)
BF12 = exp(logpriorpred1-logpriorpred2)
PMP1 = 1/(1+1/BF12+BF31)
PMP2 = 1/(1+BF12+BF32)
PMP3 = 1/(1+1/BF31+1/BF32)

cbind(c(logpriorpred1,logpriorpred2,logpriorpred3),
c(PMP1,PMP2,PMP3),
c(DIC1,DIC2,DIC3))

c(BF31,BF32,BF12)


pdf(file="midterm-predictives.pdf",width=10,height=10)
par(mfrow=c(1,1))
plot(yy,exp(logpostpred3),xlab="",ylab="Density",type="l",lwd=2)
lines(yy,exp(logpostpred2),col=2,lwd=2)
lines(yy,exp(logpostpred1),col=3,lwd=2)
legend(600,0.008,legend=c("M1","M2","M3"),col=3:1,lwd=2,bty="n")
dev.off()

pdf(file="midterm-posterior-model2.pdf",width=10,height=12)
par(mfrow=c(3,3))
for (i in 1:3){
  ts.plot(draws2[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws2[,i],main="")
  hist(draws2[,i],xlab="",main="",prob=TRUE)
}
dev.off()

pdf(file="midterm-posterior-model3.pdf",width=10,height=12)
par(mfrow=c(3,3))
for (i in 1:3){
  ts.plot(draws3[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws3[,i],main="")
  hist(draws3[,i],xlab="",main="",prob=TRUE)
}
dev.off()