Extracting and visualizing tidy draws from brms models

Introduction

This vignette describes how to use the tidybayes package to extract tidy data frames of draws from posterior distributions of model variables, fits, and predictions from brms::brm. 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(tidybayes)
library(ggplot2)
library(ggstance)
library(ggridges)
library(cowplot)
library(rstan)
library(brms)
library(ggrepel)
library(RColorBrewer)
library(gganimate)

theme_set(theme_tidybayes() + panel_border() + background_grid())

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 = brm(response ~ (1|condition), data = ABC, control = list(adapt_delta = .99),
  prior = c(
    prior(normal(0, 1), class = Intercept),
    prior(student_t(3, 0, 1), class = sd),
    prior(student_t(3, 0, 1), class = sigma)
  ))

## Compiling the C++ model

## Start sampling

The results look like this:

m

##  Family: gaussian 
##   Links: mu = identity; sigma = identity 
## Formula: response ~ (1 | condition) 
##    Data: ABC (Number of observations: 50) 
## Samples: 4 chains, each with iter = 2000; warmup = 1000; thin = 1;
##          total post-warmup samples = 4000
## 
## Group-Level Effects: 
## ~condition (Number of levels: 5) 
##               Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
## sd(Intercept)     1.15      0.42     0.60     2.26       1027 1.00
## 
## Population-Level Effects: 
##           Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
## Intercept     0.48      0.47    -0.51     1.37       1019 1.00
## 
## Family Specific Parameters: 
##       Estimate Est.Error l-95% CI u-95% CI Eff.Sample Rhat
## sigma     0.56      0.06     0.46     0.70       2191 1.00
## 
## Samples were drawn using sampling(NUTS). For each parameter, Eff.Sample 
## is a crude measure of effective sample size, and Rhat is the potential 
## scale reduction factor on split chains (at convergence, Rhat = 1).

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 variable names so that we know what variables we can extract from the model:

get_variables(m)

##  [1] "b_Intercept"              "sd_condition__Intercept"  "sigma"                    "r_condition[A,Intercept]"
##  [5] "r_condition[B,Intercept]" "r_condition[C,Intercept]" "r_condition[D,Intercept]" "r_condition[E,Intercept]"
##  [9] "lp__"                     "accept_stat__"            "stepsize__"               "treedepth__"             
## [13] "n_leapfrog__"             "divergent__"              "energy__"

Here, b_Intercept is the global mean, and the r_condition[] variables are offsets from that mean for each condition. Given these variables:

• r_condition[A,Intercept]
• r_condition[B,Intercept]
• r_condition[C,Intercept]
• r_condition[D,Intercept]
• r_condition[E,Intercept]

We might want a data frame where each row is a draw from either r_condition[A,Intercept], r_condition[B,Intercept], ...[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 simultaneously.

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 variables and their indices into tidy-format data frames.

m %>%
  spread_draws(r_condition[condition,term]) %>%
  head(10)

m %>%
  spread_draws(r_condition[c,t]) %>%
  head(10)

m %>%
  spread_draws(r_condition[condition,]) %>%
  head(10)

Point summaries and intervals

m %>%
  spread_draws(b_Intercept, sigma) %>%
  head(10)

m %>%
  spread_draws(b_Intercept, sigma) %>%
  median_qi(b_Intercept, sigma)

m %>%
  spread_draws(b_Intercept, sigma) %>%
  median_qi()

m %>%
  gather_draws(b_Intercept, sigma) %>%
  median_qi()

m %>%
  spread_draws(r_condition[condition,]) %>%
  median_qi()

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

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 (b_Intercept) plus the effect for a given condition (r_condition).

We can gather draws from b_Intercept and r_condition together in a single data frame:

m %>% 
  spread_draws(b_Intercept, r_condition[condition,]) %>%
  head(10)

Within each draw, b_Intercept is repeated as necessary to correspond to every index of r_condition. 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(`b_Intercept`, r_condition[condition,]) %>%
  mutate(condition_mean = b_Intercept + r_condition) %>%
  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(b_Intercept, r_condition[condition,]) %>%
  median_qi(condition_mean = b_Intercept + r_condition)

Plotting point summaries and intervals

Plotting point summaries and with one interval is straightforward using the ggplot2::geom_pointrange or ggstance::geom_pointrangeh geoms:

m %>%
  spread_draws(b_Intercept, r_condition[condition,]) %>%
  median_qi(condition_mean = b_Intercept + r_condition) %>%
  ggplot(aes(y = condition, x = condition_mean, xmin = .lower, xmax = .upper)) +
  geom_pointrangeh()

Intervals with multiple probability levels

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

m %>%
  spread_draws(b_Intercept, r_condition[condition,]) %>%
  median_qi(condition_mean = b_Intercept + r_condition, .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 geom_pointintervalh geom, provided by tidybayes, is a shorthand for a geom_pointrangeh with xmin, xmax, and size set appropriately based on the .lower, .upper, and .width columns in the data to produce plots of point summaries with multiple probability levels:

m %>%
  spread_draws(b_Intercept, r_condition[condition,]) %>%
  median_qi(condition_mean = b_Intercept + r_condition, .width = c(.95, .66)) %>%
  ggplot(aes(y = condition, x = condition_mean)) +
  geom_pointintervalh() 

Intervals with densities

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

m %>%
  spread_draws(b_Intercept, r_condition[condition,]) %>%
  mutate(condition_mean = b_Intercept + r_condition) %>%
  ggplot(aes(y = condition, x = condition_mean)) +
  geom_halfeyeh()

Posterior fits

Rather than calculating conditional means manually as in the previous example, we could use add_fitted_draws, which is analogous to brms::fitted.brmsfit or brms::posterior_linpred (giving posterior draws from the model’s linear predictor, in this case, 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 fits we want, then transform that grid into a long-format data frame of draws from posterior fits:

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

To plot this example, we’ll also show the use of stat_pointintervalh instead of geom_pointintervalh, which summarizes draws into points and intervals within ggplot:

ABC %>%
  data_grid(condition) %>%
  add_fitted_draws(m) %>%
  ggplot(aes(x = .value, y = condition)) +
  stat_pointintervalh(.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, which also allow precise estimation of arbitrary intervals (down to the dot resolution of the plot, here 100):

ABC %>%
  data_grid(condition) %>%
  add_fitted_draws(m) %>%
  do(tibble(.value = quantile(.$.value, ppoints(100)))) %>%
  ggplot(aes(x = .value)) +
  geom_dotplot(binwidth = .04) +
  facet_grid(fct_rev(condition) ~ .) +
  scale_y_continuous(breaks = NULL)

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_fitted_draws is analogous to brms::fitted.brmsfit (or brms::posterior_linpred), add_predicted_draws is analogous to brms::predict.brmsfit (brms::posterior_predict), giving draws from the posterior predictive distribution.

Here is an example of posterior predictive distributions plotted using ggridges::geom_density_ridges:

ABC %>%
  data_grid(condition) %>%
  add_predicted_draws(m) %>%
  ggplot(aes(x = .prediction, y = condition)) +
  geom_density_ridges()

## Picking joint bandwidth of 0.0993

We could also use tidybayes::stat_intervalh to plot predictive bands alongside the data:

ABC %>%
  data_grid(condition) %>%
  add_predicted_draws(m) %>%
  ggplot(aes(y = condition, x = .prediction)) +
  stat_intervalh(.width = c(.50, .80, .95, .99)) +
  geom_point(aes(x = response), data = ABC) +
  scale_color_brewer()

Altogether, data, posterior predictions, and posterior distributions of the means:

grid = ABC %>%
  data_grid(condition)

fits = grid %>%
  add_fitted_draws(m)

preds = grid %>%
  add_predicted_draws(m)

ABC %>%
  ggplot(aes(y = condition, x = response)) +
  stat_intervalh(aes(x = .prediction), data = preds) +
  stat_pointintervalh(aes(x = .value), data = fits, .width = c(.66, .95), position = position_nudge(y = -0.2)) +
  geom_point() +
  scale_color_brewer()

Posterior predictions, Kruschke-style

The above approach to posterior predictions integrates over the parameter uncertainty to give a single posterior predictive distribution. Another approach, often used by John Kruschke in his book Doing Bayesian Data Analysis, is to attempt to show both the predictive uncertainty and the parameter uncertainty simultaneously by showing several possible predictive distributions implied by the posterior.

We can do this pretty easily by asking for the distributional parameters for a given prediction implied by the posterior. We’ll do it explicitly here by setting dpar = c("mu", "sigma") in add_fitted_draws. Rather than specifying the parameters explicitly, you can also just set dpar = TRUE to get draws from all distributional parameters in a model, and this will work for any response distribution supported by brms.

For a more detailed description of how these charts (and some useful variations on them), see Solomon Kurz’s excellent blog post on the topic.

ABC %>%
  data_grid(condition) %>%
  add_fitted_draws(m, dpar = c("mu", "sigma"), n = 100) %>%
  mutate(
    lower = qnorm(.001, mu, sigma),
    upper = qnorm(.999, mu, sigma),
    response = map2(lower, upper, seq, length.out = 101),
    density = pmap(list(response, mu, sigma), dnorm)
  ) %>%
  unnest() %>%
  ggplot(aes(x = response, y = condition)) +
  geom_ridgeline(aes(height = density, group = interaction(condition, .draw)), 
    fill = NA, color = adjustcolor("black", alpha.f = 1/20)
  ) +
  geom_point(data = ABC, shape = 21, fill = brewer.pal(3, "Blues")[[2]], size = 2)

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 = brm(mpg ~ hp * cyl, data = mtcars)

We can draw fit curves with probability bands:

mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 51)) %>%
  add_fitted_draws(m_mpg) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl))) +
  stat_lineribbon(aes(y = .value)) +
  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)) %>%
  add_fitted_draws(m_mpg, n = 100) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl))) +
  geom_line(aes(y = .value, group = paste(cyl, .draw)), alpha = .1) +
  geom_point(data = mtcars) +
  scale_color_brewer(palette = "Dark2")

Or we can create animated hypothetical outcome plots (HOPs) of fit lines:

set.seed(123456)
ndraws = 50

p = mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 101)) %>%
  add_fitted_draws(m_mpg, n = ndraws) %>%
  ggplot(aes(x = hp, y = mpg, color = ordered(cyl))) +
  geom_line(aes(y = .value, group = paste(cyl, .draw))) +
  geom_point(data = mtcars) +
  scale_color_brewer(palette = "Dark2") +
  transition_states(.draw, 0, 1) +
  shadow_mark(future = TRUE, color = "gray50", alpha = 1/20)

animate(p, nframes = ndraws, fps = 2.5, width = 576, height = 384, res = 96, type = "cairo")

Or, for posterior predictions (instead of fits), we can go back to probability 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")

This gets difficult to judge by group, so probably better to facet into multiple plots. Fortunately, since we are using ggplot, that functionality is built in:

mtcars %>%
  group_by(cyl) %>%
  data_grid(hp = seq_range(hp, n = 101)) %>%
  add_predicted_draws(m_mpg) %>%
  ggplot(aes(x = hp, y = mpg)) +
  stat_lineribbon(aes(y = .prediction), .width = c(.99, .95, .8, .5), color = brewer.pal(5, "Blues")[[5]]) +
  geom_point(data = mtcars) +
  scale_fill_brewer() +
  facet_grid(. ~ cyl, space = "free_x", scales = "free_x")

set.seed(1234)
AB = tibble(
  group = rep(c("a", "b"), each = 20),
  response = rnorm(40, mean = rep(c(1, 5), each = 20), sd = rep(c(1, 3), each = 20))
)

AB %>%
  ggplot(aes(x = response, y = group)) +
  geom_point()

m_ab = brm(
  bf(
    response ~ group,
    sigma ~ group
  ),
  data = AB
)

## Compiling the C++ model

## Start sampling

grid = AB %>%
  data_grid(group)

fits = grid %>%
  add_fitted_draws(m_ab)

preds = grid %>%
  add_predicted_draws(m_ab)

AB %>%
  ggplot(aes(x = response, y = group)) +
  geom_halfeyeh(aes(x = .value), relative_scale = 0.7, position = position_nudge(y = 0.1), data = fits) +
  stat_intervalh(aes(x = .prediction), data = preds) +
  geom_point(data = AB) +
  scale_color_brewer()

grid %>%
  add_fitted_draws(m_ab, dpar = TRUE) %>%
  ggplot(aes(x = sigma, y = group)) +
  geom_halfeyeh() +
  geom_vline(xintercept = 0, linetype = "dashed")

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 another plotting geom, geom_halfeyeh, which gives horizontal “half-eye” plots, combining intervals with a density plot:

#N.B. the syntax for compare_levels is experimental and may change
m %>%
  spread_draws(r_condition[condition,]) %>%
  compare_levels(r_condition, by = condition) %>%
  ggplot(aes(y = condition, x = r_condition)) +
  geom_halfeyeh()

If you prefer “caterpillar” plots, ordered by something like the mean of the difference, you can reorder the factor before plotting:

#N.B. the syntax for compare_levels is experimental and may change
m %>%
  spread_draws(r_condition[condition,]) %>%
  compare_levels(r_condition, by = condition) %>%
  ungroup() %>%
  mutate(condition = reorder(condition, r_condition)) %>%
  ggplot(aes(y = condition, x = r_condition)) +
  geom_halfeyeh() +
  geom_vline(xintercept = 0, linetype = "dashed") 

Ordinal models

The brms::fitted.brmsfit function for ordinal and multinomial regression models in brms returns multiple variables for each draw: one for each outcome category (in contrast to rstanarm::stan_polr models, which return draws from the latent linear predictor). The philosophy of tidybayes is to tidy whatever format is output by a model, so in keeping with that philosophy, when applied to ordinal and multinomial brms models, add_fitted_draws adds an additional column called .category and a separate row containing the variable for each category is output for every draw and predictor.

mtcars_clean = mtcars %>%
  mutate(cyl = ordered(cyl))

head(mtcars_clean)

m_cyl = brm(cyl ~ mpg, data = mtcars_clean, family = cumulative, seed = 58393)

## Compiling the C++ model

## Start sampling

tibble(mpg = 21) %>%
  add_fitted_draws(m_cyl) %>%
  median_qi(.value)

data_plot = mtcars_clean %>%
  ggplot(aes(x = mpg, y = cyl, color = cyl)) +
  geom_point() +
  scale_color_brewer(palette = "Dark2", name = "cyl")

fit_plot = mtcars_clean %>%
  data_grid(mpg = seq_range(mpg, n = 101)) %>%
  # we can use the `value` argument to give the column with values of 
  # transformed linear predictors a more precise name and the 
  # `category` argument to give the column with category labels
  # a more precise name
  add_fitted_draws(m_cyl, value = "P(cyl | mpg)", category = "cyl") %>%
  ggplot(aes(x = mpg, y = `P(cyl | mpg)`, color = cyl)) +
  stat_lineribbon(aes(fill = cyl), alpha = 1/5) +
  scale_color_brewer(palette = "Dark2") +
  scale_fill_brewer(palette = "Dark2")

plot_grid(ncol = 1, align = "v",
  data_plot,
  fit_plot
)

ndraws = 50

p = mtcars_clean %>%
  data_grid(mpg = seq_range(mpg, n = 101)) %>%
  add_fitted_draws(m_cyl, value = "P(cyl | mpg)", category = "cyl") %>%
  ggplot(aes(x = mpg, y = `P(cyl | mpg)`, color = cyl)) +
  # we remove the `.draw` column from the data for stat_lineribbon so that the same ribbons
  # are drawn on every frame (since we use .draw to determine the transitions below)
  stat_lineribbon(aes(fill = cyl), alpha = 1/5, color = NA, data = . %>% select(-.draw)) +
  # we use sample_draws to subsample at the level of geom_line (rather than for the full dataset
  # as in previous HOPs examples) because we need the full set of draws for stat_lineribbon above
  geom_line(aes(group = paste(.draw, cyl)), size = 1, data = . %>% sample_draws(ndraws)) +
  scale_color_brewer(palette = "Dark2") +
  scale_fill_brewer(palette = "Dark2") +
  transition_manual(.draw)

animate(p, nframes = ndraws, fps = 2.5, width = 576, height = 192, res = 96, type = "cairo")