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

## Scaled Wishart prior on variance-covariance matrix
W <- diag(2)
jData <- list("y", "x", "cid", "n", "J", "u", "W")
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])/xi[1]
    mu.B[j,2] <- (gb[1] + gb[2]*u[j])/xi[2]
    a[j] <- xi[1]*B[j,1]
    b[j] <- xi[2]*B[j,2]
  }
  for (j in 1:2) {
    ga[j] ~ dnorm(0, 0.001)
    gb[j] ~ dnorm(0, 0.001)
    xi[j] ~ dunif(0, 100)
  }
  Tau.B[1:2,1:2] ~ dwish(W, 3)
  Sigma.B <- inverse(Tau.B)
  sigma[1] <- sqrt(Sigma.B[1,1])*xi[1]
  sigma[2] <- sqrt(Sigma.B[2,2])*xi[2]
  sigma[3] ~ dunif(0,100)
  rho <- Sigma.B[1,2]*xi[1]*xi[2]/(sigma[1]*sigma[2])
}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, param=c("ga", "gb", "sigma", "rho", "a", "b"), data=jData, n.iter=10000, 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
quantile(rho, c(.5, .025, .975))