mcmc_trace: officially support highlight and alpha arguments

NEWS.md

@@ -13,6 +13,7 @@
 * Replace `apply()` with `storage.mode()` for integer-to-numeric matrix conversion in `validate_predictions()`.
 * Fixed `is_chain_list()` to correctly reject empty lists instead of silently returning `TRUE`.
 * Added unit tests for `mcmc_areas_ridges_data()`, `mcmc_parcoord_data()`, and `mcmc_trace_data()`.
+* `mcmc_trace()` now supports highlighting a chain with lines using the `highlight` and `alpha` arguments. Previously this was only available via `mcmc_trace_highligh()` and with points instead of lines. (#552)
 * Added unit tests for `ppc_error_data()` and `ppc_loo_pit_data()` covering output structure, argument handling, and edge cases.
 * Added vignette sections demonstrating `*_data()` companion functions for building custom ggplot2 visualizations (#435)
 * Extract `drop_singleton_values()` helper in `mcmc_nuts_treedepth()` to remove duplicated filtering logic.

R/mcmc-traces.R

@@ -13,9 +13,6 @@
 #' @param ... Currently ignored.
 #' @param size An optional value to override the default line size
 #'   for `mcmc_trace()` or the default point size for `mcmc_trace_highlight()`.
-#' @param alpha For `mcmc_trace_highlight()`, passed to
-#'   [ggplot2::geom_point()] to control the transparency of the points
-#'   for the chains not highlighted.
 #' @param n_warmup An integer; the number of warmup iterations included in
 #'   `x`. The default is `n_warmup = 0`, i.e. to assume no warmup
 #'   iterations are included. If `n_warmup > 0` then the background for
@@ -26,6 +23,11 @@
 #'   if `n_warmup` is also set to a positive value.
 #' @param window An integer vector of length two specifying the limits of a
 #'   range of iterations to display.
+#' @param highlight For `mcmc_trace()`, `NULL` (the default) or an integer
+#'   specifying one chain to emphasize. For `mcmc_trace_highlight()`, an integer
+#'   specifying one chain to emphasize.
+#' @param alpha For `mcmc_trace()` and `mcmc_trace_highlight()`, controls the
+#'   transparency of the lines or points for the chains not highlighted.
 #' @param np For models fit using [NUTS] (more generally, any
 #'   [symplectic integrator](https://en.wikipedia.org/wiki/Symplectic_integrator)),
 #'   an optional data frame providing NUTS diagnostic information. The data
@@ -47,6 +49,7 @@
 #'   \item{`mcmc_trace()`}{
 #'    Standard trace plots of MCMC draws. For models fit using [NUTS],
 #'    the `np` argument can be used to also show divergences on the trace plot.
+#'    One chain can be emphasized using the `highlight` argument.
 #'   }
 #'   \item{`mcmc_trace_highlight()`}{
 #'    Traces are plotted using points rather than lines and the opacity of all
@@ -102,6 +105,7 @@
 #' # mix color schemes
 #' color_scheme_set("mix-blue-red")
 #' mcmc_trace(x, regex_pars = "beta")
+#' mcmc_trace(x, regex_pars = "beta", highlight = 2)
 #'
 #' # use traditional ggplot discrete color scale
 #' mcmc_trace(x, pars = c("alpha", "sigma")) +
@@ -183,6 +187,8 @@ mcmc_trace <-
            iter1 = 0,
            window = NULL,
            size = NULL,
+           alpha = 0.4,
+           highlight = NULL,
            np = NULL,
            np_style = trace_style_np(),
            divergences = NULL) {
@@ -213,18 +219,17 @@ mcmc_trace <-
     n_warmup = n_warmup,
     window = window,
     size = size,
+    alpha = alpha,
+    highlight = highlight,
     style = "line",
     np = np,
     np_style = np_style,
-    iter1 = iter1,
-    ...
+    iter1 = iter1
   )
 }
 
 #' @rdname MCMC-traces
 #' @export
-#' @param highlight For `mcmc_trace_highlight()`, an integer specifying one
-#'   of the chains that will be more visible than the others in the plot.
 mcmc_trace_highlight <- function(x,
                                  pars = character(),
                                  regex_pars = character(),
@@ -248,8 +253,7 @@ mcmc_trace_highlight <- function(x,
     size = size,
     alpha = alpha,
     highlight = highlight,
-    style = "point",
-    ...
+    style = "point"
   )
 }
 
@@ -657,8 +661,7 @@ mcmc_trace_data <- function(x,
                         alpha = 0.2,
                         np = NULL,
                         np_style = trace_style_np(),
-                        iter1 = 0,
-                        ...) {
+                        iter1 = 0) {
   style <- match.arg(style)
   data <- mcmc_trace_data(
     x,
@@ -725,17 +728,17 @@ mcmc_trace_data <- function(x,
       labels = c("Other chains", paste("Chain", highlight)))
   } else {
     scale_color <- scale_color_manual("Chain", values = chain_colors(n_chain))
+  }
 
-    if (!is.null(np)) {
-      div_rug <- divergence_rug(np, np_style, n_iter, n_chain)
-      if (!is.null(div_rug)) {
-        div_guides <- guides(
-          color = guide_legend(order = 1),
-          linetype = guide_legend(
-            order = 2, title = NULL, keywidth = rel(1/2),
-            override.aes = list(linewidth = rel(1/2)))
-        )
-      }
+  if (!is.null(np)) {
+    div_rug <- divergence_rug(np, np_style, n_iter, n_chain)
+    if (!is.null(div_rug)) {
+      div_guides <- guides(
+        color = guide_legend(order = 1),
+        linetype = guide_legend(
+          order = 2, title = NULL, keywidth = rel(1/2),
+          override.aes = list(linewidth = rel(1/2)))
+      )
     }
   }
 

tests/testthat/test-mcmc-traces.R

@@ -14,6 +14,25 @@ test_that("mcmc_trace returns a ggplot object", {
   expect_gg(mcmc_trace(chainlist1))
 })
 
+test_that("mcmc_trace highlights a chain with lines", {
+  expect_no_warning(
+    p <- mcmc_trace(arr, pars = "sigma", highlight = 2, alpha = 0.4)
+  )
+
+  expect_s3_class(p$layers[[1]]$geom, "GeomLine")
+  expect_equal(p$data$highlight, p$data$chain == 2)
+  expect_equal(p$scales$get_scales("alpha")$palette(2), c(0.4, 1))
+})
+
+test_that("mcmc_trace shows divergences when highlighting a chain", {
+  np <- rep_len(c(0, 1), nrow(arr))
+  p <- mcmc_trace(arr, pars = "sigma", highlight = 2, np = np)
+
+  expect_s3_class(p$layers[[1]]$geom, "GeomLine")
+  expect_s3_class(p$layers[[2]]$geom, "GeomRug")
+  expect_named(p$layers[[2]]$data, "Divergent")
+})
+
 # functions that require multiple chains ----------------------------------
 test_that("mcmc_trace_highlight returns a ggplot object", {
   expect_gg(mcmc_trace_highlight(arr, regex_pars = c("beta", "x\\:")))