Skip to content

Commit 8d5df1f

Browse files
committed
Add split-chain option to rank ecdf plots
Related to #333
1 parent 3ac550c commit 8d5df1f

File tree

4 files changed

+109
-5
lines changed

4 files changed

+109
-5
lines changed

R/mcmc-traces.R

+29-3
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,9 @@ mcmc_rank_hist <- function(x,
485485
#' @param plot_diff For `mcmc_rank_ecdf()`, a boolean specifying if the
486486
#' difference between the observed rank ECDFs and the theoretical expectation
487487
#' should be drawn instead of the unmodified rank ECDF plots.
488+
#' @param split_chains Logical indicating whether to split each chain into two parts.
489+
#' If TRUE, each chain is split into first and second half with "_1" and "_2" suffixes.
490+
#' Defaults to `FALSE`.
488491
#' @export
489492
mcmc_rank_ecdf <-
490493
function(x,
@@ -496,7 +499,8 @@ mcmc_rank_ecdf <-
496499
facet_args = list(),
497500
prob = 0.99,
498501
plot_diff = FALSE,
499-
interpolate_adj = NULL) {
502+
interpolate_adj = NULL,
503+
split_chains = FALSE) {
500504
check_ignored_arguments(...,
501505
ok_args = c("K", "pit", "prob", "plot_diff", "interpolate_adj", "M")
502506
)
@@ -507,8 +511,28 @@ mcmc_rank_ecdf <-
507511
transformations = transformations,
508512
highlight = 1
509513
)
514+
515+
# Split chains if requested
516+
if (split_chains) {
517+
data$n_chains = data$n_chains/2
518+
data$n_iterations = data$n_iterations/2
519+
n_samples <- length(unique(data$iteration))
520+
midpoint <- n_samples/2
521+
522+
data <- data %>%
523+
group_by(.data$chain) %>%
524+
mutate(
525+
chain = ifelse(
526+
iteration <= midpoint,
527+
paste0(.data$chain, "_1"),
528+
paste0(.data$chain, "_2")
529+
)
530+
) %>%
531+
ungroup()
532+
}
533+
510534
n_iter <- unique(data$n_iterations)
511-
n_chain <- unique(data$n_chains)
535+
n_chain <- length(unique(data$chain))
512536
n_param <- unique(data$n_parameters)
513537

514538
x <- if (is.null(K)) {
@@ -561,7 +585,9 @@ mcmc_rank_ecdf <-
561585
group = .data$chain
562586
)
563587

564-
scale_color <- scale_color_manual("Chain", values = chain_colors(n_chain))
588+
# Update legend title based on split_chains
589+
legend_title <- if (split_chains) "Split Chains" else "Chain"
590+
scale_color <- scale_color_manual(legend_title, values = chain_colors(n_chain))
565591

566592
facet_call <- NULL
567593
if (n_param == 1) {

0 commit comments

Comments
 (0)