## 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))
hasFFM <- by(x==1, county, any)

## Read in county-level soil uranium data
AllData2 <- read.delim("http://web.as.uky.edu/statistics/users/pbreheny/701/data/radon-county.txt")
Data2 <- subset(AllData2, st=="MN")
u <- log(Data2$Uppm[match(levels(county), Data2$cty)])

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

## Varying slopes, no correlation
jData <- list("y", "x", "cid", "n", "J", "u")
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(ga[1] + ga[2]*u[j], sigma[1]^(-2))
    b[j] ~ dnorm(gb[1] + gb[2]*u[j], sigma[2]^(-2))
  }
  for (j in 1:2) {
    ga[j] ~ dnorm(0, 0.001)
    gb[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("ga", "gb", "sigma", "a", "b"), data=jData, n.iter=30000, n.thin=1)
attach.jags(fit)

## R2jags output
plot(fit)
fit

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

## Posteriors for gamma
quantile(ga[,1], c(.5, .025, .975))
quantile(ga[,2], c(.5, .025, .975))
quantile(gb[,1], c(.5, .025, .975))
quantile(gb[,2], c(.5, .025, .975))

## County-level model: Intercepts
plot(u, apply(a, 2, mean), ylim=c(0.5, 2.1), pch=19, cex=0.5, las=1, ylab=~alpha)
abline(mean(ga[,1]), mean(ga[,2]), col=pal(2)[2], lwd=2)
for (j in 1:85) {
  lines(rep(u[j], 2), quantile(a[,j], c(.1, .9)), col="gray60")
}

## County-level model: Slopes
plot(u, apply(b, 2, mean), ylim=c(-2, 0.5), pch=19, cex=0.5, las=1, ylab=~beta)
abline(mean(gb[,1]), mean(gb[,2]), col=pal(2)[2], lwd=2)
for (j in 1:85) {
  lines(rep(u[j], 2), quantile(b[,j], c(.1, .9)), col="gray60")
}
points(u[!hasFFM], apply(b, 2, mean)[!hasFFM], col="red")

## Correlation of "residuals"
pma <- apply(ga, 2, mean)
pmb <- apply(gb, 2, mean)
aa <- apply(a, 2, mean) - pma[1] - pma[2]*u
bb <- apply(b, 2, mean) - pmb[1] - pmb[2]*u
plot(aa, bb, pch=19, xlab=~alpha*" residuals", ylab=~beta*" residuals")
plot(aa[hasFFM], bb[hasFFM], pch=19, xlab=~alpha*" residuals", ylab=~beta*" residuals")
cor(aa, bb)
cor(aa[hasFFM], bb[hasFFM])