############################################################################################
#
# Univariate stochastic volatility model
#
############################################################################################
#
# HEDIBERT FREITAS LOPES
# Associate Professor of Econometrics and Statistics
# The University of Chicago Booth School of Business
# 5807 South Woodlawn Avenue
# Chicago, Illinois, 60637
# Email : hlopes@ChicagoBooth.edu
# URL: http://faculty.chicagobooth.edu/hedibert.lopes/research/
#
############################################################################################
#-------------------------------------------------------
# Univariate FFBS: 
# y(t)     ~ N(alpha(t)+F(t)*theta(t);V(t))        
# theta(t) ~ N(gamma+G*theta(t-1);W)               
#
#-------------------------------------------------------
ffbsu = function(y,F,alpha,V,G,gamma,W,a1,R1,nd=1){
  n = length(y)
  if (length(F)==1){F = rep(F,n)}
  if (length(alpha)==1){alpha = rep(alpha,n)}
  if (length(V)==1){V = rep(V,n)}
  a = rep(0,n)
  R = rep(0,n)
  m = rep(0,n)
  C = rep(0,n)
  B = rep(0,n-1)
  H = rep(0,n-1)
  # time t=1
  a[1] = a1
  R[1] = R1
  f    = alpha[1]+F[1]*a[1]
  Q    = R[1]*F[1]**2+V[1]
  A    = R[1]*F[1]/Q
  m[1] = a[1]+A*(y[1]-f)
  C[1] = R[1]-Q*A**2
  # forward filtering
  for (t in 2:n){
    a[t] = gamma + G*m[t-1]
    R[t] = C[t-1]*G**2 + W
    f    = alpha[t]+F[t]*a[t]
    Q    = R[t]*F[t]**2+V[t]
    A    = R[t]*F[t]/Q
    m[t] = a[t]+A*(y[t]-f)
    C[t] = R[t]-Q*A**2
    B[t-1] = C[t-1]*G/R[t]
    H[t-1] = sqrt(C[t-1]-R[t]*B[t-1]**2)
  }
  # backward sampling
  like = 0.0
  theta = matrix(0,nd,n)
  theta[,n] = rnorm(nd,m[n],sqrt(C[n]))
  like = dnorm(theta[1,n],m[n],sqrt(C[n]),log=TRUE)
  for (t in (n-1):1){
    theta[,t] = rnorm(nd,m[t]+B[t]*(theta[,t+1]-a[t+1]),H[t])
    like = like + dnorm(theta[1,t],m[t]+B[t]*(theta[,t+1]-a[t+1]),H[t],log=TRUE)
  }
  if (nd==1){
    return(list(theta=theta[1,],like=like))
  }
  else{
    theta
  }
}

#-------------------------------------------------------
# y    = X*beta + u    u ~ N(0,sig2*I_n)
# beta|sig2 ~ N(b,sig2*inv(A))
#      sig2 ~ IG(v/2,v*lam/2)
#-------------------------------------------------------
fixedpar = function(y,X,A,b,v,lam){
  n     = length(y)
  k     = ncol(X)
  par1  = (v+n)/2
  var   = solve(crossprod(X,X)+A)
  mean  = matrix(var%*%(crossprod(X,y)+crossprod(A,b)),k,1)
  par2  = v*lam + sum((y-crossprod(t(X),mean))^2)
  par2  = (par2 + crossprod(t(crossprod(mean-b,A)),mean-b))/2
  sig2  = 1/rgamma(1,par1,par2)
  var   = var*sig2
  mean  = mean + crossprod(t(chol(var)),rnorm(2))
  return(c(mean,sig2))
}

#----------------------------------------------------------------------------
# Sample Z from 1,2,...,k, with P(Z=i) proportional to q[i]N(mu[i],sig2[i])
#----------------------------------------------------------------------------
ncind = function(y,mu,sig,q){
  w = dnorm(y,mu,sig)*q
  return(sample(1:length(q),size=1,prob=w/sum(w)))
}

#------------------------------------------------
# Quantile statistics
#------------------------------------------------
quant05 = function(x){quantile(x,0.05)}
quant95 = function(x){quantile(x,0.95)}

#------------------------------------------------
# Sampling the log-volatilities in a 
# standard univariate stochastic volatility model
#------------------------------------------------
svu = function(y,lambda,gamma,G,W,a1,R1){
  mu   = c(-11.40039,-5.24321,-9.83726,1.50746,-0.65098,0.52478,-2.35859)
  sig2 = c(5.795960,2.613690,5.179500,0.167350,0.640090,0.340230,1.262610)
  q    = c(0.007300,0.105560,0.000020,0.043950,0.340010,0.245660,0.257500)
  y    = log(y^2)
  sig  = sqrt(sig2)
  z    = sapply(y-lambda,ncind,mu,sig,q)
  run  = ffbsu(y,1.0,mu[z],sig2[z],G,gamma,W,a1,R1)
  return(list(lambda=run$theta,like=run$like))
}

##############################################################################
# Approximating a log-Chi^2(1) by a N(-1.270399;4.934854)
##############################################################################
x      = seq(-20,10,length=10000)
den    = exp(-(x)/2)*exp(-0.5*exp(x))*exp(x)/sqrt(2*pi)
norm   = dnorm(x,-1.270399,sqrt(4.934854))
par(mfrow=c(1,1))
plot(x,den,ylab="density",main="",type='n')
lines(x,den,col=1,lty=1,lwd=2)
lines(x,norm,col=2,lty=2,lwd=2)
legend(-15,0.2,legend=c("log chi^2_1","normal"),lty=1:2,col=1:2,lwd=c(2,2),bty="n")
title("E(log chi^2_1)=-1.270399 -- V(log chi^2_1)=4.934854")

#####################################################################################
# STOCHASTIC VOLATILITY: NORMAL APPROXIMATION + FORWARD FILTERING BACKWARD SAMPLING #            
#####################################################################################
set.seed(1243)
gamma     = -0.00645
G         = 0.99
W         = 0.15^2
lambda0   = 0.0
n         = 1000
lambda    = rep(0,n)
error     = rnorm(n,0.0,sqrt(W))
lambda[1] = gamma+G*lambda0+error[1]
for (t in 2:n)
  lambda[t] = gamma+G*lambda[t-1]+error[t]
h = exp(lambda/2)
y = rnorm(n,0,h)
plot(y,type="l")

# Prior hyperparameters
b         = c(0,0)
A         = diag(0.00000001,2)
v         = 0.00000002
lam       = 1.00000000
a1        = 0
R1        = 1000

# Initial values
lambdas   = lambda
pars      = c(gamma,G,sqrt(W))
M0        = 100
M         = 100

# Transformations
y1    = log(y^2)
alpha = -1.270399
V     = 4.934854
output  = ffbsu(y1,1,alpha,V,G,gamma,W,a1,R1,nd=1)
lambda  = output$theta
prop    = output$like
logden  = sum(dnorm(y,0,exp(lambda/2),log=TRUE))+ 
          sum(dnorm(lambda[2:n],gamma+G*lambda[1:(n-1)],sqrt(W),log=TRUE))-prop
accept  = rep(0,M0+M)
accepted = 0
for (iter in 1:(M0+M)){
  print(iter)
  output  = ffbsu(y1,1,alpha,V,G,gamma,W,a1,R1,nd=1)
  lambda1 = output$theta
  prop    = output$like
  lognum  = sum(dnorm(y,0,exp(lambda1/2),log=TRUE))+sum(dnorm(lambda1[2:n],gamma+G*lambda1[1:(n-1)],sqrt(W),log=TRUE))-prop
  accept[iter] = min(0,lognum-logden)
  if (log(runif(1))<accept[iter]){
    lambda = lambda1
    accepted = accepted + 1
    logden = lognum
  }
  X       = cbind(1,lambda[2:n])
  param   = fixedpar(lambda[1:(n-1)],X,A,b,v,lam)
  gamma   = param[1]
  G       = param[2]
  W       = param[3]
  pars    = rbind(pars,c(param[1:2],sqrt(param[3])))
  lambdas = rbind(lambdas,lambda)
}

hs   = exp(lambdas[(M0+1):(M0+M),])
mh   = apply(hs,2,mean)
h05  = apply(hs,2,quant05)
h95  = apply(hs,2,quant95)
hsn  = hs
mhn  = mh
h05n = h05
h95n = h95
parsn = pars

names = c("mu","phi","sigma2")
par(mfrow=c(3,2))
for (i in 1:3){
  plot(parsn[(M0+1):(M0+M),i],type="l",xlab="iteration",ylab="",main=names[i])
  abline(h=parsn[1,i],col=2)
  hist(parsn[(M0+1):(M0+M),i],xlab="",ylab="",main="")
  abline(v=parsn[1,i],col=2)
}

par(mfrow=c(1,1))
plot(1+h,type="l",ylim=c(0.0,max(1+h95n)),xlab="time",ylab="")
lines(1+h,col=1,lwd=1)
lines(1+mhn,col=2,lwd=1)
lines(1+h05n,col=3,lwd=1)
lines(1+h95n,col=3,lwd=1)
legend(400,4,legend=c("h(t)","E(h(t)|data)","90%CI"),col=c(1,2,3),lty=rep(1,3),lwd=rep(1,3),bty="n")
lines(abs(y)/max(abs(y)),type="h")


##############################################################################
# Defining the mixture of 7 normals that approximate                         #
# the log-chi-square with one degree of freedom                              #
##############################################################################
set.seed(1576)
mu     = c(-11.40039,-5.24321,-9.83726,1.50746,-0.65098,0.52478,-2.35859)
sig2   = c(  5.79596, 2.61369, 5.17950,0.16735, 0.64009,0.34023, 1.26261)
q      = c(  0.00730, 0.10556, 0.00002,0.04395, 0.34001,0.24566, 0.25750)
mm     = sum(q*mu)
vv     = sum(q*(sig2+mu^2))-(mm)^2
M      = 10000
x      = seq(-20,10,length=M)
den    = exp(-(x)/2)*exp(-0.5*exp(x))*exp(x)/sqrt(2*pi)
mix    = rep(0,M)
for (i in 1:M) 
  mix[i] = sum(q*dnorm(x[i],mu,sqrt(sig2)))
norm   = dnorm(x,mm,sqrt(vv))
par(mfrow=c(1,1))
plot(x,den,ylab="density",main="",type='n')
lines(x,den,col=1,lty=1,lwd=2)
lines(x,mix,col=2,lty=2,lwd=2)
legend(-15,0.2,legend=c("log chi^2_1","mixture of normals"),lty=1:2,col=1:2,lwd=c(2,2),bty="n")

###################################################################################
# Comparing the mixture of 7 normals X log-chi-square with one degree of freedom  #
###################################################################################
M   = 90000
ind = sample(1:7,replace=T,size=M,prob=q)
x1  = rnorm(M,mu[ind],sqrt(sig2[ind]))
xclass = seq(min(x1),max(x1),length=50)
hist(x1,prob=T,ylim=c(0,max(den)),xlab="",breaks=xclass,xlim=range(x),main="")
title(paste("Sample size=",M,sep=""))
lines(x,den,col=1,lty=1)
lines(x,mix,col=2,lty=1)
legend(-15,0.1,legend=c("log chi^2_1","mixture of normals"),lty=c(1,1),col=1:2,bty="n")
text(-12,0.24,paste("Sample mean=",round(mean(x1),5),sep=""))
text(-12,0.23,paste("Sample variance=",round(var(x1),5),sep=""))
text(-12,0.21,paste("Mixture mean=",round(mm,5),sep=""))
text(-12,0.20,paste("Mixture variance=",round(vv,5),sep=""))

##############################################################################
#  STOCHASTIC VOLATILITY + FORWARD FILTERING BACKWARD SAMPLING               #
##############################################################################
set.seed(1243)
gamma     = -0.00645
G         = 0.99
W         = 0.15^2
lambda0   = 0.0
n         = 1000
lambda    = rep(0,n)
error     = rnorm(n,0.0,sqrt(W))
lambda[1] = gamma+G*lambda0+error[1]
for (t in 2:n)
  lambda[t] = gamma+G*lambda[t-1]+error[t]
h = exp(lambda/2)
y = rnorm(n,0,h)
plot(y,type="l")

# Prior hyperparameters
b         = c(0,0)
A         = diag(0.00000001,2)
v         = 0.00000002
lam       = 1.00000000
a1        = 0
R1        = 1000

# Initial values
lambdas   = lambda
pars      = c(gamma,G,sqrt(W))
M0        = 100
M         = 100
output    = svu(y,lambda,gamma,G,W,a1,R1)
lambda    = output$lambda
prop      = output$like
logden    = sum(dnorm(y,0,exp(lambda/2),log=TRUE))+ 
            sum(dnorm(lambda[2:n],gamma+G*lambda[1:(n-1)],sqrt(W),log=TRUE))-prop
accept    = rep(0,M0+M)
accepted = 0
for (iter in 1:(M0+M)){
	print(iter)
  output  = svu(y,lambda,gamma,G,W,a1,R1)
  lambda1 = output$lambda
  prop    = output$like
  lognum  = sum(dnorm(y,0,exp(lambda1/2),log=TRUE))+ 
            sum(dnorm(lambda1[2:n],gamma+G*lambda1[1:(n-1)],sqrt(W),log=TRUE))-prop
  accept[iter] = min(0,lognum-logden)
  if (log(runif(1))<accept[iter]){
    lambda = lambda1
    accepted = accepted + 1
    logden = lognum
  }
  X       = cbind(1,lambda[2:n])
  param   = fixedpar(lambda[1:(n-1)],X,A,b,v,lam)
  gamma   = param[1]
  G       = param[2]
  W       = param[3]
  pars    = rbind(pars,c(param[1:2],sqrt(param[3])))
  lambdas = rbind(lambdas,lambda)
}
date()
hs  = exp(lambdas[(M0+1):(M0+M),])
mh  = apply(hs,2,mean)
h05 = apply(hs,2,quant05)
h95 = apply(hs,2,quant95)
hsm   = hs
mhm   = mh
h05m  = h05
h95m  = h95
parsm = pars

names = c("mu","phi","sigma2")
par(mfrow=c(3,2))
for (i in 1:3){
  plot(parsm[(M0+1):(M0+M),i],type="l",xlab="iteration",ylab="",main=names[i])
  abline(h=parsm[1,i],col=2)
  hist(parsm[(M0+1):(M0+M),i],xlab="",ylab="",main="")
  abline(v=parsm[1,i],col=2)
}

par(mfrow=c(1,1))
plot(1+h,type="l",ylim=c(0.0,max(1+h95m)),xlab="time",ylab="")
lines(1+h,col=1,lwd=1)
lines(1+mhm,col=2,lwd=1)
lines(1+h05m,col=3,lwd=1)
lines(1+h95m,col=3,lwd=1)
legend(400,4,legend=c("h(t)","E(h(t)|data)","90%CI"),col=c(1,2,3),lty=rep(1,3),lwd=rep(1,3),bty="n")
lines(abs(y)/max(abs(y)),type="h")

# Comparison
par(mfrow=c(1,1))
plot(h,type="l",ylim=c(0.0,max(h,mhn,mhm)),xlab="time",ylab="")
lines(h,col=1,lwd=1.5)
lines(mhn,col=2,lwd=1.5)
lines(mhm,col=3,lwd=1.5)
legend(400,2.5,legend=c("h(t)","normal approximation","mixture normals"),
col=c(1,2,3),lty=rep(1,3),lwd=rep(1.5,3),bty="n")

# Comparison
par(mfrow=c(2,2))
plot(1:250,h[1:250],type="l",ylim=c(0.0,max(h,mhn,mhm)),xlab="time",ylab="")
lines(1:250,h[1:250],col=1,lwd=1.5)
lines(1:250,mhn[1:250],col=2,lwd=1.5)
lines(1:250,mhm[1:250],col=3,lwd=1.5)
plot(251:500,h[251:500],type="l",ylim=c(0.0,max(h,mhn,mhm)),xlab="time",ylab="")
lines(251:500,h[251:500],col=1,lwd=1.5)
lines(251:500,mhn[251:500],col=2,lwd=1.5)
lines(251:500,mhm[251:500],col=3,lwd=1.5)
legend(300,2.5,legend=c("h(t)","normal approximation","mixture normals"),
col=c(1,2,3),lty=rep(1,3),lwd=rep(1.5,3),bty="n")
plot(501:750,h[501:750],type="l",ylim=c(0.0,max(h,mhn,mhm)),xlab="time",ylab="")
lines(501:750,h[501:750],col=1,lwd=1.5)
lines(501:750,mhn[501:750],col=2,lwd=1.5)
lines(501:750,mhm[501:750],col=3,lwd=1.5)
plot(501:750,h[751:n],type="l",ylim=c(0.0,max(h,mhn,mhm)),xlab="time",ylab="")
lines(501:750,h[751:n],col=1,lwd=1.5)
lines(501:750,mhn[751:n],col=2,lwd=1.5)
lines(501:750,mhm[751:n],col=3,lwd=1.5)