Ensure rank overlay plot starts at 0 even if not all bins present

fixes #331

Author: jgabry

R/mcmc-traces.R

@@ -304,17 +304,24 @@ mcmc_rank_overlay <- function(x,
     mutate(cut = cut(.data$value_rank, n_bins)) %>%
     group_by(.data$cut) %>%
     mutate(bin_start = min(.data$value_rank)) %>%
-    ungroup() %>%
-    select(-c("cut"))
+    ungroup()
 
+  # Count how many values fall into each bin per chain & parameter
   d_bin_counts <- data %>%
     left_join(histobins, by = "value_rank") %>%
     count(.data$parameter, .data$chain, .data$bin_start)
 
+  # Now ensure that all combinations of parameter, chain, and bin_start exist
+  # even if no counts are present (https://github.com/stan-dev/bayesplot/issues/331)
+  all_params_chains <- dplyr::distinct(data, .data$parameter, .data$chain)
+  all_bins <- dplyr::distinct(histobins, .data$bin_start, .data$cut)
+  combos <- dplyr::cross_join(all_params_chains, all_bins)
+  d_bin_counts <- full_join(combos, d_bin_counts, by = c("parameter", "chain", "bin_start")) %>%
+    mutate(n = dplyr::coalesce(n, 0L))
+
   # Duplicate the final bin, setting the left edge to the greatest x value, so
   # that the entire x-axis is used,
   right_edge <- max(data$value_rank)
-
   d_bin_counts <- d_bin_counts %>%
     dplyr::filter(.data$bin_start == max(.data$bin_start)) %>%
     mutate(bin_start = right_edge) %>%

tests/testthat/_snaps/mcmc-traces/mcmc-rank-overlay-not-all-bins.svg

@@ -73,4 +73,11 @@ vdiff_dframe_chains_lp <- vdiff_dframe_chains_divergences
 vdiff_dframe_chains_lp$Parameter <- NULL
 vdiff_dframe_chains_lp$Value <- runif(2000, -100, -50)
 
+vdiff_dframe_rank_overlay_bins_test <- posterior::as_draws_df(
+  list(
+    list(theta = -2 + 0.003 * 1:1000 + stats::arima.sim(list(ar = 0.7), n = 1000, sd = 0.5)),
+    list(theta = 1 + -0.01 * 1:1000 + stats::arima.sim(list(ar = 0.7), n = 1000, sd = 0.5))
+  )
+)
+
 set.seed(seed = NULL)

tests/testthat/data-for-mcmc-tests.R

@@ -154,6 +154,9 @@ test_that("mcmc_rank_overlay renders correctly", {
     n_bins = 4
   )
 
+  # https://github.com/stan-dev/bayesplot/issues/331
+  p_not_all_bins_exist <- mcmc_rank_overlay(vdiff_dframe_rank_overlay_bins_test)
+
   vdiffr::expect_doppelganger("mcmc_rank_overlay (default)", p_base)
   vdiffr::expect_doppelganger(
     "mcmc_rank_overlay (reference line)",
@@ -164,6 +167,9 @@ test_that("mcmc_rank_overlay renders correctly", {
     "mcmc_rank_overlay (wide bins)",
     p_one_param_wide_bins
   )
+
+  # https://github.com/stan-dev/bayesplot/issues/331
+  vdiffr::expect_doppelganger("mcmc_rank_overlay (not all bins)", p_not_all_bins_exist)
 })
 
 test_that("mcmc_rank_hist renders correctly", {