gganimate_brms_posterior.R

library(tidyverse)
library(brms)
library(tidybayes)
library(gganimate)
library(magick)

tb <- tibble(
  index = 1:15,
  a = runif(n = 15, min = 0, max = 2),
  b = rnorm(15, mean = 2 * a, sd = 1),
)

tb

tb %>%
  ggplot(aes(x = a, y = b)) +
  geom_point() +
  geom_smooth(method = "lm")



# rep(list(tb), nrow(tb)) %>%
#   enframe(name = "index", value = "data") %>%
#   mutate(
#     samp_data = map2(data, index, ~ dplyr::slice(.x, 1:.y))
#   )


# tb %>%
#   group_split(index) %>%
#   purrr::accumulate(bind_rows) %>%
#   enframe()



tb %>%
  group_split(index) %>%
  purrr::accumulate(bind_rows) %>%
  enframe() %>%
  filter(name > 5) %>%
  mutate(
    mod = map(value, ~ lm(b ~ a, data = .))
  ) %>%
  mutate(tidied = map(mod, broom::tidy)) %>%
  unnest(tidied) %>%
  filter(term != "(Intercept)")




a <- tb %>%
  group_split(index) %>%
  purrr::accumulate(bind_rows) %>%
  enframe() %>%
  filter(name > 5) %>%
  unnest(value) %>%
  ggplot(
    aes(x = a, y = b)
  ) +
  geom_point(size = 4) +
  geom_smooth(aes(group = factor(name)), method = "lm", se = FALSE) +
  transition_states(name) +
  theme(legend.position = "none") +
  shadow_mark(past = FALSE) +
  labs(title = "n: {closest_state}")

a_gif <- animate(a, width = 240, height = 240)



results <- tb %>%
  group_split(index) %>%
  purrr::accumulate(bind_rows) %>%
  discard(~ nrow(.x) <= 5) %>%
  brms::brm_multiple(
    bf(b ~ a),
    data = .,
    prior = c(
      set_prior("normal(0, 5)", class = "Intercept"),
      set_prior("normal(2, 1)", class = "b")
    ),
    chains = 2, iter = 4000, warmup = 2000, seed = 1024,
    combine = F
  )



# # prior 
# tibble(
#   x = seq(-1, 5, length.out = 100),
#   y = dnorm(x, mean = 2, sd = 1)
# ) %>% 
#   ggplot(aes(x = x, y = y)) +
#   geom_line() +
#   labs(title = "prior distrition: Normal(mean = 2, sd = 1)")
# 
# 
# # facet
# results %>%
#   enframe() %>%
#   mutate(
#     name = str_pad(name + 5, width = 2, pad = "0"),
#     name = str_c("sample = ", name)
#   ) %>%
#   group_by(name) %>%
#   transmute(post = map(value, posterior_samples)) %>%
#   unnest(post) %>%
#   ggplot(
#     aes(x = b_a, y = 0)
#   ) +
#   geom_line(
#     data = tibble(
#       x = seq(-1, 5, length.out = 100),
#       y = dnorm(x, mean = 2, sd = 1)
#     ),
#     aes(x = x, y = y)
#   ) +
#   tidybayes::geom_halfeyeh() +
#   facet_wrap(vars(name))




b <- results %>%
  enframe() %>%
  mutate(
    name = str_pad(name + 5, width = 2, pad = "0"),
    name = str_c("sample = ", name)
  ) %>%
  group_by(name) %>%
  transmute(post = map(value, posterior_samples)) %>%
  unnest(post) %>%
  ggplot(
    aes(x = b_a, y = 0)
  ) +
  geom_line(
    data = tibble(
      x = seq(-1, 5, length.out = 100),
      y = dnorm(x, mean = 2, sd = 1)
    ),
    aes(x = x, y = y)
  ) +
  tidybayes::geom_halfeyeh() +
  transition_states(name) +
  shadow_mark(past = FALSE) +
  enter_fade() +
  exit_shrink() +
  ease_aes("sine-in-out") +
  #view_follow(fixed_x = c(-5, 8)) +
  labs(title = "n: {closest_state}")

b_gif <- animate(b, width = 240, height = 240)


a_mgif <- image_read(a_gif)
b_mgif <- image_read(b_gif)

new_gif <- image_append(c(a_mgif[1], b_mgif[1]))
for (i in 2:100) {
  combined <- image_append(c(a_mgif[i], b_mgif[i]))
  new_gif <- c(new_gif, combined)
}

new_gif

anim_save('test2.gif', new_gif)



# Ref
# https://mjskay.github.io/tidybayes/
# https://solomonkurz.netlify.com/post/would-you-like-all-your-posteriors-in-one-plot/
# https://github.com/thomasp85/gganimate/wiki/Animation-Composition