## 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, correlated with varying intercepts
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) {
    B[j,1:2] ~ dmnorm(mu.B[j,], Tau.B)
    mu.B[j,1] <- ga[1] + ga[2]*u[j]
    mu.B[j,2] <- gb[1] + gb[2]*u[j]
    a[j] <- B[j,1]
    b[j] <- B[j,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) 
  }
  Tau.B[1:2,1:2] <- inverse(Sigma.B)
  Sigma.B[1,1] <- sigma[1]^2
  Sigma.B[2,2] <- sigma[2]^2
  rho ~ dunif(-1,1)
  Sigma.B[1,2] <- rho*sigma[1]*sigma[2]
  Sigma.B[2,1] <- rho*sigma[1]*sigma[2]
}
require(R2jags); invisible(runif(1))
## This takes a while:
fit <- jags(model=model, param=c("ga", "gb", "sigma", "rho", "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: 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")

## Draws from the posterior
ind <- round(seq(1, nrow(ga), length=50))
plot(u, apply(b, 2, mean), ylim=c(-2, 0.5), pch=19, cex=0.5, las=1, ylab=~beta)
for (j in ind) {
  abline(gb[j,1], gb[j,2], col="gray80")
}
abline(mean(gb[,1]), mean(gb[,2]), col=pal(2)[2], lwd=2)
points(u, apply(b, 2, mean), pch=19, cex=0.5)

## Correlation
quantile(rho, c(.5, .025, .975))