# ============================================================================
# Metropolis-Hastings: implementation, then diagnostics, then failure modes
# ============================================================================

overlay <- function(draws, f, main = "", xlim = range(draws)) {
  hist(draws, breaks = 60, freq = FALSE, col = "grey90", border = "grey60",
       xlim = xlim, xlab = "x", main = main)
  curve(f(x), add = TRUE, lwd = 2)
}

## ---------------------------------------------------------------------------
## 1. ONE CHAIN, WRITTEN OUT
## ---------------------------------------------------------------------------

set.seed(1)

n <- 5000
x <- numeric(n)
x[1] <- 50

for (t in 2:n) {
  y <- rnorm(1, x[t - 1], 2.4)
  if (log(runif(1)) < dnorm(y, log = TRUE) - dnorm(x[t - 1], log = TRUE)) {
    x[t] <- y
  } else {
    x[t] <- x[t - 1]
  }
}

plot(x, type = "l", xlab = "t", ylab = "x")
plot(x[1:200], type = "l", xlab = "t", ylab = "x")
overlay(x[-(1:500)], dnorm)

estimate <- function(x) c(mean = mean(x), tail = mean(x > 2.5))
truth    <- c(mean = 0, tail = pnorm(2.5, lower.tail = FALSE))

round(rbind(all_draws = estimate(x),
            post_burn = estimate(x[-(1:500)]),
            truth     = truth), 5)

## ---------------------------------------------------------------------------
## 2. THE SAME ALGORITHM, WRAPPED UP
## ---------------------------------------------------------------------------

metropolis <- function(x0, n, log_density, propose, log_ratio) {
  draws <- matrix(NA_real_, nrow = n, ncol = length(x0))
  draws[1, ] <- x0
  accepted <- 0

  for (t in 2:n) {
    current  <- draws[t - 1, ]
    proposed <- propose(current)
    if (log(runif(1)) < log_ratio(current, proposed, log_density)) {
      draws[t, ] <- proposed
      accepted <- accepted + 1
    } else {
      draws[t, ] <- current
    }
  }

  list(draws = drop(draws), rate = accepted / (n - 1))
}

standard_normal <- function(x) dnorm(x, log = TRUE)
random_walk     <- function(sd) function(x) rnorm(length(x), x, sd)
symmetric       <- function(x, y, log_density) log_density(y) - log_density(x)

## --- Three proposal scales ----------

n <- 50000
scales <- c(small = 0.1, tuned = 2.4, large = 25)

chains <- lapply(scales, function(s)
  metropolis(0, n, standard_normal, random_walk(s), symmetric))

sapply(chains, function(ch) ch$rate)

for (nm in names(chains)) {
  plot(chains[[nm]]$draws[1:2000], type = "l", main = nm, xlab = "t", ylab = "x")
  acf(chains[[nm]]$draws, lag.max = 60, main = nm)
  overlay(chains[[nm]]$draws, dnorm, nm)
}

test <- metropolis(50, 10000, standard_normal, random_walk(0.01), symmetric)
hist(test$draws)
plot(test$draws, type = 'l')

## --- What the draws are for ----------

round(cbind(sapply(chains, function(ch) estimate(ch$draws)), truth), 5)

## ---------------------------------------------------------------------------
## 3. EFFECTIVE SAMPLE SIZE
## ---------------------------------------------------------------------------

ess <- function(x) {
  rho <- acf(x, lag.max = 1000, plot = FALSE)$acf[-1]
  cut <- which(rho < 0.05)[1]
  if (is.na(cut)) cut <- length(rho)
  length(x) / (1 + 2 * sum(rho[1:cut]))
}

sapply(chains, function(ch) ess(ch$draws))

## --- Does it predict the error? ----------

replicate_means <- sapply(scales, function(s)
  replicate(30, mean(metropolis(0, 20000, standard_normal,
                                random_walk(s), symmetric)$draws)))

observed  <- apply(replicate_means, 2, sd)
predicted <- 1 / sqrt(sapply(scales, function(s)
  ess(metropolis(0, 20000, standard_normal, random_walk(s), symmetric)$draws)))

round(rbind(observed, predicted), 4)

## --- Tuning ----------

sweep <- t(sapply(c(0.25, 0.5, 1, 1.5, 2, 2.38, 3, 4, 6, 10, 20), function(s) {
  ch <- metropolis(0, 40000, standard_normal, random_walk(s), symmetric)
  c(sd = s, rate = ch$rate, ess = ess(ch$draws))
}))

sweep

plot(sweep[, "rate"], sweep[, "ess"], type = "b", pch = 19,
     xlab = "acceptance rate", ylab = "effective sample size")
abline(v = 0.44, lty = 3)

## ---------------------------------------------------------------------------
## 4. FAILURE: A WRONG ACCEPTANCE RATIO
## ---------------------------------------------------------------------------

shrink_toward_zero <- function(x) rnorm(1, x / 2, 1)

log_q <- function(from, to) dnorm(to, from / 2, 1, log = TRUE)

hastings <- function(x, y, log_density)
  log_density(y) - log_density(x) + log_q(y, x) - log_q(x, y)

correct <- metropolis(0, n, standard_normal, shrink_toward_zero, hastings)
buggy   <- metropolis(0, n, standard_normal, shrink_toward_zero, symmetric)

summarise <- function(ch) c(rate = ch$rate, ess = ess(ch$draws),
                            mean = mean(ch$draws), sd = sd(ch$draws),
                            tail = mean(ch$draws > 2.5))

rbind(correct = summarise(correct),
      buggy   = summarise(buggy),
      truth   = c(NA, NA, 0, 1, pnorm(2.5, lower.tail = FALSE)))

overlay(correct$draws, dnorm, "Hastings ratio included")
overlay(buggy$draws, dnorm, "Hastings ratio omitted")

## ---------------------------------------------------------------------------
## 5. FAILURE: A TARGET WITH TWO MODES
## ---------------------------------------------------------------------------

mixture <- function(x) log(0.5 * dnorm(x, -10) + 0.5 * dnorm(x, 10))

trapped <- metropolis(10, n, mixture, random_walk(2.4), symmetric)

rbind(trapped = c(rate = trapped$rate, ess = ess(trapped$draws),
                  mean = mean(trapped$draws), above_zero = mean(trapped$draws > 0)),
      truth   = c(NA, NA, 0, 0.5))

acf(trapped$draws, lag.max = 10, plot = FALSE)$acf[2:5]

plot(trapped$draws[1:2000], type = "l", xlab = "t", ylab = "x")
overlay(trapped$draws, function(x) 0.5 * dnorm(x, -10) + 0.5 * dnorm(x, 10),
        xlim = c(-14, 14))

second <- metropolis(-10, n, mixture, random_walk(2.4), symmetric)
mean(second$draws)

## ---------------------------------------------------------------------------
## 6. FAILURE: A CORRELATED TARGET
## ---------------------------------------------------------------------------

rho   <- 0.99
Sigma <- matrix(c(1, rho, rho, 1), 2, 2)
Omega <- solve(Sigma)

bivariate <- function(z) -0.5 * drop(t(z) %*% Omega %*% z)

bivariate_pair <- function(a, b)
  -0.5 * (Omega[1,1]*a^2 + 2*Omega[1,2]*a*b + Omega[2,2]*b^2)

ridge <- metropolis(c(0, 0), n, bivariate, random_walk(0.2), symmetric)

along  <- (ridge$draws[, 1] + ridge$draws[, 2]) / sqrt(2)
across <- (ridge$draws[, 1] - ridge$draws[, 2]) / sqrt(2)

c(rate = ridge$rate, ess_x1 = ess(ridge$draws[, 1]),
  ess_along = ess(along), ess_across = ess(across))

g <- seq(-4, 4, length.out = 120)
contour(g, g, outer(g, g, function(a, b) exp(bivariate_pair(a, b))),
        levels = exp(-qchisq(c(0.5, 0.9, 0.99), 2) / 2),
        drawlabels = FALSE, col = "grey55", xlab = "x1", ylab = "x2",
        main = "first 500 draws, 50/90/99% contours")
lines(ridge$draws[1:500, ], col = "red")

plot(along[1:3000], type = "l", xlab = "t", ylab = "along the ridge")
plot(across[1:3000], type = "l", xlab = "t", ylab = "across the ridge")

sapply(c(200, 500, 2000, 50000), function(k)
  c(n = k,
    along  = diff(range(along[1:k])) / (6 * sqrt(1 + rho)),
    across = diff(range(across[1:k])) / (6 * sqrt(1 - rho))))