## code to accompany the Data Augmentation example from
## lecture 18 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy

## data for a data augmentation examples
y <- c(125, 18, 20, 34)

## uniform (Beta) prior
alpha <- beta <- 1

## choose chain length and starting point
T <- 1000
theta <- rep(0.5, T)
z <- rep(y[1]/2, T)

## T Gibbs sampling/Data Augmentation rounds
for(t in 2:T) {

  ## samlping theta conditional on z
  theta[t] <- rbeta(1, z[t-1]+y[4]+alpha,
                    y[2]+y[3]+beta)

  ## sampling z condition on theta
  z[t] <- rbinom(1, y[1], theta[t]/(2+theta[t]))
}

par(mfrow=c(2,1), cex.main=1.5, cex.lab=1.5,
    cex.axis=1.5)

## plot a trace and histogram of the marginal
## posterior for the latent z variable
plot(z, type="l", bty="n", xlab="s")
hist(z[10:1000], bty="n", xlab="z", main="")

## plot a trace and histogram of the marginal
## posterior for the parameter theta
plot(theta, type="l", bty="n", xlab="s")
hist(theta[10:1000],  bty="n", xlab="theta", main="")