diff --git a/R/mcmc-traces.R b/R/mcmc-traces.R index ac5bb467..01f79e1f 100644 --- a/R/mcmc-traces.R +++ b/R/mcmc-traces.R @@ -485,6 +485,9 @@ mcmc_rank_hist <- function(x, #' @param plot_diff For `mcmc_rank_ecdf()`, a boolean specifying if the #' difference between the observed rank ECDFs and the theoretical expectation #' should be drawn instead of the unmodified rank ECDF plots. +#' @param split_chains Logical indicating whether to split each chain into two parts. +#' If TRUE, each chain is split into first and second half with "_1" and "_2" suffixes. +#' Defaults to `FALSE`. #' @export mcmc_rank_ecdf <- function(x, @@ -496,7 +499,8 @@ mcmc_rank_ecdf <- facet_args = list(), prob = 0.99, plot_diff = FALSE, - interpolate_adj = NULL) { + interpolate_adj = NULL, + split_chains = FALSE) { check_ignored_arguments(..., ok_args = c("K", "pit", "prob", "plot_diff", "interpolate_adj", "M") ) @@ -507,8 +511,28 @@ mcmc_rank_ecdf <- transformations = transformations, highlight = 1 ) + + # Split chains if requested + if (split_chains) { + data$n_chains = data$n_chains/2 + data$n_iterations = data$n_iterations/2 + n_samples <- length(unique(data$iteration)) + midpoint <- n_samples/2 + + data <- data %>% + group_by(.data$chain) %>% + mutate( + chain = ifelse( + iteration <= midpoint, + paste0(.data$chain, "_1"), + paste0(.data$chain, "_2") + ) + ) %>% + ungroup() + } + n_iter <- unique(data$n_iterations) - n_chain <- unique(data$n_chains) + n_chain <- length(unique(data$chain)) n_param <- unique(data$n_parameters) x <- if (is.null(K)) { @@ -561,7 +585,9 @@ mcmc_rank_ecdf <- group = .data$chain ) - scale_color <- scale_color_manual("Chain", values = chain_colors(n_chain)) + # Update legend title based on split_chains + legend_title <- if (split_chains) "Split Chains" else "Chain" + scale_color <- scale_color_manual(legend_title, values = chain_colors(n_chain)) facet_call <- NULL if (n_param == 1) { diff --git a/tests/testthat/_snaps/mcmc-traces/mcmc-rank-ecdf-split-chain.svg b/tests/testthat/_snaps/mcmc-traces/mcmc-rank-ecdf-split-chain.svg new file mode 100644 index 00000000..d16c9ca3 --- /dev/null +++ b/tests/testthat/_snaps/mcmc-traces/mcmc-rank-ecdf-split-chain.svg @@ -0,0 +1,70 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +-0.4 +-0.2 +0.0 +0.2 + + + + + + + + + + + +0.0 +0.2 +0.4 +0.6 +0.8 +1.0 +theta +Split Chains + + + + +1_1 +1_2 +2_1 +2_2 +mcmc_rank_ecdf (split chain) + + diff --git a/tests/testthat/data-for-mcmc-tests.R b/tests/testthat/data-for-mcmc-tests.R index 17f5b8d0..fe892579 100644 --- a/tests/testthat/data-for-mcmc-tests.R +++ b/tests/testthat/data-for-mcmc-tests.R @@ -80,7 +80,7 @@ vdiff_dframe_rank_overlay_bins_test <- posterior::as_draws_df( ) ) -vdiff_dframe_rank_overlay_split_chain_test <- posterior::as_draws_df( +vdiff_dframe_rank_split_chain_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.003 * 1:1000 + stats::arima.sim(list(ar = 0.7), n = 1000, sd = 0.5)) diff --git a/tests/testthat/test-mcmc-traces.R b/tests/testthat/test-mcmc-traces.R index 7c6da533..f79b4c2e 100644 --- a/tests/testthat/test-mcmc-traces.R +++ b/tests/testthat/test-mcmc-traces.R @@ -158,7 +158,7 @@ test_that("mcmc_rank_overlay renders correctly", { p_not_all_bins_exist <- mcmc_rank_overlay(vdiff_dframe_rank_overlay_bins_test) # https://github.com/stan-dev/bayesplot/issues/333 - p_split_chains <- mcmc_rank_overlay(vdiff_dframe_rank_overlay_split_chain_test, + p_split_chains <- mcmc_rank_overlay(vdiff_dframe_rank_split_chain_test, split_chains = TRUE) vdiffr::expect_doppelganger("mcmc_rank_overlay (default)", p_base) @@ -261,6 +261,11 @@ test_that("mcmc_rank_ecdf renders correctly", { plot_diff = TRUE ) + # https://github.com/stan-dev/bayesplot/issues/333 + p_split_chains <- mcmc_rank_ecdf(vdiff_dframe_rank_split_chain_test, + plot_diff = TRUE, + split_chains = TRUE) + vdiffr::expect_doppelganger("mcmc_rank_ecdf (default)", p_base) vdiffr::expect_doppelganger("mcmc_rank_ecdf (one parameter)", p_one_param) vdiffr::expect_doppelganger("mcmc_rank_ecdf (diff)", p_diff) @@ -268,6 +273,9 @@ test_that("mcmc_rank_ecdf renders correctly", { "mcmc_rank_ecdf (one param, diff)", p_diff_one_param ) + + # https://github.com/stan-dev/bayesplot/issues/333 + vdiffr::expect_doppelganger("mcmc_rank_ecdf (split chain)", p_split_chains) }) test_that("mcmc_trace with 'np' renders correctly", {