## 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))

## How many counties have first-floor measurements
table(by(x==0, county, any)) ## All have basements
hasFFM <- by(x==1, county, any)
table(hasFFM)

## Outcome
radon <- Data$activity
y <- log(ifelse(radon==0,.1, radon))

## For comparison
fit1 <- lm(y~floor, Data)
fit3 <- lm(y~0+county*floor, Data)

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

## R2jags summaries
plot(fit)
fit

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

## 3 counties
ind <- c(grep("Aitkin", levels(county)), grep("Ramsey", levels(county)), grep("St Louis", levels(county)))
par(oma=c(0,0,1,0), mfrow=c(1,3))
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(coef(fit1)[1], coef(fit1)[2], lwd=3, col=pal(3)[1])
  abline(coef(fit3)[ind[j]], coef(fit3)[ind[j]+J], lwd=3, col=pal(3)[2])
  abline(a=median(a[,ind[j]]), b=median(b[,ind[j]]), 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
nnj <- as.numeric(table(county))*exp(runif(J, -.1, .1))
col <- c("gray60", "#008DFFFF")[hasFFM+1]
plot(nnj, apply(b, 2, median), ylim=c(-3, 2), pch=19, cex=0.4, log="x", las=1, ylab=~beta, col=col, xlab="n")
for (j in 1:85) {
  lines(rep(nnj[j], 2), quantile(b[,j], c(.025, .975)), col=col[j])
}
abline(h=mean(mu[,2]), col="gray")
points(nnj, coef(fit3)[-(1:85)], pch=19, cex=0.4)
