R/er-plot-style-summary.R

In erplots: Model-Agnostic Exposure-Response Plots

#' Summary annotation builders for exposure-response plots
#'
#' Builder functions for the `summary` layer ([er_plot_add_summary()]),
#' drawing a text/label annotation from a model's p-value, coefficients,
#' goodness-of-fit statistics, or observation counts.
#'
#' @include er-plot-style.R er-style-registry.R
#' @param data The original data frame.
#' @param config Configuration for the specific plot.
#' @param stratify Logical: whether to stratify.
#' @param exposure Exposure variable.
#' @param response Response variable.
#' @param strata Stratification variable.
#' @param theme Theme components.
#' @param inset Distance from the panel edge for the annotation label.
#'   Defaults to `0.05`.
#' @param label_size Label text size. Defaults to `NULL` ([ggplot2::geom_label()]'s own default).
#' @param label_colour Label text colour. Defaults to `NULL` ([ggplot2::geom_label()]'s own default).
#' @param label_fill Label background fill. Defaults to `NULL` ([ggplot2::geom_label()]'s own default).
#' @param fields Fields from `glance` to include for `er_style_summary_gof()`,
#'   and the order they're shown in: one or more of `"n"`, `"aic"`, `"bic"`, or
#'   `"r_squared"`. Defaults to all four, in that order. A field is shown
#'   only when both present and non-`NA` in the model's `glance` result.
#' @param ... Additional named arguments forwarded from [er_plot_add_model()]'s own `...`.
#'
#' @details
#' See [er_style()] for the shared builder interface these functions
#' implement, including how to write a custom builder of your own.
#'
#' @section Choosing a builder:
#' Each builder draws a different kind of annotation, with its own
#' data requirements:
#'
#' * `er_style_summary_pvalue()` (the default) -- a formatted p-value
#'   from the model's [er_summary()] result.
#' * `er_style_summary_n()` -- observation counts. Doesn't have to
#'   originate from a fitted model at all.
#' * `er_style_summary_coefficients()` -- one line per row of the
#'   model's `coefficients` table (see [er_summary()]'s `coefficients`
#'   field), useful for models with several parameters and no single
#'   privileged p-value (e.g. a multi-parameter nonlinear model). Draws
#'   nothing if `coefficients` wasn't supplied, or if the layer is
#'   stratified.
#' * `er_style_summary_gof()` -- a single-line, comma-separated
#'   goodness-of-fit annotation from the model's `glance` field (see
#'   [er_summary()]) -- a curated subset (N, AIC, BIC, R-squared)
#'   rather than every reserved `glance` column, showing only whichever
#'   of those four are actually present and non-`NA`. Same restrictions
#'   as `er_style_summary_coefficients()`: draws nothing if none of
#'   those fields are available, or if the layer is stratified.
#'
#' @section Tags:
#' All four builders are tagged `er_style_tag(fn, layer = "plot_summary")`,
#' so [er_plot_add_summary()] errors informatively if a builder tagged
#' for a different layer is passed to it instead.
#'
#' @returns A geom, or a list of geoms; see [er_style()].
#'
#' @examples
#' if (requireNamespace("erglm", quietly = TRUE)) {
#'   library(erglm)
#'   mod <- erglm_model(ae1 ~ aucss, erglm_data, family = binomial())
#'
#'   # er_style_summary_pvalue(): the default, drawn from the model's own
#'   # er_summary()
#'   erglm_data |>
#'     er_plot(aucss, ae1) |>
#'     er_plot_add_model(mod) |>
#'     er_plot_add_summary(model = mod, style = er_style_summary_pvalue) |>
#'     plot()
#'
#'   # er_style_summary_n(): model-agnostic observation count
#'   erglm_data |>
#'     er_plot(aucss, ae1) |>
#'     er_plot_add_model(mod) |>
#'     er_plot_add_summary(style = er_style_summary_n) |>
#'     plot()
#' }
#'
#' @name er_style_summary
#' @seealso [er_style()]
NULL


# internal helper to construct a geom_label with optional style args
.summary_label_geom <- function(summary_data, x, y, hjust, vjust,
                                label_size = NULL, label_colour = NULL, label_fill = NULL) {
  # geom_label's parameters: label.size, text.colour, fill, size.unit
  args <- list(
    data = summary_data,
    mapping = ggplot2::aes(x = I(x), y = I(y), label = lbl),
    hjust = hjust,
    vjust = vjust,
    show.legend = FALSE,
    inherit.aes = FALSE
  )
  if (!is.null(label_size)) args$size <- label_size
  if (!is.null(label_colour)) args$colour <- label_colour
  if (!is.null(label_fill)) args$fill <- label_fill

  do.call(ggplot2::geom_label, args)
}

#' @rdname er_style_summary
#' @export
er_style_summary_pvalue <- function(data, config, stratify, exposure, response, strata, theme,
                                     inset = 0.05, label_size = NULL, label_colour = NULL, label_fill = NULL, ...) {

  if (is.null(config$p_value) || stratify) return(list())

  corner <- names(sort(config$corner_distance)[4])
  summary_data <- tibble::tibble(lbl = theme$format_p(config$p_value))
  x_left  <- inset
  x_right <- 1 - inset
  y_top   <- 1 - inset
  y_bot   <- inset

  if (corner == "top_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_top, 0, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "top_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_top, 1, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_bot, 0, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_bot, 1, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  return(geoms)
}
er_style_summary_pvalue <- er_style_tag(er_style_summary_pvalue, layer = "plot_summary", label = "pvalue")

#' @rdname er_style_summary
#' @export
er_style_summary_n <- function(data, config, stratify, exposure, response, strata, theme,
                                inset = 0.05, label_size = NULL, label_colour = NULL, label_fill = NULL, ...) {

  if (stratify && !is.null(strata$name)) {
    counts <- data |>
      dplyr::mutate(strata_value = .get_strata_values(data, strata$name)) |>
      dplyr::count(strata_value) |>
      dplyr::mutate(lbl = paste0(strata_value, ": N=", n))
    lbl <- paste(counts$lbl, collapse = "\n")
  } else {
    lbl <- paste0("N=", nrow(data))
  }

  corner <- names(sort(config$corner_distance)[4])
  summary_data <- tibble::tibble(lbl = lbl)
  x_left  <- inset
  x_right <- 1 - inset
  y_top   <- 1 - inset
  y_bot   <- inset

  if (corner == "top_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_top, 0, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "top_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_top, 1, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_bot, 0, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_bot, 1, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  return(geoms)
}
er_style_summary_n <- er_style_tag(er_style_summary_n, layer = "plot_summary", label = "n")

#' @rdname er_style_summary
#' @export
er_style_summary_coefficients <- function(data, config, stratify, exposure, response, strata, theme,
                                           inset = 0.05, label_size = NULL, label_colour = NULL, label_fill = NULL, ...) {

  coefs <- config$summary$coefficients
  if (is.null(coefs) || stratify) return(list())

  # `label` falls back to `term`; `p_value` is optional per row -- see
  # `?er_model_interface`'s `coefficients` contract. Checked via `%in%`
  # names() rather than `$` directly, since tibble's `$` warns on access
  # to a column that isn't there.
  term_label <- if ("label" %in% names(coefs)) coefs$label else coefs$term
  row_p_value <- if ("p_value" %in% names(coefs)) coefs$p_value else rep(NA_real_, nrow(coefs))
  line <- ifelse(
    is.na(row_p_value),
    paste0(term_label, ": ", theme$format_number(coefs$estimate)),
    paste0(term_label, ": ", theme$format_number(coefs$estimate), " (", theme$format_p(row_p_value), ")")
  )

  corner <- names(sort(config$corner_distance)[4])
  summary_data <- tibble::tibble(lbl = paste(line, collapse = "\n"))
  x_left  <- inset
  x_right <- 1 - inset
  y_top   <- 1 - inset
  y_bot   <- inset

  if (corner == "top_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_top, 0, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "top_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_top, 1, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_bot, 0, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_bot, 1, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  return(geoms)
}
er_style_summary_coefficients <- er_style_tag(er_style_summary_coefficients, layer = "plot_summary", label = "coefficients")

#' @rdname er_style_summary
#' @export
er_style_summary_gof <- function(data, config, stratify, exposure, response, strata, theme,
                                 inset = 0.05, fields = c("n", "aic", "bic", "r_squared"),
                                 label_size = NULL, label_colour = NULL, label_fill = NULL, ...) {

  glance <- config$summary$glance
  if (is.null(glance) || stratify) return(list())

  # a curated, compact subset of `glance`'s reserved columns (see
  # `?er_model_interface`) -- `df_residual`/`logLik`/`deviance`/
  # `converged` are part of the contract but deliberately left out of
  # this compact annotation; a model package wanting to show those can
  # write its own builder reading `config$summary$glance` directly.
  # Each field is shown only if the column is both present and non-`NA`,
  # so a model that only populates some of `glance` (e.g. `aic` but not
  # `r_squared`) still gets a sensible, partial annotation rather than a
  # blank or an error. The `fields` argument controls which of the four
  # recognised fields to show and in what order.
  field_specs <- list(
    n         = list(label = "N",   format = function(x) as.character(as.integer(x))),
    aic       = list(label = "AIC", format = theme$format_number),
    bic       = list(label = "BIC", format = theme$format_number),
    r_squared = list(label = "R\u00b2", format = theme$format_number)
  )

  line <- character(0)
  for (col in fields) {
    if (col %in% names(field_specs) && col %in% names(glance) && !is.na(glance[[col]])) {
      spec <- field_specs[[col]]
      line <- c(line, paste0(spec$label, " = ", spec$format(glance[[col]])))
    }
  }
  if (length(line) == 0) return(list())

  corner <- names(sort(config$corner_distance)[4])
  summary_data <- tibble::tibble(lbl = paste(line, collapse = ", "))
  x_left  <- inset
  x_right <- 1 - inset
  y_top   <- 1 - inset
  y_bot   <- inset

  if (corner == "top_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_top, 0, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "top_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_top, 1, 1,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_left") {
    geoms <- .summary_label_geom(summary_data, x_left, y_bot, 0, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  if (corner == "bottom_right") {
    geoms <- .summary_label_geom(summary_data, x_right, y_bot, 1, 0,
                                 label_size = label_size, label_colour = label_colour, label_fill = label_fill)
  }

  return(geoms)
}
er_style_summary_gof <- er_style_tag(er_style_summary_gof, layer = "plot_summary", label = "gof")