R/gglasagna.R

In tidyfun: Tidy Functional Data Wrangling and Visualization

# from https://github.com/r-lib/scales/blob/master/R/scale-discrete.r
is.discrete <- function(x) {
  is.factor(x) || is.character(x) || is.logical(x)
}

#' Lasagna plots for `tf`s using `ggplot2`
#'
#' Lasagna plots show one color bar for each function.
#'
#' The vertical order of the lasagna layers is **increasing** in
#'
#' - `order` (if provided),
#' - the values returned by `order_by` (if provided),
#' - and the row number of the observations.
#'
#' i.e., lowest values are on top so that by default the first layer is the first
#' observation in `data` and the vertical order of the layers is the
#' ordering of observations obtained by `dplyr::arrange(data, order, order_by(value), row_number())`.
#'
#' @param data A data frame containing the `tf` column to visualize.
#' @param tf bare name of the `tf` column to visualize
#' @param order (optional) bare name of a column in `data` to define vertical
#'   order of lasagna layers.
#' @param label (optional) bare name of a column in `data` to define labels for
#'   lasagna layers. Defaults to names of `y`, if present, or row numbers.
#' @param arg `arg` to evaluate `y` on
#' @param order_by a function applied to each row in `y[, arg]` that must
#'   return a scalar value to define the order of lasagna layers.
#' @param order_ticks add horizontal lines indicating borders between levels of
#'   `order` (if it is a discrete variable) and labels for its levels?
#'   Defaults to TRUE. Supply a named list to override tick appearance, including
#'   label styling, line type, alpha, rotation, and label placement.
#'   Disable this when faceting; the tick annotations are not designed for faceted layouts.
#' @returns a `ggplot2` object
#' @export
#' @importFrom grid unit grobTree textGrob gpar
#' @family tidyfun visualization
#' @examples
#' \donttest{
#' library(ggplot2)
#' set.seed(1221)
#' data <- expand.grid(group = factor(1:5), rep = 1:10)
#' data <- dplyr::mutate(data,
#'   id = paste(group, rep, sep = "-"),
#'   f = tf_rgp(50),
#'   fb = tfb(f)
#' )
#'
#' gglasagna(data, f, label = id)
#' gglasagna(data, fb, label = id, order = group)
#' # order is lowest first / on top by default
#' gglasagna(data, f, label = id, order = tf_depth(f))
#' gglasagna(data, f, label = id, order_by = dplyr::first) +
#'   facet_wrap(~group, scales = "free")
#' # order of layers is by "order_by" within "order":
#' gglasagna(data, fb, label = id, order = group, order_by = dplyr::first)
#' }
gglasagna <- function(
  data,
  tf,
  order = NULL,
  label = NULL,
  arg = NULL,
  order_by = NULL,
  order_ticks = TRUE
) {
  order_ticks_args <- list(
    color = "black",
    linetype = 2,
    alpha = 0.5,
    fontsize = theme_get()$text$size * 0.8,
    rot = 90,
    label_offset = grid::unit(0.02, "npc")
  )
  if (is.list(order_ticks)) {
    order_ticks_args <- utils::modifyList(order_ticks_args, order_ticks)
    order_ticks <- TRUE
  }

  # FIXME: render errors for weird arg lengths (e.g. 93)
  if (!is_tf(pull(data, {{ tf }}))) {
    cli::cli_abort(
      "{.arg tf} must be a {.cls tf} object, not {.obj_type_friendly {pull(data, {{ tf }})}}."
    )
  }
  order <- enexpr(order)
  has_order <- !is.null(order)
  order_by_label <- enexpr(order_by)
  has_order_by <- !is.null(order_by_label)
  if (has_order) {
    order_label <- as_label(order)
  } else {
    order_label <- NULL
    order <- expr(..row)
    order_ticks <- FALSE
  }
  label <- enexpr(label)
  has_label <- !is.null(label)
  if (!has_label) {
    label <- expr(names(!!enexpr(tf)) %||% row_number())
    labelname <- ""
  } else {
    labelname <- as_label(label)
  }
  y_name <- as_name(enexpr(tf))
  data <- mutate(
    data,
    ..label = !!label,
    ..row = row_number(),
    ..order = !!order
  )
  tf_eval <- data |>
    # TODO: add .preserve for all list columns not being plotted
    mutate(..y = !!label) |> # vertical position variable
    tf_unnest(y_name, .arg = arg, names_sep = "___", try_dropping = FALSE) |>
    rename(..x = matches("___arg"), ..fill = matches("___value"))
  if (has_order_by) {
    order_by_label <- as_label(order_by_label)
    if (!is.function(order_by)) {
      cli::cli_abort(
        "{.arg order_by} must be a function, not {.obj_type_friendly {order_by}}."
      )
    }
    order_by_value <- tf_eval |>
      group_by(..y) |>
      summarise(..order_by_value = order_by(..fill)) |>
      ungroup() |>
      mutate(..order_by_value = rank(..order_by_value))
    tf_eval <- left_join(tf_eval, order_by_value, by = "..y")
    # override previous ordering
    if (!has_order) {
      tf_eval <- tf_eval |> mutate(..order = ..order_by_value)
    }
  } else {
    order_by_label <- NULL
    tf_eval <- tf_eval |> mutate(..order_by_value = ..row)
  }
  # create order of rows
  tf_eval <- tf_eval |>
    arrange(..order, ..order_by_value, ..row) |>
    mutate(..y = ordered(..y, levels = rev(unique(..y)))) |>
    # mutate(..y = as.numeric(..y)) |>
    rename(!!y_name := ..fill)
  labeldata <- with(
    tf_eval,
    tibble(
      breaks = unique(..y),
      labels = ..label[!duplicated(..y)]
    )
  )

  p <- ggplot(tf_eval) +
    geom_tile(aes(
      y = ..y,
      x = ..x,
      fill = .data[[y_name]],
      colour = .data[[y_name]]
    )) +
    labs(x = "", y = "")
  if (!is.null(order_label) || !is.null(order_by_label)) {
    p <- p +
      labs(
        caption = paste(
          "ordered by:",
          paste0(
            order_label,
            ifelse(has_order_by & has_order, "; ", ""),
            order_by_label
          )
        )
      )
  }
  if (!is.null(order_ticks) && is.discrete(pull(tf_eval, ..order))) {
    order_ticks_data <- data |>
      arrange(desc(!!order)) |>
      mutate(ticks = row_number()) |>
      group_by(!!order) |>
      summarize(tick_hi = max(ticks), tick_lo = min(ticks)) |>
      mutate(
        label_pos = (tick_hi + tick_lo) / 2,
        tick_pos = lead(tick_hi) + 0.5
      )
    p <- p +
      geom_hline(
        data = order_ticks_data,
        aes(yintercept = tick_pos),
        col = order_ticks_args$color,
        alpha = order_ticks_args$alpha,
        linetype = order_ticks_args$linetype
      )
    # adding a secondary axis with labels for the order-variable only works
    # if we make the y-axis continuous -- but a) it isn't,really and
    # b)scales="free" is needed for good looking facets but won't
    # close gaps for pseudo-continuous y-scales,
    # so we use this annotation_custom weirdness instead
    order_ticks_labelpos <- with(
      order_ticks_data,
      grid::unit(label_pos / max(tick_hi), "npc")
    )
    order_ticks_labels <- grid::grobTree(
      grid::textGrob(
        order_ticks_data |> pull(!!order),
        x = order_ticks_args$label_offset,
        y = order_ticks_labelpos,
        rot = order_ticks_args$rot,
        gp = grid::gpar(fontsize = order_ticks_args$fontsize)
      )
    )
    p <- p + annotation_custom(order_ticks_labels)
  }
  p <- p +
    scale_y_discrete(
      labelname,
      breaks = labeldata$breaks,
      labels = labeldata$labels
    )
  p
}