R/er-plot-compose.R

In erplots: Model-Agnostic Exposure-Response Plots

# composition/polishing steps -------------------------------------------------

.polish_margins <- function(object) {

  p <- object$plot

  margins <- ggplot2::margin(t = 5.5, r = 5.5, b = 5.5, l = 5.5, unit = "pt")
  zero_pt <- ggplot2::unit(0, "pt")

  base_mar <- margins
  panel_position <- object$layer$data$config$panel_position %||% character(0)

  for (panel_name in names(p$data)) {
    panel_mar <- margins
    position <- panel_position[[panel_name]]
    if (identical(position, "above")) {
      base_mar[1] <- zero_pt
      panel_mar[3] <- zero_pt
    }
    if (identical(position, "below")) {
      base_mar[3] <- zero_pt
      panel_mar[1] <- zero_pt
    }
    p$data[[panel_name]] <- p$data[[panel_name]] + ggplot2::theme(margins = panel_mar)
  }

  # `p$base` is only built when at least one of the model/summary/
  # quantile/overlay layers is present (see `er_plot_build()`) -- a
  # group-only or panel-layout-data-only plot has no base panel to
  # margin-adjust.
  if (!is.null(p$base)) {
    p$base <- p$base + ggplot2::theme(margins = base_mar)
  }
  if (!is.null(p$group)) {
    for(g in seq_along(p$group)) {
      p$group[[g]] + ggplot2::theme(margins = margins)
    }
  }

  return(p)
}

.polish_labels <- function(object) {
  p <- object$plot

  # `p$base` is only built when at least one of the model/summary/
  # quantile/overlay layers is present (see `er_plot_build()`) -- a
  # group-only or panel-layout-data-only plot has no base panel to
  # label. `ggplot2::get_labs()` errors on `NULL`, so this whole block
  # is skipped rather than guarded piecemeal.
  if (!is.null(p$base)) {
    p$base <- p$base + ggplot2::labs(
      x = object$exposure$label,
      y = object$response$label
    )
    ll <- names(ggplot2::get_labs(p$base))

    # `fill` on the base plot almost always means strata (e.g.
    # `er_style_model_ribbonline()`'s ribbon), but an "overlay"-layout data
    # builder can claim `fill` for something else entirely --
    # `er_style_data_hex()` uses it for bin density, and tags itself with
    # `er_style_tag(builder, fill_role = "density")` to say so (mirroring
    # `er_style_group_histogram()`'s `y_role` tag). Such a builder can
    # only coexist with other `fill`-mapped layers if they don't map
    # `fill` themselves (a discrete `fill = strata` ribbon and a
    # continuous density `fill` collide as two scales for one aesthetic,
    # and ggplot2 errors) -- so if `fill` is present at all alongside a
    # density-tagged overlay builder, it's safe to assume the density is
    # the sole source and label it accordingly rather than as strata.
    overlay_style <- object$layer$overlay$config$style
    fill_is_density <- identical(.style_fill_role(overlay_style), "density")

    if ("fill" %in% ll) {
      p$base <- p$base + ggplot2::labs(fill = if (fill_is_density) "Count" else object$strata$label)
    }
    if ("colour" %in% ll) p$base <- p$base + ggplot2::labs(color = object$strata$label)
  }

  # the data layer's `colour` aesthetic means strata everywhere except
  # when `config$color_role == "response"` (continuous/count response;
  # there, `colour` is the response value itself, so its label is the
  # response's, not the strata's. When that response-coloured layer is also
  # faceted by stratum (more than one panel), each panel is tagged with its
  # stratum level via a plot title -- not the y-axis label, which patchwork's
  # `axes = "collect"` merges across all stacked panels (see
  # `er_plot_build()`), so a per-panel y-axis label would visually
  # overlap with the others rather than sit next to its own panel.
  data_color_role <- object$layer$data$config$color_role %||% "strata"
  data_color_label <- if (identical(data_color_role, "response")) {
    object$response$label
  } else {
    object$strata$label
  }
  data_panel_names <- names(p$data)
  data_is_faceted <- identical(data_color_role, "response") && length(data_panel_names) > 1

  for (panel_name in data_panel_names) {
    p$data[[panel_name]] <- p$data[[panel_name]] + ggplot2::labs(
      x = object$exposure$label,
      y = NULL,
      title = if (data_is_faceted) panel_name else NULL
    )
    ll <- names(ggplot2::get_labs(p$data[[panel_name]]))
    if ("fill" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + ggplot2::labs(fill = data_color_label)
    if ("colour" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + ggplot2::labs(color = data_color_label)
  }

  if (!is.null(p$group)) {
    for(g in names(p$group)) {
      # most group builders (e.g. `er_style_group_boxplot()`/
      # `er_style_group_violin()`) put the group variable itself on the
      # y-axis, so the group variable's own label is the right y-axis
      # title. A histogram-style builder instead needs its y-axis free
      # for counts (with group levels shown via facet strips), and tags
      # itself with `er_style_tag(builder, y_role = "count")` to say so --
      # see `er_style_group_histogram()`.
      group_style <- object$layer$group$config[[g]]$style
      y_label <- if (identical(.style_y_role(group_style), "count")) {
        "Count"
      } else {
        object$layer$group$config[[g]]$y$label
      }
      p$group[[g]] <- p$group[[g]] + ggplot2::labs(
        x = object$exposure$label,
        y = y_label
      )
      ll <- names(ggplot2::get_labs(p$group[[g]]))
      if ("fill" %in% ll) p$group[[g]] <- p$group[[g]] + ggplot2::labs(fill = object$strata$label)
      if ("colour" %in% ll) p$group[[g]] <- p$group[[g]] + ggplot2::labs(color = object$strata$label)
    }
  }

  return(p)
}

.polish_scales <- function(object) {
  p <- object$plot
  color_discrete <- object$theme$color_discrete
  fill_discrete <- object$theme$fill_discrete
  color_continuous <- object$theme$color_continuous
  fill_continuous <- object$theme$fill_continuous
  if (is.null(color_discrete) && is.null(fill_discrete) &&
      is.null(color_continuous) && is.null(fill_continuous)) {
    return(p)
  }

  # mirrors `.polish_labels()`'s own eligibility logic: `color_discrete`/
  # `fill_discrete` only ever override `colour`/`fill` where it's
  # genuinely mapped to strata (discrete); `color_continuous`/
  # `fill_continuous` are the symmetric counterpart, only ever overriding
  # where it's mapped to something else continuous instead (density,
  # or -- for a future custom builder -- a response-coloured data layer)
  if (!is.null(p$base)) {
    overlay_style <- object$layer$overlay$config$style
    fill_is_density <- identical(.style_fill_role(overlay_style), "density")
    ll <- names(ggplot2::get_labs(p$base))
    if (!is.null(color_discrete) && "colour" %in% ll) p$base <- p$base + color_discrete
    if (!is.null(fill_discrete) && !fill_is_density && "fill" %in% ll) p$base <- p$base + fill_discrete
    if (!is.null(fill_continuous) && fill_is_density && "fill" %in% ll) p$base <- p$base + fill_continuous
  }

  data_color_role <- object$layer$data$config$color_role %||% "strata"
  data_is_discrete <- identical(data_color_role, "strata")
  data_is_continuous <- identical(data_color_role, "response")
  if (data_is_discrete) {
    for (panel_name in names(p$data)) {
      ll <- names(ggplot2::get_labs(p$data[[panel_name]]))
      if (!is.null(color_discrete) && "colour" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + color_discrete
      if (!is.null(fill_discrete) && "fill" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + fill_discrete
    }
  }
  if (data_is_continuous) {
    for (panel_name in names(p$data)) {
      ll <- names(ggplot2::get_labs(p$data[[panel_name]]))
      if (!is.null(color_continuous) && "colour" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + color_continuous
      if (!is.null(fill_continuous) && "fill" %in% ll) p$data[[panel_name]] <- p$data[[panel_name]] + fill_continuous
    }
  }

  if (!is.null(p$group)) {
    for (g in names(p$group)) {
      ll <- names(ggplot2::get_labs(p$group[[g]]))
      if (!is.null(color_discrete) && "colour" %in% ll) p$group[[g]] <- p$group[[g]] + color_discrete
      if (!is.null(fill_discrete) && "fill" %in% ll) p$group[[g]] <- p$group[[g]] + fill_discrete
    }
  }

  return(p)
}

.polish_arrangement <- function(object) {

  plot_list <- list()
  plot_info <- tibble::tibble(
    id = integer(),
    size = numeric(),
    plot = character(),
    name = character()
  )
  ind <- 0L

  data_panels <- names(object$plot$data)
  panel_position <- object$layer$data$config$panel_position %||% character(0)
  above_panels <- data_panels[panel_position[data_panels] == "above"]
  below_panels <- data_panels[panel_position[data_panels] == "below"]

  # divide the data layer's total height budget evenly across however
  # many panels it has -- 2 for the binary upper/lower split (unchanged
  # from before), 1 for an unstratified continuous/count panel, or N for
  # an N-stratum continuous/count facet fallback
  data_panel_height <- object$theme$height$data / max(length(data_panels), 1)

  for (panel_name in above_panels) {
    ind <- ind + 1L
    plot_list[[ind]] <- object$plot$data[[panel_name]]
    plot_info <- plot_info |>
      tibble::add_row(
        id = ind,
        size = data_panel_height,
        plot = "data",
        name = paste0("data_", panel_name)
      )
  }

  # `object$plot$base` is only built when at least one of the model/
  # summary/quantile/overlay layers is present (see `er_plot_build()`)
  # -- a group-only or panel-layout-data-only plot has no base panel to
  # place, so it's omitted from the arrangement entirely rather than
  # inserting a `NULL` into `plot_list` (which `patchwork::wrap_plots()`
  # can't render).
  if (!is.null(object$plot$base)) {
    ind <- ind + 1L
    plot_list[[ind]] <- object$plot$base
    plot_info <- plot_info |>
      tibble::add_row(
        id = ind,
        size = object$theme$height$base,
        plot = "base",
        name = "base"
      )
  }

  for (panel_name in below_panels) {
    ind <- ind + 1L
    plot_list[[ind]] <- object$plot$data[[panel_name]]
    plot_info <- plot_info |>
      tibble::add_row(
        id = ind,
        size = data_panel_height,
        plot = "data",
        name = paste0("data_", panel_name)
      )
  }

  if (!is.null(object$plot$group)) {
    group_n <- purrr::map_dbl(object$layer$group$config, \(vv) vv$n_groups)
    group_prop <- group_n / sum(group_n)
    for(g in seq_along(object$plot$group)) {
      ind <- ind + 1L
      plot_list[[ind]] <- object$plot$group[[g]]
      plot_info <- plot_info |>
        tibble::add_row(
          id = ind,
          size = object$theme$height$group * group_prop[g],
          plot = "group",
          name = paste("group", g, sep = "_")
        )
    }
  }

  return(list(plots = plot_list, info = plot_info))
}

.polish_legends <- function(object, composition) {
  if (is.null(object$strata$name)) return(composition)
  has_strata <- purrr::map_lgl(object$layer, \(x) x$stratify %||% FALSE)

  # the data layer's `stratify` flag drives per-stratum faceting (not a
  # shared colour legend) whenever its colour channel is already spoken
  # for by the response value (`color_role == "response"`, continuous/
  # count response). Exclude it from strata-legend deduplication
  # in that case so each stratum panel keeps its own response colourbar.
  if (!is.null(object$layer$data) && identical(object$layer$data$config$color_role, "response")) {
    has_strata["data"] <- FALSE
  }

  if (!any(has_strata)) return(composition)
  stratified_parts <- names(has_strata[has_strata])
  # `model`/`summary`/`quantile`/an `"overlay"`-layout data builder all draw
  # into the single base panel (`composition$info`'s `"base"` row), not a
  # panel of their own -- so all four need mapping onto `"base"` here.
  # Missing one of these means `stratified_plots` can end up naming a plot
  # that isn't actually a row in `composition$info` (e.g. a plot with only
  # a stratified overlay data layer and nothing else stratified), leaving
  # `has_legend` empty and crashing the `for()` loop below on `2:0`.
  stratified_plots <- dplyr::case_when(
    stratified_parts == "quantile" ~ "base",
    stratified_parts == "model" ~ "base",
    stratified_parts == "summary" ~ "base",
    stratified_parts == "overlay" ~ "base",
    TRUE ~ stratified_parts
  )
  stratified_plots <- unique(stratified_plots)
  has_legend <- composition$info |>
    dplyr::filter(plot %in% stratified_plots) |>
    dplyr::pull(id)
  # Fewer than two legend-bearing plots means there's nothing to
  # deduplicate against (zero can happen if a future layer's `stratify`
  # flag is set but never mapped above; one means a single shared legend
  # already, nothing to strip).
  if (length(has_legend) <= 1L) return(composition)
  for(ind in has_legend[-1]) {
    composition$plots[[ind]] <- composition$plots[[ind]] +
      ggplot2::guides(
        color = ggplot2::guide_none(),
        fill = ggplot2::guide_none()
      )
  }
  return(composition)
}

.polish_theme <- function(object, composition) {
  for (ind in seq_along(composition$plots)) {
    composition$plots[[ind]] <- composition$plots[[ind]] + object$theme$theme_extra
  }
  return(composition)
}