##################################################################################################
# 
# Simple linear regression with Student-t errors
#
#         y(i) = alpha + beta*x(i) + sig*e(i)      
#
# where e(1),...,e(n) are iid t(v)
#
# Prior distribution:
#
# alpha ~ N(0,Va)
# beta  ~ N(0,Vb)
#
# (sig,v) are fixed
#
##################################################################################################
#
# Hedibert Freitas Lopes
# Booth School of Business 
# University of Chicago
#
##################################################################################################
logpost = function(alpha,beta,sig2){
-0.5*(alpha-mua)^2/Va-0.5*(beta-mub)^2/Vb-(a+1)*log(sig2)-b/sig2-
0.5*(v+1)*sum(log(1+(y-alpha-beta*x)^2/(v*sig2)))-0.5*n*log(sig2)
}
logpostv = function(v){
n*lgamma((v+1)/2)-n*lgamma(v/2)-0.5*n*log(v)-0.5*(v+1)*sum(log(1+(y-alpha-beta*x)^2/(v*sig2)))
}
logpostv1 = function(v){
0.5*log(v/(v+3))+0.5*log(trigamma(v/2)-trigamma((v+1)/2)-2*(v+3)/(v*(v+1)^2))+
n*lgamma((v+1)/2)-n*lgamma(v/2)-0.5*n*log(v)-0.5*(v+1)*sum(log(1+(y-alpha-beta*x)^2/(v*sig2)))
}


pdf(file="hw2-2013-solution.pdf",width=12,height=10)

# Reading the data
# ----------------
data = read.table("school-spending.txt",header=TRUE)
y = data[,2]
x = data[,3]/10000
n = length(y)

# OLS estimation
# --------------
par(mfrow=c(1,1))
plot(x,y,main="OLS fit")
abline(lm(y~x),col=2,lwd=2)

# Fixed parameters and prior hyperparameters
# ------------------------------------------
v      = 4.46
sig2   = 46.7^2
mua    = -70
mub    = 600
Va     = 10000
Vb     = 10000
a      = 5
b      = 15000

par(mfrow=c(1,3))
aaa = seq(mua-4*sqrt(Va),mua+4*sqrt(Va),length=1000)
plot(aaa,dnorm(aaa,mua,sqrt(Va)),type="l",xlab="",ylab="Density",main="Prior for alpha")
bbb = seq(mub-4*sqrt(Vb),mub+4*sqrt(Vb),length=1000)
plot(bbb,dnorm(bbb,mub,sqrt(Vb)),type="l",xlab="",ylab="Density",main="Prior for beta")
sss = seq(0,20000,length=1000)
plot(sss,dgamma(1/sss,a,b)/sss^2,type="l",xlab="",ylab="Density",main="Prior for sig2")

# a) MH random walk 
# Roughly 190 seconds to run on my Mac OS X version 10.5.8 
# 2x3.2 GHz Quad-Core Intel Xeon 8 GB 800 MHz DDR2 FB-DIMM computer
# -----------------------------------------------------------------

# Initial values
alpha = mua
beta  = mub
sig2  = b/(a-1)

# MCMC setup
sda = 5
sdb = 20
dof = 10
burnin = 1000
thin   = 1000
M      = 2000
niter  = burnin+thin*M
draws  = matrix(0,niter,3)
set.seed(164548)
rw.time = system.time({
for (iter in 1:niter){
  alpha1  = rnorm(1,alpha,sda)
  beta1   = rnorm(1,beta,sdb)
  sig21   = 1/rgamma(1,dof,(dof-1)*sig2)
  logratio = min(0,logpost(alpha1,beta1,sig21)-logpost(alpha,beta,sig2)+
  dgamma(1/sig2,dof,(dof-1)*sig21,log=TRUE)-2*log(sig2)-dgamma(1/sig21,dof,(dof-1)*sig2,log=TRUE)+2*log(sig21))
  if (log(runif(1))<logratio){
    alpha = alpha1
    beta  = beta1
    sig2  = sig21  
  }
  draws[iter,] = c(alpha,beta,sig2)
}
})
draws = draws[seq(burnin+1,niter,by=thin),]

names = c("alpha","beta","sig2")
par(mfrow=c(3,3))
for (i in 1:3){
  ts.plot(draws[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws[,i],main="")
  hist(draws[,i],xlab="",main="",prob=TRUE,breaks=seq(min(draws[,i]),max(draws[,i]),length=15))
}

write(t(draws),"draws-rw.txt",ncolumns=3)
summary(cbind(draws,sqrt(draws[,3])))
#           alpha            beta           sig2             sig
# Min.   :-240.17   Min.   :383.6   Min.   :1090   Min.   :33.01  
# 1st Qu.:-108.64   1st Qu.:542.7   1st Qu.:2028   1st Qu.:45.03  
# Median : -77.52   Median :585.8   Median :2372   Median :48.70  
# Mean   : -77.46   Mean   :584.8   Mean   :2427   Mean   :48.97  
# 3rd Qu.: -46.37   3rd Qu.:627.6   3rd Qu.:2747   3rd Qu.:52.41  
# Max.   :  81.15   Max.   :806.1   Max.   :5347   Max.   :73.12 


# b) Gibbs sampler
# Roughly 17 seconds to run on my Mac OS X version 10.5.8 
# 2x3.2 GHz Quad-Core Intel Xeon 8 GB 800 MHz DDR2 FB-DIMM computer
# -----------------------------------------------------------------

# Initial values
alpha = mua
beta  = mub
sig2  = b/(a-1)
lambda = rep(1,n)

# MCMC set up
burnin = 1000
thin   = 100
M      = 2000
niter  = burnin+thin*M
draws1 = matrix(0,niter,3)
set.seed(164548)
gs.time = system.time({
for (iter in 1:niter){
	var = 1/(1/Va+sum(1/lambda)/sig2)
	mean = var*(mua/Va + sum((y-beta*x)/lambda)/sig2)
	alpha = rnorm(1,mean,sqrt(var))
	var = 1/(1/Vb+sum(x^2/lambda)/sig2)
	mean = var*(mub/Vb + sum((y-alpha)*x/lambda)/sig2)
	beta = rnorm(1,mean,sqrt(var))
	sig2 = 1/rgamma(1,(a+n)/2,(b+sum((y-alpha-beta*x)^2/lambda))/2)
	lambda = 1/rgamma(n,(v+1)/2,(v+(y-alpha-beta*x)^2/sig2)/2)
  draws1[iter,] = c(alpha,beta,sig2)
}
})
draws1 = draws1[seq(burnin+1,niter,by=thin),]

names = c("alpha","beta","sig2")
par(mfrow=c(3,3))
for (i in 1:3){
  ts.plot(draws1[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws1[,i],main="")
  hist(draws1[,i],xlab="",main="",prob=TRUE,breaks=seq(min(draws1[,i]),max(draws1[,i]),length=15))
}

write(t(draws1),"draws-gs.txt",ncolumns=3)
summary(cbind(draws1,sqrt(draws1[,3])))
#           alpha            beta           sig2             sig
# Min.   :-228.02   Min.   :334.9   Min.   :1006   Min.   :31.72  
# 1st Qu.:-109.81   1st Qu.:545.3   1st Qu.:1952   1st Qu.:44.18  
# Median : -79.45   Median :586.8   Median :2258   Median :47.52  
# Mean   : -79.43   Mean   :587.0   Mean   :2342   Mean   :48.07  
# 3rd Qu.: -49.15   3rd Qu.:628.4   3rd Qu.:2637   3rd Qu.:51.36  
# Max.   : 110.03   Max.   :784.0   Max.   :5425   Max.   :73.66  


# d) Gibbs sampler for (alph,beta,sig2,v) using uniform prior for v
# -----------------------------------------------------------------

# Initial values
alpha  = mua
beta   = mub
sig2   = b/(a-1)
lambda = rep(1,n)
v      = 4.46
vs     = 1:100

# MCMC set up
burnin = 10000
thin   = 100
M      = 1000
niter  = burnin+thin*M
draws2 = matrix(0,niter,4)
set.seed(164548)
gs1.time = system.time({
for (iter in 1:niter){
	var = 1/(1/Va+sum(1/lambda)/sig2)
	mean = var*(mua/Va + sum((y-beta*x)/lambda)/sig2)
	alpha = rnorm(1,mean,sqrt(var))
	var = 1/(1/Vb+sum(x^2/lambda)/sig2)
	mean = var*(mub/Vb + sum((y-alpha)*x/lambda)/sig2)
	beta = rnorm(1,mean,sqrt(var))
	sig2 = 1/rgamma(1,(a+n)/2,(b+sum((y-alpha-beta*x)^2/lambda))/2)
  pv = rep(0,n)	
	for (i in 1:100) 
	  pv[i] = logpostv(i)
	pv = exp(pv-max(pv))
	v = sample(1:100,size=1,prob=pv)
	lambda = 1/rgamma(n,(v+1)/2,(v+(y-alpha-beta*x)^2/sig2)/2)
  draws2[iter,] = c(alpha,beta,sig2,v)
}
})
draws2 = draws2[seq(burnin+1,niter,by=thin),]

names = c("alpha","beta","sig2","v")
par(mfrow=c(4,3))
for (i in 1:4){
  ts.plot(draws2[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws2[,i],main="")
  hist(draws2[,i],xlab="",main="",prob=TRUE,breaks=seq(min(draws2[,i]),max(draws2[,i]),length=15))
}

par(mfrow=c(1,1))
plot(table(draws2[,4]),ylab="Frequency",xlab="Degrees of freedom")

write(t(draws2),"draws-gs1.txt",ncolumns=4)
summary(cbind(draws2,sqrt(draws2[,3])))
#           alpha            beta             sig2               nu
# Min.   :-223.90   Min.   :367.8   Min.   : 977.4   Min.   :  2.00  
# 1st Qu.:-122.47   1st Qu.:570.2   1st Qu.:2484.2   1st Qu.:  8.00  
# Median : -91.92   Median :607.5   Median :2969.3   Median : 17.00  
# Mean   : -93.20   Mean   :609.3   Mean   :3039.3   Mean   : 29.42  
# 3rd Qu.: -64.42   3rd Qu.:649.8   3rd Qu.:3504.6   3rd Qu.: 47.00  
# Max.   :  84.97   Max.   :777.3   Max.   :6693.8   Max.   :100.00  

#           sig
# Min.   :31.26  
# 1st Qu.:49.84  
# Median :54.49  
# Mean   :54.64  
# 3rd Qu.:59.20  
# Max.   :81.82  


# d.2)  Using Fonseca et al's (2008) independence Jeffreys prior for v
# --------------------------------------------------------------------

# Initial values
alpha  = mua
beta   = mub
sig2   = b/(a-1)
lambda = rep(1,n)
v      = 4.46
vs     = 1:100

# MCMC set up
burnin = 10000
thin   = 100
M      = 1000
niter  = burnin+thin*M
draws3 = matrix(0,niter,4)
set.seed(164548)
gs1.time = system.time({
for (iter in 1:niter){
	var = 1/(1/Va+sum(1/lambda)/sig2)
	mean = var*(mua/Va + sum((y-beta*x)/lambda)/sig2)
	alpha = rnorm(1,mean,sqrt(var))
	var = 1/(1/Vb+sum(x^2/lambda)/sig2)
	mean = var*(mub/Vb + sum((y-alpha)*x/lambda)/sig2)
	beta = rnorm(1,mean,sqrt(var))
	sig2 = 1/rgamma(1,(a+n)/2,(b+sum((y-alpha-beta*x)^2/lambda))/2)
  pv = rep(0,n)	
	for (i in 1:100) 
	  pv[i] = logpostv1(i)
	pv = exp(pv-max(pv))
	v = sample(1:100,size=1,prob=pv)
	lambda = 1/rgamma(n,(v+1)/2,(v+(y-alpha-beta*x)^2/sig2)/2)
  draws3[iter,] = c(alpha,beta,sig2,v)
}
})
draws3 = draws3[seq(burnin+1,niter,by=thin),]

names = c("alpha","beta","sig2","v")
par(mfrow=c(4,3))
for (i in 1:4){
  ts.plot(draws3[,i],xlab="Iterations",ylab="",main=names[i])
  acf(draws3[,i],main="")
  hist(draws3[,i],xlab="",main="",prob=TRUE,breaks=seq(min(draws3[,i]),max(draws3[,i]),length=15))
}

par(mfrow=c(1,1))
plot(table(draws3[,4]),ylab="Frequency",xlab="Degrees of freedom")

write(t(draws3),"draws-gs2.txt",ncolumns=4)
summary(cbind(draws3,sqrt(draws2[,3])))
#           alpha            beta             sig2               nu
# Min.   :-231.97   Min.   :364.8   Min.   : 977.5   Min.   : 1.000  
# 1st Qu.:-111.13   1st Qu.:548.4   1st Qu.:1975.3   1st Qu.: 3.000  
# Median : -81.69   Median :590.8   Median :2330.1   Median : 5.000  
# Mean   : -81.04   Mean   :589.3   Mean   :2458.7   Mean   : 7.161  
# 3rd Qu.: -50.38   3rd Qu.:628.4   3rd Qu.:2846.5   3rd Qu.: 7.000  
# Max.   : 106.40   Max.   :781.5   Max.   :6821.3   Max.   :91.000  
#           sig       
# Min.   :31.26  
# 1st Qu.:49.84  
# Median :54.49  
# Mean   :54.64  
# 3rd Qu.:59.20  
# Max.   :81.82  


par(mfrow=c(1,1))
plot(table(draws2[,4])/M,ylab="Probability",xlab="Degrees of freedom",ylim=c(0,0.4),type="b",pch=16)
lines(table(draws3[,4])/M,col=2,type="b",pch=16)
vv=1:100
prior = exp(0.5*log(vv/(vv+3))+0.5*log(trigamma(vv/2)-trigamma((vv+1)/2)-2*(vv+3)/(vv*(vv+1)^2)))
prior = prior/sum(prior)
lines(prior,col=2)
abline(h=1/100,col=1)
abline(h=0,col=1,lty=3)
title("Prior (solid lines) x Posterior (dotted lines)")

dev.off()