## code to accompany the sparrow Poisson regression example
## from lecture 16 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy, adapted from Peter Hoff

library(mvtnorm) ## for sampling from MVNs
library(coda) ## for ESS

## load necessary libraries & data
sparrow <- read.table("sparrow.txt", header=TRUE)

## present a plot of the data
boxplot(sparrow$fledged[sparrow$age == 1],
        sparrow$fledged[sparrow$age == 2],
        sparrow$fledged[sparrow$age == 3],
        sparrow$fledged[sparrow$age == 4],
        sparrow$fledged[sparrow$age == 5],
        sparrow$fledged[sparrow$age == 6],
        names=1:6, xlab="age", ylab="offspring")

## set up the design matrix with intercept and quadratic
X <- cbind(rep(1,nrow(sparrow)), sparrow$age, sparrow$age^2)
p <- ncol(X)

## set up the response
y <- sparrow$fledged
n <- length(y)

## priors mean and variance
beta.0 <- rep(0, p)
Sdiag.0 <- rep(10^2, p)

## set up the proposal variance
Sq <- var(log(y+1/2))*solve(t(X) %*% X)

## allocate space for MCMC samples
S <- 10000
beta <- matrix(NA, nrow=S, ncol=p)

## initial value
beta[1,] <- rep(0,p)

## MH iterations
for(s in 2:S) {

  ## propose a new value of beta centered at the old one
  beta.p <- drop(rmvnorm(1, beta[s-1,], Sq))

  ## posterior of proposal in log space
  lpost.p <- sum(dpois(y, exp(X %*% beta.p), log=TRUE))
  lpost.p <- lpost.p + sum(dnorm(beta.p, beta.0, sqrt(Sdiag.0), log=TRUE))

  ## posterior of s-1^th sample in log space
  lpost <- sum(dpois(y, exp(X %*% beta[s-1,]), log=TRUE))
  lpost <- lpost + sum(dnorm(beta[s-1,], beta.0, sqrt(Sdiag.0), log=TRUE))

  ## accept or reject
  if(runif(1) < exp(lpost.p-lpost)) beta[s,] <- beta.p ## accpet
  else beta[s,] <- beta[s-1,] ## reject
}


## plot the chains
par(mfrow=c(3,1), mar=c(5,4,1,1))
plot(beta[,1], type="l", xlab="b1")
plot(beta[,2], type="l", xlab="b2")
plot(beta[,3], type="l", xlab="b3")

## calculating ESSs
effectiveSize(beta[-(1:1000),1])
effectiveSize(beta[-(1:1000),2])
effectiveSize(beta[-(1:1000),3])

## plot the posterior densities
par(mfrow=c(2,2))
plot(density(beta[-(1:1000),1]),
     main="", bty="n", lwd=2, xlab="b1", ylab="")
plot(density(beta[-(1:1000),2]),
     main="", bty="n", lwd=2, xlab="b2", ylab="")
plot(density(beta[-(1:1000),3]),
     main="", bty="n", lwd=2, xlab="b3", ylab="")

## set up a design matrix for the posterior predictive
agep <- seq(0,6,length=100)
Xp <- cbind(rep(1, length(agep)), agep, agep^2)

## plot the posterior predictive median and 95% quantiles
lambdap <-  t(exp(Xp %*% t(beta[-(1:1000),])))
q1p <- apply(qpois(0.025, lambdap), 2, mean)
q2p <- apply(qpois(0.5, lambdap), 2, mean)
q3p <- apply(qpois(0.975, lambdap), 2, mean)
matplot(agep, cbind(q1p,q2p,q3p), type="l", col=1,
        lty=c(2,1,2), lwd=2, bty="n", xlab="age",
        ylab="number of offspring")

## add on the posterior predictive median and 95% quantiles
## of the mean
q1pm <- apply(lambdap, 2, quantile, probs=0.025)
q2pm <- apply(lambdap, 2, quantile, probs=0.5)
q3pm <- apply(lambdap, 2, quantile, probs=0.975)
matplot(agep, cbind(q1pm,q2pm,q3pm), type="l", col=2,
        lty=c(2,1,2), lwd=2, add=TRUE)
legend("topright", c("pred", "mean"), col=1:2, lwd=2, bty="n")