Maintain chain information from prediction functions

[n-kall]

Currently posterior_predict and similar functions return an array of draws (not a draws_array from posterior), which is S x N dimensions. I would like to combine this array of predictions with the posterior draws (in draws format), but as this array does not have any record of the chains, posterior::bind_draws does not work ("chain_ids do not match").

Example

fit <- brm(yield ~ N * P * K, npk)
posterior::bind_draws(as_draws(fit), as_draws(posterior_predict(fit)))
#> Error: 'chain_ids' of bound objects do not match.

I'd like to avoid having to merge the chains of the posterior draws, so if it would be possible for the prediction functions from brms to return chain information (e.g. using the draws_array format), that would be very useful.

[paul-buerkner]

Agreed. However, I cannot just change the output format of such a frequently used function. This has to be part of brms 3.0. I am definitely planning on outputting posterior draws formats. Which of them it will become, I am not yet sure.

[n-kall]

That makes total sense. In the meantime, here is a workaround function that I've been using which might be useful for others. It should also work for multivariate responses:

##' predictions as draws
##'
##' @param x brmsfit object
##' @param predict_fn function for predictions
##' @param prediction_names optional names of the predictions
##' @param ... further arguments passed to predict_fn
##' @return draws array of predictions
predictions_as_draws <- function(x, predict_fn, prediction_names = NULL, ...) {
  terms <- brms::brmsterms(x$formula)
  if(inherits(terms, "mvbrmsterms")) {
    responses <- brms::brmsterms(x$formula)$responses
    mv <- TRUE
  } else {
    responses <- ""
    mv <- FALSE
  }
  pred_draws <- list()
  predictions <- predict_fn(x, ...)
  if (!(mv)) {
    # add additional dimension in univariate case
    dim(predictions) <- c(dim(predictions), 1)
  }
  for (resp in seq_along(responses)) {
    # create draws array of predictions for each response variable
    predicted_draws <- posterior::as_draws_array(
      array(
        predictions[, , resp],
        dim = c(
          posterior::ndraws(x) / posterior::nchains(x),
          posterior::nchains(x), dim(predictions)[2]
        )
      )
    )
    # name predicted variables
    posterior::variables(predicted_draws) <-  c(
      paste0(
        responses[[resp]],
        "_pred[",
        seq_along(posterior::variables(predicted_draws)),
        "]")
    )
    pred_draws[[resp]] <- predicted_draws
  }
  # bind draws from different responses
  out <- posterior::bind_draws(pred_draws)
  if (!(is.null(prediction_names))) {
    posterior::variables(out) <- prediction_names
  }
  out
}

[avehtari]

Would it be possible to add an option to choose the output format, and keep the current format as the default?