###############################################################################################
#
#              Bayesian simple linear regression with mixture of normal errors
#
###############################################################################################
#
#  Author: Hedibert Freitas Lopes
#  Affiliation: Graduate School of Business, University of Chicago
#
#  This version: May, 13th 2005
#
###############################################################################################
#
#  A few steps before you "source" this file into R:
#
#  [1] In R change the directory to the one you currently have the executable file
#
#                  "mixtureofnormal-regression.exe"
#
#
#  [2] Change the simulation specification accordinlgy.
#
#
#  [3] Once [1] and [2] are done, simply type 
#
#  > source("Bayesianregressionwithmixtureerrors.R")
#
#
#  [4] Several files will be created plus graphical summaries in a PDF file called
#
#                  "Bayesianregressionwithmixtureerrors.pdf"
#
###############################################################################################
quantab  = function(x){return(x[(x>quantile(x,0.05))&(x<quantile(x,0.95))])}
quant025 = function(x){quantile(x,0.025)}
quant975 = function(x){quantile(x,0.975)}
set.seed(932747)

#  Simulation spefification
# -------------------------
n     = 500                 # sample size
k     = 3                   # number of components
w     = c(0.20,0.6,0.20)    # weights
mu    = c(-5,2,7)           # means 
sig2  = c(10,4,10)          # variances
beta  = 4                   # regression coefficient
M00   = 100                 # Burn-in (MCMC)
MM    = 1000                # draws (MCMC)
LAG   = 1                   # step (MCMC)


MMM   = 100
sig   = sqrt(sig2)          
sigt  = sqrt(sum(w*sig2))
X     = rnorm(n,5,1)
X1    = seq(min(X),max(X),length=MMM)
z     = sample(1:k,size=n,replace=T,prob=w)
error = rnorm(n,mu[z],sig[z])
y     = beta*X + error
xx = seq(min(error),max(error),length=30)
xxx = seq(min(error),max(error),length=1000)

pdf(file="Bayesianregressionwithmixtureerrors.pdf")
par(mfrow=c(2,2))
hist(error,xlab="error",ylab="",prob=T,breaks=xx,main="")
f = NULL
for (i in 1:1000)
  f = c(f,sum(w*dnorm(xxx[i],mu,sig)))
lines(xxx,f,col=2)
plot(error,col=0,xlab="observation",ylab="error",main="")
for (i in 1:n)
text(i,error[i],z[i],col=z[i])
plot(X,y,xlab="regressor",ylab="dependent variable")
abline(mu[1],beta,col=2)
abline(mu[2],beta,col=3)
abline(mu[3],beta,col=4)
abline(lm(y~X),col=1,lwd=3)
plot(error,lm(y~X)$resi,xlab="true error",ylab="OLS error")
x     = seq(-20,10,length=MMM)
f     = NULL
for (i in 1:MMM)
  f = c(f,sum(w*dnorm(beta*x[i],mu,sig)))


# Creating file "input.dat" to be read by "mixtureofnormal-regression.exe"
# ------------------------------------------------------------------------
write(c(n,k),"input.dat",ncolumns=2)
write(c(M00,MM,LAG),"input.dat",ncolumns=3,append=T)
write(c(beta,w,mu,sig2),"input.dat",ncolumns=3*k+1,append=T)
write(c(0.01,0.01,0.01,0.01,0,100),"input.dat",ncolumns=6,append=T)
write(t(cbind(y,X)),"dataset.dat",ncolumns=2)
write(z,"z.dat",ncolumns=1)

# Running "mixtureofnormal-regression.exe"
# ----------------------------------------
system.time(system("mixtureofnormal-regression.exe"))


# Graphical summaries
# -------------------
mus   = matrix(scan("mu.dat"),MM,k,byrow=T)
sig2s = matrix(scan("sigma2.dat"),MM,k,byrow=T)
sigs  = sqrt(sig2s)
ws    = matrix(scan("w.dat"),MM,k,byrow=T)
taus  = scan("tau.dat")
ms    = scan("m.dat")
betas = scan("beta.dat")

par(mfrow=c(k+1,3))
for (i in 1:k){
  ts.plot(mus[,i],xlab="iterations",ylab="",main=paste("mu",i,sep=""))
  abline(h=mu[i],col=2)
  ts.plot(sig2s[,i],xlab="iterations",ylab="",main=paste("sigma2",i,sep=""))
  abline(h=sig2[i],col=2)
  ts.plot(ws[,i],xlab="iterations",ylab="",main=paste("w",i,sep=""))
  abline(h=w[i],col=2)
}
ts.plot(betas,xlab="iterations",ylab="",main="beta")
abline(h=beta,col=2)
ts.plot(ms,xlab="iterations",ylab="",main="m")
ts.plot(taus,xlab="iterations",ylab="",main="tau")

par(mfrow=c(k+1,3))
for (i in 1:k){
  hist(quantab(mus[,i]),xlab="",ylab="",main=paste("mu",i,sep=""))
  abline(v=mu[i],col=2)
  hist(quantab(sig2s[,i]),xlab="",ylab="",main=paste("sigma2",i,sep=""))
  abline(v=sig2[i],col=2)
  hist(quantab(ws[,i]),xlab="",ylab="",main=paste("w",i,sep=""))
  abline(v=w[i],col=2)
}
hist(quantab(betas),xlab="",ylab="",main="beta")
abline(v=beta,col=2)
hist(quantab(ms),xlab="",ylab="",main="m")
hist(quantab(taus),xlab="",ylab="",main="tau")


LLL = 100
xxx = seq(-20,25,length=LLL)
f = matrix(0,MM,LLL)
for (i in 1:LLL)
  for (l in 1:k)
    f[,i] = f[,i] + ws[,l]*dnorm(xxx[i],mus[,l],sqrt(sig2s[l]))
par(mfrow=c(1,2))
plot(xxx,f[1,],xlab="error",ylab="predictive density",type="l",ylim=c(0,0.08))
for (i in seq(2,MM,by=10))
  lines(xxx,f[i,])
f025 = apply(f,2,quant025)
f050 = apply(f,2,median)
f975 = apply(f,2,quant975)
plot(xxx,f050,xlab="error",ylab="predictive density",type="l",ylim=c(0,0.08))
lines(xxx,f025,col=2)
lines(xxx,f975,col=2)

zs    = matrix(scan("zs.dat"),MM,n,byrow=T)
mz    = matrix(0,n,k)
zz    = diag(1,k)
for (i in 1:n)
  for (j in 1:k)
    mz[i,j] = sum(zs[,i]==j)
mz    = mz/MM
error = NULL
maxz  = NULL
ind = 1:k
for (i in 1:n){
  error = rbind(error,zz[z[i],]-mz[i,])
  maxz  = c(maxz,ind[mz[i,]==max(mz[i,])])
}
error = apply(error^2,1,sum)
count = matrix(0,k,k)
for (i in 1:n)
  count[z[i],maxz[i]]=count[z[i],maxz[i]]+1 
par(mfrow=c(1,2))
plot(error,xlab="observation",ylab="S(z(i))")
title(paste("S(z)=",round(mean(error),4),sep=""))
plot(1:k,1:k,xlab="ESTIMATED Z",ylab="TRUE Z",main="",col=0,axes=F,xlim=c(-0.5,k+0.5),ylim=c(-0.5,k+0.5))
axis(1,at=1:k,labels=1:k)
axis(2,at=1:k,labels=1:k)
for (i in 1:k)
  for (j in 1:k)
    text(i,j,count[i,j]/n)
title(paste("Proportion of correct match=",sum(diag(count))/n,sep=""))
dev.off()