Compute (Bayesian) R2 for Turing models

Specific Domains Probabilistic programming
,
1

Somehow my google search didn’t yield many results for how to compute R2 for Turing models.
StatsBase.jl defines a few generic methods, but I was in particular interested in the so-called Bayesian R2 (Gelman et al., 2019), defined as:

image
image907×405 17.8 KB

Would that be easy to implement in Julia/Turing? Its R implementation seems pretty straightforward:

bayes_R2 <- function(fit) {
  y <- rstanarm::get_y(fit)
  ypred <- rstanarm::posterior_linpred(fit, transform = TRUE)
  if (family(fit)$family == "binomial" && NCOL(y) == 2) {
    trials <- rowSums(y)
    y <- y[, 1]
    ypred <- ypred %*% diag(trials)
  }
  e <- -1 * sweep(ypred, 2, y)
  var_ypred <- apply(ypred, 1, var)
  var_e <- apply(e, 1, var)
  var_ypred / (var_ypred + var_e)
}
2

Yes, this is easily possible in Turing.
The simplest approach is probably to return the R² value from the model directly,

@model function mymodel()
    # ...model definition
    r2 = compute_r2(args)
    return r2
end

where compute_r2 is a function that computes the R² values given some parameters of the model.
Then you can call generated_quantities on the sampled model to get the R² values for each iteration.

model = mymodel()
chain = sample(model, NUTS(), 1000)
generated_quantities(model, chain)
3

PosteriorStats.jl implements Gelman’s Bayesian R^2. See API · PosteriorStats.jl.

5 Likes