Extracting and visualizing tidy draws from rstanarm models

Introduction

This vignette describes how to use the tidybayes and ggdist packages to extract and visualize tidy data frames of draws from posterior distributions of model variables, means, and predictions from rstanarm. For a more general introduction to tidybayes and its use on general-purpose Bayesian modeling languages (like Stan and JAGS), see vignette("tidybayes").

Setup

The following libraries are required to run this vignette:

library(magrittr)
library(dplyr)
library(purrr)
library(forcats)
library(tidyr)
library(modelr)
library(ggdist)
library(tidybayes)
library(ggplot2)
library(cowplot)
library(rstan)
library(rstanarm)
library(RColorBrewer)

theme_set(theme_tidybayes() + panel_border())

These options help Stan run faster:

rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())

Example dataset

To demonstrate tidybayes, we will use a simple dataset with 10 observations from 5 conditions each:

set.seed(5)
n = 10
n_condition = 5
ABC =
  tibble(
    condition = rep(c("A","B","C","D","E"), n),
    response = rnorm(n * 5, c(0,1,2,1,-1), 0.5)
  )

A snapshot of the data looks like this:

head(ABC, 10)

This is a typical tidy format data frame: one observation per row. Graphically:

ABC %>%
  ggplot(aes(y = condition, x = response)) +
  geom_point()

Model

Let’s fit a hierarchical model with shrinkage towards a global mean:

m = stan_lmer(response ~ (1|condition), data = ABC, 
  prior = normal(0, 1, autoscale = FALSE),
  prior_aux = student_t(3, 0, 1, autoscale = FALSE),
  adapt_delta = .99)

The results look like this:

m

## stan_lmer
##  family:       gaussian [identity]
##  formula:      response ~ (1 | condition)
##  observations: 50
## ------
##             Median MAD_SD
## (Intercept) 0.6    0.5   
## 
## Auxiliary parameter(s):
##       Median MAD_SD
## sigma 0.6    0.1   
## 
## Error terms:
##  Groups    Name        Std.Dev.
##  condition (Intercept) 1.13    
##  Residual              0.56    
## Num. levels: condition 5 
## 
## ------
## * For help interpreting the printed output see ?print.stanreg
## * For info on the priors used see ?prior_summary.stanreg

Extracting draws from a fit in tidy-format using spread_draws

Now that we have our results, the fun begins: getting the draws out in a tidy format! First, we’ll use the get_variables() function to get a list of raw model variables names so that we know what variables we can extract from the model:

get_variables(m)

##  [1] "(Intercept)"                              "b[(Intercept) condition:A]"              
##  [3] "b[(Intercept) condition:B]"               "b[(Intercept) condition:C]"              
##  [5] "b[(Intercept) condition:D]"               "b[(Intercept) condition:E]"              
##  [7] "sigma"                                    "Sigma[condition:(Intercept),(Intercept)]"
##  [9] "accept_stat__"                            "stepsize__"                              
## [11] "treedepth__"                              "n_leapfrog__"                            
## [13] "divergent__"                              "energy__"

Here, (Intercept) is the global mean, and the b parameters are offsets from that mean for each condition. Given these parameters:

• b[(Intercept) condition:A]
• b[(Intercept) condition:B]
• b[(Intercept) condition:C]
• b[(Intercept) condition:D]
• b[(Intercept) condition:E]

We might want a data frame where each row is a draw from either b[(Intercept) condition:A], b[(Intercept) condition:B], ...:C], ...:D], or ...:E], and where we have columns indexing which chain/iteration/draw the row came from and which condition (A to E) it is for. That would allow us to easily compute quantities grouped by condition, or generate plots by condition using ggplot, or even merge draws with the original data to plot data and posteriors.

The workhorse of tidybayes is the spread_draws() function, which does this extraction for us. It includes a simple specification format that we can use to extract model variables and their indices into tidy-format data frames.

m %>%
  spread_draws(b[term,group]) %>%
  head(10)

m %>%
  spread_draws(b[t,g]) %>%
  head(10)

m %>%
  spread_draws(b[,group]) %>%
  head(10)

m %>% 
  spread_draws(b[,group]) %>%
  separate(group, c("group", "condition"), ":") %>%
  head(10)

m %>% 
  spread_draws(b[,group,condition], sep = "[, :]") %>%
  head(10)

Point summaries and intervals

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  head(10)

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  median_qi(`(Intercept)`, sigma)

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  median_qi()

m %>%
  gather_draws(`(Intercept)`, sigma) %>%
  median_qi()

m %>%
  spread_draws(b[,group]) %>%
  median_qi()

m %>%
  spread_draws(b[,group]) %>%
  group_by(group) %>%       # this line not necessary (done by spread_draws)
  median_qi(b)                # b is not necessary (it is the only non-group column)

m %>%
  spread_draws(b[,group]) %>%
  summarise_draws()

Combining variables with different indices in a single tidy format data frame

spread_draws() and gather_draws() support extracting variables that have different indices into the same data frame. Indices with the same name are automatically matched up, and values are duplicated as necessary to produce one row per all combination of levels of all indices. For example, we might want to calculate the mean within each condition (call this condition_mean). In this model, that mean is the intercept ((Intercept)) plus the effect for a given condition (b).

We can gather draws from (Intercept) and b together in a single data frame:

m %>% 
  spread_draws(`(Intercept)`, b[,group]) %>%
  head(10)

Within each draw, (Intercept) is repeated as necessary to correspond to every index of b. Thus, the mutate function from dplyr can be used to find their sum, condition_mean (which is the mean for each condition):

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  mutate(condition_mean = `(Intercept)` + b) %>%
  median_qi(condition_mean)

median_qi() uses tidy evaluation (see vignette("tidy-evaluation", package = "rlang")), so it can take column expressions, not just column names. Thus, we can simplify the above example by moving the calculation of condition_mean from mutate into median_qi():

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  median_qi(condition_mean = `(Intercept)` + b)

Plotting intervals with multiple probability levels

median_qi() and its sister functions can produce an arbitrary number of probability intervals by setting the .width = argument:

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  median_qi(condition_mean = `(Intercept)` + b, .width = c(.95, .8, .5))

The results are in a tidy format: one row per group and uncertainty interval width (.width). This facilitates plotting. For example, assigning -.width to the size aesthetic will show all intervals, making thicker lines correspond to smaller intervals. The ggdist::geom_pointinterval() geom automatically sets the size aesthetic appropriately based on the .width columns in the data to produce plots of points with multiple probability levels:

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  median_qi(condition_mean = `(Intercept)` + b, .width = c(.95, .66)) %>%
  ggplot(aes(y = group, x = condition_mean, xmin = .lower, xmax = .upper)) +
  geom_pointinterval()

Intervals with densities

To see the density along with the intervals, we can use ggdist::stat_eye() (“eye plots”, which combine intervals with violin plots), or ggdist::stat_halfeye() (interval + density plots):

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  mutate(condition_mean = `(Intercept)` + b) %>%
  ggplot(aes(y = group, x = condition_mean)) +
  stat_halfeye()

Or say you want to annotate portions of the densities in color; the fill aesthetic can vary within a slab in all geoms and stats in the ggdist::geom_slabinterval() family, including ggdist::stat_halfeye(). For example, if you want to annotate a domain-specific region of practical equivalence (ROPE), you could do something like this:

m %>%
  spread_draws(`(Intercept)`, b[,group]) %>%
  mutate(condition_mean = `(Intercept)` + b) %>%
  ggplot(aes(y = group, x = condition_mean, fill = stat(abs(x) < .8))) +
  stat_halfeye() +
  geom_vline(xintercept = c(-.8, .8), linetype = "dashed") +
  scale_fill_manual(values = c("gray80", "skyblue"))

Other visualizations of distributions: stat_slabinterval

There are a variety of additional stats for visualizing distributions in the ggdist::geom_slabinterval() family of stats and geoms:

See vignette("slabinterval", package = "ggdist") for an overview.

Posterior means and predictions

Rather than calculating conditional means manually as in the previous example, we could use add_epred_draws(), which is analogous to rstanarm::posterior_epred() (giving posterior draws from the expectation of the posterior predictive; i.e. posterior distributions of conditional means), but uses a tidy data format. We can combine it with modelr::data_grid() to first generate a grid describing the predictions we want, then transform that grid into a long-format data frame of draws from conditional means:

ABC %>%
  data_grid(condition) %>%
  add_epred_draws(m) %>%
  head(10)

To plot this example, we’ll also show the use of ggdist::stat_pointinterval() instead of ggdist::geom_pointinterval(), which summarizes draws into point summaries and intervals within ggplot:

ABC %>%
  data_grid(condition) %>%
  add_epred_draws(m) %>%
  ggplot(aes(x = .epred, y = condition)) +
  stat_pointinterval(.width = c(.66, .95))

Quantile dotplots

Intervals are nice if the alpha level happens to line up with whatever decision you are trying to make, but getting a shape of the posterior is better (hence eye plots, above). On the other hand, making inferences from density plots is imprecise (estimating the area of one shape as a proportion of another is a hard perceptual task). Reasoning about probability in frequency formats is easier, motivating quantile dotplots (Kay et al. 2016, Fernandes et al. 2018), which also allow precise estimation of arbitrary intervals (down to the dot resolution of the plot, 100 in the example below).

Within the slabinterval family of geoms in tidybayes is the dots and dotsinterval family, which automatically determine appropriate bin sizes for dotplots and can calculate quantiles from samples to construct quantile dotplots. ggdist::stat_dotsinterval() is the horizontal variant designed for use on samples:

ABC %>%
  data_grid(condition) %>%
  add_epred_draws(m) %>%
  ggplot(aes(x = .epred, y = condition)) +
  stat_dotsinterval(quantiles = 100)

The idea is to get away from thinking about the posterior as indicating one canonical point or interval, but instead to represent it as (say) 100 approximately equally likely points.

Posterior predictions

Where add_epred_draws() is analogous to rstanarm::posterior_epred(), add_predicted_draws() is analogous to rstanarm::posterior_predict(), giving draws from the posterior predictive distribution.

We could use tidybayes::stat_interval() to plot predictive bands alongside the data and posterior distributions of the means:

grid = ABC %>%
  data_grid(condition)

means = grid %>%
  add_epred_draws(m)

preds = grid %>%
  add_predicted_draws(m)

ABC %>%
  ggplot(aes(y = condition, x = response)) +
  stat_interval(aes(x = .prediction), data = preds) +
  stat_pointinterval(aes(x = .epred), data = means, .width = c(.66, .95), position = position_nudge(y = -0.3)) +
  geom_point() +
  scale_color_brewer()

Fit/prediction curves

To demonstrate drawing fit curves with uncertainty, let’s fit a slightly naive model to part of the mtcars dataset:

m_mpg = stan_glm(mpg ~ hp * cyl, data = mtcars)

We can plot fit curves with probability bands:

mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 51)) %>%
  add_epred_draws(m_mpg) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl))) +
  stat_lineribbon(aes(y = .epred)) +
  geom_point(data = mtcars) +
  scale_fill_brewer(palette = "Greys") +
  scale_color_brewer(palette = "Set2")

Or we can sample a reasonable number of fit lines (say 100) and overplot them:

mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 101)) %>%
  # NOTE: this shows the use of ndraws to subsample within add_epred_draws()
  # ONLY do this IF you are planning to make spaghetti plots, etc.
  # NEVER subsample to a small sample to plot intervals, densities, etc.
  add_epred_draws(m_mpg, ndraws = 100) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl))) +
  geom_line(aes(y = .epred, group = paste(cyl, .draw)), alpha = .1) +
  geom_point(data = mtcars) +
  scale_color_brewer(palette = "Dark2")

Or we could plot posterior predictions (instead of means). For this example we’ll also use alpha to make it easier to see overlapping bands:

mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 101)) %>%
  add_predicted_draws(m_mpg) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl), fill = ordered(cyl))) +
  stat_lineribbon(aes(y = .prediction), .width = c(.95, .80, .50), alpha = 1/4) +
  geom_point(data = mtcars) +
  scale_fill_brewer(palette = "Set2") +
  scale_color_brewer(palette = "Dark2")

See vignette("tidy-brms") for additional examples of fit lines, including animated hypothetical outcome plots (HOPs).

Comparing levels of a factor

If we wish compare the means from each condition, compare_levels() facilitates comparisons of the value of some variable across levels of a factor. By default it computes all pairwise differences.

Let’s demonstrate compare_levels() with ggdist::stat_halfeye(). We’ll also re-order by the mean of the difference:

m %>%
  spread_draws(b[,,condition], sep = "[, :]") %>%
  compare_levels(b, by = condition) %>%
  ungroup() %>%
  mutate(condition = reorder(condition, b)) %>%
  ggplot(aes(y = condition, x = b)) +
  stat_halfeye() +
  geom_vline(xintercept = 0, linetype = "dashed") 

Ordinal models

data(esoph)
m_esoph_rs = stan_polr(tobgp ~ agegp, data = esoph, prior = R2(0.25), prior_counts = rstanarm::dirichlet(1))

esoph %>%
  data_grid(agegp) %>%
  add_linpred_draws(m_esoph_rs) %>%
  ggplot(aes(x = as.numeric(agegp), y = .linpred)) +
  stat_lineribbon() +
  scale_fill_brewer(palette = "Greys")

get_variables(m_esoph_rs)

##  [1] "agegp.L"        "agegp.Q"        "agegp.C"        "agegp^4"        "agegp^5"        "0-9g/day|10-19"
##  [7] "10-19|20-29"    "20-29|30+"      "accept_stat__"  "stepsize__"     "treedepth__"    "n_leapfrog__"  
## [13] "divergent__"    "energy__"

thresholds = m_esoph_rs %>%
  gather_draws(`.*[|].*`, regex = TRUE) %>%
  group_by(.draw) %>%
  select(.draw, threshold = .value) %>%
  summarise_all(list) %>%
  mutate(threshold = map(threshold, ~ c(., Inf)))

head(thresholds, 10)

thresholds[1,]$threshold

## [[1]]
## [1] -0.8200002  0.6020005  1.4201284        Inf

esoph %>%
  data_grid(agegp) %>%
  add_linpred_draws(m_esoph_rs) %>%
  inner_join(thresholds, by = ".draw") %>%
  mutate(`P(Y = category)` = map2(threshold, .linpred, function(alpha, beta_x)
      # this part is logit^-1(alpha_j - beta*x) - logit^-1(alpha_j-1 - beta*x)
      plogis(alpha - beta_x) - 
      plogis(lag(alpha, default = -Inf) - beta_x)
    )) %>%
  mutate(.category = list(levels(esoph$tobgp))) %>%
  unnest(c(threshold, `P(Y = category)`, .category)) %>%
  ggplot(aes(x = agegp, y = `P(Y = category)`, color = .category)) +
  stat_pointinterval(position = position_dodge(width = .4)) +
  scale_size_continuous(guide = "none") +
  scale_color_manual(values = brewer.pal(6, "Blues")[-c(1,2)]) 

esoph_plot = esoph %>%
  data_grid(agegp) %>%
  add_linpred_draws(m_esoph_rs) %>%
  inner_join(thresholds, by = ".draw") %>%
  mutate(`P(Y = category)` = map2(threshold, .linpred, function(alpha, beta_x)
      # this part is logit^-1(alpha_j - beta*x) - logit^-1(alpha_j-1 - beta*x)
      plogis(alpha - beta_x) - 
      plogis(lag(alpha, default = -Inf) - beta_x)
    )) %>%
  mutate(.category = list(levels(esoph$tobgp))) %>%
  unnest(c(threshold, `P(Y = category)`, .category)) %>%
  ggplot(aes(x = `P(Y = category)`, y = .category)) +
  coord_cartesian(expand = FALSE) +
  facet_grid(. ~ agegp, switch = "x") +
  theme_classic() +
  theme(strip.background = element_blank(), strip.placement = "outside") +
  ggtitle("P(tobacco consumption category | age group)") +
  xlab("age group")

esoph_plot +
  stat_summary(fun = median, geom = "bar", fill = "gray65", width = 1, color = "white") +
  stat_pointinterval()

esoph_plot +
  stat_ccdfinterval() +
  expand_limits(x = 0) #ensure bars go to 0