Improvements to sample_chain interface (#43)

* Copy initial state to avoid mutating

* Don't require passing proposal to adapter initializers

* Update adapter initialise syntax in vignette

* Remove unused parameter docstring

* Cover default target_accept_prob logic in tests

* Add test for adapter interface

Author: matt-graham

R/adaptation.R

@@ -1,16 +1,14 @@
 #' Create object to adapt proposal scale to coerce average acceptance rate.
 #'
-#' @param proposal Proposal object to adapt. Must define an `update` function
-#'   which accepts a parameter `scale` for setting scale parameter of proposal.
 #' @param initial_scale Initial value to use for scale parameter. If not set
 #'   explicitly a proposal and dimension dependent default will be used.
 #' @param target_accept_prob Target value for average accept probability for
 #'   chain. If not set a proposal dependent default will be used.
 #' @param kappa Decay rate exponent in `[0.5, 1]` for adaptation learning rate.
 #'
 #' @return List of functions with entries
-#' * `initialize`, a function for initializing adapter state at beginning of
-#'   chain,
+#' * `initialize`, a function for initializing adapter state and proposal
+#'   parameters at beginning of chain,
 #' * `update` a function for updating adapter state and proposal parameters on
 #'   each chain iteration,
 #' * `finalize` a function for performing any final updates to adapter state and
@@ -27,24 +25,22 @@
 #'   grad_log_density = function(x) -x
 #' )
 #' proposal <- barker_proposal(target_distribution)
-#' adapter <- scale_adapter(
-#'   proposal,
-#'   initial_scale = 1., target_accept_prob = 0.4
-#' )
+#' adapter <- scale_adapter(initial_scale = 1., target_accept_prob = 0.4)
+#' adapter$initialize(proposal, chain_state(c(0, 0)))
 scale_adapter <- function(
-    proposal, initial_scale = NULL, target_accept_prob = NULL, kappa = 0.6) {
+    initial_scale = NULL, target_accept_prob = NULL, kappa = 0.6) {
   log_scale <- NULL
-  if (is.null(target_accept_prob)) {
-    target_accept_prob <- proposal$default_target_accept_prob()
-  }
-  initialize <- function(initial_state) {
+  initialize <- function(proposal, initial_state) {
     if (is.null(initial_scale)) {
       initial_scale <- proposal$default_initial_scale(initial_state$dimension())
     }
     log_scale <<- log(initial_scale)
     proposal$update(scale = initial_scale)
   }
-  update <- function(sample_index, state_and_statistics) {
+  update <- function(proposal, sample_index, state_and_statistics) {
+    if (is.null(target_accept_prob)) {
+      target_accept_prob <- proposal$default_target_accept_prob()
+    }
     gamma <- sample_index^(-kappa)
     accept_prob <- state_and_statistics$statistics$accept_prob
     log_scale <<- log_scale + gamma * (accept_prob - target_accept_prob)
@@ -53,16 +49,14 @@ scale_adapter <- function(
   list(
     initialize = initialize,
     update = update,
-    finalize = function() {},
+    finalize = NULL,
     state = function() list(log_scale = log_scale)
   )
 }
 
 #' Create object to adapt proposal with per dimension scales based on estimates
 #' of target distribution variances.
 #'
-#' @param proposal Proposal object to adapt. Must define an `update` function
-#'   which accepts a parameter `shape` for setting shape parameter of proposal.
 #' @param kappa Decay rate exponent in `[0.5, 1]` for adaptation learning rate.
 #'
 #' @inherit scale_adapter return
@@ -74,15 +68,16 @@ scale_adapter <- function(
 #'   grad_log_density = function(x) -x
 #' )
 #' proposal <- barker_proposal(target_distribution)
-#' adapter <- variance_adapter(proposal)
-variance_adapter <- function(proposal, kappa = 0.6) {
+#' adapter <- variance_adapter()
+#' adapter$initialize(proposal, chain_state(c(0, 0)))
+variance_adapter <- function(kappa = 0.6) {
   mean_estimate <- NULL
   variance_estimate <- NULL
-  initialize <- function(initial_state) {
+  initialize <- function(proposal, initial_state) {
     mean_estimate <<- initial_state$position()
     variance_estimate <<- rep(1., initial_state$dimension())
   }
-  update <- function(sample_index, state_and_statistics) {
+  update <- function(proposal, sample_index, state_and_statistics) {
     gamma <- sample_index^(-kappa)
     position <- state_and_statistics$state$position()
     mean_estimate <<- mean_estimate + gamma * (position - mean_estimate)
@@ -124,20 +119,23 @@ variance_adapter <- function(proposal, kappa = 0.6) {
 #'   grad_log_density = function(x) -x
 #' )
 #' proposal <- barker_proposal(target_distribution)
-#' adapter <- robust_shape_adapter(
-#'   proposal,
-#'   initial_scale = 1.,
-#'   target_accept_prob = 0.4
-#' )
+#' adapter <- robust_shape_adapter(initial_scale = 1., target_accept_prob = 0.4)
+#' adapter$initialize(proposal, chain_state(c(0, 0)))
 robust_shape_adapter <- function(
-    proposal, initial_scale, target_accept_prob = 0.4, kappa = 0.6) {
+    initial_scale = NULL, target_accept_prob = NULL, kappa = 0.6) {
   rlang::check_installed("ramcmc", reason = "to use this function")
   shape <- NULL
-  initialize <- function(initial_state) {
+  initialize <- function(proposal, initial_state) {
+    if (is.null(initial_scale)) {
+      initial_scale <- proposal$default_initial_scale(initial_state$dimension())
+    }
     shape <<- initial_scale * diag(initial_state$dimension())
     proposal$update(shape = shape)
   }
-  update <- function(sample_index, state_and_statistics) {
+  update <- function(proposal, sample_index, state_and_statistics) {
+    if (is.null(target_accept_prob)) {
+      target_accept_prob <- proposal$default_target_accept_prob()
+    }
     momentum <- state_and_statistics$proposed_state$momentum()
     accept_prob <- state_and_statistics$statistics$accept_prob
     shape <<- ramcmc::adapt_S(

R/chains.R

@@ -72,7 +72,7 @@ sample_chain <- function(
   if (is.vector(initial_state) && is.atomic(initial_state)) {
     state <- chain_state(initial_state)
   } else if (is.vector(initial_state) && "position" %in% names(initial_state)) {
-    state <- initial_state
+    state <- initial_state$copy()
   } else {
     stop("initial_state must be a vector or list with an entry named position.")
   }
@@ -163,7 +163,7 @@ chain_loop <- function(
     statistic_names) {
   progress_bar <- get_progress_bar(use_progress_bar, n_iteration, stage_name)
   for (adapter in adapters) {
-    adapter$initialize(state)
+    adapter$initialize(proposal, state)
   }
   if (record_traces_and_statistics) {
     trace_names <- names(unlist(trace_function(state)))
@@ -176,25 +176,25 @@ chain_loop <- function(
     traces <- NULL
     statistics <- NULL
   }
-  for (s in seq_len(n_iteration)) {
+  for (chain_iteration in seq_len(n_iteration)) {
     state_and_statistics <- sample_metropolis_hastings(
       state, target_distribution, proposal
     )
     for (adapter in adapters) {
-      adapter$update(s + 1, state_and_statistics)
+      adapter$update(proposal, chain_iteration + 1, state_and_statistics)
     }
     state <- state_and_statistics$state
     if (record_traces_and_statistics) {
-      traces[s, ] <- unlist(trace_function(state))
+      traces[chain_iteration, ] <- unlist(trace_function(state))
       adapter_states <- lapply(adapters, function(a) a$state())
-      statistics[s, ] <- unlist(
+      statistics[chain_iteration, ] <- unlist(
         c(state_and_statistics$statistics, adapter_states)
       )
     }
     if (!is.null(progress_bar)) progress_bar$tick()
   }
   for (adapter in adapters) {
-    if (!is.null(adapter$finalize)) adapter$finalize()
+    if (!is.null(adapter$finalize)) adapter$finalize(proposal)
   }
   list(final_state = state, traces = traces, statistics = statistics)
 }

README.Rmd

@@ -59,7 +59,7 @@ results <- sample_chain(
   initial_state = rnorm(dimension),
   n_warm_up_iteration = 1000,
   n_main_iteration = 1000,
-  adapters = list(scale_adapter(proposal), variance_adapter(proposal))
+  adapters = list(scale_adapter(), variance_adapter())
 )
 mean_accept_prob <- mean(results$statistics[, "accept_prob"])
 adapted_shape <- proposal$parameters()$shape

README.md

@@ -56,7 +56,7 @@ results <- sample_chain(
   initial_state = rnorm(dimension),
   n_warm_up_iteration = 1000,
   n_main_iteration = 1000,
-  adapters = list(scale_adapter(proposal), variance_adapter(proposal))
+  adapters = list(scale_adapter(), variance_adapter())
 )
 mean_accept_prob <- mean(results$statistics[, "accept_prob"])
 adapted_shape <- proposal$parameters()$shape