Add split-chain option to rank overlay plots

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,28 @@ mcmc_rank_overlay <- function(x,
     transformations = transformations
   )
 
-  n_chains <- unique(data$n_chains)
+  # Split chains if requested
+  if (split_chains) {
+    data$n_chains = data$n_chains/2
+    data$n_iterations = data$n_iterations/2
+    # 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(
+          .data$iteration <= midpoint,
+          paste0(.data$chain, "_1"),
+          paste0(.data$chain, "_2")
+        )
+      ) %>%
+      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 +344,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 +357,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 +380,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 +
@@ -457,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,
@@ -468,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")
   )
@@ -479,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(
+          .data$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)) {
@@ -533,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) {