## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

## ----setup--------------------------------------------------------------------
library(rmcmc)

## -----------------------------------------------------------------------------
dimension <- 10
scales <- c(0.01, rep(1, dimension - 1))

## -----------------------------------------------------------------------------
target_distribution <- list(
  log_density = function(x) -sum((x / scales)^2) / 2,
  gradient_log_density = function(x) -x / scales^2
)

## -----------------------------------------------------------------------------
proposal <- barker_proposal()

## -----------------------------------------------------------------------------
adapters <- list(
  scale_adapter(
    algorithm = "stochastic_approximation",
    initial_scale = dimension^(-1 / 6),
    target_accept_prob = 0.574,
    kappa = 0.6
  ),
  shape_adapter(type = "variance", kappa = 0.6)
)

## -----------------------------------------------------------------------------
set.seed(791285301L)
initial_state <- chain_state(10 * rnorm(dimension))

## -----------------------------------------------------------------------------
n_warm_up_iteration <- 10000
n_main_iteration <- 10000

## -----------------------------------------------------------------------------
barker_results <- sample_chain(
  target_distribution = target_distribution,
  proposal = proposal,
  initial_state = initial_state,
  n_warm_up_iteration = n_warm_up_iteration,
  n_main_iteration = n_main_iteration,
  adapters = adapters,
  trace_warm_up = TRUE
)

## -----------------------------------------------------------------------------
mean_accept_prob <- mean(barker_results$statistics[, "accept_prob"])
cat(sprintf("Average acceptance probability is %.2f", mean_accept_prob))

## -----------------------------------------------------------------------------
clipped_dimension <- min(5, dimension)
final_shape <- proposal$parameters()$shape
cat(
  sprintf("Adapter scale est.: %s", toString(final_shape[1:clipped_dimension])),
  sprintf("True target scales: %s", toString(scales[1:clipped_dimension])),
  sep = "\n"
)

## -----------------------------------------------------------------------------
library(posterior)

## -----------------------------------------------------------------------------
summarize_draws(barker_results$traces)

## -----------------------------------------------------------------------------
draws <- as_draws_matrix(barker_results$traces)
summary(draws)

## -----------------------------------------------------------------------------
cat(
  sprintf(
    "Effective sample size of mean(target_log_density) is %.0f",
    ess_mean(extract_variable(draws, "target_log_density"))
  )
)

## -----------------------------------------------------------------------------
mala_results <- sample_chain(
  target_distribution = target_distribution,
  proposal = langevin_proposal(),
  initial_state = initial_state,
  n_warm_up_iteration = n_warm_up_iteration,
  n_main_iteration = n_main_iteration,
  adapters = list(
    scale_adapter(algorithm = "stochastic_approximation", kappa = 0.6),
    shape_adapter(type = "variance", kappa = 0.6)
  ),
  trace_warm_up = TRUE
)

## -----------------------------------------------------------------------------
cat(
  sprintf(
    "Average acceptance probability is %.2f",
    mean(mala_results$statistics[, "accept_prob"])
  )
)

## -----------------------------------------------------------------------------
cat(
  sprintf(
    "Effective sample size of mean(target_log_density) is %.0f",
    ess_mean(
      extract_variable(
        as_draws_matrix(mala_results$traces), "target_log_density"
      )
    )
  )
)

## -----------------------------------------------------------------------------
visualize_scale_adaptation <- function(warm_up_statistics, label) {
  n_warm_up_iteration <- nrow(warm_up_statistics)
  old_par <- par(mfrow = c(1, 2))
  on.exit(par(old_par))
  plot(
    exp(warm_up_statistics[, "log_scale"]),
    type = "l",
    xlab = expression(paste("Chain iteration ", t)),
    ylab = expression(paste("Scale ", sigma[t]))
  )
  plot(
    cumsum(warm_up_statistics[, "accept_prob"]) / 1:n_warm_up_iteration,
    type = "l",
    xlab = expression(paste("Chain iteration ", t)),
    ylab = expression(paste("Average acceptance rate ", alpha[t])),
    ylim = c(0, 1)
  )
  mtext(
    sprintf("Scale adaptation for %s", label),
    side = 3, line = -2, outer = TRUE
  )
}

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_scale_adaptation(barker_results$warm_up_statistics, "Barker proposal")

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_scale_adaptation(mala_results$warm_up_statistics, "Langevin proposal")

## -----------------------------------------------------------------------------
visualize_shape_adaptation <- function(warm_up_statistics, dimensions, label) {
  matplot(
    sqrt(warm_up_statistics[, paste0("variance_estimate", dimensions)]),
    type = "l",
    xlab = expression(paste("Chain iteration ", t)),
    ylab = expression(paste("Shape ", diag(Sigma[t]^(1 / 2)))),
    log = "y"
  )
  legend(
    "right",
    paste0("coordinate ", dimensions),
    lty = dimensions,
    col = dimensions,
    bty = "n"
  )
  mtext(
    sprintf("Shape adaptation for %s", label),
    side = 3, line = -2, outer = TRUE
  )
}

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_shape_adaptation(
  barker_results$warm_up_statistics, 1:clipped_dimension, "Barker proposal"
)

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_shape_adaptation(
  mala_results$warm_up_statistics, 1:clipped_dimension, "Langevin proposal"
)

## -----------------------------------------------------------------------------
visualize_traces <- function(traces, dimensions, label) {
  matplot(
    traces[, paste0("position", dimensions)],
    type = "l",
    xlab = expression(paste("Chain iteration ", t)),
    ylab = expression(paste("Position ", X[t])),
  )
  legend(
    "topright",
    paste0("coordinate ", dimensions),
    lty = dimensions,
    col = dimensions,
    bty = "n"
  )
  mtext(sprintf("Traces for %s", label), side = 3, line = -2, outer = TRUE)
}

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_traces(
  barker_results$warm_up_traces, 1:clipped_dimension, "Barker proposal"
)

## ----fig.width=7, fig.height=4------------------------------------------------
visualize_traces(
  mala_results$warm_up_traces, 1:clipped_dimension, "Langevin proposal"
)

