R/funnel.R

In baggr: Bayesian Aggregate Treatment Effects

#' Funnel plots for baggr models
#'
#' @param bg a [baggr()] model
#' @param show whether to plot raw study-level inputs (`"inputs"`) or
#'   posterior summaries (`"posterior"`)
#' @param level confidence level for reference lines
#' @param label logical: add study/group labels?
#' @param covariate optional name of a column in the model input data used to
#'   colour points.
#'
#' @return A ggplot funnel plot for the supplied model
#' @details
#' Funnel plots provide a visual check of how study-level effects vary with
#' precision. Apparent asymmetry can indicate small-study effects, but can also
#' arise due to unexplained heterogeneity between studies.
#'
#' For models with group-level covariates, colouring points by a covariate can
#' help inspect whether asymmetry is partly explained by meta-regression effects.
#' In Rubin summary-data meta-regression models (`model = "rubin"` with
#' `covariates`), `show = "posterior"` plots posterior study effects from
#' [group_effects()], which include the fitted covariate contribution. By
#' contrast, `show = "inputs"` plots the original study-level estimates.
#'
#' @export
#' @importFrom ggrepel geom_text_repel
#'
#' @examples
#' bg <- baggr(schools, iter = 500, refresh = 0)
#' funnel_plot(bg, label = TRUE)
#'
funnel_plot <- function(bg,
                        show = c("inputs", "posterior"),
                        level = 0.95,
                        label = FALSE,
                        covariate = NULL) {
  stopifnot(inherits(bg, "baggr"))
  if(length(bg$effects) != 1)
    stop("Funnel plot currently defined for 1-dimensional effects.")
  show <- match.arg(show)

  # raw study-level data (tau/se/group)
  if(bg$model %in% c("rubin_full", "logit")) {
    studies <- bg$summary_data
  } else {
    studies <- bg$data
  }
  if(bg$model %in% c("rubin_full", "mutau"))
    studies$se <- studies$se.tau
  if(is.null(studies$group))
    studies$group <- group_names(bg)

  raw_df <- studies[, c("group", "tau", "se")]
  names(raw_df) <- c("group", "effect", "se")

  if(!is.null(covariate)) {
    if(!(covariate %in% names(studies))) {
      warning("Requested covariate column was not found; points are shown without colouring")
    } else {
      raw_df$.covariate <- studies[[covariate]]
    }
  }

  # posterior summaries (mean & SD) for partial/full pooling outputs
  post_df <- as.data.frame(group_effects(bg, summary = TRUE,
                                         interval = level)[,,1])
  post_df$group <- group_names(bg)
  post_df <- post_df[, c("group", "mean", "sd")]
  names(post_df) <- c("group", "effect", "se")
  if(!is.null(raw_df$.covariate))
    post_df$.covariate <- raw_df$.covariate

  plot_df <- if(show == "inputs") raw_df else post_df

  # hyper-mean + heterogeneity
  if(bg$pooling == "none")
    stop("Need a pooled model to plot average treatment effect for the funnel.")
  hypermean <- hypermean(bg)[["mean"]]
  hetero <- if(bg$pooling == "partial") hypersd(bg)[["mean"]] else 0

  crit <- stats::qnorm(1 - (1 - level)/2)
  se_grid <- seq(0, max(plot_df$se), length.out = 200)
  fan <- data.frame(
    se = se_grid,
    lower = hypermean - crit * sqrt(se_grid^2 + hetero^2),
    upper = hypermean + crit * sqrt(se_grid^2 + hetero^2)
  )

  effect <- group <- head <- lower <- se <- tail <- upper <- NULL
  .covariate <- NULL

  point_mapping <- if(!is.null(plot_df$.covariate)) {
    ggplot2::aes(colour = .covariate)
  }

  ggplot2::ggplot(plot_df, ggplot2::aes(x = effect, y = se)) +
    ggplot2::geom_point(mapping = point_mapping) +
    ggplot2::geom_line(data = fan,
                       ggplot2::aes(x = lower, y = se),
                       linetype = "dashed") +
    ggplot2::geom_line(data = fan,
                       ggplot2::aes(x = upper, y = se),
                       linetype = "dashed") +
    ggplot2::geom_vline(xintercept = hypermean, linewidth = 0.4) +
    {if(label) ggrepel::geom_text_repel(ggplot2::aes(label = group),
                                        min.segment.length = 0)} +
    ggplot2::scale_y_reverse() +
    ggplot2::labs(x = bg$effects,
                  y = "Standard error") +
    baggr_theme_get()
}

#' @rdname funnel_plot
#' @param ... arguments passed to [funnel_plot()]
#' @export
funnel <- function(...) {
  funnel_plot(...)
}