diff --git a/NEWS.md b/NEWS.md index fb765a60..4138aba6 100644 --- a/NEWS.md +++ b/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. diff --git a/R/mcmc-traces.R b/R/mcmc-traces.R index ddf9c549..699c2fe4 100644 --- a/R/mcmc-traces.R +++ b/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))) + ) } } diff --git a/man/MCMC-traces.Rd b/man/MCMC-traces.Rd index ff30246c..2d884a05 100644 --- a/man/MCMC-traces.Rd +++ b/man/MCMC-traces.Rd @@ -22,6 +22,8 @@ mcmc_trace( iter1 = 0, window = NULL, size = NULL, + alpha = 0.4, + highlight = NULL, np = NULL, np_style = trace_style_np(), divergences = NULL @@ -160,6 +162,13 @@ range of iterations to display.} \item{size}{An optional value to override the default line size for \code{mcmc_trace()} or the default point size for \code{mcmc_trace_highlight()}.} +\item{alpha}{For \code{mcmc_trace()} and \code{mcmc_trace_highlight()}, controls the +transparency of the lines or points for the chains not highlighted.} + +\item{highlight}{For \code{mcmc_trace()}, \code{NULL} (the default) or an integer +specifying one chain to emphasize. For \code{mcmc_trace_highlight()}, an integer +specifying one chain to emphasize.} + \item{np}{For models fit using \link{NUTS} (more generally, any \href{https://en.wikipedia.org/wiki/Symplectic_integrator}{symplectic integrator}), an optional data frame providing NUTS diagnostic information. The data @@ -174,13 +183,6 @@ divergences (if the \code{np} argument is specified).} \item{divergences}{Deprecated. Use the \code{np} argument instead.} -\item{alpha}{For \code{mcmc_trace_highlight()}, passed to -\code{\link[ggplot2:geom_point]{ggplot2::geom_point()}} to control the transparency of the points -for the chains not highlighted.} - -\item{highlight}{For \code{mcmc_trace_highlight()}, an integer specifying one -of the chains that will be more visible than the others in the plot.} - \item{div_color, div_size, div_alpha}{Optional arguments to the \code{trace_style_np()} helper function that are eventually passed to \code{\link[ggplot2:geom_rug]{ggplot2::geom_rug()}} if the \code{np} argument is also specified. They control @@ -237,6 +239,7 @@ section, below, for details. \item{\code{mcmc_trace()}}{ Standard trace plots of MCMC draws. For models fit using \link{NUTS}, the \code{np} argument can be used to also show divergences on the trace plot. +One chain can be emphasized using the \code{highlight} argument. } \item{\code{mcmc_trace_highlight()}}{ Traces are plotted using points rather than lines and the opacity of all @@ -291,6 +294,7 @@ mcmc_trace(x, regex_pars = "beta") # 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")) + diff --git a/tests/testthat/test-mcmc-traces.R b/tests/testthat/test-mcmc-traces.R index 5b8715f4..a4e8df64 100644 --- a/tests/testthat/test-mcmc-traces.R +++ b/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\\:")))