## Read in data
AllData <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/radon.txt")
Data <- subset(AllData, state=="MN")
county <- droplevels(Data$county)
x <- Data$floor
cid <- as.numeric(county) ## County ID
n <- nrow(Data)
J <- length(levels(county))

## Outcome
radon <- Data$activity
hist(radon, col="gray", border="white")
y <- log(ifelse(radon==0,.1, radon))
hist(y, col="gray", border="white")

## Identical parameters
fit1 <- lm(y~floor, Data)

## Independent parameters
fit2 <- lm(y~0+floor+county, Data)

## Hierarchical
jData <- list("y", "x", "cid", "n", "J")
model <- function() {
  ## Likelihood
  for (i in 1:n) {
    y[i] ~ dnorm(a[cid[i]] + b*x[i], sigma[1]^(-2))
  }
  ## Priors
  b ~ dnorm(0, 0.0001)
  for (j in 1:2) {
    sigma[j] ~ dgamma(0.001, 0.001)
  }
  for (j in 1:J) {
    a[j] ~ dnorm(mu, sigma[2]^(-2))
  }
  mu ~ dnorm(0, 0.0001)
}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, param=c("mu", "sigma", "a", "b"), data=jData, n.iter=10000, n.thin=1)
plot(fit)
fit
attach.jags(fit)

## Convergence
max(gelman.diag(as.mcmc(fit))$psrf)
min(effectiveSize(as.mcmc(fit)))

## Mu
source("http://web.as.uky.edu/statistics/users/pbreheny/760/S13/notes/estimate.R")
estimate(c(1,0), fit1)
estimate(c(0, rep(1/J, J)), fit2)
quantile(mu, c(.025, 0.5, .975))
B <- rbind(estimate(c(1,0), fit1)[1:3],
           estimate(c(0, rep(1/J,J)), fit2)[1:3],
           quantile(mu, c(.025, 0.5, .975))[c(2,1,3)])
rownames(B) <- c("Identical", "Independent", "Hierarchical")
CIplot(B, xlab=~mu)

## Beta
estimate(c(0,1), fit1)
estimate(c(1, rep(0, J)), fit2)
quantile(b, c(.025, 0.5, .975))
B <- rbind(estimate(c(0,1), fit1)[1:3],
           estimate(c(1, rep(0, J)), fit2)[1:3],
           quantile(b, c(.025, 0.5, .975))[c(2,1,3)])
rownames(B) <- c("Identical", "Independent", "Hierarchical")
CIplot(B, xlab=~beta)

## 3 counties
ind <- c(grep("COOK", levels(county)), grep("CARLTON", levels(county)), grep("ST LOUIS", levels(county)))
L <- matrix(0, 3, J+1)
L[1, ind[1]+1] <- 1
L[2, ind[2]+1] <- 1
L[3, ind[3]+1] <- 1
for (j in 1:3) {
  B <- rbind(estimate(c(1,0), fit1)[1:3],
             estimate(L[j,], fit2)[1:3],
             quantile(a[,ind[j]], c(.025, 0.5, .975))[c(2,1,3)])
  rownames(B) <- c("Identical", "Independent", "Hierarchical")
  CIplot(B, xlab=~beta, mar=c(5, 6, 2, 1))
  mtext(paste(capwords(levels(county)[ind[j]], strict=TRUE), "(n=", table(county)[ind[j]], ")", sep=""))
}

## Plot regression lines
par(mar=c(5,4,4,2))
for (j in 1:3) {
  ii <- cid==ind[j]
  plot(jitter(x[ii], amount=.05), y[ii], xlim=c(-.05,1.05), ylim=range(y), pch=19, col="gray", xlab="Floor", ylab="Log(Radon)")
  abline(fit1, col=pal(3)[1], lwd=3)
  abline(a=coef(fit2)[ind[j]+1], b=coef(fit2)[1], col=pal(3)[2], lwd=3)
  abline(a=median(a[,ind[j]]), b=median(b), col=pal(3)[3], lwd=3)
  mtext(paste(capwords(levels(county)[ind[j]], strict=TRUE), "(n=", table(county)[ind[j]], ")", sep=""))
  toplegend(c("Identical", "Independent", "Hierarchical"), col=pal(3), lwd=3, ncol=3)  
}

## CI vs. n plot
n <- as.numeric(table(county))*exp(runif(J, -.1, .1))
plot(n, coef(fit2)[-1], ylim=c(-1,4), pch=19, cex=0.3, log="x", las=1, ylab=~alpha)
mtext("Independent")
L <- cbind(0, diag(85))
for (j in 1:85) {
  lines(rep(n[j], 2), estimate(L[j,], fit2)[2:3])
}
abline(h=coef(fit1)[1], col="gray")
plot(n, apply(a, 2, median), ylim=c(-1, 4), pch=19, cex=0.3, log="x", las=1, ylab=~alpha)
mtext("Hierarchical")
for (j in 1:85) {
  lines(rep(n[j], 2), quantile(a[,j], c(.025, .975)))
}
abline(h=coef(fit1)[1], col="gray")

## Variance components
apply(sigma, 2, mean)
quantile(sigma[,2]^2/sigma[,1]^2, c(.025, .5, .975))
icc <- sigma[,2]^2/(sigma[,1]^2+sigma[,2]^2)
quantile(icc, c(.025, .5, .975))

## A hypothetical quantity of interest
mean(a[,ind[2]] > a[,ind[3]])

## Prediction: Carlton county
nn <- length(mu)
yy <- rnorm(nn, a[,9], sigma[,1])
quantile(yy, c(.025, .5, .975))
## Note that the median of yy is equivalent to the mean of a[,9]
## However, mean(a[,9]) is preferable b/c it has lower MC error

## Lower MC error for quantiles, method 1 (root-finding)
f <- function(x, q) mean(pnorm(x, a[,9], sigma[,1])) - q
uniroot(f, c(-10,10), q=.025)$root
uniroot(f, c(-10,10), q=.975)$root

## Lower MC error for quantiles, method 2 (more samples)
yy <- rnorm(nn*50, a[,9], sigma[,1])
quantile(yy, c(.025, .5, .975))

## Prediction: New county
aa <- rnorm(nn, mu, sigma[,2])
yy <- rnorm(nn, aa, sigma[,1])
quantile(yy, c(.025, .5, .975))