Add Langevin proposal implementation (#25)

R/langevin.R

@@ -0,0 +1,84 @@
+#' Sample new state from Langevin proposal.
+#'
+#' @inherit barker_proposal return params
+#'
+#' @keywords internal
+sample_langevin <- function(
+    state,
+    target_distribution,
+    scale_and_shape,
+    sample_auxiliary = stats::rnorm) {
+  state$update(momentum = sample_auxiliary(state$dimension()))
+  involution_langevin(state, scale_and_shape, target_distribution)
+}
+
+#' Apply involution underlying Langevin proposal to a chain state.
+#'
+#' @inheritParams sample_langevin
+#'
+#' @return Chain state after involution is applied.
+#'
+#' @keywords internal
+involution_langevin <- function(state, scale_and_shape, target_distribution) {
+  gradient <- state$gradient_log_density(target_distribution)
+  momentum <- state$momentum() + Matrix::drop(
+    Matrix::t(scale_and_shape) %*% gradient
+  ) / 2
+  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
+  state$update(momentum = -momentum)
+  state
+}
+
+#' Compute logarithm of Langevin proposal density ratio.
+#'
+#' @inherit log_density_ratio_barker return params
+#'
+#' @keywords internal
+log_density_ratio_langevin <- function(
+    state,
+    proposed_state,
+    target_distribution,
+    scale_and_shape) {
+  sum(proposed_state$momentum()^2 - state$momentum()^2) / 2
+}
+
+#' Create a new Langevin proposal object.
+#'
+#' @inherit barker_proposal return params description
+#'
+#' @export
+#'
+#' @examples
+#' target_distribution <- list(
+#'   log_density = function(x) -sum(x^2) / 2
+#' )
+#' proposal <- langevin_proposal(target_distribution, scale = 1.)
+#' 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)
+langevin_proposal <- function(
+    target_distribution,
+    scale = NULL,
+    shape = NULL,
+    sample_auxiliary = stats::rnorm) {
+  scale_and_shape_proposal(
+    sample = function(state, scale_and_shape) {
+      sample_langevin(
+        state, target_distribution, scale_and_shape, sample_auxiliary
+      )
+    },
+    log_density_ratio = function(state, proposed_state, scale_and_shape) {
+      log_density_ratio_langevin(
+        state, proposed_state, target_distribution, scale_and_shape
+      )
+    },
+    scale = scale,
+    shape = shape
+  )
+}

tests/testthat/test-langevin.R

@@ -0,0 +1,34 @@
+langevin_proposal_with_standard_normal_target <- function(
+    scale = NULL, shape = NULL) {
+  target_distribution <- standard_normal_target_distribution()
+  langevin_proposal(target_distribution, scale = scale, shape = shape)
+}
+
+test_scale_and_shape_proposal(
+  langevin_proposal_with_standard_normal_target,
+  proposal_name = "Langevin proposal",
+  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.)) {
+    test_that(
+      sprintf(
+        "Langevin involution is an involution (dimension %i, scale %.1f)",
+        dimension, scale
+      ), {
+        target_distribution <- standard_normal_target_distribution()
+        withr::with_seed(seed = default_seed(), {
+          state <- chain_state(
+            position = rnorm(dimension), momentum = rnorm(dimension)
+          )
+        })
+        inv_state <- involution_langevin(state, scale, target_distribution)
+        inv_inv_state <- involution_langevin(inv_state, scale, target_distribution)
+        expect_equal(state$position(), inv_inv_state$position())
+        expect_equal(state$momentum(), inv_inv_state$momentum())
+      }
+    )
+  }
+}