From 04ef3504e62ce1e6bcc8e01b0c405a3c8a25e0fc Mon Sep 17 00:00:00 2001
From: Matt Graham <matthew.m.graham@gmail.com>
Date: Wed, 23 Oct 2024 17:45:45 +0100
Subject: [PATCH] Provide per proposal defaults for target acceptance rate and
 initial scale (#42)

* Provide proposal dependent defaults for target acceptance rate and initial scale

* Update test to cover case using default initial_scale
---
 R/adaptation.R                   | 13 ++++++++++---
 R/barker.R                       |  8 +++++++-
 R/langevin.R                     |  4 +++-
 R/proposal.R                     | 11 +++++++++--
 R/random_walk.R                  |  4 +++-
 README.Rmd                       | 17 +++++------------
 README.md                        | 23 ++++++++---------------
 man/barker_proposal.Rd           |  4 ++++
 man/langevin_proposal.Rd         |  4 ++++
 man/random_walk_proposal.Rd      |  4 ++++
 man/robust_shape_adapter.Rd      |  5 +++--
 man/sample_langevin.Rd           |  4 ++++
 man/sample_random_walk.Rd        |  4 ++++
 man/scale_adapter.Rd             | 12 +++++++++---
 tests/testthat/test-adaptation.R | 25 +++++++++++++++++++++++--
 tests/testthat/test-proposal.R   |  6 +++++-
 16 files changed, 105 insertions(+), 43 deletions(-)

diff --git a/R/adaptation.R b/R/adaptation.R
index bae97ba..2b81ced 100644
--- a/R/adaptation.R
+++ b/R/adaptation.R
@@ -2,9 +2,10 @@
 #'
 #' @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.
+#' @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.
+#'   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
@@ -31,9 +32,15 @@
 #'   initial_scale = 1., target_accept_prob = 0.4
 #' )
 scale_adapter <- function(
-    proposal, initial_scale, target_accept_prob = 0.4, kappa = 0.6) {
+    proposal, 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) {
+    if (is.null(initial_scale)) {
+      initial_scale <- proposal$default_initial_scale(initial_state$dimension())
+    }
     log_scale <<- log(initial_scale)
     proposal$update(scale = initial_scale)
   }
diff --git a/R/barker.R b/R/barker.R
index e51ac27..52f5fa1 100644
--- a/R/barker.R
+++ b/R/barker.R
@@ -87,6 +87,10 @@ log_density_ratio_barker <- function(
 #'   for a given pair of current and proposed chain states,
 #' * `update`: a function to update parameters of proposal,
 #' * `parameters`: a function to return list of current parameter values.
+#' * `default_target_accept_prob`: a function returning the default target
+#'   acceptance rate to use for any scale adaptation.
+#' * `default_initial_scale`: a function which given a dimension gives a default
+#'   value to use for the initial proposal scale parameter.
 #'
 #' @export
 #'
@@ -118,6 +122,8 @@ barker_proposal <- function(
       )
     },
     scale = scale,
-    shape = shape
+    shape = shape,
+    default_target_accept_prob = 0.4,
+    default_initial_scale = function(dimension) 2.38 / (dimension)^(1 / 3)
   )
 }
diff --git a/R/langevin.R b/R/langevin.R
index 853cd9e..0d2dd34 100644
--- a/R/langevin.R
+++ b/R/langevin.R
@@ -80,6 +80,8 @@ langevin_proposal <- function(
       )
     },
     scale = scale,
-    shape = shape
+    shape = shape,
+    default_target_accept_prob = 0.574,
+    default_initial_scale = function(dimension) 2.38 / (dimension)^(1 / 3)
   )
 }
diff --git a/R/proposal.R b/R/proposal.R
index 0a54a97..171e009 100644
--- a/R/proposal.R
+++ b/R/proposal.R
@@ -20,7 +20,12 @@ get_shape_matrix <- function(scale, shape) {
   }
 }
 
-scale_and_shape_proposal <- function(sample, log_density_ratio, scale, shape) {
+scale_and_shape_proposal <- function(
+    sample,
+    log_density_ratio,
+    scale, shape,
+    default_target_accept_prob,
+    default_initial_scale) {
   scale <- scale
   shape <- shape
   list(
@@ -38,6 +43,8 @@ scale_and_shape_proposal <- function(sample, log_density_ratio, scale, shape) {
         shape <<- shape
       }
     },
-    parameters = function() list(scale = scale, shape = shape)
+    parameters = function() list(scale = scale, shape = shape),
+    default_target_accept_prob = function() default_target_accept_prob,
+    default_initial_scale = default_initial_scale
   )
 }
diff --git a/R/random_walk.R b/R/random_walk.R
index cee8cdd..a3b0dd9 100644
--- a/R/random_walk.R
+++ b/R/random_walk.R
@@ -42,6 +42,8 @@ random_walk_proposal <- function(
     },
     log_density_ratio = function(state, proposed_state, scale_and_shape) 0,
     scale = scale,
-    shape = shape
+    shape = shape,
+    default_target_accept_prob = 0.234,
+    default_initial_scale = function(dimension) 2.38 / (dimension)^(1 / 2)
   )
 }
diff --git a/README.Rmd b/README.Rmd
index e3ee7c6..424d19a 100644
--- a/README.Rmd
+++ b/README.Rmd
@@ -47,26 +47,19 @@ library(rmcmc)
 
 set.seed(876287L)
 dimension <- 3
-scales <- exp(2 * rnorm(dimension))
+scales <- exp(rnorm(dimension))
 target_distribution <- list(
   log_density = function(x) -sum((x / scales)^2) / 2,
   gradient_log_density = function(x) -x / scales^2
 )
 proposal <- barker_proposal(target_distribution)
-adapters <- list(
-  scale_adapter(proposal, initial_scale = 1., target_accept_prob = 0.4),
-  variance_adapter(proposal)
-)
-n_warm_up_iteration <- 1000
-n_main_iteration <- 1000
-initial_state <- chain_state(rnorm(dimension))
 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
+  initial_state = rnorm(dimension),
+  n_warm_up_iteration = 1000,
+  n_main_iteration = 1000,
+  adapters = list(scale_adapter(proposal), variance_adapter(proposal))
 )
 mean_accept_prob <- mean(results$statistics[, "accept_prob"])
 adapted_shape <- proposal$parameters()$shape
diff --git a/README.md b/README.md
index 5b675d3..6c07526 100644
--- a/README.md
+++ b/README.md
@@ -44,26 +44,19 @@ library(rmcmc)
 
 set.seed(876287L)
 dimension <- 3
-scales <- exp(2 * rnorm(dimension))
+scales <- exp(rnorm(dimension))
 target_distribution <- list(
   log_density = function(x) -sum((x / scales)^2) / 2,
   gradient_log_density = function(x) -x / scales^2
 )
 proposal <- barker_proposal(target_distribution)
-adapters <- list(
-  scale_adapter(proposal, initial_scale = 1., target_accept_prob = 0.4),
-  variance_adapter(proposal)
-)
-n_warm_up_iteration <- 1000
-n_main_iteration <- 1000
-initial_state <- chain_state(rnorm(dimension))
 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
+  initial_state = rnorm(dimension),
+  n_warm_up_iteration = 1000,
+  n_main_iteration = 1000,
+  adapters = list(scale_adapter(proposal), variance_adapter(proposal))
 )
 mean_accept_prob <- mean(results$statistics[, "accept_prob"])
 adapted_shape <- proposal$parameters()$shape
@@ -73,7 +66,7 @@ cat(
   sprintf("Adapter scale est.: %s", toString(adapted_shape)),
   sep = "\n"
 )
-#> Average acceptance probability is 0.46
-#> True target scales: 2.26617033226883, 1.89818769776724, 0.0767505506297473
-#> Adapter scale est.: 1.77277384748788, 1.71554065105575, 0.0804144979270686
+#> Average acceptance probability is 0.40
+#> True target scales: 1.50538046096953, 1.37774732725824, 0.277038897322645
+#> Adapter scale est.: 1.35010920408606, 1.5140138215658, 0.248974800274054
 ```
diff --git a/man/barker_proposal.Rd b/man/barker_proposal.Rd
index 6414887..85ba9e0 100644
--- a/man/barker_proposal.Rd
+++ b/man/barker_proposal.Rd
@@ -45,6 +45,10 @@ current chain state,
 for a given pair of current and proposed chain states,
 \item \code{update}: a function to update parameters of proposal,
 \item \code{parameters}: a function to return list of current parameter values.
+\item \code{default_target_accept_prob}: a function returning the default target
+acceptance rate to use for any scale adaptation.
+\item \code{default_initial_scale}: a function which given a dimension gives a default
+value to use for the initial proposal scale parameter.
 }
 }
 \description{
diff --git a/man/langevin_proposal.Rd b/man/langevin_proposal.Rd
index 810848e..76b8084 100644
--- a/man/langevin_proposal.Rd
+++ b/man/langevin_proposal.Rd
@@ -41,6 +41,10 @@ current chain state,
 for a given pair of current and proposed chain states,
 \item \code{update}: a function to update parameters of proposal,
 \item \code{parameters}: a function to return list of current parameter values.
+\item \code{default_target_accept_prob}: a function returning the default target
+acceptance rate to use for any scale adaptation.
+\item \code{default_initial_scale}: a function which given a dimension gives a default
+value to use for the initial proposal scale parameter.
 }
 }
 \description{
diff --git a/man/random_walk_proposal.Rd b/man/random_walk_proposal.Rd
index d9a6c87..78dea8c 100644
--- a/man/random_walk_proposal.Rd
+++ b/man/random_walk_proposal.Rd
@@ -41,6 +41,10 @@ current chain state,
 for a given pair of current and proposed chain states,
 \item \code{update}: a function to update parameters of proposal,
 \item \code{parameters}: a function to return list of current parameter values.
+\item \code{default_target_accept_prob}: a function returning the default target
+acceptance rate to use for any scale adaptation.
+\item \code{default_initial_scale}: a function which given a dimension gives a default
+value to use for the initial proposal scale parameter.
 }
 }
 \description{
diff --git a/man/robust_shape_adapter.Rd b/man/robust_shape_adapter.Rd
index 8a737b5..9f8cac4 100644
--- a/man/robust_shape_adapter.Rd
+++ b/man/robust_shape_adapter.Rd
@@ -16,10 +16,11 @@ robust_shape_adapter(
 \item{proposal}{Proposal object to adapt. Must define an \code{update} function
 which accepts a parameter \code{scale} for setting scale parameter of proposal.}
 
-\item{initial_scale}{Initial value to use for scale parameter.}
+\item{initial_scale}{Initial value to use for scale parameter. If not set
+explicitly a proposal and dimension dependent default will be used.}
 
 \item{target_accept_prob}{Target value for average accept probability for
-chain.}
+chain. If not set a proposal dependent default will be used.}
 
 \item{kappa}{Decay rate exponent in \verb{[0.5, 1]} for adaptation learning rate.}
 }
diff --git a/man/sample_langevin.Rd b/man/sample_langevin.Rd
index e2aef69..d388cd7 100644
--- a/man/sample_langevin.Rd
+++ b/man/sample_langevin.Rd
@@ -34,6 +34,10 @@ current chain state,
 for a given pair of current and proposed chain states,
 \item \code{update}: a function to update parameters of proposal,
 \item \code{parameters}: a function to return list of current parameter values.
+\item \code{default_target_accept_prob}: a function returning the default target
+acceptance rate to use for any scale adaptation.
+\item \code{default_initial_scale}: a function which given a dimension gives a default
+value to use for the initial proposal scale parameter.
 }
 }
 \description{
diff --git a/man/sample_random_walk.Rd b/man/sample_random_walk.Rd
index 07b64e7..c9b600c 100644
--- a/man/sample_random_walk.Rd
+++ b/man/sample_random_walk.Rd
@@ -34,6 +34,10 @@ current chain state,
 for a given pair of current and proposed chain states,
 \item \code{update}: a function to update parameters of proposal,
 \item \code{parameters}: a function to return list of current parameter values.
+\item \code{default_target_accept_prob}: a function returning the default target
+acceptance rate to use for any scale adaptation.
+\item \code{default_initial_scale}: a function which given a dimension gives a default
+value to use for the initial proposal scale parameter.
 }
 }
 \description{
diff --git a/man/scale_adapter.Rd b/man/scale_adapter.Rd
index 2a80c1b..9943a85 100644
--- a/man/scale_adapter.Rd
+++ b/man/scale_adapter.Rd
@@ -4,16 +4,22 @@
 \alias{scale_adapter}
 \title{Create object to adapt proposal scale to coerce average acceptance rate.}
 \usage{
-scale_adapter(proposal, initial_scale, target_accept_prob = 0.4, kappa = 0.6)
+scale_adapter(
+  proposal,
+  initial_scale = NULL,
+  target_accept_prob = NULL,
+  kappa = 0.6
+)
 }
 \arguments{
 \item{proposal}{Proposal object to adapt. Must define an \code{update} function
 which accepts a parameter \code{scale} for setting scale parameter of proposal.}
 
-\item{initial_scale}{Initial value to use for scale parameter.}
+\item{initial_scale}{Initial value to use for scale parameter. If not set
+explicitly a proposal and dimension dependent default will be used.}
 
 \item{target_accept_prob}{Target value for average accept probability for
-chain.}
+chain. If not set a proposal dependent default will be used.}
 
 \item{kappa}{Decay rate exponent in \verb{[0.5, 1]} for adaptation learning rate.}
 }
diff --git a/tests/testthat/test-adaptation.R b/tests/testthat/test-adaptation.R
index d37cd60..9eed6c8 100644
--- a/tests/testthat/test-adaptation.R
+++ b/tests/testthat/test-adaptation.R
@@ -15,7 +15,9 @@ check_adapter <- function(adapter) {
 dummy_proposal_with_scale_parameter <- function(scale = NULL) {
   list(
     update = function(scale) scale <<- scale,
-    parameters = function() list(scale = scale)
+    parameters = function() list(scale = scale),
+    default_target_accept_prob = function() 0.234,
+    default_initial_scale = function(dimension) 1 / sqrt(dimension)
   )
 }
 
@@ -46,7 +48,7 @@ for (target_accept_prob in c(0.2, 0.4, 0.6)) {
             kappa = kappa
           )
           check_adapter(adapter)
-          adapter$initialize(initial_state = NULL)
+          adapter$initialize(initial_state = chain_state(rep(0, dimension)))
           adapter_state <- adapter$state()
           expect_named(adapter_state, "log_scale")
           expect_length(adapter_state$log_scale, 1)
@@ -87,6 +89,25 @@ for (target_accept_prob in c(0.2, 0.4, 0.6)) {
   }
 }
 
+for (dimension in c(1L, 2L, 5L)) {
+  test_that(
+    sprintf(
+      "Scale adapter with only proposal specified works in dimension %i",
+      dimension
+    ),
+    {
+      proposal <- dummy_proposal_with_scale_parameter()
+      adapter <- scale_adapter(proposal)
+      check_adapter(adapter)
+      adapter$initialize(initial_state = chain_state(rep(0, dimension)))
+      adapter_state <- adapter$state()
+      expect_named(adapter_state, "log_scale")
+      expect_length(adapter_state$log_scale, 1)
+      expect_equal(adapter_state$log_scale, -log(dimension) / 2)
+    }
+  )
+}
+
 for (dimension in c(1L, 2L, 5L)) {
   for (kappa in c(0.5, 0.6, 0.8)) {
     for (correlation in c(0.0, 0.5, 0.8)) {
diff --git a/tests/testthat/test-proposal.R b/tests/testthat/test-proposal.R
index 73018a0..6b4567b 100644
--- a/tests/testthat/test-proposal.R
+++ b/tests/testthat/test-proposal.R
@@ -28,7 +28,9 @@ for (dimension in c(1, 2)) {
       chain_state(position = state$position() + offset)
     }
     log_density_ratio <- function(state, proposed_state, scale_and_shape) 0
-    proposal <- scale_and_shape_proposal(sample, log_density_ratio, NULL, NULL)
+    proposal <- scale_and_shape_proposal(
+      sample, log_density_ratio, NULL, NULL, NULL, function(d) NULL
+    )
     state <- chain_state(position = rnorm(dimension))
     expect_error(proposal$sample(state), "must be set")
     expect_identical(proposal$parameters(), list(scale = NULL, shape = NULL))
@@ -40,5 +42,7 @@ for (dimension in c(1, 2)) {
     check_chain_state(proposed_state)
     expect_all_different(state$position(), proposed_state$position())
     expect_identical(proposal$log_density(state, proposed_state), 0)
+    expect_identical(proposal$default_target_accept_prob(), NULL)
+    expect_identical(proposal$default_initial_scale(dimension), NULL)
   })
 }