code for comparing HMC to Adaptive HMC

```r
# compare adaptive HMC to HMC on some knotty problems

library(greta)

n_chains <- 10
warmup <- 2000

# very hard model
correlation <- 0.99
rel_sd_range <- 1e5

# # very easy model
# correlation <- 0
# rel_sd_range <- 1

# sample from an correlated multivariate Gaussian with different marginal
# variances
dim <- 4
C <- matrix(correlation, dim, dim)
diag(C) <- 1
sd <- seq(1, rel_sd_range, length.out = dim)
Sigma <- diag(sd) %*% C %*% diag(sd)
x <- multivariate_normal(mean = zeros(1, dim),
Sigma = Sigma)
m <- model(x)
init <- initials(x = t(rep(0, dim)))

# default HMC
hmc_time <- system.time(
draws_hmc <- mcmc(m,
warmup = warmup,
chains = n_chains,
sampler = hmc())
)

# HMC with insufficient leapfrog steps (representative of more complex
# posteriors, when using the default step sizes)
hmc_l2_time <- system.time(
draws_hmc_l2 <- mcmc(m,
warmup = warmup,
chains = n_chains,
sampler = hmc(Lmin = 1, Lmax = 3))
)

# adaptive HMC, default initialisation
ahmc_time <- system.time(
draws_ahmc <- mcmc(m,
warmup = warmup,
chains = n_chains,
sampler = adaptive_hmc())
)

# adaptive HMC, initialised well, at a single point
ahmc_init_time <- system.time(
draws_ahmc_init <- mcmc(m,
warmup = warmup,
chains = n_chains,
initial_values = init,
sampler = adaptive_hmc())
)

# convergence/sampling statistics

# run times are similar (hmc_l2 expected to be faster)
hmc_time["elapsed"]
hmc_l2_time["elapsed"]
ahmc_time["elapsed"]
ahmc_init_time["elapsed"]

# what's the sampler efficiency (effective numbers of samples, per second
# overall sampling) for the worst-sampled variable
worst_efficiency <- function(draws, time) {
neffs <- coda::effectiveSize(draws)
min(neffs) / time["elapsed"]
}

# higher is better
worst_efficiency(draws_hmc, hmc_time)
worst_efficiency(draws_hmc_l2, hmc_l2_time)
worst_efficiency(draws_ahmc, ahmc_time)
worst_efficiency(draws_ahmc_init, ahmc_init_time)

# what's the upper bound for the univariate potential scale reduction factor,
# for the worst-sampled variable
worst_rhat <- function(draws) {
rhats <- coda::gelman.diag(draws,
autoburnin = FALSE,
multivariate = FALSE)
uppers <- rhats$psrf[, 2]
max(uppers)
}

# lower is better, <1.1 is ideal
worst_rhat(draws_hmc)
worst_rhat(draws_hmc_l2)
worst_rhat(draws_ahmc)
worst_rhat(draws_ahmc_init)

bayesplot::mcmc_trace(draws_hmc)
bayesplot::mcmc_trace(draws_ahmc)

plot(draws_hmc)
plot(draws_ahmc)

par(mfrow = c(1, 1))
plot(as.matrix(draws_ahmc)[, 1:2],
cex = 0.5,
pch = 16,
col = grey(0.8))
points(as.matrix(draws_hmc)[, 1:2],
cex = 0.5,
pch = 16,
col = "red")
```