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

## -----------------------------------------------------------------------------
# library(tensorflow)
# # assume it's version 1.14, with eager not yet being the default
# tf$compat$v1$enable_v2_behavior()
# 
# library(tfprobability)
# library(keras)
# 
# library(dplyr)
# library(tidyr)
# library(ggplot2)
# 
# # generate the data
# x_min <- -40
# x_max <- 60
# n <- 150
# w0 <- 0.125
# b0 <- 5
# 
# normalize <- function(x) (x - x_min) / (x_max - x_min)
# 
# # training data; predictor
# x <- x_min + (x_max - x_min) * runif(n) %>% as.matrix()
# 
# # training data; target
# eps <- rnorm(n) * (3 * (0.25 + (normalize(x)) ^ 2))
# y <- (w0 * x * (1 + sin(x)) + b0) + eps
# 
# # test data (predictor)
# x_test <- seq(x_min, x_max, length.out = n) %>% as.matrix()

## -----------------------------------------------------------------------------
# ggplot(data.frame(x = x, y = y), aes(x, y)) + geom_point()

## ----eval=TRUE, echo=FALSE, layout="l-body-outset", fig.cap = "Simulated data"----
knitr::include_graphics("images/uncertainty_data.png")

## -----------------------------------------------------------------------------
# prior_trainable <-
#   function(kernel_size,
#            bias_size = 0,
#            dtype = NULL) {
#     n <- kernel_size + bias_size
#     keras_model_sequential() %>%
#       layer_variable(n, dtype = dtype, trainable = TRUE) %>%
#       layer_distribution_lambda(function(t) {
#         tfd_independent(tfd_normal(loc = t, scale = 1),
#                         reinterpreted_batch_ndims = 1)
#       })
#   }
# 

## -----------------------------------------------------------------------------
# posterior_mean_field <-
#   function(kernel_size,
#            bias_size = 0,
#            dtype = NULL) {
#     n <- kernel_size + bias_size
#     c <- log(expm1(1))
#     keras_model_sequential(list(
#       layer_variable(shape = 2 * n, dtype = dtype),
#       layer_distribution_lambda(
#         make_distribution_fn = function(t) {
#           tfd_independent(tfd_normal(
#             loc = t[1:n],
#             scale = 1e-5 + tf$nn$softplus(c + t[(n + 1):(2 * n)])
#             ), reinterpreted_batch_ndims = 1)
#         }
#       )
#     ))
#   }

## -----------------------------------------------------------------------------
# model <- keras_model_sequential() %>%
#   layer_dense_variational(
#     units = 2,
#     make_posterior_fn = posterior_mean_field,
#     make_prior_fn = prior_trainable,
#     # scale by the size of the dataset
#     kl_weight = 1 / n
#   ) %>%
#   layer_distribution_lambda(function(x)
#     tfd_normal(loc = x[, 1, drop = FALSE],
#                scale = 1e-3 + tf$math$softplus(0.01 * x[, 2, drop = FALSE])
#                )
#     )
# 

## -----------------------------------------------------------------------------
# negloglik <- function(y, model) - (model %>% tfd_log_prob(y))
# model %>% compile(optimizer = optimizer_adam(0.01), loss = negloglik)
# model %>% fit(x, y, epochs = 1000)

## -----------------------------------------------------------------------------
# # each time we ask the model to predict, we get a different line
# yhats <- purrr::map(1:100, function(x) model(tf$constant(x_test)))
# means <-
#   purrr::map(yhats, purrr::compose(as.matrix, tfd_mean)) %>% abind::abind()
# sds <-
#   purrr::map(yhats, purrr::compose(as.matrix, tfd_stddev)) %>% abind::abind()
# 
# means_gathered <- data.frame(cbind(x_test, means)) %>%
#   gather(key = run, value = mean_val,-X1)
# sds_gathered <- data.frame(cbind(x_test, sds)) %>%
#   gather(key = run, value = sd_val,-X1)
# 
# lines <-
#   means_gathered %>% inner_join(sds_gathered, by = c("X1", "run"))
# mean <- apply(means, 1, mean)
# 
# ggplot(data.frame(x = x, y = y, mean = as.numeric(mean)), aes(x, y)) +
#   geom_point() +
#   theme(legend.position = "none") +
#   geom_line(aes(x = x_test, y = mean), color = "violet", size = 1.5) +
#   geom_line(
#     data = lines,
#     aes(x = X1, y = mean_val, color = run),
#     alpha = 0.6,
#     size = 0.5
#   ) +
#   geom_ribbon(
#     data = lines,
#     aes(
#       x = X1,
#       ymin = mean_val - 2 * sd_val,
#       ymax = mean_val + 2 * sd_val,
#       group = run
#     ),
#     alpha = 0.05,
#     fill = "grey",
#     inherit.aes = FALSE
#   )

## ----eval=TRUE, echo=FALSE, layout="l-body-outset", fig.cap = "Displaying both epistemic and aleatoric uncertainty on the simulated dataset."----
knitr::include_graphics("images/uncertainty.png")