plot_rhats.R

plot_rhats <- function(dgirt_out) {
  # Save to disk plots of split R-hats for theta-bars (from a particular model)
  #
  # This function produces two plots showing the average R-hat in each year by
  # 1) state and 2) party, each averaging over the other. It assumes that
  # particular variable names exist in dgirt_out as the grouping variable (D_pid3), 
  # time variable (D_year), and geographic variable (D_abb).

  # From the stanfit object extract R-hats and effective sample sizes
  tb_sum <- summary(dgirt_out, par = "theta_bar", verbose = TRUE)$summary
  tb_sum <- data.frame(group = rownames(tb_sum),
                       rhat = tb_sum[, "Rhat"],
                       n_eff = tb_sum[, "n_eff"])
  tb_sum$t <- as.integer(gsub(".*\\[(\\d+),.*", "\\1", tb_sum$group))
  tb_sum$g <- as.integer(gsub(".*,(\\d+).*", "\\1", tb_sum$group))
  data.table::setDT(tb_sum)
  nrow(tb_sum[rhat > 1.1, ]) / nrow(tb_sum)

  fnames = dgirt:::flatnames(dgirt_out) %>%
    dplyr::filter(param == "theta_bar") %>%
    dplyr::select(-hier_param, -item, -param, state = D_abb, pid3 = D_pid3,
           year = D_year) %>%
    dplyr::mutate(t = as.integer(gsub(".*\\[(\\d+),.*", "\\1", fname)),
           g = as.integer(gsub(".*,(\\d+).*", "\\1", fname))) %>%
    dplyr::select(-fname)

  tb_sum = left_join(tb_sum, fnames, by = c("t", "g")) %>%
    dplyr::select(-t, -g, -group) %>%
    # these values vary over versions of the model
    dplyr::mutate(pid3 = ifelse(pid3 == "democrat" | pid3 == "D_pid31", "D", pid3),
           pid3 = ifelse(pid3 == "independent" | pid3 == "D_pid32", "I", pid3),
           pid3 = ifelse(pid3 == "republican" | pid3 == "D_pid33", "R", pid3),
           state = sub("D_abb", "", state))

  p = tb_sum %>%
    group_by(year, state) %>%
    summarize(rhat = mean(rhat)) %>%
    ggplot() +
    geom_line(aes(year, rhat)) +
    facet_wrap(~ state) +
    geom_hline(yintercept = 1.1, alpha = 0.5) +
    theme_few() +
    theme(axis.text.x = element_text(angle = 90, hjust = 1))
  p

  ggsave(paste0("mean_state_rhats-", Sys.Date(), ".pdf"))

  p = tb_sum %>%
    group_by(year, pid3) %>%
    summarize(rhat = mean(rhat)) %>%
    ggplot() +
    geom_line(aes(year, rhat)) +
    facet_wrap(~ pid3) +
    geom_hline(yintercept = 1.1, alpha = 0.5) +
    theme_few() +
    theme(axis.text.x = element_text(angle = 90, hjust = 1))
  ggsave(paste0("mean_party_rhats-", Sys.Date(), ".pdf"))
}