Make diagnostic color matching explicit for neff

Author: utkarshpawade

NEWS.md

@@ -2,6 +2,7 @@
 
 * New functions `mcmc_dots` and `mcmc_dots_by_chain` for dot plots of MCMC draws by @behramulukir (#402)
 * Default to `quantiles=100` for all dot plots by @behramulukir (#402)
+* Make diagnostic color scale helpers handle `"neff"` and `"neff_ratio"` explicitly, avoiding reliance on partial matching.
 
 # bayesplot 1.15.0
 

R/mcmc-diagnostics.R

@@ -429,7 +429,10 @@ scale_fill_diagnostic <- function(diagnostic = c("rhat", "neff")) {
 
 diagnostic_color_scale <- function(diagnostic = c("rhat", "neff_ratio"),
                                    aesthetic = c("color", "fill")) {
-  diagnostic <- match.arg(diagnostic)
+  diagnostic <- match.arg(diagnostic, choices = c("rhat", "neff", "neff_ratio"))
+  if (diagnostic == "neff") {
+    diagnostic <- "neff_ratio"
+  }
   aesthetic <- match.arg(aesthetic)
   dc <- diagnostic_colors(diagnostic, aesthetic)
   do.call(
@@ -445,7 +448,10 @@ diagnostic_color_scale <- function(diagnostic = c("rhat", "neff_ratio"),
 
 diagnostic_colors <- function(diagnostic = c("rhat", "neff_ratio"),
                               aesthetic = c("color", "fill")) {
-  diagnostic <- match.arg(diagnostic)
+  diagnostic <- match.arg(diagnostic, choices = c("rhat", "neff", "neff_ratio"))
+  if (diagnostic == "neff") {
+    diagnostic <- "neff_ratio"
+  }
   aesthetic <- match.arg(aesthetic)
   color_levels <- c("light", "mid", "dark")
   if (diagnostic == "neff_ratio") {

tests/testthat/test-mcmc-diagnostics.R

@@ -79,6 +79,13 @@ test_that("'description' & 'rating' columns are correct (#176)", {
   expect_equal(df$description, expected_descriptions)
 })
 
+test_that("diagnostic color helpers handle neff names explicitly", {
+  expect_no_error(scale_color_diagnostic("neff"))
+  expect_no_error(scale_fill_diagnostic("neff"))
+  expect_no_error(diagnostic_color_scale("neff", aesthetic = "color"))
+  expect_no_error(diagnostic_color_scale("neff_ratio", aesthetic = "color"))
+})
+
 test_that("mcmc_acf & mcmc_acf_bar return a ggplot object", {
   expect_gg(mcmc_acf(arr, pars = "beta[1]", regex_pars = "x\\:[2,5]"))
   expect_gg(mcmc_acf_bar(arr, pars = "beta[1]", regex_pars = "x\\:[2,5]"))