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

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

## Group level predictor
jData <- list("y", "x", "cid", "n", "J", "u")
model <- function() {
  ## House level
  for (i in 1:n) {
    y[i] ~ dnorm(a[cid[i]] + b*x[i], sigma[1]^(-2))
    r1[i] <- y[i] - a[cid[i]] - b*x[i]
  }
  RSS[1] <- sum(r1^2)
  
  ## County level
  for (j in 1:J) {
    a[j] ~ dnorm(g[1] + g[2]*u[j], sigma[2]^(-2))
    r2[j] <- a[j] - g[1] - g[2]*u[j]
  }
  RSS[2] <- sum(r2^2)
  
  b ~ dnorm(0, 0.001)
  for (j in 1:2) {
    g[j] ~ dnorm(0, 0.001)
    sigma[j] ~ dunif(0, 100)
  }  
}
require(R2jags); invisible(runif(1))
fit <- jags(model=model, param=c("g", "sigma", "a", "b", "RSS"), data=jData, n.iter=10000, n.thin=1)
attach.jags(fit)

## Classical model(s)
summary(lm(y~x))$r.sq
summary(lm(y~x+county))$r.sq
summary(lm(y~x+county))$adj.r.sq

rss0 <- function(x) crossprod(x-mean(x))

## Data level
1-mean(RSS[,1])/rss0(y)

## Group level
1-mean(RSS[,2])/mean(apply(a, 1, rss0))

## Shrinkage factors
lam <- pmin(apply(a, 2, var)/mean(sigma[,2])^2, 1)
nn <- as.numeric(table(county))
nnj <- nn*exp(runif(J, -.1, .1))
plot(nnj, lam, pch=19, cex=0.7, log="x", ylim=c(0,1), xlab="Sample size", ylab="Shrinkage factor", las=1)

## Summarizing pooling
ahat <- g[,1] + tcrossprod(g[,2], u)
e.a <- a-ahat
1-var(apply(e.a, 2, mean))/mean(apply(e.a, 1, var))