Extracting and visualizing tidy draws from rstanarm models

Matthew Kay

2020-11-01

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, fits, 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)
library(gganimate)

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)
condition response
A -0.4204277
B 1.6921797
C 1.3722541
D 1.0350714
E -0.1442796
A -0.3014540
B 0.7639168
C 1.6823143
D 0.8571132
E -0.9309459

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:

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.

Gathering variable indices into a separate column in a tidy format data frame

Given a parameter like this:

b[(Intercept) condition:D]

We can provide spread_draws() with a column specification like this:

b[term,group]

Where term corresponds to (Intercept) and group to condition:D. There is nothing too magical about what spread_draws() does with this specification: under the hood, it splits the parameter indices by commas and spaces (you can split by other characters by changing the sep argument). It lets you assign columns to the resulting indices in order. So b[(Intercept) condition:D] has indices (Intercept) and condition:D, and spread_draws() lets us extract these indices as columns in the resulting tidy data frame of draws from b:

m %>%
  spread_draws(b[term,group]) %>%
  head(10)
term group b .chain .iteration .draw
(Intercept) condition:A -1.2677915 1 1 1
(Intercept) condition:A -0.2132916 1 2 2
(Intercept) condition:A 0.0192337 1 3 3
(Intercept) condition:A -0.4226583 1 4 4
(Intercept) condition:A -0.4007204 1 5 5
(Intercept) condition:A 0.0275432 1 6 6
(Intercept) condition:A -0.3184563 1 7 7
(Intercept) condition:A -0.4145287 1 8 8
(Intercept) condition:A -0.5637111 1 9 9
(Intercept) condition:A -1.1462066 1 10 10

We can choose whatever names we want for the index columns; e.g.:

m %>%
  spread_draws(b[t,g]) %>%
  head(10)
t g b .chain .iteration .draw
(Intercept) condition:A -1.2677915 1 1 1
(Intercept) condition:A -0.2132916 1 2 2
(Intercept) condition:A 0.0192337 1 3 3
(Intercept) condition:A -0.4226583 1 4 4
(Intercept) condition:A -0.4007204 1 5 5
(Intercept) condition:A 0.0275432 1 6 6
(Intercept) condition:A -0.3184563 1 7 7
(Intercept) condition:A -0.4145287 1 8 8
(Intercept) condition:A -0.5637111 1 9 9
(Intercept) condition:A -1.1462066 1 10 10

But the more descriptive and less cryptic names from the previous example are probably preferable.

In this particular model, there is only one term ((Intercept)), thus we could omit that index altogether to just get each group and the value of b for the corresponding condition:

m %>%
  spread_draws(b[,group]) %>%
  head(10)
group b .chain .iteration .draw
condition:A -1.2677915 1 1 1
condition:A -0.2132916 1 2 2
condition:A 0.0192337 1 3 3
condition:A -0.4226583 1 4 4
condition:A -0.4007204 1 5 5
condition:A 0.0275432 1 6 6
condition:A -0.3184563 1 7 7
condition:A -0.4145287 1 8 8
condition:A -0.5637111 1 9 9
condition:A -1.1462066 1 10 10

Since all the groups in this case are from the condition factor, we may also want to separate out a column just containing the corresponding condition (A, B, C, etc). We can do that using tidyr::separate:

m %>% 
  spread_draws(b[,group]) %>%
  separate(group, c("group", "condition"), ":") %>%
  head(10)
group condition b .chain .iteration .draw
condition A -1.2677915 1 1 1
condition A -0.2132916 1 2 2
condition A 0.0192337 1 3 3
condition A -0.4226583 1 4 4
condition A -0.4007204 1 5 5
condition A 0.0275432 1 6 6
condition A -0.3184563 1 7 7
condition A -0.4145287 1 8 8
condition A -0.5637111 1 9 9
condition A -1.1462066 1 10 10

Alternatively, we could change the sep argument to spread_draws() to also split on : (sep is a regular expression). Note: This works in this example, but will not work well on rstanarm models where interactions between factors are used as grouping levels in a multilevel model, thus : is not included in the default separators.

m %>% 
  spread_draws(b[,group,condition], sep = "[, :]") %>%
  head(10)
group condition b .chain .iteration .draw
condition A -1.2677915 1 1 1
condition A -0.2132916 1 2 2
condition A 0.0192337 1 3 3
condition A -0.4226583 1 4 4
condition A -0.4007204 1 5 5
condition A 0.0275432 1 6 6
condition A -0.3184563 1 7 7
condition A -0.4145287 1 8 8
condition A -0.5637111 1 9 9
condition A -1.1462066 1 10 10

Note: If you have used spread_draws() with a raw sample from Stan or JAGS, you may be used to using recover_types() before spread_draws() to get index column values back (e.g. if the index was a factor). This is not necessary when using spread_draws() on rstanarm models, because those models already contain that information in their variable names. For more on recover_types(), see vignette("tidybayes").

Point summaries and intervals

With simple model variables

tidybayes provides a family of functions for generating point summaries and intervals from draws in a tidy format. These functions follow the naming scheme [median|mean|mode]_[qi|hdi], for example, median_qi(), mean_qi(), mode_hdi(), and so on. The first name (before the _) indicates the type of point summary, and the second name indicates the type of interval. qi yields a quantile interval (a.k.a. equi-tailed interval, central interval, or percentile interval) and hdi yields a highest (posterior) density interval. Custom point or interval functions can also be applied using the point_interval() function.

For example, we might extract the draws corresponding to the posterior distributions of the overall mean and standard deviation of observations:

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  head(10)
.chain .iteration .draw (Intercept) sigma
1 1 1 1.4707274 0.7488978
1 2 2 0.2947973 0.5082004
1 3 3 0.3301366 0.5640631
1 4 4 0.5871531 0.5641068
1 5 5 0.5221652 0.5328644
1 6 6 0.6041039 0.5441214
1 7 7 0.3356784 0.6227631
1 8 8 0.5821142 0.5585193
1 9 9 1.0652540 0.6601671
1 10 10 1.1857010 0.5061982

Like with b[term,group], this gives us a tidy data frame. If we want the median and 95% quantile interval of the variables, we can apply median_qi():

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  median_qi(`(Intercept)`, sigma)
(Intercept) (Intercept).lower (Intercept).upper sigma sigma.lower sigma.upper .width .point .interval
0.6086138 -0.4855255 1.584869 0.5606433 0.4568133 0.6975851 0.95 median qi

We can specify the columns we want to get medians and intervals from, as above, or if we omit the list of columns, median_qi() will use every column that is not a grouping column or a special column (like .chain, .iteration, or .draw). Thus in the above example, (Intercept) and sigma are redundant arguments to median_qi() because they are also the only columns we gathered from the model. So we can simplify this to:

m %>%
  spread_draws(`(Intercept)`, sigma) %>%
  median_qi()
(Intercept) (Intercept).lower (Intercept).upper sigma sigma.lower sigma.upper .width .point .interval
0.6086138 -0.4855255 1.584869 0.5606433 0.4568133 0.6975851 0.95 median qi

If you would rather have a long-format list of intervals, use gather_draws() instead:

m %>%
  gather_draws(`(Intercept)`, sigma) %>%
  median_qi()
.variable .value .lower .upper .width .point .interval
(Intercept) 0.6086138 -0.4855255 1.5848694 0.95 median qi
sigma 0.5606433 0.4568133 0.6975851 0.95 median qi

For more on gather_draws(), see vignette("tidybayes").

With indexed variables

When we have a model variable with one or more indices, such as b, we can apply median_qi() (or other functions in the point_interval() family) as we did before:

m %>%
  spread_draws(b[,group]) %>%
  median_qi()
group b .lower .upper .width .point .interval
condition:A -0.4099590 -1.4241662 0.6630131 0.95 median qi
condition:B 0.3853878 -0.6377729 1.4480078 0.95 median qi
condition:C 1.2029810 0.2439743 2.3109966 0.95 median qi
condition:D 0.4012586 -0.5931526 1.4861856 0.95 median qi
condition:E -1.4849292 -2.5190234 -0.4131678 0.95 median qi

How did median_qi() know what to aggregate? Data frames returned by spread_draws() are automatically grouped by all index variables you pass to it; in this case, that means spread_draws() groups its results by group. median_qi() respects those groups, and calculates the point summaries and intervals within all groups. Then, because no columns were passed to median_qi(), it acts on the only non-special (.-prefixed) and non-group column, b. So the above shortened syntax is equivalent to this more verbose call:

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)
group b .lower .upper .width .point .interval
condition:A -0.4099590 -1.4241662 0.6630131 0.95 median qi
condition:B 0.3853878 -0.6377729 1.4480078 0.95 median qi
condition:C 1.2029810 0.2439743 2.3109966 0.95 median qi
condition:D 0.4012586 -0.5931526 1.4861856 0.95 median qi
condition:E -1.4849292 -2.5190234 -0.4131678 0.95 median qi

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)
.chain .iteration .draw (Intercept) group b
1 1 1 1.4707274 condition:A -1.2677915
1 1 1 1.4707274 condition:B -0.1775179
1 1 1 1.4707274 condition:C 0.4365900
1 1 1 1.4707274 condition:D -0.4126736
1 1 1 1.4707274 condition:E -2.1516102
1 2 2 0.2947973 condition:A -0.2132916
1 2 2 0.2947973 condition:B 0.8856697
1 2 2 0.2947973 condition:C 1.6586191
1 2 2 0.2947973 condition:D 0.7126382
1 2 2 0.2947973 condition:E -1.2611341

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)
group condition_mean .lower .upper .width .point .interval
condition:A 0.1987924 -0.1490188 0.5322955 0.95 median qi
condition:B 0.9948920 0.6540764 1.3561489 0.95 median qi
condition:C 1.8290093 1.4753124 2.1830999 0.95 median qi
condition:D 1.0154955 0.6644085 1.3666874 0.95 median qi
condition:E -0.8807627 -1.2302815 -0.5128633 0.95 median qi

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)
group condition_mean .lower .upper .width .point .interval
condition:A 0.1987924 -0.1490188 0.5322955 0.95 median qi
condition:B 0.9948920 0.6540764 1.3561489 0.95 median qi
condition:C 1.8290093 1.4753124 2.1830999 0.95 median qi
condition:D 1.0154955 0.6644085 1.3666874 0.95 median qi
condition:E -0.8807627 -1.2302815 -0.5128633 0.95 median qi

Plotting points and intervals

Plotting point summaries and with one interval is straightforward using ggdist::geom_pointinterval():

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

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(`(Intercept)`, b[,group]) %>%
  median_qi(condition_mean = `(Intercept)` + b, .width = c(.95, .8, .5))
group condition_mean .lower .upper .width .point .interval
condition:A 0.1987924 -0.1490188 0.5322955 0.95 median qi
condition:B 0.9948920 0.6540764 1.3561489 0.95 median qi
condition:C 1.8290093 1.4753124 2.1830999 0.95 median qi
condition:D 1.0154955 0.6644085 1.3666874 0.95 median qi
condition:E -0.8807627 -1.2302815 -0.5128633 0.95 median qi
condition:A 0.1987924 -0.0243299 0.4099232 0.80 median qi
condition:B 0.9948920 0.7687734 1.2267481 0.80 median qi
condition:C 1.8290093 1.5985601 2.0610815 0.80 median qi
condition:D 1.0154955 0.7812895 1.2365038 0.80 median qi
condition:E -0.8807627 -1.1069924 -0.6509838 0.80 median qi
condition:A 0.1987924 0.0816338 0.3077564 0.50 median qi
condition:B 0.9948920 0.8745594 1.1189284 0.50 median qi
condition:C 1.8290093 1.7087691 1.9516226 0.50 median qi
condition:D 1.0154955 0.8918738 1.1290132 0.50 median qi
condition:E -0.8807627 -1.0014610 -0.7577554 0.50 median qi

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:

The slabinterval family of geoms and stats

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

Posterior fits

Rather than calculating conditional means manually as in the previous example, we could use add_fitted_draws(), which is analogous to rstanarm::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)
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.
condition .row .chain .iteration .draw .value
A 1 NA NA 1 0.2029359
A 1 NA NA 2 0.0815057
A 1 NA NA 3 0.3493703
A 1 NA NA 4 0.1644948
A 1 NA NA 5 0.1214448
A 1 NA NA 6 0.6316470
A 1 NA NA 7 0.0172222
A 1 NA NA 8 0.1675855
A 1 NA NA 9 0.5015429
A 1 NA NA 10 0.0394944

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_fitted_draws(m) %>%
  ggplot(aes(x = .value, y = condition)) +
  stat_pointinterval(.width = c(.66, .95))
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.

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_dots() is the horizontal variant designed for use on samples:

ABC %>%
  data_grid(condition) %>%
  add_fitted_draws(m) %>%
  ggplot(aes(x = .value, y = condition)) +
  stat_dots(quantiles = 100)
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.

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 rstanarm::posterior_linpred(), add_predicted_draws() is analogous to rstanarm::posterior_predict(), giving draws from the posterior predictive distribution.

Here is an example of posterior predictive distributions plotted using ggdist::stat_slab():

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

We could also use tidybayes::stat_interval() to plot predictive bands alongside the data:

ABC %>%
  data_grid(condition) %>%
  add_predicted_draws(m) %>%
  ggplot(aes(y = condition, x = .prediction)) +
  stat_interval() +
  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)
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.
preds = grid %>%
  add_predicted_draws(m)

ABC %>%
  ggplot(aes(y = condition, x = response)) +
  stat_interval(aes(x = .prediction), data = preds) +
  stat_pointinterval(aes(x = .value), data = fits, .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_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")
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.

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")
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.

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)
## Instead of posterior_linpred(..., transform=TRUE) please call posterior_epred(), which provides equivalent functionality.
animate(p, nframes = ndraws, fps = 2.5, width = 432, height = 288, res = 96, dev = "png", 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 can get difficult to judge by group, so could be 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")

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():

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

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(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

Ordinal model with categorical predictor

Here’s an ordinal model with a categorical predictor:

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

The rstanarm::posterior_linpred() function for ordinal regression models in rstanarm returns only the link-level prediction for each draw (in contrast to brms::fitted.brmsfit(), which returns one prediction per category for ordinal models, see the ordinal regression examples in vignette("tidy-brms")). The philosophy of tidybayes is to tidy whatever format is output by a model, so in keeping with that philosophy, when applied to ordinal rstanarm models, add_fitted_draws() just returns the link-level prediction (Note: setting scale = "response" for such models will not usually make sense).

For example, here is a plot of the link-level fit:

esoph %>%
  data_grid(agegp) %>%
  add_fitted_draws(m_esoph_rs, scale = "linear") %>%
  ggplot(aes(x = as.numeric(agegp), y = .value)) +
  stat_lineribbon(color = "red") +
  scale_fill_brewer(palette = "Greys")

This can be hard to interpret. To turn this into predicted probabilities on a per-category basis, we have to use the fact that an ordinal logistic regression defines the probability of an outcome in category \(j\) or less as:

\[ \textrm{logit}\left[Pr(Y\le j)\right] = \alpha_j - \beta x \]

Thus, the probability of category \(j\) is:

\[ \begin{align} Pr(Y = j) &= Pr(Y \le j) - Pr(Y \le j - 1)\\ &= \textrm{logit}^{-1}(\alpha_j - \beta x) - \textrm{logit}^{-1}(\alpha_{j-1} - \beta x) \end{align} \]

To derive these values, we need two things:

The thresholds in rstanarm are coefficients with names containing |, indicating which categories they are thresholds between. We can see those parameters in the list of variables in the model:

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__"

We can extract those automatically by using the regex = TRUE argument to gather_draws() to find all variables containing a | character. We will then use dplyr::summarise_all(list) to turn these thresholds into a list column, and add a final threshold equal to \(+\infty\) (to represent the highest category):

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)
.draw threshold
1 -0.7923445, 0.4245173, 1.2704626, Inf
2 -1.17748852, 0.01450121, 1.30498864, Inf
3 -1.1495349, 0.1566035, 1.2123842, Inf
4 -1.4922447, 0.0523582, 1.2962055, Inf
5 -0.4998088, 0.5846447, 1.6627990, Inf
6 -0.9618395, 0.3006262, 1.1391999, Inf
7 -0.8411123, 0.2065357, 1.4143414, Inf
8 -1.1155892, 0.1232009, 1.2170391, Inf
9 -0.8299133, 0.3095756, 1.3044277, Inf
10 -1.04741544, 0.02811229, 1.00848691, Inf

For example, the threshold vector from one row of this data frame (i.e., from one draw from the posterior) looks like this:

thresholds[1,]$threshold
## [[1]]
## [1] -0.7923445  0.4245173  1.2704626        Inf

We can combine those thresholds (the \(\alpha_j\) values from the above formula) with the .value from add_fitted_draws() (\(\beta x\) from the above formula) to calculate per-category probabilities:

esoph %>%
  data_grid(agegp) %>%
  add_fitted_draws(m_esoph_rs, scale = "linear") %>%
  inner_join(thresholds, by = ".draw") %>%
  mutate(`P(Y = category)` = map2(threshold, .value, 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 = FALSE) +
  scale_color_manual(values = brewer.pal(6, "Blues")[-c(1,2)]) 

It is hard to see the changes in categories in the above plot; let’s try something that gives a better gist of the distribution within each year:

esoph_plot = esoph %>%
  data_grid(agegp) %>%
  add_fitted_draws(m_esoph_rs, scale = "linear") %>%
  inner_join(thresholds, by = ".draw") %>%
  mutate(`P(Y = category)` = map2(threshold, .value, 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 = "gray75", width = 1, color = "white") +
  stat_pointinterval()

The bars in this case might present a false sense of precision, so we could also try CCDF barplots instead:

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

This output should be very similar to the output from the corresponding m_esoph_brm model in vignette("tidy-brms") (modulo different priors), though it takes a bit more work to produce in rstanarm compared to brms.