## ----setup, include = FALSE---------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  eval = FALSE
)

## -----------------------------------------------------------------------------
# # assume it's version 1.14, with eager not yet being the default
# library(tensorflow)
# tf$enable_v2_behavior()
# 
# library(tfprobability)
# library(rethinking)
# library(zeallot)
# library(purrr)
# 
# data("reedfrogs")
# d <- reedfrogs
# str(d)

## -----------------------------------------------------------------------------
# n_tadpole_tanks <- nrow(d)
# n_surviving <- d$surv
# n_start <- d$density
# 
# model <- tfd_joint_distribution_sequential(
#   list(
#     # a_bar, the prior for the mean of the normal distribution of per-tank logits
#     tfd_normal(loc = 0, scale = 1.5),
#     # sigma, the prior for the variance of the normal distribution of per-tank logits
#     tfd_exponential(rate = 1),
#     # normal distribution of per-tank logits
#     # parameters sigma and a_bar refer to the outputs of the above two distributions
#     function(sigma, a_bar)
#       tfd_sample_distribution(
#         tfd_normal(loc = a_bar, scale = sigma),
#         sample_shape = list(n_tadpole_tanks)
#       ),
#     # binomial distribution of survival counts
#     # parameter l refers to the output of the normal distribution immediately above
#     function(l)
#       tfd_independent(
#         tfd_binomial(total_count = n_start, logits = l),
#         reinterpreted_batch_ndims = 1
#       )
#   )
# )

## -----------------------------------------------------------------------------
# s <- model %>% tfd_sample(2)
# s

## -----------------------------------------------------------------------------
# model %>% tfd_log_prob(s)

## -----------------------------------------------------------------------------
# logprob <- function(a, s, l)
#   model %>% tfd_log_prob(list(a, s, l, n_surviving))

## -----------------------------------------------------------------------------
# # number of steps after burnin
# n_steps <- 500
# # number of chains
# n_chain <- 4
# # number of burnin steps
# n_burnin <- 500
# 
# hmc <- mcmc_hamiltonian_monte_carlo(
#   target_log_prob_fn = logprob,
#   num_leapfrog_steps = 3,
#   # one step size for each parameter
#   step_size = list(0.1, 0.1, 0.1),
# ) %>%
#   mcmc_simple_step_size_adaptation(target_accept_prob = 0.8,
#                                    num_adaptation_steps = n_burnin)
# 

## -----------------------------------------------------------------------------
# # initial values to start the sampler
# c(initial_a, initial_s, initial_logits, .) %<-% (model %>% tfd_sample(n_chain))
# 
# # optionally retrieve metadata such as acceptance ratio and step size
# trace_fn <- function(state, pkr) {
#   list(pkr$inner_results$is_accepted,
#        pkr$inner_results$accepted_results$step_size)
# }
# 
# run_mcmc <- function(kernel) {
#   kernel %>% mcmc_sample_chain(
#     num_results = n_steps,
#     num_burnin_steps = n_burnin,
#     current_state = list(initial_a, tf$ones_like(initial_s), initial_logits),
#     trace_fn = trace_fn
#   )
# }
# 
# run_mcmc <- tf_function(run_mcmc)
# res <- run_mcmc(hmc)

## -----------------------------------------------------------------------------
# mcmc_trace <- res$all_states

## -----------------------------------------------------------------------------
# map(mcmc_trace, ~ compose(dim, as.array)(.x))

## -----------------------------------------------------------------------------
# mcmc_potential_scale_reduction(mcmc_trace)
# mcmc_effective_sample_size(mcmc_trace)