#######################################################
## You can copy and paste this code into your R GUI  ##
## to reproduce Figure 1 of the paper, "On the proof ##
## of posterior contraction for sparse generalized   ##
## linear models with multivariate responses" by     ##
## Shao-Hsuan Wang, Ray Bai, and Hsin-Hsiung Huang   ##
#######################################################


##################
# Load libraries #
##################
library(pracma)
library(rmutil)


#########################################
#########################################
## This code plots the central region  ##
## and the tails for various shrinkage ## 
## priors.                             ##
#########################################
#########################################

# Grid points for central region
beta.central <- seq(-3,3,0.02)
null.index <- which(beta.central == 0) # Removes the value of zero from the central region
beta.central <- beta.central[-null.index]

# Grid points for tails
beta.tail <- seq(3,8,0.1)


######################################
# Code to plot the Student's t prior #
# with 3 degrees of freedom          #
######################################
studentst.central.points = dt(beta.central, df=3)
studentst.tail.points = dt(beta.tail, df=3)

####################################
# Code to plot the horseshoe prior #   
####################################
hs.prior <- function(beta){
  K = 1/(pi*sqrt(2*pi))
  return(K*exp(beta^2/2)*expint(beta^2/2))
}

hs.central.points = hs.prior(beta.central)
hs.tail.points = hs.prior(beta.tail)


#############################################
# Code to plot the normal-exponential-gamma #
# (NEG) prior                               #
#############################################
NEG.prior <- function(beta, a){
  # Joint of f(beta, zeta).
  f <- function(zeta){
    (1/(2*pi*zeta))*exp(-(beta^2)/(2*zeta))*(gamma(a+1)/gamma(a))*(1+zeta)^(-(1+a))
  }
  # Integrate out zeta
  return(pracma::quadinf(f,1e-12,Inf,tol=1e-6)$Q)
}
NEG.prior.vec= Vectorize(NEG.prior, c("beta")) # Handles a vector as an input

a = 1
NEG.central.points = NEG.prior.vec(beta.central,a)
NEG.tail.points = NEG.prior.vec(beta.tail,a)


################################
# Code to plot the generalized #
# double Pareto prior          #
################################
GDP.prior <- function(beta, alpha, eta){
  
  return((alpha/(2*eta))*(1+abs(beta)/eta)^(-(alpha+1)))
}
GDP.prior.vec= Vectorize(GDP.prior, c("beta")) # Handles a vector as an input

alpha=eta=1
GDP.central.points = GDP.prior.vec(beta.central,alpha,eta)
GDP.tail.points = GDP.prior.vec(beta.tail,alpha,eta)



###################################
# Code to plot the HS+ prior with #
# sparsity parameter tau=1        #
###################################
hsplus.prior <- function(beta){
  
  # Joint of f(beta, zeta). See Supplementary material S.1 of
  # the Bhadra et al. (2017) paper
  f <- function(zeta){
    (1/(pi^2*sqrt(2*pi)))*exp(-(zeta*beta^2)/2)*(log(zeta)/(zeta-1))
  }
  # Integrate out zeta
  return(pracma::quadinf(f,1e-12,Inf,tol=1e-6)$Q)
}
hsplus.prior.vec <- Vectorize(hsplus.prior, c("beta")) # Handles a vector as an input

hsplus.central.points = hsplus.prior.vec(beta.central)
hsplus.tail.points = hsplus.prior.vec(beta.tail)


################################
################################
## Plot the shrinkage priors' ##
## central region and tails   ##
################################
################################
pdf(file="shrinkageillustration.pdf",height=7,width=7)

par(mfrow=c(2,1), mai = c(0.5, 0.5, 0.5, 0.5))
plot(c(0,0), xlim=c(-3,3),ylim=c(0,0.8),type="n",ylab="", xlab="", main="Comparison of priors: central region")
abline(v=0)
lines(beta.central,studentst.central.points,col="darkorchid",lty=4,lwd=4,cex=0.3)
lines(beta.central,hs.central.points, col="darkorange2",lty=2,lwd=4,cex=0.3)
lines(beta.central,NEG.central.points,col="darkgreen",type="o",cex=0.6,pch=0)
lines(beta.central,GDP.central.points,col="red",type="o",lty=6,lwd=1,cex=0.5,pch=24)
lines(beta.central,hsplus.central.points,col="blue",lty=1,lwd=3,cex=0.3)

legend("topright",legend=c(expression("t"[3]), "HS","NEG","GDP","HS+"),
       lty=c(4,2,1,6,1),lwd=c(2,2,2,2,2), pch=c(NA,NA,0,24,NA),
       col=c("darkorchid","darkorange2","darkgreen","red","blue"))

plot(c(0,0), xlim=c(3,8),ylim=c(0,0.035),type="n",ylab="", xlab="", main="Comparison of priors: tails")
abline(v=0)
lines(beta.tail,studentst.tail.points,col="darkorchid",lty=4,lwd=4,cex=0.3)
lines(beta.tail,hs.tail.points, col="darkorange2",lty=2,lwd=4,cex=0.3)
lines(beta.tail,NEG.tail.points,col="darkgreen",type="o",lty=1,lwd=1.5,cex=0.75,pch=0)
lines(beta.tail,GDP.tail.points,col="red",type="o",lty=6,lwd=1.5,cex=0.75,pch=24)
lines(beta.tail,hsplus.tail.points,col="blue",lty=1,lwd=3,cex=0.3)

legend("topright",legend=c(expression("t"[3]), "HS","NEG","GDP","HS+"),
       lty=c(4,2,1,6,1),lwd=c(2,2,2,2,2), pch=c(NA,NA,0,24,NA),
       col=c("darkorchid","darkorange2","darkgreen","red","blue"))
dev.off()