## code to accompany the mice intestinal tumor example from
## lecture 18 of PSTAT 215A, Bayesian Inference, UCSB
##
## Robert B. Gramacy, adapted from Peter Hoff


## load libraries and other files
library(mvtnorm)
library(coda)
source("rwish.R")

## read in the data file, and construct Y
mice <- read.table("mice.txt", header=TRUE)
Y <- t(mice)
n <- nrow(Y)
m <- ncol(Y)

## Plotting the data
matplot(Y, col=1:21, type="l", bty="l",
        ylab="number of tumors", xlab="location")
lines(apply(Y, 1, mean), lwd=3)

## create the design matrix (common to all mice)
x <- (1:n)/n
X <- cbind(rep(1, n), x, x^2, x^3, x^4)
p <- ncol(X)

## priors
ols <- matrix(NA, nrow=m, ncol=p)
for(j in 1:m) ols[j,] <- lm(log(Y[,j] + 1/n)~-1+X)$coef
mu.0 <- apply(ols, 2, mean)
Lambda.0 <- S.0 <- cov(ols)
Si.0 <- Lambdai.0 <- solve(S.0)
eta.0 <- p + 2

## allocate memory for samples
S <- 50000
beta.0 <- matrix(NA, nrow=S, ncol=p)
Sigma.0 <- array(NA, dim=c(S,p,p))

## initial values
beta <- array(NA, dim=c(m, S, p))
for(j in 1:m) beta[j,1,] <- ols[j,]
beta.0[1,] <- mu.0
Sigma.0[1,,] <- S.0
Sigmai.0 <- solve(Sigma.0[1,,])

## begin MCMC iterations
for(s in 2:S) {

  ## update beta.0
  Lambda.m <- solve(Lambdai.0 + m*Sigmai.0)
  mu.m <- Lambda.m %*% (Lambdai.0 %*% mu.0 + Sigmai.0 %*% apply(beta[,s-1,], 2, sum))
  beta.0[s,] <- drop(rmvnorm(1, mu.m, Lambda.m))

  ## update Sigma.0
  S.beta0 <- matrix(0, nrow=p, ncol=p)
  for(j in 1:m) ## calculate S.beta0 for Sigma.0
    S.beta0 <- S.beta0 + (beta[j,s-1,] - beta.0[s,]) %*% t(beta[j,s-1,] - beta.0[s,])
  Sigmai.0 <- rwish(eta.0 + m, solve(S.0 + S.beta0))
  Sigma.0[s,,] <- solve(Sigmai.0)

  ## update each beta
  for(j in 1:m) {

    ## propose a new beta vector
    beta.p <- drop(rmvnorm(1, beta[j,s-1,], 0.5*Sigma.0[s,,]))

    ## calculate the log posterior probability of the proposal
    lpost.p <- sum(dpois(Y[,j], exp(X %*% beta.p), log=TRUE))
    lpost.p <- lpost.p + dmvnorm(beta.p, beta.0[s,], Sigma.0[s,,], log=TRUE)

    ## calculate the log posterior probability of the previous beta
    lpost <- sum(dpois(Y[,j], exp(X %*% beta[j,s-1,]), log=TRUE))
    lpost <- lpost + dmvnorm(beta[j,s-1,], beta.0[s,], Sigma.0[s,,], log=TRUE)

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

## save the image of this long run
save.image(file="mice.Rsave")


## caluclate the ESS of the components of beta.0
apply(beta.0[-(1:1000),], 2, effectiveSize)

## can plot traces of the samples of components of beta
plot(beta[4,-(1:1000),5], type="l", main="", xlab="t",
     ylab="beta[4,5]", bty="n")


## obtain samples from the predictive
lambda.0 <- lambda <- Ytilde <- matrix(NA, nrow=S, ncol=n)
for(s in 1:S) {
  lambda.0[s,] <- exp(X %*% beta.0[s,])
  beta.new <- drop(rmvnorm(1, beta.0[s,], Sigma.0[s,,]))
  lambda[s,] <- exp(X %*% beta.new)
  Ytilde[s,] <- rpois(n, lambda[s,])
}


## plotting lambda.0
matplot(t(apply(lambda.0[-(1:1000),], 2, quantile, c(0.025, 0.5, 0.975))),
        type="l", lwd=2, lty=c(2,1,2), col=c(2,1,2), ylab="number of tumors",
        xlab="location", ylim=c(0,20), bty="n", main="lambda.0")

## plotting lambda
matplot(t(apply(lambda[-(1:1000),], 2, quantile, c(0.025, 0.5, 0.975))),
        type="l", lwd=2, lty=c(2,1,2), col=c(2,1,2), ylab="number of tumors",
        xlab="location", ylim=c(0,20), bty="n", main="lambda")

## plotting Ytilde
matplot(t(apply(Ytilde[-(1:1000),], 2, quantile, c(0.025, 0.5, 0.975))),
        type="l", lwd=2, lty=c(2,1,2), col=c(2,1,2), ylab="number of tumors",
        xlab="location", ylim=c(0,20), bty="n", main="Ytilde")