From 8a4f9edc810a4a19623646338185163b6b4ae25d Mon Sep 17 00:00:00 2001 From: Maximilian Scholz Date: Mon, 16 Dec 2024 14:35:49 +0100 Subject: [PATCH] Add split-chain option to rank overlay plots Related to #333 --- R/mcmc-traces.R | 34 ++++++++-- .../mcmc-rank-overlay-split-chains.svg | 67 +++++++++++++++++++ tests/testthat/data-for-mcmc-tests.R | 7 ++ tests/testthat/test-mcmc-traces.R | 7 ++ 4 files changed, 111 insertions(+), 4 deletions(-) create mode 100644 tests/testthat/_snaps/mcmc-traces/mcmc-rank-overlay-split-chains.svg diff --git a/R/mcmc-traces.R b/R/mcmc-traces.R index 571a6ce2..e5743ba3 100644 --- a/R/mcmc-traces.R +++ b/R/mcmc-traces.R @@ -277,6 +277,9 @@ trace_style_np <- function(div_color = "red", div_size = 0.25, div_alpha = 1) { #' of rank-normalized MCMC samples. Defaults to `20`. #' @param ref_line For the rank plots, whether to draw a horizontal line at the #' average number of ranks per bin. Defaults to `FALSE`. +#' @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_overlay <- function(x, pars = character(), @@ -285,7 +288,8 @@ mcmc_rank_overlay <- function(x, facet_args = list(), ..., n_bins = 20, - ref_line = FALSE) { + ref_line = FALSE, + split_chains = FALSE) { check_ignored_arguments(...) data <- mcmc_trace_data( x, @@ -294,7 +298,26 @@ mcmc_rank_overlay <- function(x, transformations = transformations ) - n_chains <- unique(data$n_chains) + # Split chains if requested + if (split_chains) { + # Calculate midpoint for each chain + n_samples <- length(unique(data$iteration)) + midpoint <- n_samples/2 + + # Create new data frame with split chains + data <- data %>% + group_by(.data$chain) %>% + mutate( + chain = ifelse( + iteration <= midpoint, + paste0(.data$chain, "_1"), # First half + paste0(.data$chain, "_2") # Second half + ) + ) %>% + ungroup() + } + + n_chains <- length(unique(data$chain)) n_param <- unique(data$n_parameters) # We have to bin and count the data ourselves because @@ -319,6 +342,7 @@ mcmc_rank_overlay <- function(x, bin_start = unique(histobins$bin_start), stringsAsFactors = FALSE )) + d_bin_counts <- all_combos %>% left_join(d_bin_counts, by = c("parameter", "chain", "bin_start")) %>% mutate(n = dplyr::if_else(is.na(n), 0L, n)) @@ -331,7 +355,9 @@ mcmc_rank_overlay <- function(x, mutate(bin_start = right_edge) %>% dplyr::bind_rows(d_bin_counts) - scale_color <- scale_color_manual("Chain", values = chain_colors(n_chains)) + # 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_chains)) layer_ref_line <- if (ref_line) { geom_hline( @@ -352,7 +378,7 @@ mcmc_rank_overlay <- function(x, } ggplot(d_bin_counts) + - aes(x = .data$bin_start, y = .data$n, color = .data$chain) + + aes(x = .data$bin_start, y = .data$n, color = .data$chain) + geom_step() + layer_ref_line + facet_call + diff --git a/tests/testthat/_snaps/mcmc-traces/mcmc-rank-overlay-split-chains.svg b/tests/testthat/_snaps/mcmc-traces/mcmc-rank-overlay-split-chains.svg new file mode 100644 index 00000000..0079d95e --- /dev/null +++ b/tests/testthat/_snaps/mcmc-traces/mcmc-rank-overlay-split-chains.svg @@ -0,0 +1,67 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + +0 +20 +40 +60 + + + + + + + + + + +0 +500 +1000 +1500 +2000 +Rank +Split Chains + + + + +1_1 +1_2 +2_1 +2_2 +mcmc_rank_overlay (split chains) + + diff --git a/tests/testthat/data-for-mcmc-tests.R b/tests/testthat/data-for-mcmc-tests.R index 1136cefe..17f5b8d0 100644 --- a/tests/testthat/data-for-mcmc-tests.R +++ b/tests/testthat/data-for-mcmc-tests.R @@ -80,4 +80,11 @@ vdiff_dframe_rank_overlay_bins_test <- posterior::as_draws_df( ) ) +vdiff_dframe_rank_overlay_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)) + ) +) + set.seed(seed = NULL) diff --git a/tests/testthat/test-mcmc-traces.R b/tests/testthat/test-mcmc-traces.R index 62d46c88..7c6da533 100644 --- a/tests/testthat/test-mcmc-traces.R +++ b/tests/testthat/test-mcmc-traces.R @@ -157,6 +157,10 @@ test_that("mcmc_rank_overlay renders correctly", { # https://github.com/stan-dev/bayesplot/issues/331 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, + split_chains = TRUE) + vdiffr::expect_doppelganger("mcmc_rank_overlay (default)", p_base) vdiffr::expect_doppelganger( "mcmc_rank_overlay (reference line)", @@ -170,6 +174,9 @@ test_that("mcmc_rank_overlay renders correctly", { # https://github.com/stan-dev/bayesplot/issues/331 vdiffr::expect_doppelganger("mcmc_rank_overlay (not all bins)", p_not_all_bins_exist) + + # https://github.com/stan-dev/bayesplot/issues/333 + vdiffr::expect_doppelganger("mcmc_rank_overlay (split chains)", p_split_chains) }) test_that("mcmc_rank_hist renders correctly", {