## Chapter 7
##
## Code to support Video Vignette
##
## Videos and supporting code are not a complete portrayal
## of Chapter content.

##
## a: correlation tests
##

## frat dash data
x <- c(12.3, 16.1, 18.6, 13.5, 16.8, 14.1, 18.1, 17.8, 15.2, 13.1)
y <- c(13.6, 17.1, 19.3, 13.9, 16.3, 16.9, 17.9, 15.9, 16.1, 16.1)
plot(x, y, xlab="time before beer", ylab="time after beer")

## formula
n <- length(y)
ybar <- mean(y)
xbar <- mean(x)
sum((y - ybar)*(x - xbar))/(n-1)

## same
syx <- cov(y, x)
syx

## follows the formula for estimated correlation
s2y <- var(y)
s2x <- var(x)
syx / sqrt(s2y*s2x)
r <- cor(y, x)
r

## simulate under H0: rho = 0, i.e., two separate Gaussians for Y and X
## it doesn't matter what muy and mux you use, so I use ybar and xbar

N <- 100000
Rs <- rep(NA, N)
for(i in 1:N) {
  sigma2ys <- (n - 1)*s2y/rchisq(1, n - 1)
  Ys <- rnorm(n, ybar, sqrt(sigma2ys))
  sigma2xs <- (n - 1)*s2x/rchisq(1, n - 1)
  Xs <- rnorm(n, xbar, sqrt(sigma2xs))
  Rs[i] <- cor(Xs, Ys)
}

## visualize
hist(Rs, xlim=c(-1.15, 1.15), main="")
abline(v=c(r, -r), lty=1:2, col=2, lwd=2)
legend("topright", c("r", "reflect"), col=2, lty=1:2, lwd=2, bty="n")

## p-value calculation
pval.mc <- 2*mean(Rs > r)
pval.mc

## the math way

u = r*sqrt(n - 2)/sqrt(1 - r^2)
pval.t <- 2*pt(-abs(u), df=n - 2)
pval.t

##
## b: simple linear model
##

## MLE/OLS coefficients
b1 <- r * sqrt(s2y/s2x)
b0 <- ybar - b1*xbar

## visualize
plot(x, y, xlab="time before beer", ylab="time after beer")
abline(b0, b1)

## bias-corrected MLE for sigma2
yhat <- b0 + b1*x  ## (better than muhat)
s2 <- sum((y - yhat)^2)/(n - 2)
s2

## simulation-based testing
beta1 <- 0
B1s <- rep(NA, N)
for(i in 1:N) {
  sigma2s <- (n - 2)*s2/rchisq(1, n - 2)         ## Chapter 3 or 6
  Ys <- rnorm(n, b0 + beta1*x, sqrt(sigma2s))    ## Gaussian SLR under H0
  B1s[i] <- cor(Ys, x)*sd(Ys)/sd(x)              ## or coef(lm(Ys ~ x))[2]
}

## visualizing
hist(B1s, main="")
abline(v=c(b1, -b1), lty=1:2, col=2, lwd=2)
legend("topright", c("b1", "reflect"), col=2, lty=1:2, lwd=2, bty="n")

## p-value calculation
b1.pval.mc <- 2*mean(B1s > b1)
b1.pval.mc

## CIs switching beta1 and b1 with quantile as usual

## the math way

## standard error
sb1 <- sqrt(s2/((n - 1)*s2x))

## test stat with beta1 = 0
t1 <- b1/sb1

## p-value
b1.pval.t <- 2*pt(-abs(t1), df=n - 2)
b1.pval.t

##
## c: prediction
##

## predictive grid
np <- 100
xp <- seq(10, 20, length=np)

## best line, but not the only reasonable line
yp <- b0 + b1*xp

## prediction on that grid is similar to abline
plot(x, y, xlab="time before beer", ylab="time after beer")
lines(xp, yp, lwd=2)

## simulation
YYmeans <- YYs <- matrix(NA, nrow=N, ncol=np)
for(i in 1:N) {
  sigma2s <- (n - 2)*s2/rchisq(1, n - 2)              ## noise variance
  Ys <- rnorm(n, b0 + b1*x, sqrt(sigma2s))            ## synthetic data
  B1s <- cor(Ys, x)*sd(Ys)/sd(x)                      ## re-fit slope
  B0s <- mean(Ys) - B1s*mean(x)                       ## re-fit intercept
  YYmeans[i,] <- B0s + B1s*xp                         ## predict line
  YYs[i,] <- rnorm(np, YYmeans[i,], sqrt(sigma2s))    ## noise around line
}

## visualization of the "CI", via YYmeans
sel <- seq(1, N, by=100)  ## subset important because too many lines
matplot(xp, t(YYmeans[sel,]), col="gray", type="l", lty=1,
        xlab="time before beer", ylab="time after beer",)
points(x, y)
legend("bottomright", "selection of YYmeans", col="gray", lty=1, bty="n")

## extract quantiles for each xp
low.PI.mc <- apply(YYs, 2, quantile, prob=0.025)
high.PI.mc <- apply(YYs, 2, quantile, prob=0.975)
low.CI.mc <- apply(YYmeans, 2, quantile, prob=0.025)
high.CI.mc <- apply(YYmeans, 2, quantile, prob=0.975)

## visual of the summary of prediction, CI and PI
plot(x, y, xlim=range(xp), ylim=range(c(low.PI.mc, high.PI.mc)))
lines(xp, colMeans(YYmeans), lwd=2)
lines(xp, low.CI.mc, lwd=2, lty=2)
lines(xp, high.CI.mc, lwd=2, lty=2)
lines(xp, low.PI.mc, lwd=2, lty=2, col=2)
lines(xp, high.PI.mc, lwd=2, lty=2, col=2)
legend("bottomright", c("mean", "95% CI", "95% PI"), 
       col=c(1, 1, 2),  lty=c(1, 2, 2), lwd=2, bty="n")

## math way

## CI-based uncertainty
s2fit <- s2*(1/n + (xp - xbar)^2/((n - 1)*s2x))
sfit <- sqrt(s2fit)
low.CI.t <- yp + qt(0.025, n - 2)*sfit
high.CI.t <- yp + qt(0.975, n - 2)*sfit

## augmenting visual, with green covering black
lines(xp, low.CI.t, lwd=2, lty=2, col=3)
lines(xp, high.CI.t, lwd=2, lty=2, col=3)

## full PI-based uncertainty
spred <- sqrt(s2fit + s2)
low.PI.t <- yp + qt(0.025, n - 2)*spred
high.PI.t <- yp + qt(0.975, n - 2)*spred

## augmenting visual, with blue covering red
lines(xp, low.PI.t, lwd=2, lty=2, col=4)
lines(xp, high.PI.t, lwd=2, lty=2, col=4)