## code to accompany the rejection sampling example from
## lecture 6 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy

## set up the distribution we want to sample from, f
a <- 2
x <- seq(0,10,length=1000)
fx <- dgamma(x, a, 1)

## set up the optimal rejection sampler, g
bstar <- 1/a
Mstar <- a^a * exp(-(a-1))/gamma(a)
Mgx <- Mstar*dgamma(x, 1, bstar)

## plot the f density and the optimal M*g function
ylim <- range(Mgx)
par(mfrow=c(1,1), cex.main=1.5, cex.lab=1.5, cex.axis=1.5)
main <- paste("a=", a, ", 1/Mstar=", signif(1/Mstar,3), sep="")
fname <- paste("Gamma[", a, ",1]", sep="")
gname <- paste("M*Gamma[1,", bstar, "]", sep="")
plot(x, fx, type="l", lwd=2, col="blue", ylim=ylim, main=main)
lines(x, Mgx, lwd=2, lty=2, col="red")
legend("topright", c(fname, gname), col=c("blue", "red"),
       lty=c(1,2), lwd=2, cex=1.5)

## wait until we're ready to start
readline("press RETURN to start rejection sampling: ")

## start taking n samples
n <- 50
xs <- ys <- NULL
par(mfrow=c(2,1), cex.main=1.5, cex.lab=1.5, cex.axis=1.5)

for(i in 1:n) {

  ## sample from g
  y <- rgamma(1, 1, bstar)
  ys <- c(ys, y)

  ## calculate the acceptance ratio
  ratio <- dgamma(y, a, 1)/(Mstar*dgamma(y, 1, bstar))

  ## accept or reject
  u <- runif(1)
  if(u < ratio) {
    col <- "blue"
    xs <- c(xs, y)
  } else col <- "red"

  ## re-plot the f and M*g lines
  plot(x, fx, type="l", lwd=2, col="blue", ylim=ylim, main=main,
       cex.lab=1.5, cex.axis=1.5, cex.main=1.5)
  lines(x, Mgx, lwd=2, lty=2, col="red")
  legend("topright", c(fname, gname), col=c("blue", "red"),
         lty=c(1,2))

  ## add info about the proposed y value
  text(x=4.5, y=0.27,
       labels=paste(signif(y,2), "~g    f/Mg=", signif(ratio,2),
         "    ", signif(u, 2), "~U", sep=""), pos=4)
  abline(v=y, col=col, lty=3, lwd=2)

  ## put all proposed y values in red, and accepted x values in blue
  points(ys, rep(0, length(ys)), col="red", cex=1.5)
  points(xs, rep(0, length(xs)), col="blue", cex=1.5)

  ## make histograms of proposed y samples and accepted x samples
  hmain <- paste("acceptance rate: ", signif(length(xs)/length(ys),4), sep="")
  hist(ys, col="red", xlim=range(x), main=hmain)
  if(length(xs) >= 1) {
    xh <- hist(xs, plot=FALSE)
    plot(xh, add=TRUE, col="blue")
  }

  ## wait until we're ready for the next sample
  readline(paste("i=", i, ", press RETURN to continue: ", sep=""))
}