mcmc_sample_chain: Implements Markov chain Monte Carlo via repeated...
In tfprobability: Interface to 'TensorFlow Probability'

mcmc_sample_chain

R Documentation

Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

Description

This function samples from an Markov chain at current_state and whose stationary distribution is governed by the supplied TransitionKernel instance (kernel).

Usage

mcmc_sample_chain(
  kernel = NULL,
  num_results,
  current_state,
  previous_kernel_results = NULL,
  num_burnin_steps = 0,
  num_steps_between_results = 0,
  trace_fn = NULL,
  return_final_kernel_results = FALSE,
  parallel_iterations = 10,
  seed = NULL,
  name = NULL
)

Arguments

`kernel`	An instance of `tfp$mcmc$TransitionKernel` which implements one step of the Markov chain.
`num_results`	Integer number of Markov chain draws.
`current_state`	`Tensor` or `list` of `Tensor`s representing the current state(s) of the Markov chain(s).
`previous_kernel_results`	A `Tensor` or a nested collection of `Tensor`s representing internal calculations made within the previous call to this function (or as returned by `bootstrap_results`).
`num_burnin_steps`	Integer number of chain steps to take before starting to collect results. Default value: 0 (i.e., no burn-in).
`num_steps_between_results`	Integer number of chain steps between collecting a result. Only one out of every `num_steps_between_samples + 1` steps is included in the returned results. The number of returned chain states is still equal to `num_results`. Default value: 0 (i.e., no thinning).
`trace_fn`	A function that takes in the current chain state and the previous kernel results and return a `Tensor` or a nested collection of `Tensor`s that is then traced along with the chain state.
`return_final_kernel_results`	If `TRUE`, then the final kernel results are returned alongside the chain state and the trace specified by the `trace_fn`.
`parallel_iterations`	The number of iterations allowed to run in parallel. It must be a positive integer. See `tf$while_loop` for more details.
`seed`	Optional, a seed for reproducible sampling.
`name`	string prefixed to Ops created by this function. Default value: `NULL`, (i.e., "mcmc_sample_chain").

Details

This function can sample from multiple chains, in parallel. (Whether or not there are multiple chains is dictated by the kernel.)

The current_state can be represented as a single Tensor or a list of Tensors which collectively represent the current state. Since MCMC states are correlated, it is sometimes desirable to produce additional intermediate states, and then discard them, ending up with a set of states with decreased autocorrelation. See Owen (2017). Such "thinning" is made possible by setting num_steps_between_results > 0. The chain then takes num_steps_between_results extra steps between the steps that make it into the results. The extra steps are never materialized (in calls to sess$run), and thus do not increase memory requirements.

Warning: when setting a seed in the kernel, ensure that sample_chain's parallel_iterations=1, otherwise results will not be reproducible. In addition to returning the chain state, this function supports tracing of auxiliary variables used by the kernel. The traced values are selected by specifying trace_fn. By default, all kernel results are traced but in the future the default will be changed to no results being traced, so plan accordingly. See below for some examples of this feature.

Value

list of:

checkpointable_states_and_trace: if return_final_kernel_results is TRUE. The return value is an instance of CheckpointableStatesAndTrace.
all_states: if return_final_kernel_results is FALSE and trace_fn is NULL. The return value is a Tensor or Python list of Tensors representing the state(s) of the Markov chain(s) at each result step. Has same shape as input current_state but with a prepended num_results-size dimension.
states_and_trace: if return_final_kernel_results is FALSE and trace_fn is not NULL. The return value is an instance of StatesAndTrace.

References

Art B. Owen. Statistically efficient thinning of a Markov chain sampler. Technical Report, 2017.

Examples


  dims <- 10
  true_stddev <- sqrt(seq(1, 3, length.out = dims))
  likelihood <- tfd_multivariate_normal_diag(scale_diag = true_stddev)

  kernel <- mcmc_hamiltonian_monte_carlo(
    target_log_prob_fn = likelihood$log_prob,
    step_size = 0.5,
    num_leapfrog_steps = 2
  )

  states <- kernel %>% mcmc_sample_chain(
    num_results = 1000,
    num_burnin_steps = 500,
    current_state = rep(0, dims),
    trace_fn = NULL
  )

  sample_mean <- tf$reduce_mean(states, axis = 0L)
  sample_stddev <- tf$sqrt(
    tf$reduce_mean(tf$math$squared_difference(states, sample_mean), axis = 0L))

tfprobability documentation built on Sept. 1, 2022, 5:07 p.m.