Don't require passing proposal to adapter initializers

R/adaptation.R

@@ -9,8 +9,8 @@
 #' @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 +27,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,7 +51,7 @@ scale_adapter <- function(
   list(
     initialize = initialize,
     update = update,
-    finalize = function() {},
+    finalize = NULL,
     state = function() list(log_scale = log_scale)
   )
 }
@@ -74,15 +72,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 +123,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

@@ -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