Add Hamiltonian proposal (#45)

NAMESPACE

@@ -3,6 +3,7 @@
 export(barker_proposal)
 export(chain_state)
 export(example_gaussian_stan_model)
+export(hamiltonian_proposal)
 export(langevin_proposal)
 export(random_walk_proposal)
 export(robust_shape_adapter)

R/hamiltonian.R

@@ -0,0 +1,103 @@
+#' Sample new state from Hamiltonian proposal.
+#'
+#' @keywords internal
+sample_hamiltonian <- function(
+    state,
+    target_distribution,
+    n_step,
+    scale_and_shape,
+    sample_auxiliary = stats::rnorm) {
+  state$update(momentum = sample_auxiliary(state$dimension()))
+  involution_hamiltonian(state, n_step, scale_and_shape, target_distribution)
+}
+
+#' Apply involution underlying Hamiltonian proposal to a chain state.
+#'
+#' @inheritParams sample_hamiltonian
+#'
+#' @return Chain state after involution is applied.
+#'
+#' @keywords internal
+involution_hamiltonian <- function(
+    state, n_step, scale_and_shape, target_distribution) {
+  # Initial half-step
+  gradient <- state$gradient_log_density(target_distribution)
+  momentum <- state$momentum() + Matrix::drop(
+    Matrix::t(scale_and_shape) %*% gradient
+  ) / 2
+  # Inner full 'leapfrog' steps
+  for (step in seq_len(n_step - 1)) {
+    position <- state$position() + Matrix::drop(scale_and_shape %*% momentum)
+    state <- chain_state(position = position, momentum = momentum)
+    gradient <- state$gradient_log_density(target_distribution)
+    momentum <- state$momentum() + Matrix::drop(
+      Matrix::t(scale_and_shape) %*% gradient
+    )
+  }
+  # Final half-step
+  position <- state$position() + Matrix::drop(scale_and_shape %*% momentum)
+  state <- chain_state(position = position, momentum = momentum)
+  gradient <- state$gradient_log_density(target_distribution)
+  momentum <- state$momentum() + Matrix::drop(
+    Matrix::t(scale_and_shape) %*% gradient
+  ) / 2
+  # Negate momentum to make an involution
+  state$update(momentum = -momentum)
+  state
+}
+
+#' Compute logarithm of Hamiltonian proposal density ratio.
+#'
+#' @inherit log_density_ratio_barker return params
+#'
+#' @keywords internal
+log_density_ratio_hamiltonian <- function(
+    state,
+    proposed_state,
+    target_distribution,
+    scale_and_shape) {
+  sum(state$momentum()^2 - proposed_state$momentum()^2) / 2
+}
+
+#' Create a new Hamiltonian proposal object.
+#'
+#' @inherit barker_proposal return params description
+#'
+#' @param n_step Number of leapfrog steps to simulate Hamiltonian dynamics for
+#'   in each proposed move.
+#'
+#' @export
+#'
+#' @examples
+#' target_distribution <- list(
+#'   log_density = function(x) -sum(x^2) / 2,
+#'   gradient_log_density = function(x) -x
+#' )
+#' proposal <- hamiltonian_proposal(target_distribution, scale = 1., n_step = 5)
+#' state <- chain_state(c(0., 0.))
+#' withr::with_seed(876287L, proposed_state <- proposal$sample(state))
+#' log_density_ratio <- proposal$log_density_ratio(state, proposed_state)
+#' proposal$update(scale = 0.5)
+hamiltonian_proposal <- function(
+    target_distribution,
+    n_step,
+    scale = NULL,
+    shape = NULL,
+    sample_auxiliary = stats::rnorm) {
+  scale_and_shape_proposal(
+    sample = function(state, scale_and_shape) {
+      sample_hamiltonian(
+        state, target_distribution, n_step, scale_and_shape, sample_auxiliary
+      )
+    },
+    log_density_ratio = function(state, proposed_state, scale_and_shape) {
+      log_density_ratio_hamiltonian(
+        state, proposed_state, target_distribution, scale_and_shape
+      )
+    },
+    scale = scale,
+    shape = shape,
+    default_target_accept_prob = 0.8,
+    default_initial_scale = function(dimension) 2.38 / (dimension)^(1 / 4)
+  )
+}

tests/testthat/test-hamiltonian.R

@@ -0,0 +1,51 @@
+hamiltonian_proposal_with_standard_normal_target <- function(n_step) {
+  function(scale = NULL, shape = NULL) {
+    hamiltonian_proposal(
+      target_distribution = standard_normal_target_distribution(),
+      n_step = n_step,
+      scale = scale,
+      shape = shape
+    )
+  }
+}
+
+for (n_step in c(1L, 2L, 5L)) {
+  test_scale_and_shape_proposal(
+    hamiltonian_proposal_with_standard_normal_target(n_step),
+    proposal_name = sprintf("Hamiltonian proposal with n_step = %i", n_step),
+    dimensions = c(1L, 2L),
+    scales = c(0.5, 1., 2.)
+  )
+}
+
+for (dimension in c(1L, 2L, 3L)) {
+  for (scale in c(0.5, 1., 2.)) {
+    for (n_step in c(1L, 2L, 5L)) {
+      test_that(
+        sprintf(
+          paste0(
+            "Hamiltonian involution is an involution ",
+            "(dimension %i, scale %.1f, n_step %i)"
+          ),
+          dimension, scale, n_step
+        ),
+        {
+          target_distribution <- standard_normal_target_distribution()
+          withr::with_seed(seed = default_seed(), {
+            state <- chain_state(
+              position = rnorm(dimension), momentum = rnorm(dimension)
+            )
+          })
+          inv_state <- involution_hamiltonian(
+            state, n_step, scale, target_distribution
+          )
+          inv_inv_state <- involution_hamiltonian(
+            inv_state, n_step, scale, target_distribution
+          )
+          expect_equal(state$position(), inv_inv_state$position())
+          expect_equal(state$momentum(), inv_inv_state$momentum())
+        }
+      )
+    }
+  }
+}

vignettes/barker-proposal.Rmd

@@ -50,6 +50,8 @@ target_distribution <- list(
 
 - `barker_proposal`: The robust gradient-based Barker proposal proposed by @livingstone2022barker.
 - `langevin_proposal`: A gradient-based proposal based on a discretization of Langevin dynamics.
+- `hamiltonian_proposal`: A gradient-based proposal based on a discretization of Hamiltonian dynamics, simulated for a fixed number of integrator steps. 
+  With a single integrator step equivalent to `langevin_proposal`.
 - `random_walk_proposal`: A Gaussian random-walk proposal.
 
 Each function requires the first argument to specify the target distribution the proposal is to be constructed for.