R/summary.R

In paasim/stanutils: Make stan output workable within tidyverse

#' Get the stanfit-summary as a tibble
#'
#' Uses [summary()] for `stanfit` and converts it into a tibble. Also separates
#' the parameter name and its dimensions into separate columns similar to
#' [extract_tidy()].
#'
#' @param f An object of class `stanfit`.
#' @param pars The names of the paremters to be extracted. If empty, all the
#' paremters are extracted.
#' @param probs The quantiles paremters to be extracted. If empty, the 0.1, 0.5
#' and 0.9-quantiles are extracted.
#'
#' @return A tibble with the following columns columns:
#' * `par` The name of the parameter.
#' * `d1:dn` The dimensions of the paremeter. `NA` if the parameter is of lower dimension.
#' * `mean` Posterior mean for the parameter.
#' * `se_mean` Standard error for the mean.
#' * `sd` Posterior standard deviation for the parameter.
#' * `q_xx` User-specified quantiles for the posterior distribution of the parameter.
#' * `n_eff` Effective sample-size.
#' * `rhat` Rhat.
#'
#' @export
summary_tidy <- function(f, pars = character(0L), probs = c(0.1, 0.5, 0.9)) {
  # extract the summary
  sum_pars <- if (is_empty(pars)) summary else partial(summary, pars = pars)
  sum_probs <- if (is_empty(probs)) sum_pars else partial(sum_pars, probs = probs)
  result <- sum_probs(f)$summary

  result_tbl <- as_tibble(result)
  pars_raw <- rownames(result)

  # extract dimensions
  dim_regex <- "(?<=\\[)[\\d,]+(?=\\])"
  dims <- str_extract(pars_raw, dim_regex)
  max_dim <- get_max_dim(dims)

  if (max_dim > 0L) {
    dim_tbl <- str_split_fixed(dims, ",", max_dim) %>%
      set_colnames(str_c("d", 1:max_dim)) %>%
      as_tibble() %>%
      mutate_all(as.integer)
  } else {
    dim_tbl <- tibble()
  }

  # extract parameters
  par_tbl <- tibble(par = str_remove(pars_raw, "\\[[\\d,]+\\]") %>% factor())

  # format quantiles
  pct_vals <- colnames(result_tbl) %>% str_subset("^\\d+(\\.\\d+)?\\%$")
  name_map <- c(set_names(pct_vals, format_pct(pct_vals)), rhat = "Rhat")

  bind_cols(par_tbl, dim_tbl, result_tbl) %>%
    rename(!!!name_map)
}

get_max_dim <- function(dims) {
  str_count(dims, ",") %>% add(1L) %>% coalesce(0L) %>% max()
}

format_pct <- function(str) str_c("q_", str_sub(str, end = -2L))