Fix missing legends for unobserved levels in rhat and neff plots

closes #327

R/mcmc-diagnostics.R

@@ -142,7 +142,8 @@ mcmc_rhat <- function(rhat, ..., size = NULL) {
       mapping = aes(
         yend = .data$parameter,
         xend = ifelse(min(.data$value) < 1, 1, -Inf)),
-      na.rm = TRUE) +
+      na.rm = TRUE,
+      show.legend = TRUE) +
       bayesplot_theme_get()
 
   if (min(data$value) < 1) {
@@ -238,7 +239,8 @@ mcmc_neff <- function(ratio, ..., size = NULL) {
       fill = .data$rating)) +
     geom_segment(
       aes(yend = .data$parameter, xend = -Inf),
-      na.rm = TRUE) +
+      na.rm = TRUE,
+      show.legend = TRUE) +
     diagnostic_points(size) +
     vline_at(
       c(0.1, 0.5, 1),
@@ -408,7 +410,7 @@ zero_pad_int <- function(xs) {
 }
 
 diagnostic_points <- function(size = NULL) {
-  args <- list(shape = 21, na.rm = TRUE)
+  args <- list(shape = 21, na.rm = TRUE, show.legend = TRUE)
   do.call("geom_point", c(args, size = size))
 }
 
@@ -454,7 +456,6 @@ diagnostic_colors <- function(diagnostic = c("rhat", "neff_ratio"),
   }
 
   color_labels <- diagnostic_color_labels[[diagnostic]]
-
   list(diagnostic = diagnostic,
        aesthetic = aesthetic,
        color_levels = color_levels,

tests/testthat/test-mcmc-diagnostics.R

@@ -146,6 +146,17 @@ test_that("mcmc_neff renders correctly", {
   vdiffr::expect_doppelganger("mcmc_neff (default)", p_base)
 })
 
+test_that("mcmc_neff renders legend correctly even if some levels missing", {
+  testthat::skip_on_cran()
+  testthat::skip_if_not_installed("vdiffr")
+  skip_on_r_oldrel()
+
+  neffs <- c(0.1, 0.2, 0.3, 0.4) # above 0.5 is missing but should still appear in legend
+
+  p_base <- mcmc_neff(neffs)
+  vdiffr::expect_doppelganger("mcmc_neff (missing levels)", p_base)
+})
+
 test_that("mcmc_neff_hist renders correctly", {
   testthat::skip_on_cran()
   testthat::skip_if_not_installed("vdiffr")