vsva = function(y,X,Z=X,backward=TRUE,restrict=TRUE,a=1+1e-4,b=1+1e-4,tol=1e-10)
{
# Details: 	Variable Selection with the Vriational Approximation ranking algorithm
#		for the heteroscedastic linear regression models. 
#		One can restrict (if restrict = T) the search for inclusion 
#		in the variance model to those that have been included in the mean model.
#		The search can combine both forward and backward moves (if backward=T).
# 
# References: "Variational approximation for heteroscedastic models and matching 
#		pursuit algorithms" by David Nott et al.
# Author:	Minh-Ngoc Tran (ngoctm@nus.edu.sg) with a helper function from David Nott
# Inputs:
#
# y			vector of responses
# X			design matrix for the mean model (including the intercept, if desired)
# Z			design matrix for the variance model (including the intercept, if desired). By default, Z:=X.  
# restrict		whether or not restrict the search in the variance model. If restrict=TRUE, Z:=X by default
# backward		whether or not to do backward selection besides forward selection
# a, b		hyperparameters in the inverse gamma priors for sigma2
# tol			stopping tolerance for variational lower bound
#
# Outputs:
# 
# C		logical vector indicating variables selected in the mean model
# V		logical vector indicating variables selected in the variance model
# mub		vector of estimated coefficients in the mean model
# muv		vector of estimated coefficients in the variance model
#
#==============================================

if (restrict) Z=X
p = dim(X)[2]; q = dim(Z)[2]; n = dim(X)[1]
C = c(TRUE,rep(FALSE,p-1)) # initialize the active set for mean model
V = c(TRUE,rep(FALSE,q-1)) # initialize the active set for variance model
fit = vbhetlm1(y,as.matrix(X[,C]),as.matrix(Z[,V]),a,b,tol)
Lopt = fit$lb[length(fit$lb)]

lmprior = function(C,V) {							# model prior 
lmprior = logfac(sum(C))+logfac(p-sum(C))+logfac(sum(V))+logfac(q-sum(V))
return(lmprior)
}
logfac <- function(x) {	if (x==0) return(0) else return(sum(log(1:x))) }

# forward search
stoprule = FALSE
while (!stoprule) {
	oldC = C; oldV = V # store the mean and variance model at the beginning of a circile

	mub = fit$mub
	muv = fit$muv
	sv = fit$sv
	sig2b = fit$sig2b
	XC = as.matrix(X[,C])
	ZV = as.matrix(Z[,V])	
	L = rep(0,p)
	for (j in 1:p) {
		if (C[j]) L[j] = -Inf
		else {
			xj = X[,j]
			sig2bjhat = (1/sig2b+sum(xj*xj/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum))))
			sig2bjhat = 1/sig2bjhat
			mubjhat = sum(xj*(y-XC%*%mub)/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum)))
			mubjhat = mubjhat*sig2bjhat
			L[j] = .5*log(sig2bjhat/sig2b)+.5*(1/sig2bjhat-1/sig2b)*mubjhat^2	#+lmprior(Ctry,V)
		}
	}
	for (j in 1:p) if (L[j]==max(L)) {jhat = j; break }
	Ctry = C; Ctry[jhat] = TRUE
	fittry = vbhetlm1(y,as.matrix(X[,Ctry]),as.matrix(Z[,V]),a,b,tol)
	Ltry = fittry$lb[length(fittry$lb)]
	if (Ltry+lmprior(Ctry,V)>Lopt+lmprior(C,V)) { C = Ctry; Lopt = Ltry; fit = fittry}
	
	mub = fit$mub
	muv = fit$muv
	sb = fit$sb
	sv = fit$sv
	sig2v = fit$sig2v
	XC = as.matrix(X[,C])
	L = rep(0,q)
	if (restrict) L[V|!C] = -Inf else L[V] = -Inf
	for (j in 1:q) {
		if (L[j]==0) {
			zj = Z[,j]
 			v = (y-XC%*%mub)*(y-XC%*%mub)+apply((XC%*%sb)*XC,1,sum)
			v = v/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum))
			mu0 = .5*sum(zj*(v-1))/(1/sig2v+.5*sum(zj^2*v))
			muvjhat = Newton(zj,sig2v,v,mu0,tol) 		 
			sig2vjhat = 1/sig2v+.5*sum((zj^2*v)/exp(zj*muvjhat))
			sig2vjhat = 1/sig2vjhat
			vec = 1/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum)+zj*muvjhat-.5*zj^2*sig2vjhat)
			vec = vec - 1/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum))
			vec = vec*((y-XC%*%mub)*(y-XC%*%mub)+apply((XC%*%sb)*XC,1,sum))
			L[j] = .5+.5*log(sig2vjhat/sig2v)-.5*sig2vjhat/sig2v
			L[j] = L[j] -.5*muvjhat^2/sig2v-.5*sum(zj*muvjhat) - .5*sum(vec)	#+lmprior(C,Vtry)			
		}
	}
	for (j in 1:q) if (L[j]==max(L)) {jhat = j; break }
	Vtry = V; Vtry[jhat] = TRUE
	fittry = vbhetlm1(y,as.matrix(X[,C]),as.matrix(Z[,Vtry]),a,b,tol)
	Ltry = fittry$lb[length(fittry$lb)]
	if (Ltry+lmprior(C,Vtry)>Lopt+lmprior(C,V)) { V = Vtry; Lopt = Ltry; fit = fittry}

	if ((sum(oldC == C)==p)&&(sum(oldV == V)==q)) stoprule = TRUE
}

# backward search
if (backward) {
stoprule = FALSE
while (!stoprule) {
	oldC = C; oldV = V # store the mean and variance model at the beginning of a circile

	if ((restrict)&&(sum(C&V)>1)) {
		mub = as.matrix(rep(0,p))
		mub[C] = as.numeric(fit$mub)
		muv = as.matrix(rep(0,q))
		muv[V] = as.numeric(fit$muv)
		sv = matrix(rep(0,q*q),ncol=q)	
		sv[V,V] = as.numeric(fit$sv)
		sig2v = fit$sig2v
		sig2b = fit$sig2b
		L = rep(0,p)
		for (j in 1:p) {
			if (!C[j]) {L[j] = Inf}
			else {
				Cnew = C; Cnew[j] = FALSE
				Vnew = V; Vnew[j] = FALSE
				XC = as.matrix(X[,Cnew]); ZV = as.matrix(Z[,Vnew])
				xj = X[,j]
				mub.new = mub[Cnew]; muv.new = muv[Vnew]; sv.new = sv[Vnew,Vnew]
				sig2bjhat = (1/sig2b+sum(xj*xj/exp(ZV%*%muv.new-1/2*apply((ZV%*%sv.new)*ZV,1,sum))))
				sig2bjhat = 1/sig2bjhat
				mubjhat = sum(xj*(y-XC%*%mub.new)/exp(ZV%*%muv.new-1/2*apply((ZV%*%sv.new)*ZV,1,sum)))
				mubjhat = mubjhat*sig2bjhat
				L[j] = .5*log(sig2bjhat/sig2b)+.5*(1/sig2bjhat-1/sig2b)*mubjhat^2	#-lmprior(Cnew,Vnew)
			}
		}
		for (j in 1:p) if (L[j]==min(L)) {jhat = j; break }
		Ctry = C; Ctry[jhat] = FALSE; Vtry = V; Vtry[jhat] = FALSE
		fittry = vbhetlm1(y,as.matrix(X[,Ctry]),as.matrix(Z[,Vtry]),a,b,tol)
		Ltry = fittry$lb[length(fittry$lb)]
		if (Ltry+lmprior(Ctry,Vtry)>Lopt+lmprior(C,V)) { C = Ctry; V = Vtry; Lopt = Ltry; fit = fittry}
	} 
	if ((!restrict)&&(sum(C)>1)) {
		mub = as.matrix(rep(0,p))
		mub[C] = as.numeric(fit$mub)
		muv = fit$muv
		sv = fit$sv
		sig2b = fit$sig2b
		ZV = as.matrix(Z[,V])	
		L = rep(0,p)
		for (j in 1:p) {
			if (!C[j]) {L[j] = Inf} 
			else {
				Cnew = C; Cnew[j] = FALSE
				XC = as.matrix(X[,Cnew])
				xj = X[,j]
				sig2bjhat = (1/sig2b+sum(xj*xj/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum))))
				sig2bjhat = 1/sig2bjhat
				mub.new = mub[Cnew]
				mubjhat = sum(xj*(y-XC%*%mub.new)/exp(ZV%*%muv-1/2*apply((ZV%*%sv)*ZV,1,sum)))
				mubjhat = mubjhat*sig2bjhat
				L[j] = .5*log(sig2bjhat/sig2b)+.5*(1/sig2bjhat-1/sig2b)*mubjhat^2	#-lmprior(Cnew,V)
			}
		}
		for (j in 1:p) if (L[j]==min(L)) {jhat = j; break }
		Ctry = C; Ctry[jhat] = FALSE
		fittry = vbhetlm1(y,as.matrix(X[,Ctry]),ZV,a,b,tol)
		Ltry = fittry$lb[length(fittry$lb)]
		if (Ltry+lmprior(Ctry,V)>Lopt+lmprior(C,V)) { C = Ctry; Lopt = Ltry; fit = fittry}
	}
	if (sum(V)>1) {	
	mub = fit$mub
	sb = fit$sb
	muv = as.matrix(rep(0,q))
	muv[V] = as.numeric(fit$muv)
	sv = matrix(rep(0,q*q),ncol=q)	
	sv[V,V] = as.numeric(fit$sv)
	sig2v = fit$sig2v
	XC = as.matrix(X[,C])
	L = rep(0,q)
	for (j in 1:q) {
		if (!V[j]) L[j] = Inf
		else {
			Vnew = V; Vnew[j] = FALSE
			ZV = as.matrix(Z[,Vnew])
			zj = Z[,j]
			muv.new = muv[Vnew]
			sv.new = sv[Vnew,Vnew]
 			v = (y-XC%*%mub)*(y-XC%*%mub)+apply((XC%*%sb)*XC,1,sum)
			v = v/exp(ZV%*%muv.new-1/2*apply((ZV%*%sv.new)*ZV,1,sum))
			mu0 = .5*sum(zj*(v-1))/(1/sig2v+.5*sum(zj^2*v))
			muvjhat = Newton(zj,sig2v,v,mu0,tol) 		 
			sig2vjhat = 1/sig2v+.5*sum((zj^2*v)/exp(zj*muvjhat))
			sig2vjhat = 1/sig2vjhat
			vec = 1/exp(ZV%*%muv.new-1/2*apply((ZV%*%sv.new)*ZV,1,sum)+zj*muvjhat-.5*zj^2*sig2vjhat)
			vec = vec - 1/exp(ZV%*%muv.new-1/2*apply((ZV%*%sv.new)*ZV,1,sum))
			vec = vec*((y-XC%*%mub)*(y-XC%*%mub)+apply((XC%*%sb)*XC,1,sum))
			L[j] = .5+.5*log(sig2vjhat/sig2v)-.5*sig2vjhat/sig2v
			L[j] = L[j] -.5*muvjhat^2/sig2v-.5*sum(zj*muvjhat) - .5*sum(vec)	#-lmprior(C,Vnew)			
		}
	}
	for (j in 1:q) if (L[j]==min(L)) {jhat = j; break }
	Vtry = V
	Vtry[jhat] = FALSE
	fittry = vbhetlm1(y,as.matrix(X[,C]),as.matrix(Z[,Vtry]),a,b,tol)
	Ltry = fittry$lb[length(fittry$lb)]
	if (Ltry+lmprior(C,Vtry)>Lopt+lmprior(C,V)) { V = Vtry; Lopt = Ltry; fit = fittry}
	}
	if ((sum(oldC == C)==p)&(sum(oldV == V)==q)) stoprule = TRUE
}
}

mub = rep(0,p)
mub[C]=as.numeric(fit$mub)
muv = rep(0,q)
muv[V]=as.numeric(fit$muv)
list(C=C,V=V,mub=mub,muv=muv)
}

Newton = function(zj,sig2v,v,mu0,tol) {
# helper function to estimate the mode by the Newton method

xold = mu0
fvalold = .5*xold^2/sig2v+.5*sum(zj)*xold+.5*sum(v*exp(-xold*zj))
Df = xold/sig2v+.5*sum(zj)-.5*sum(zj*v*exp(-xold*zj))
D2f = 1/sig2v+.5*sum(zj*zj*v*exp(-xold*zj))
xnew = xold-Df/D2f
fvalnew = .5*xnew^2/sig2v+.5*sum(zj)*xnew+.5*sum(v*exp(-xnew*zj))

while((fvalold-fvalnew)>tol) {
  xold = xnew
  fvalold = fvalnew
  Df = xold/sig2v+.5*sum(zj)-.5*sum(zj*v*exp(-xold*zj))
  D2f = 1/sig2v+.5*sum(zj*zj*v*exp(-xold*zj))
  xnew = xold-Df/D2f
  fvalnew = .5*xnew^2/sig2v+.5*sum(zj)*xnew+.5*sum(v*exp(-xnew*zj))
}
return(xnew)
}

vbhetlm1<-function(y,X,Z,a=1+1e-5,b=1+1e-5,tol=1e-10) {
# Details:	Maximizing the lower bound for the heteroscedastic linear regression model.
# 		The prior mean vectors are assumed to be 0 and the prior covariance matrices
# 		are assumed to be the identity matrices, up to unknow scale parameters.
# References: "Variational approximation for heteroscedastic models and matching 
#		pursuit algorithms" by David Nott et al.
# Author:	Minh-Ngoc Tran (ngoctm@nus.edu.sg) with a helper function from David Nott
# Inputs:
# y		vector of responses
# X		design matrix for the mean model (you need to include the intercept, 
#   		if desired).
# Z		design matrix for the variance model (you need to include the intercept, 
#   		if desired).
# a,b		hyperparameters of the inverse gamma priors 
# tol		stopping tolerance for variational lower bound
# 
# Outputs:
# mub		mean in variational normal posterior for mean coefficients
# muv		mean in variational normal posterior for variance coefficients
# sb		covariance matrix in variational normal posterior for mean
#  		coefficients
# sv		covariance matrix in variational normal posterior for variance
#  		coefficients
# lb		vector of lower bound values over iterations
#
######################################################################

n<-length(y)
p<-dim(X)[2]
q<-dim(Z)[2]

vlb<-function(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv) {
  t1<-solve(psb,sb)
  t2<-solve(psv,sv)
  lb<-(p+q)/2-n/2*log(2*pi)+1/2*log(det(t1))+1/2*log(det(t2))
  lb<-lb-1/2*sum(diag(t1))-1/2*sum(diag(t2))
  v1<-mub-pmub
  lb<-lb-1/2*sum(v1*solve(psb,v1))
  v2<-muv-pmuv
  lb<-lb-1/2*sum(v2*solve(psv,v2))
  lv<-Z%*%muv
  lb<-lb-1/2*sum(lv)
  lm<-X%*%mub
  ln<-(y-lm)^2
  ln<-ln+apply((X%*%sb)*X,1,sum)
  ld<-exp(lv-1/2*apply((Z%*%sv)*Z,1,sum))
  lb<-lb-1/2*sum(ln/ld)
  return(lb)
}

lmfit = lm(y ~ X-1)
r = residuals(lmfit)
r<-log(r*r)
init.lm<-lm(r ~ Z-1)
muv<-coef(init.lm)
sv<-vcov(init.lm)
dd<-1/exp(Z%*%muv-1/2*apply((Z%*%sv)*Z,1,sum))
Xtilde<-matrix(rep(sqrt(dd),times=p),nrow=n)*X
ytilde<-sqrt(dd)*y
weighted.lm<-lm(ytilde ~ Xtilde-1)
mub<-weighted.lm$coef
sb<-vcov(weighted.lm)/summary(weighted.lm)$sigma^2
sig2b = (b+.5*sum(mub*mub)+.5*sum(diag(sb)))/(a+1+p/2)
psb = sig2b*diag(p)
sig2v = (b+.5*sum(muv*muv)+.5*sum(diag(sv)))/(a+1+q/2)
psv = sig2v*diag(q)
pmub = rep(0,p)
pmuv = rep(0,q)

lb<-vlb(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv)
lball<-lb
lbold<-lb-abs(lb)

while((lb-lbold)>tol) {
  lbold<-lb
  w<-(y-X%*%mub)
  w<-w*w+apply((X%*%sb)*X,1,sum)
  working.glm<-bgamglm(w,Z,pmuv,psv,muv,tol=tol)
  muv<-working.glm$coef
  sv = working.glm$variance

  sig2v = (b+.5*sum(muv*muv)+.5*sum(diag(sv)))/(a+1+q/2)
  psv = sig2v*diag(q)

  dd<-1/exp(Z%*%muv-1/2*apply((Z%*%sv)*Z,1,sum))
  D = Diagonal(x=as.numeric(dd))
  sb = solve(t(X)%*%D%*%X+solve(psb))
  mub = sb%*%(solve(psb,pmub)+t(X)%*%D%*%y)
  
  sig2b = (b+.5*sum(mub*mub)+.5*sum(diag(sb)))/(a+1+p/2)
  psb = sig2b*diag(p)

  lb<-vlb(y,X,Z,p,q,n,pmub,pmuv,psb,psv,mub,muv,sb,sv)
  lball<-c(lball,lb)
}
list(mub=mub,muv=muv,sb=sb,sv=sv,sig2b=sig2b,sig2v=sig2v,lb=lball)

} # end of function vbhetlm1


bgamglm<-function(w,Z,pm,pv,p0,tol) {
# Helper function for the function vbhetlm1.
# This function is to compute posterior mode for gamma generalized
# linear model with log link, known shape parameters
# and normal prior on coefficients
#
# w - vector of responses
# Z - design matrix
# pm, pv - prior mean and covariance matrix
# p0 - starting guess for mode
# tolerance for stopping

xold = p0
V = .5*w*exp(-Z%*%xold)
fvalold = .5*sum(Z%*%xold)+sum(V)+0.5*sum((xold-pm)*solve(pv,(xold-pm)))
Df = .5*colSums(Z)-t(Z)%*%V+solve(pv,xold-pm)
W = Diagonal(x = as.numeric(V))
D2f = t(Z)%*%W%*%Z+solve(pv)
xnew = xold - solve(D2f,Df)
V = .5*w*exp(-Z%*%xnew)
fvalnew = .5*sum(Z%*%xnew)+sum(V)+0.5*sum((xnew-pm)*solve(pv,(xnew-pm)))
fvalall = fvalold

while((fvalold-fvalnew)>tol) {
  xold = xnew
  fvalold = fvalnew
  fvalall = c(fvalall,fvalnew)
  Df = .5*colSums(Z)-t(Z)%*%V+solve(pv,xold-pm)
  W = Diagonal(x = as.numeric(V))
  D2f = t(Z)%*%W%*%Z+solve(pv)
  xnew = xold - solve(D2f,Df)
  V = .5*w*exp(-Z%*%xnew)
  fvalnew = .5*sum(Z%*%xnew)+sum(V)+0.5*sum((xnew-pm)*solve(pv,(xnew-pm)))
}
list(coef=xnew,variance=solve(D2f),fval=fvalall)
}


