rm(list=ls())


pdf(file="takehome-exam-solution.pdf",width=12,height=10)


# my functions
ldig = function(x,a,b){
  dgamma(1/x,a,b,log=TRUE)-2*log(x)
}
logprior = function(beta,sig2,b0,sB0,a,b){
  dnorm(beta,b0,sB0,log=TRUE)+ldig(sig2,a,b)
}
loglike = function(y,x,beta,sig2,n,nu){
  sum(dt((y-beta*x)/sqrt(sig2),df=nu,log=TRUE))-0.5*n*log(sig2)
}                   
logproposal = function(beta,sig2,betahat1,sd1,a1,b1){
  dnorm(beta,betahat1,sd1,log=TRUE)+ldig(sig2,a1,b1)
}
pnu.f = function(nu){
 sqrt(nu/(nu+3))*sqrt(trigamma(nu/2)+trigamma((nu+1)/2)-2*(nu+3)/(nu*(nu+1)^2))
}
summary = function(x){
  c(mean(x),var(x),quantile(x,0.025),quantile(x,0.975))
}

# Reading the data
data = read.table("spending.txt",header=TRUE)
y    = data[,2]
x    = data[,3]
n    = nrow(data)

# OLS solution
reg = lm(y~x-1)
betahat = reg$coef[1]
sig2hat  = sum(reg$res^2)/(n-1)

###############################################
# a) Assuming nu=5 - SIR
###############################################
nu = 5
b0 = 0
B0 = 1
sB0 = sqrt(B0)
nu0 = 1
modesig2 = 1
# x ~ IG(a,b) => mode(x)=b/(a+1)
sig20 = 2*modesig2*(nu0/2+1)/nu0
a = nu0/2
b = nu0*sig20/2


# Prior distribution
N = 200
sig2s = seq(0.001,15,length=N)
betas = seq(-5,5,length=N)
lprior = matrix(0,N,N)
for (i in 1:N)
  for (j in 1:N)
    lprior[i,j] = logprior(betas[i],sig2s[j],b0,sB0,a,b)
prior = exp(lprior-max(lprior))

# Likelihood function
sig2s1 = seq(0.001,1.0,length=N)
betas1 = seq(0,1.5,length=N)
llike     = matrix(0,N,N)
for (i in 1:N)
  for (j in 1:N)
    llike[i,j] = loglike(y,x,betas1[i],sig2s1[j],n,nu)                   
                   
like   = exp(llike-max(llike))
likebeta = apply(like,1,sum)
likesig2 = apply(like,2,sum)
betahat1 = betas1[likebeta==max(likebeta)]
sig2hat1 = sig2s1[likesig2==max(likesig2)]

par(mfrow=c(1,2))
contour(betas,sig2s,prior,xlab=expression(beta),
        ylab=expression(sigma^2),drawlabels=FALSE,main="Prior")
            
contour(betas1,sig2s1,like,xlab=expression(beta),
        ylab=expression(sigma^2),drawlabels=FALSE,main="Likelihood")
points(betahat,sig2hat,col=2,pch=16,cex=2)
points(betahat1,sig2hat1,col=4,pch=16,cex=2)
text(0.5,0.8,paste("OLS: (",round(betahat,3),",",round(sig2hat,5),")",sep=""),col=2)
text(0.5,0.75,paste("MLE: (",round(betahat1,3),",",round(sig2hat1,5),")",sep=""),col=4)

# Proposal distribution
lprop = matrix(0,N,N)
sd1 = 0.25
a1   = 10
b1   = (a1+1)*sig2hat1
for (i in 1:N)
  for (j in 1:N)
    lprop[i,j] = logproposal(betas1[i],sig2s1[j],betahat1,sd1,a1,b1)
prop = exp(lprop-max(lprop))

par(mfrow=c(1,1))
contour(betas1,sig2s1,prop,xlab=expression(beta),
        ylab=expression(sigma^2),drawlabels=FALSE,main="Likelihood x proposal")
contour(betas1,sig2s1,like,xlab=expression(beta),
        ylab=expression(sigma^2),drawlabels=FALSE,add=TRUE,col=2)
             
# SIR in action
set.seed(123235)
M = 100000
N  = 10000
draws = matrix(0,M,2)
draws[,1] = rnorm(M,betahat1,sd1)
draws[,2] = 1/rgamma(M,a1,b1)
plot(draws,xlab=expression(beta),ylab=expression(sigma^2),col=grey(0.8),pch=16)
contour(betas1,sig2s1,prop,,drawlabels=FALSE,add=TRUE)
w = rep(0,M)
for (i in 1:M)
 w[i] = loglike(y,x,draws[i,1],draws[i,2],n,nu)+logprior(draws[i,1],draws[i,2],b0,sB0,a,b)-logproposal(draws[i,1],draws[i,2],betahat1,sd1,a1,b1)
w = exp(w-max(w))          
k  = sample(1:M,size=N,replace=TRUE,prob=w)      
draws1 = draws[k,]  
     
# Posterior inference         
plot(draws1,xlab=expression(beta),ylab=expression(sigma^2),col=grey(0.8),pch=16)
points(betahat1,sig2hat1,col=4,pch=16,cex=2)

par(mfrow=c(1,2))
hist(draws1[,1],xlab="",prob=TRUE,main=expression(beta))
points(betahat1,0,col=4,pch=16,cex=2)
breaks = seq(min(draws1[,2]),max(draws1[,2]),length=50)
hist(draws1[,2],xlab="",prob=TRUE,main=expression(sigma^2),breaks=breaks)
points(sig2hat1,0,col=4,pch=16,cex=2)

t(apply(draws1,2,summary))         
#          Posterior moments  Posterior 95% C.I.
#            Mean    Variance      2.5%     97.5%
#beta   0.6502281 0.010956795 0.4448216 0.8599034
#sigma2 0.3436215 0.007056576 0.2136839 0.5419747

###############################################
# b) Assuming nu=5 - GIBBS SAMPLER
###############################################
set.seed(3244365)
M = 10000
M0 = 10000
LAG = 1
niter = M0+M*LAG
draws2 = matrix(0,niter,2)
beta = betahat
sig2 = sig2hat
for (iter in 1:niter){
  e2 = (y-beta*x)^2
  lambda = 1/rgamma(n,(nu+1)/2,(nu+e2/sig2)/2)
  y1 = y/sqrt(lambda)
  x1  = x/sqrt(lambda)
  sig2 = 1/rgamma(1,(nu0+n)/2,(nu0*sig20+sum((y1-beta*x1)^2))/2)
  V = 1/(1/B0+sum(x1^2)/sig2)
  m = V*(b0/B0+sum(x1*y1)/sig2)
  beta = rnorm(1,m,sqrt(V))
  draws2[iter,] = c(beta,sig2)
}
draws2 = draws2[seq(M0+1,niter,by=LAG),]


# Posterior inference         
plot(draws2,xlab=expression(beta),ylab=expression(sigma^2),col=grey(0.8),pch=16)
points(betahat1,sig2hat1,col=4,pch=16,cex=2)

par(mfrow=c(1,2))
hist(draws2[,1],xlab="",prob=TRUE,main=expression(beta))
points(betahat1,0,col=4,pch=16,cex=2)
breaks = seq(min(draws2[,2]),max(draws2[,2]),length=50)
hist(draws2[,2],xlab="",prob=TRUE,main=expression(sigma^2),breaks=breaks)
points(sig2hat1,0,col=4,pch=16,cex=2)

t(apply(draws2,2,summary))   
#          Posterior moments  Posterior 95% C.I.
#            Mean    Variance      2.5%     97.5%
#beta   0.6509079 0.010992207 0.4470766 0.8583139
#sigma2 0.3417835 0.006686252 0.2138918 0.5354451


par(mfrow=c(1,2))
plot(density(draws1[,1]),xlab="",main=expression(beta))
lines(density(draws2[,1]),col=2)
plot(density(draws1[,2]),xlab="",main=expression(sigma^2))
lines(density(draws2[,2]),col=2)


###############################################
# c) Learning nu - discrete uniform prior
###############################################
set.seed(3244365)
maxnu = 200
M = 10000
M0 = 10000
LAG = 1
niter = M0+M*LAG
draws3 = matrix(0,niter,3)
beta = betahat
sig2 = sig2hat
pnu  = rep(0,maxnu)
for (iter in 1:niter){
  e2 = (y-beta*x)^2
  lambda = 1/rgamma(n,(nu+1)/2,(nu+e2/sig2)/2)
  y1 = y/sqrt(lambda)
  x1  = x/sqrt(lambda)
  sig2 = 1/rgamma(1,(nu0+n)/2,(nu0*sig20+sum((y1-beta*x1)^2))/2)
  V = 1/(1/B0+sum(x1^2)/sig2)
  m = V*(b0/B0+sum(x1*y1)/sig2)
  beta = rnorm(1,m,sqrt(V))
  for (i in 1:maxnu)
    pnu[i] = loglike(y,x,beta,sig2,n,i)
  pnu = exp(pnu-max(pnu))
  nu = sample(1:maxnu,size=1,prob=pnu)
  draws3[iter,] = c(beta,sig2,nu)
}
draws3 = draws3[seq(M0+1,niter,by=LAG),]

probnu = rep(0,maxnu)
for (i in 1:maxnu)
  probnu[i] = 100*mean(draws3[,3]==i)

# Posterior inference         
par(mfrow=c(1,2))
hist(draws3[,1],xlab="",prob=TRUE,main=expression(beta))
points(betahat1,0,col=4,pch=16,cex=2)
breaks = seq(min(draws3[,2]),max(draws3[,2]),length=50)
hist(draws3[,2],xlab="",prob=TRUE,main=expression(sigma^2),breaks=breaks)
points(sig2hat1,0,col=4,pch=16,cex=2)

par(mfrow=c(1,1))
plot(1:maxnu,probnu,xlab="",ylab="Percentage",main=expression(nu))

t(apply(draws3,2,summary))   
#          Posterior moments  Posterior 95% C.I.
#            Mean    Variance      2.5%     97.5%
# beta  0.7276333 1.109268e-02 0.5181887   0.9338805
# sig2  0.4489441 1.082242e-02 0.2770887   0.6826232
# nu   85.9372000 3.519922e+03 6.0000000 194.0000000

par(mfrow=c(1,2))
plot(density(draws1[,1]),xlab="",main=expression(beta))
lines(density(draws3[,1]),col=2)

plot(density(draws1[,2]),xlab="",main=expression(sigma^2))
lines(density(draws3[,2]),col=2)

###############################################
# d) Learning nu - Fonseca et al's prior
###############################################
set.seed(3244365)
maxnu = 200
M = 10000
M0 = 10000
LAG = 1
niter = M0+M*LAG
draws4 = matrix(0,niter,3)
beta = betahat
sig2 = sig2hat
pnu  = rep(0,maxnu)
for (iter in 1:niter){
  e2 = (y-beta*x)^2
  lambda = 1/rgamma(n,(nu+1)/2,(nu+e2/sig2)/2)
  y1 = y/sqrt(lambda)
  x1  = x/sqrt(lambda)
  sig2 = 1/rgamma(1,(nu0+n)/2,(nu0*sig20+sum((y1-beta*x1)^2))/2)
  V = 1/(1/B0+sum(x1^2)/sig2)
  m = V*(b0/B0+sum(x1*y1)/sig2)
  beta = rnorm(1,m,sqrt(V))
  for (i in 1:maxnu)
    pnu[i] = loglike(y,x,beta,sig2,n,i)+log(pnu.f(i))
  pnu = exp(pnu-max(pnu))
  nu = sample(1:maxnu,size=1,prob=pnu)
  draws4[iter,] = c(beta,sig2,nu)
}
draws4 = draws4[seq(M0+1,niter,by=LAG),]

probnu1 = rep(0,maxnu)
for (i in 1:maxnu)
  probnu1[i] = 100*mean(draws4[,3]==i)

# Posterior inference         
par(mfrow=c(1,2))
hist(draws4[,1],xlab="",prob=TRUE,main=expression(beta))
points(betahat1,0,col=4,pch=16,cex=2)
breaks = seq(min(draws4[,2]),max(draws4[,2]),length=50)
hist(draws4[,2],xlab="",prob=TRUE,main=expression(sigma^2),breaks=breaks)
points(sig2hat1,0,col=4,pch=16,cex=2)

par(mfrow=c(1,1))
plot(1:maxnu,probnu1,xlab="",ylab="Percentage",main=expression(nu))

t(apply(draws4,2,summary))   
#          Posterior moments  Posterior 95% C.I.
#            Mean    Variance      2.5%     97.5%#
#[1,]  0.7131771 1.168804e-02 0.5013975   0.9252924
#[2,]  0.4290057 1.100811e-02 0.2551988   0.6728589
#[3,] 59.7909000 3.060796e+03 4.0000000 187.0000000

par(mfrow=c(1,2))
plot(density(draws1[,1]),xlab="",main=expression(beta))
lines(density(draws3[,1]),col=2)
lines(density(draws4[,1]),col=4)

plot(density(draws1[,2]),xlab="",main=expression(sigma^2))
lines(density(draws3[,2]),col=2)
lines(density(draws4[,2]),col=4)

par(mfrow=c(1,2))
plot(1:maxnu,probnu,xlab="",ylab="Percentage",main=expression(nu),type="b",ylim=range(probnu,probnu1))
lines(1:maxnu,probnu1,col=2,type="b")

plot(1:maxnu,cumsum(probnu),type="l",xlab=expression(nu),ylab="CDF")
lines(1:maxnu,cumsum(probnu1),col=2)

dev.off()

write(t(draws1),"draws1.txt",ncolumns=2)
write(t(draws2),"draws2.txt",ncolumns=2)
write(t(draws3),"draws3.txt",ncolumns=3)
write(t(draws4),"draws4.txt",ncolumns=3)