---
title: "Budworms with Parameter-Expanded Data Augmentation"
author: "Luke Tierney"
output:
  html_document:
     toc: yes
---

<style type="text/css">
/* Whole document: */
body{
  font-family: Helvetica;
  font-size: 16pt;
}
/* Headers */
h1,h2{
  font-size: 24pt;
  font-weight: 800;
}
h3,h4,h5,h6{
  font-size: 20pt;
  font-weight: 800;
}
pre {
  font-size: 0.9em;
}
</style>

```{r, include = FALSE}
set.seed(1234)
```

```{r global_options, include=FALSE}
knitr::opts_chunk$set(collapse=TRUE)
options(rgl.useNULL=TRUE)
knitr::knit_hooks$set(webgl = rgl::hook_webgl)
```

## The original data augmentation approach

```{r}
dose <- c(1, 2, 4, 8, 16, 32)
died <- c(1, 4, 9, 13, 18, 20)
x <- log2(dose)
xx <- rep(x - mean(x), each = 20)

z <- unlist(lapply(died,
                   function(x) c(rep(1, x), rep(0, 20 - x))))
genAlpha <- function(y, sigma = 1)
    rnorm(1, mean(y), sigma / sqrt(length(y)))
genBeta <- function(y, xx, sxx2, sigma = 1)
    rnorm(1, sum(xx * y) / sxx2, sigma / sqrt(sxx2))
genY <- function(z, mu, sigma = 1) {
    p <- pnorm(-mu)
    u <- runif(length(z))
    ifelse(z == 1,
           qnorm(log(1 - u) + 
                 pnorm(0, mu, sigma, lower.tail = FALSE, log.p = TRUE),
                 mu, sigma, lower.tail = FALSE, log.p = TRUE),
           qnorm(log(u) + pnorm(0, mu, sigma, log.p = TRUE),
                 mu, sigma, log.p = TRUE))
}

da <- function(z, alpha, beta, xx, N) {
    val <- matrix(0, nrow = N, ncol = 2)
    colnames(val) <- c("alpha", "beta")
    sxx2 <- sum(xx^2)
    for (i in 1 : N) {
        y <- genY(z, alpha + beta * xx)
        alpha <- genAlpha(y)
        beta <- genBeta(y, xx, sxx2)
        val[i,1] <- alpha
        val[i,2] <- beta
    }
    val
}

alpha0 <- qnorm(mean(z))
beta0 <- 0
N <- 100000

v <- da(z, alpha0, beta0, xx, N)
```


## Parameter expanded data augmentation

Adding an unidentifiable `sigma` parameter to improve mixing:

```{r}
rinvgamma <- function(n, a, b)
    1 / rgamma(n, a, scale = b)

rinvgammaT <- function(n, a, b, T)
    1 / qgamma((1-runif(n)) *
               pgamma(1 / T, a, scale = b, lower.tail = FALSE),
               a, scale = b, lower.tail = FALSE)

## using log.p is a bit better numerically
rinvgammaT <- function(n, a, b, T)
    1 / qgamma(log(1-runif(n)) +
               pgamma(1 / T, a, scale = b, lower.tail = FALSE,
                      log.p = TRUE),
               a, scale = b, lower.tail = FALSE, log.p = TRUE)

a0 <- 0.1
v0 <- 0.1
SM2 <- 10000 # upper bound on sigma^2

genSigma <- function(y, mu, T, extra = 0) {
    sy2 <- sum((y - mu)^2)
    sqrt(rinvgammaT(1, (length(y) + v0 + extra) / 2, 2 / (sy2 + a0), T))
}
```

Generate from the full conditional distributions:

```{r}
da1 <- function(z, alpha, beta, xx, N, sigma = 1) {
    val <- matrix(0, nrow = N, ncol = 3)
    colnames(val) <- c("alpha", "beta", "sigma")
    sxx2 <- sum(xx^2)
    alpha.tilde <- alpha
    beta.tilde <- beta
    for (i in 1 : N) {
        y <- genY(z, alpha.tilde + beta.tilde * xx, sigma)
        sigma <- genSigma(y, alpha.tilde + beta.tilde * xx, SM2, 2)
        alpha.tilde <- genAlpha(y, sigma)
        beta.tilde <- genBeta(y, xx, sxx2, sigma)
        val[i,1] <- alpha.tilde / sigma
        val[i,2] <- beta.tilde / sigma
        val[i,3] <- sigma
    }
    val
}

v1 <- da1(z, alpha0, beta0, xx, N)
```

- Generate sigma from its conditional given `y`.
- Generate others from full conditional distributions.

```{r}
da2 <- function(z, alpha, beta, xx, N, sigma = 1) {
    val <- matrix(0, nrow = N, ncol = 3)
    colnames(val) <- c("alpha", "beta", "sigma")
    sxx2 <- sum(xx^2)
    alpha.tilde <- alpha
    beta.tilde <- beta
    for (i in 1 : N) {
        y <- genY(z, alpha.tilde + beta.tilde * xx, sigma)
        sigma <- genSigma(y, mean(y) + sum(xx * y) / sxx2 * xx, SM2)
        alpha.tilde <- genAlpha(y, sigma)
        beta.tilde <- genBeta(y, xx, sxx2, sigma)
        val[i,1] <- alpha.tilde / sigma
        val[i,2] <- beta.tilde / sigma
        val[i,3] <- sigma
    }
    val
}

v2 <- da2(z, alpha0, beta0, xx, N)
```

- Generate `sigma.star` from its prior distribution.
- Generate `y` from the conditional distribution given observable `alpha`, `beta` and multiply by `sigma.star`.
- Generate  `sigma`, `alpha`, `beta` from their full conditionals.

```{r}
da3 <- function(z, alpha, beta, xx, N, sigma = 1) {
    val <- matrix(0, nrow = N, ncol = 3)
    colnames(val) <- c("alpha", "beta", "sigma")
    sxx2 <- sum(xx^2)
    alpha.tilde <- alpha
    beta.tilde <- beta
    for (i in 1 : N) {
        sigma.star <- sqrt(rinvgammaT(1, v0 / 2, 2 / a0, SM2))
        alpha <- alpha.tilde / sigma
        beta <- beta.tilde / sigma
        y <-  sigma.star * genY(z, alpha + beta * xx, 1)
        sigma <- genSigma(y, mean(y) + sum(xx * y) / sxx2 * xx, SM2)
        alpha.tilde <- genAlpha(y, sigma)
        beta.tilde <- genBeta(y, xx, sxx2, sigma)
        val[i,1] <- alpha.tilde / sigma
        val[i,2] <- beta.tilde / sigma
        val[i,3] <- sigma
    }
    val
}

v3 <- da3(z, alpha0, beta0, xx, N)
```

Density plots:

```{r}
plot(density(v[,"beta"]))
lines(density(v1[,"beta"]), col = "red")
lines(density(v2[,"beta"]), col = "blue")
lines(density(v3[,"beta"]), col = "forestgreen")
```

Autocorrelation plots:

```{r, echo = FALSE}
opar <- par(mfrow=c(2,2))
acf(v[,"alpha"])
acf(v1[,"alpha"])
acf(v2[,"alpha"])
acf(v3[,"alpha"])
acf(v[,"beta"])
acf(v1[,"beta"])
acf(v2[,"beta"])
acf(v3[,"beta"])
par(opar)
```

