Streamlined way to use powerscale with predictions from brms

[jflournoy]

Amazing package and paper. Thank you! I suppose this is either a feature request or a cry for help. :)

I'm trying to look at sensitivity of the predictions from a brms model and my current implementation seems rather convoluted. I think a lot of this has to do first with the fact that brms does not provide chain and iteration info or that the functions in the posterior package don't extract this information properly. I have reproducible code below and though it works for the purpose, even after getting the prediction draws properly bound to the draws from the parameters, it's still a bit of rough work to get the data into shape for other purposes.

Only the last powerscale call works

Package versions:

> packageVersion('priorsense')
[1] ‘0.0.0.9000’
> packageVersion('brms')
[1] ‘2.17.7’
> packageVersion('tidybayes')
[1] ‘3.0.4’

library(priorsense)
library(brms)
library(tidybayes)
packageVersion('brms')
packageVersion('tidybayes')

d <- data.table::fread('https://raw.githubusercontent.com/dill/mgcv-workshop/master/data/time_series/Florida_climate.csv')
fit <- brm(bf(air_temp ~ 1 + s(year) + s(month)), 
           data = d, chains = 4, cores = 4, 
           file = 'climate', file_refit = 'on_change',
           backend = 'rstan', control = list(adapt_delta = .999, max_treedepth = 20))

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = \(x) tidybayes::add_epred_draws(x, newdata = unique(x$data[, c('year', 'month')])))

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = \(x) brms::posterior_epred(x, newdata = unique(x$data[, c('year', 'month')])))

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = \(x) as_draws_df(brms::posterior_epred(x, newdata = unique(x$data[, c('year', 'month')]))))

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = \(x) as_draws_rvars(brms::posterior_epred(x, newdata = unique(x$data[, c('year', 'month')]))))

custom_prediction <- function(x){
  pred_draws <- brms::posterior_epred(x, newdata = unique(x$data[, c('year', 'month')]))
  pred_draws_df <- posterior::as_draws_df(pred_draws)
  iter_per_chain <- brms::ndraws(x) / brms::nchains(x)
  pred_draws_df$.chain <- (pred_draws_df$.draw - 1) %/% iter_per_chain + 1
  pred_draws_df$.iteration <- (pred_draws_df$.draw - 1) %% iter_per_chain + 1
  return(pred_draws_df)
}

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = custom_prediction)

[n-kall]

Thanks for opening this issue. I completely agree that current way predictions are handled is far from ideal. If there's only one predicted variable (for example R2), it should work like prediction = function(x) draws_df(R2 = bayes_R2(x, summary = F), .nchains = 4)). But in your case the prediction function is returning many variables, so this won't work as is.

I don't have an immediate solution for you as the way it is currently implemented is using posterior::bind_draws to join the predictions to the posterior draws, and I suppose this requires the same iters and chains in both draws objects (as you manually match in your custom_prediction function).

I have been planning to transition from functions that operate on fit objects to functions that operate directly on draws objects. This should allow more freedom to change the draws (such as adding predictions) before performing power-scaling. So this issue gives me some motivation to do that sooner.

[n-kall]

This is related to paul-buerkner/brms#1534. I will likely incorporate the workaround posted there into priorsense as the change to brms will be incorporated in a later brms version.

[n-kall]

This is now improved in the separate-scaling branch, with the helper function predictions_as_draws.

@jflournoy your example should now work like this:

priorsense::powerscale(
  fit, 
  alpha = .9, 
  component = 'prior', 
  prediction = \(x) priorsense::predictions_as_draws(x, brms::posterior_epred, newdata = unique(x$data[, c('year', 'month')])))

The predictions are named _pred[1] ... _pred[n] by default, but can be changed with the prediction_names argument.

[jflournoy]

I will test this out and report back. Thank you!

[jflournoy]

works very nicely. Thanks for the fix!