Nothing
.add_aesthetics <- function(dots, to_add) {
# Add some default aesthetics if not included in ...
# @param dots list(...)
# @param to_add = character vector
add_names <- gsub("pt_", "", to_add)
for (j in seq_along(to_add)) {
if (!add_names[j] %in% names(dots))
dots[[add_names[j]]] <- rstanvis_aes_ops(to_add[j])
}
dots
}
`%ifNULL%` <- function(x, replacement) {
if (!is.null(x)) x
else replacement
}
is.stanreg <- function(x) inherits(x, "stanreg")
is.stanfit <- function(x) inherits(x, "stanfit")
.get_stanreg_parnames <- function(x) {
stopifnot(is.stanreg(x))
pars <- rownames(x$stan_summary)
sel <- grepl("mean_ppd|lp__|log-posterior", pars, ignore.case = TRUE)
c(pars[!sel], pars[sel])
}
.check_object <- function(object, unconstrain = FALSE) {
if (is.stanreg(object)) {
if (object$algorithm == "optimizing")
stop("Plots not yet available for optimization (algorithm='optimizing')",
call. = FALSE)
if (unconstrain)
stop("Option 'unconstrain' not yet available for stanreg objects.",
call. = FALSE)
.mode_check(object$stanfit)
}
else .mode_check(object)
}
.mode_check <- function(object) {
stopifnot(is.stanfit(object))
if (object@mode == 1L) {
stop("Stan model '", object@model_name, "' is of mode 'test_grad';\n",
"sampling is not conducted.\n",
call. = FALSE)
} else if (object@mode == 2L) {
stop("Stan model '", object@model_name, "' does not contain samples.\n",
call. = FALSE)
}
invisible(TRUE)
}
.vb_check <- function(x) {
if (is.stanreg(x)) x <- x$stanfit
if (isTRUE(x@stan_args[[1]]$method == "variational"))
stop("Plot only available for models estimated using MCMC", call. = FALSE)
}
.reshape_sample <- function(x) {
res <- lapply(seq_along(x), function(i) {
data.frame(value = unlist(x[[i]], use.names = FALSE),
parameter = rep(names(x[[i]]), each = length(x[[i]][[1L]])),
chain = i)
})
res <- do.call(rbind, res)
res$parameter <- as.character(res$parameter)
res
}
.make_plot_data <- function(object, pars, include = TRUE,
inc_warmup = FALSE, unconstrain = FALSE) {
window <- NULL
if (is.stanreg(object)) {
sim <- object$stanfit@sim
}
else sim <- object@sim
nopars <- missing(pars)
if (is.stanfit(object) && !nopars) {
if ("log-posterior" %in% pars)
pars[which(pars == "log-posterior")] <- "lp__"
}
if (!include) {
if (nopars) stop("pars must be specified if include=FALSE.", call. = FALSE)
else {
if (is.stanreg(object))
pars <- setdiff(sim$fnames_oi, pars)
else
pars <- setdiff(sim$pars_oi, pars)
}
}
if (nopars) {
if (is.stanreg(object))
pars <- names(object$coefficients)
else
pars <- setdiff(sim$pars_oi, c("lp__", "log-posterior"))
}
else {
if (!is.stanreg(object))
pars <- check_pars_second(sim, pars)
}
pars <- remove_empty_pars(pars, sim$dims_oi)
if (unconstrain && "lp__" %in% pars) {
if (length(pars) == 1L) stop("Can't unconstrain lp__", call. = FALSE)
pars <- pars[-which(pars == "lp__")]
}
tidx <- pars_total_indexes(sim$pars_oi,
sim$dims_oi,
sim$fnames_oi,
pars)
tidx <- lapply(tidx, function(x) attr(x, "row_major_idx"))
tidx <- unlist(tidx, use.names = FALSE)
num_plots <- length(tidx)
if (nopars && num_plots > 10) {
# if pars is not specified then default to showing 10 of the parameters
tidx <- tidx[1:10]
message("'pars' not specified. Showing first 10 parameters by default.")
}
if (!is.null(window)) {
window <- sort(window)
if (window[1] < 1) window[1] <- 1
if (window[1] > sim$iter[1])
stop("wrong specification of argument window", call. = FALSE)
if (is.na(window[2]) || window[2] > sim$iter[1])
window[2] <- sim$iter[1]
} else {
window <- c(1, sim$iter[1])
}
if ((all(sim$warmup2 == 0) || !inc_warmup) && window[1] <= sim$warmup[1]) {
window[1] <- sim$warmup[1] + 1
}
if (window[1] > window[2]) {
stop("the given window does not include sample")
}
if (window[1] > sim$warmup[1]) inc_warmup <- FALSE
thin <- sim$thin
warmup2 <- sim$warmup2[1]
warmup <- sim$warmup
start_i <- window[1]
window_size <- (window[2] - start_i) %/% thin
id <- seq.int(start_i, by = thin, length.out = window_size)
start_idx <- (if (warmup2 == 0) (start_i - warmup) else start_i) %/% thin
if (start_idx < 1) start_idx <- 1
idx <- seq.int(start_idx, by = 1, length.out = window_size)
if (unconstrain) {
sim$samples <- .upars(object)
sel <- grep(paste(pars, collapse ="|"), names(sim$samples[[1]]), value = TRUE)
samp_use <- lapply(sim$samples, function(chain) {
out <- lapply(chain[names(chain) %in% sel], function(x) x[idx])
names(out) <- sel
out
})
} else {
samp_use <- lapply(sim$samples, function(chain) {
out <- lapply(chain[tidx], function(x) x[idx])
names(out) <- sim$fnames_oi[tidx]
out
})
}
nchains <- length(samp_use)
if (unconstrain) {
if (is.stanreg(object))
object <- object$stanfit
pblock <- .get_stan_params(object)
pars2 <- unlist(lapply(strsplit(pars, "\\["), "[[", 1))
not_pblock <- length(setdiff(unique(pars2), pblock))
if (not_pblock)
stop("If 'unconstrain' is TRUE only variables declared in the ",
"'parameters' block can be included in 'pars'.",
call. = FALSE)
}
dat <- .reshape_sample(samp_use)
dat$iteration <- idx * thin
dat$chain <- factor(dat$chain)
fnames <- if (unconstrain)
names(samp_use[[1]]) else sim$fnames_oi[tidx]
lp <- which(dat$parameter == "lp__")
if (!identical(lp, integer(0))) {
dat$parameter[lp] <- "log-posterior"
fnames[which(fnames == "lp__")] <- "log-posterior"
}
dat$parameter <- factor(dat$parameter, levels = fnames)
list(samp = dat,
nchains = nchains,
nparams = length(fnames),
warmup = warmup)
}
# get parameter names for parameters block only
.get_stan_params <- function(object) {
stopifnot(is.stanfit(object))
params <- grep("context__.vals_r", fixed = TRUE, value = TRUE,
x = strsplit(get_cppcode(get_stanmodel(object)), "\n")[[1]])
params <- sapply(strsplit(params, "\""), FUN = function(x) x[[2]])
intersect(params, object@model_pars)
}
# unconstrain
.upars <- function(object, permuted = FALSE, inc_warmup = TRUE) {
if (is.stanreg(object))
object <- object$stanfit
pars <- extract(object, permuted = permuted, inc_warmup = inc_warmup)
nchains <- ncol(pars)
pn <- dimnames(pars)$parameters
param_names <- pn[pn != "lp__"]
sel <- object@model_pars != "lp__"
skeleton <- create_skeleton(object@model_pars, object@par_dims)[sel]
upars <- apply(pars, 1:2, FUN = function(theta) {
unconstrain_pars(object, rstan_relist(theta, skeleton))
})
if (length(dim(upars)) == 2) {
upars <- array(upars, dim = c(dim(upars)[1], dim(upars)[2], 1))
upars <- aperm(upars, c(2, 3, 1))
}
else upars <- aperm(upars, c(3, 1, 2))
pblock <- .get_stan_params(object)
mark <- c()
for (p in pblock) {
patt <- paste0("^", p, "|^", p, "\\[")
sel <- grep(patt, param_names)
if (length(sel))
mark <- c(mark, sel)
}
param_names <- param_names[sort(mark)]
cpp_code <- strsplit(get_cppcode(get_stanmodel(object)), "\n")[[1]]
param_names <- .remove_udiag_pars(cpp_code, pblock, param_names)
lapply(seq_len(nchains), function(chain) {
x <- upars[chain,, ]
if (NCOL(x) == 1) x <- rbind(x)
plist <- lapply(seq_len(nrow(x)), function(param) x[param, ])
names(plist) <- param_names
plist
})
}
.remove_udiag_pars <- function(cpp_code, pblock, param_names) {
patts <- c("cholesky_corr", "cholesky_cov", "corr_matrix", "cov_matrix")
for (patt in patts) {
to_drop <- c()
for (p in pblock) {
par_mentions = grep(paste0("(",p,")"), x = cpp_code,
fixed = TRUE, value = TRUE)
if (length(grep(patt, par_mentions))) {
for (v in grep(p, param_names)) {
ij <- strsplit(x = param_names[v], split = c("\\[|,|\\]"))[[1]][2:3]
if (diff(as.numeric(ij)) > -1)
to_drop = c(v,to_drop)
}
}
}
if (length(to_drop))
param_names <- param_names[-to_drop]
}
return(param_names)
}
color_vector <- function(n) {
hues = seq(15, 375, length=n+1)
hcl(h=hues, l=50, c=50)[1:n]
}
color_vector_chain <- function(n) {
hues = seq(15, 375, length=n+1)
hcl(h=hues, l=80, c=50)[1:n]
}
# autocorrelation ---------------------------------------------------------
.ac_fun <- function(x, lag.max, partial = FALSE) {
if (!partial)
acf(x, lag.max = lag.max, plot = FALSE)$acf[,, 1L]
else
pacf(x, lag.max = lag.max, plot = FALSE)$acf[,, 1L]
}
.ac_plot_data <- function(dat, lags, partial = FALSE) {
ch <- dat[, grep("chain", colnames(dat))]
nc <- length(unique(ch))
ac_list <- tapply(dat$value, INDEX = ch, FUN = .ac_fun, lag.max = lags,
partial = partial, simplify = FALSE)
nl <- if (partial) lags else lags + 1
ch <- factor(rep(1:nc, each = nl), labels = paste0("chain:", 1:nc))
ll <- rep(seq(if (partial) 1 else 0, lags), nc)
data.frame(chains = ch, ac = do.call(c, ac_list), lag = ll)
}
.ac_plot_data_multi <- function(dat, lags, partial = FALSE) {
ch <- dat[, grep("chain", colnames(dat))]
nc <- length(unique(ch))
pa <- factor(dat[, grep("parameter", colnames(dat))])
np <- length(unique(pa))
ac_list <- tapply(dat$value, INDEX = list(ch, pa),
FUN = .ac_fun, lag.max = lags,
partial = partial, simplify = FALSE)
nl <- if (partial) lags else lags + 1
ch <- factor(rep(rep(1:nc, each = nl), np), labels = paste0("chain:", 1:nc))
ll <- rep(seq(if (partial) 1 else 0, lags), nc * np)
pp <- factor(rep(1:np, each = nc * nl), labels = levels(pa))
data.frame(parameters = pp, chains = ch, ac = do.call(c, ac_list), lag = ll)
}
# rhat, ess, mcse ---------------------------------------------------------
.rhat_neff_mcse_hist <- function(which = c("rhat", "n_eff_ratio", "mcse_ratio"),
object, pars=NULL, ...) {
if (is.stanreg(object)) object <- object$stanfit
thm <- rstanvis_hist_theme()
dots <- .add_aesthetics(list(...), c("fill", "color"))
if (which == "n_eff_ratio") {
lp <- suppressWarnings(get_logposterior(object, inc_warmup = FALSE))
SS <- prod(length(lp), length(lp[[1L]]))
}
if (!is.null(pars)) smry <- summary(object, pars = pars)$summary
else smry <- summary(object)$summary
xlab <- switch(which,
rhat = "Rhat statistic",
n_eff_ratio = "Effective sample size / Sample size",
mcse_ratio = "Monte Carlo SE / Posterior SD"
)
stat <- switch(which,
rhat = smry[, "Rhat"],
n_eff_ratio = smry[, "n_eff"] / SS,
mcse_ratio = smry[, "se_mean"] / smry[, "sd"]
)
df <- data.frame(stat)
plot_labs <- ggplot2::labs(y = "", x = xlab)
base <- ggplot2::ggplot(df, ggplot2::aes(x = stat))
base +
do.call(ggplot2::geom_histogram, dots) +
plot_labs +
thm
}
# nuts stuff --------------------------------------------------------------
.LP_NAME <- "log-posterior"
.LP_LAB <- "Log Posterior"
.METROP_LAB <- "Mean Metrop. Acceptance"
.STEPSIZE_LAB <- "Sampled Step Size"
.TREEDEPTH_LAB <- "Treedepth"
.NDIVERGENT_LAB <- "Divergent"
.NUTS_VLINE_CLR <- "#222222"
.NUTS_FILL <- "#66a7e0"
.NUTS_CLR <- "#006dcc"
.NUTS_PT_CLR <- "#328ad6"
.NDIVERGENT_FILL <- "#ae0001"
.MAXTD_FILL <- "#eeba30"
.NDIVERGENT_CLR <- "black"
.MAXTD_CLR <- "black"
.DIV_AND_MAXTD_SHAPE <- 21
.nuts_return <- function(graphs, ...) {
gtable <- do.call(gridExtra::arrangeGrob, c(graphs, ...))
plot(gtable)
invisible(graphs)
}
.nuts_args_check <- function(...) {
if ("par" %in% names(dots <- list(...)))
stop("'par' argument should not be specified.", call. = FALSE)
}
.max_td <- function(x) {
if (is.stanreg(x))
x <- x$stanfit
cntrl <- x@stan_args[[1L]]$control
if (is.null(cntrl)) 11
else {
max_td <- cntrl$max_treedepth
if (is.null(max_td)) 11
else max_td
}
}
.sampler_params_post_warmup <- function(object, which = "stepsize__", as.df = FALSE) {
if (which == "divergent__" && utils::packageVersion("rstan") < "2.10")
which <- "n_divergent__"
if (is.stanreg(object))
object <- object$stanfit
sampler_params <- suppressWarnings(get_sampler_params(object))
warmup_val <- floor(object@sim$warmup / object@sim$thin)
warmup_saved <- object@stan_args[[1L]][["save_warmup"]]
if (!is.null(warmup_saved)) {
if (!warmup_saved) warmup_val <- 0
}
sp <-lapply(1:length(sampler_params), function(i) {
out <- sampler_params[[i]]
out <- if (warmup_val == 0) out[, which] else out[-(1:warmup_val), which]
names(out) <- (warmup_val + 1):(warmup_val + NROW(out))
out
})
if (length(which) == 1L && as.df) .sampler_param_df(sp, warmup_val)
else sp
}
.sampler_param_df <- function(sp, warmup_val) {
sp_mat <- do.call("cbind", sp)
colnames(sp_mat) <- paste0("chain:", 1:ncol(sp_mat))
sp_mat <- cbind(iterations = (warmup_val+1):(warmup_val + nrow(sp_mat)), sp_mat)
as.data.frame(sp_mat)
}
.reshape_df <- function(df) {
iter_col <- grep("iteration", colnames(df))
mdf <- reshape(df, direction = "long",
v.names = "value",
varying = list(colnames(df)[-iter_col]),
timevar = "variable",
times = colnames(df)[-iter_col],
idvar = colnames(df)[iter_col])
rownames(mdf) <- 1:nrow(mdf)
mdf$variable <- as.factor(mdf$variable)
attr(mdf, "reshapeLong") <- NULL
mdf
}
.sampler_param_vs_param <- function(p, sp, divergent = NULL, hit_max_td = NULL,
p_lab, sp_lab, chain = 0, violin = FALSE,
smoother = FALSE, ...) {
dots <- .add_aesthetics(list(...), c("alpha", "shape"))
dots$alpha <- 0.5 * dots$alpha
dots$color <- .NUTS_PT_CLR
dots$fill <- .NUTS_FILL
xy_labs <- ggplot2::labs(
y = if (missing(p_lab)) NULL else p_lab,
x = if (missing(sp_lab)) NULL else sp_lab
)
df <- data.frame(sp = do.call("c", sp), p = c(p))
if (violin) df$sp <- as.factor(round(df$sp, 4))
if (!is.null(divergent)) df$divergent <- do.call("c", divergent)
if (!is.null(hit_max_td)) df$hit_max_td <- do.call("c", hit_max_td)
base <- ggplot2::ggplot(df, ggplot2::aes(x = .data$sp, y = .data$p)) + xy_labs
if (chain == 0) {
if (violin)
graph <- base + ggplot2::geom_violin(color = .NUTS_CLR, fill = .NUTS_FILL)
else {
graph <- base + do.call(ggplot2::geom_point, dots)
if (smoother) graph <- graph + ggplot2::stat_smooth(se = FALSE)
if (!is.null(divergent))
graph <-
graph + ggplot2::geom_point(
data = subset(df, divergent == 1),
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = .NDIVERGENT_CLR,
fill = .NDIVERGENT_FILL,
alpha = 0.8,
size = 3,
shape = .DIV_AND_MAXTD_SHAPE
)
if (!is.null(hit_max_td))
graph <-
graph + ggplot2::geom_point(
data = subset(df, hit_max_td == 1),
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = .MAXTD_CLR,
fill = .MAXTD_FILL,
alpha = 0.8,
size = 3,
shape = .DIV_AND_MAXTD_SHAPE
)
}
return(graph)
}
chain_data <- data.frame(sp = sp[, chain], p = p[, chain])
if (!is.null(divergent)) chain_data$div <- divergent[, chain]
if (!is.null(hit_max_td)) chain_data$hit <- hit_max_td[, chain]
chain_clr <- color_vector_chain(ncol(sp))[chain]
chain_fill <- chain_clr
if (violin) {
chain_data$sp <- as.factor(round(chain_data$sp, 4))
graph <- base +
ggplot2::geom_violin(color = .NUTS_CLR, fill = .NUTS_FILL, ...) +
ggplot2::geom_violin(
data = chain_data,
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = chain_clr,
fill = chain_fill,
...
)
return(graph)
}
graph <- base + do.call(ggplot2::geom_point, dots)
if (smoother) graph <- graph + ggplot2::stat_smooth(se = FALSE)
graph <- graph +
ggplot2::geom_point(
data = chain_data,
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = chain_fill,
...
)
if (smoother) graph <- graph +
ggplot2::stat_smooth(
data = chain_data,
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = chain_fill,
se = FALSE
)
if (!is.null(divergent))
graph <- graph +
ggplot2::geom_point(
data = chain_data[chain_data$div == 1, , drop = FALSE],
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = .NDIVERGENT_CLR,
fill = .NDIVERGENT_FILL,
size = 3,
shape = .DIV_AND_MAXTD_SHAPE
)
if (!is.null(hit_max_td))
graph <-
graph + ggplot2::geom_point(
data = chain_data[chain_data$hit == 1, , drop = FALSE],
mapping = ggplot2::aes(x = .data$sp, y = .data$p),
color = .MAXTD_CLR,
fill = .MAXTD_FILL,
size = 3,
shape = .DIV_AND_MAXTD_SHAPE
)
graph
}
.sampler_param_vs_sampler_param_violin <- function(df_x, df_y, lab_x, lab_y,
chain = 0) {
xy_labs <- ggplot2::labs(y = lab_y, x = lab_x)
df <- data.frame(x = do.call("c", df_x), y = do.call("c", df_y))
df$x <- as.factor(df$x)
base <- ggplot2::ggplot(df, ggplot2::aes(x = .data$x, y = .data$y)) + xy_labs
graph <- base + ggplot2::geom_violin(color = .NUTS_CLR, fill = .NUTS_FILL)
if (chain == 0) return(graph)
chain_clr <- color_vector_chain(ncol(df_x))[chain]
chain_fill <- chain_clr
chain_data <- data.frame(x = as.factor(df_x[, chain]), y = df_y[, chain])
graph + ggplot2::geom_violin(
data = chain_data,
mapping = ggplot2::aes(x = .data$x, y = .data$y),
color = chain_clr,
fill = chain_fill,
alpha = 0.5
)
}
.p_hist <- function(df, lab, chain = 0, ...) {
mdf <- .reshape_df(df) # reshape2::melt(df, id.vars = grep("iteration", colnames(df), value = TRUE))
dots <- .add_aesthetics(list(...), "linewidth")
dots$binwidth <- diff(range(mdf$value))/30
dots$fill <- .NUTS_FILL
dots$color <- .NUTS_CLR
base <- ggplot2::ggplot(mdf, ggplot2::aes(x = .data$value)) +
do.call(ggplot2::geom_histogram, dots) +
ggplot2::labs(x = if (missing(lab)) NULL else lab, y = "")
if (chain == 0) {
graph <- base +
ggplot2::geom_vline(xintercept = mean(mdf$value), color = .NUTS_VLINE_CLR, linewidth = .8) +
ggplot2::geom_vline(xintercept = median(mdf$value), color = .NUTS_VLINE_CLR, lty = 2, linewidth = 1)
return(graph)
}
chain_data <- mdf[mdf$variable == paste0("chain:",chain), ]
chain_clr <- color_vector_chain(ncol(df) - 1)[chain]
chain_fill <- chain_clr
base + ggplot2::geom_histogram(data = chain_data,
binwidth = diff(range(chain_data$value))/30,
fill = chain_fill, alpha = 0.5) +
ggplot2::geom_vline(xintercept = mean(chain_data$value), color = chain_clr, linewidth = .8) +
ggplot2::geom_vline(xintercept = median(chain_data$value),
color = chain_clr, lty = 2, linewidth = 1)
}
.treedepth_ndivergent_hist <- function(df_td, df_nd, chain = 0,
divergent = c("All", 0, 1), ...) {
x_lab <- if (divergent == "All")
"Treedepth" else paste0("Treedepth (Divergent = ", divergent,")")
plot_labs <- ggplot2::labs(x = x_lab, y = "")
mdf_td <- .reshape_df(df_td) #reshape2::melt(df_td, id.vars = grep("iteration", colnames(df_td), value = TRUE))
mdf_nd <- .reshape_df(df_nd) #reshape2::melt(df_nd, id.vars = grep("iteration", colnames(df_nd), value = TRUE))
mdf <- cbind(mdf_td, div = mdf_nd$value)
plot_data <- if (divergent == "All") mdf else mdf[mdf$div == divergent, , drop = FALSE]
if (nrow(plot_data) == 0) return(NULL)
count <- NaN
graph <- ggplot2::ggplot(
data = plot_data,
mapping = ggplot2::aes(x = factor(.data$value)),
na.rm = TRUE
) +
ggplot2::geom_bar(
mapping = ggplot2::aes(
y = ggplot2::after_stat(count) / sum(ggplot2::after_stat(count))
),
width = 1, fill = .NUTS_FILL, color = .NUTS_CLR
) +
plot_labs
if (chain == 0) return(graph)
chain_clr <- color_vector_chain(ncol(df_td) - 1)[chain]
chain_fill <- chain_clr
chain_data <- plot_data[plot_data$variable == paste0("chain:", chain), , drop = FALSE]
graph + ggplot2::geom_bar(
data = chain_data,
mapping = ggplot2::aes(
y = ggplot2::after_stat(count) / sum(ggplot2::after_stat(count))
),
fill = chain_fill, width = 1
)
}
Any scripts or data that you put into this service are public.
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.