guide_axis_combmatrix.R

library(grid)
library(scales)
library(gtable)
library(rlang)

# Constructor -------------------------------------------------------------

#' Combination matrix axis
#'
#' @inheritParams guide_axis
#' @param sep A `character(1)` to split label strings.
#' @param levels A `character()` with the order of labels.
#' @param zebra_even,zebra_odd An `element_rect` object to determine the
#'   look of alternating bands.
#' @param connector An `element_line` object to determine the look of the
#'   lines connecting positive points.
#' @param pos_shape A symbol to use for positive points.
#' @param neg_shape A symbol to use for negative points.
#' @param pos_gp A `gpar` object giving graphical parameters for positive
#'   points.
#' @param neg_gp A `gpar` object giving graphical parameters for negative
#'   points.
#'
#' @return A `guide` object that can be given to a position scale or the
#'   `guides()` functions.
#' @export
#' @md
#'
#' @examples
#' df <- data.frame(
#'   x = c("A,B,C", "A,B", "A,C", "C", ""),
#'   y = 1:5
#' )
#'
#' ggplot(df, aes(x, y)) +
#'   geom_col() +
#'   guides(x = guide_axis_combmatrix())
guide_axis_combmatrix <- function(
  title = waiver(),
  check.overlap = FALSE,
  angle = NULL,
  n.dodge = 1,
  order = 0,
  position = waiver(),
  sep = "[^[:alnum:]]+",
  levels = NULL,
  zebra_even = element_rect(),
  zebra_odd  = element_rect(),
  connector  = element_line(),
  pos_shape = 19,
  neg_shape = 19,
  pos_gp = gpar(col = "black"),
  neg_gp = gpar(col = "#D0D0D0")
) {
  # Validate some inputs
  if (!(inherits(zebra_even, c("element_rect", "element_blank")))) {
    stop("The `zebra_even` argument is expected to be an 'element_rect'.")
  }
  if (!(inherits(zebra_odd,  c("element_rect", "element_blank")))) {
    stop("The `zebra_odd`  argument is expected to be an 'element_rect'")
  }
  if (!(inherits(connector, c("element_line", "element_blank")))) {
    stop("The `connector` argument is expected to be an 'element_line'")
  }

  # Make attempt to match connector colour to point colour if left unspecified
  if (!inherits(connector, "element_blank") && is.null(connector$colour)) {
    if (!is.null(pos_gp$col)) {
      connector$colour <- pos_gp$col[1]
    }
  }

  structure(list(
    title = title,
    check.overlap = check.overlap,
    angle = angle,
    n.dodge = n.dodge,
    order = order,
    position = position,
    available_aes = c("x", "y"),
    name = "axis",
    sep = sep,
    levels = levels,
    zebra_even = zebra_even,
    zebra_odd  = zebra_odd,
    connector  = connector,
    pos_shape = pos_shape,
    neg_shape = neg_shape,
    pos_gp = pos_gp,
    neg_gp = neg_gp
  ), class = c("guide", "axis_combmatrix", "axis"))
}

# Methods -----------------------------------------------------------------

#' @export
guide_train.axis_combmatrix <- function(guide, scale, aesthetic = NULL) {
  guide  <- NextMethod()
  labs <- strsplit(guide$key$.label, guide$sep)
  uniq <- Reduce(union, labs)
  if (!is.null(guide$levels)) {
    uniq <- intersect(guide$levels, uniq)
  }
  if (length(uniq) == 0) {
    warning("No appropriate labels have been found.")
  }
  guide$levels <- rev(uniq)
  guide
}

#' @export
guide_gengrob.axis_combmatrix <- function(guide, theme) {
  key <- guide$key
  aes <- names(key)[!grepl("^\\.", names(key))][1]
  position  <- match.arg(guide$position, c("top", "bottom", "right", "left"))

  # Calculate theme elements
  line_elem  <- calc_element(paste("axis.line",  aes, position, sep = "."), theme)
  text_elem  <- calc_element(paste("axis.text",  aes, position, sep = "."), theme)
  tick_len   <- calc_element(paste("axis.ticks.length", aes, position, sep = "."), theme)
  connector  <- calc_element("line", theme)
  zebra_even <- calc_element("plot.background",  theme)
  zebra_odd  <- calc_element("panel.background", theme)

  # Inheritance of specific elements
  zebra_even <- inherit_element(guide$zebra_even, zebra_even)
  zebra_odd  <- inherit_element(guide$zebra_odd,  zebra_odd)
  connector  <- inherit_element(guide$connector,  connector)

  # Orientation parameters
  is_vertical <- position %in% c("left", "right")
  is_second   <- position %in% c("top",  'right')
  init_gtable <- if (is_vertical) gtable_row else gtable_col
  alt_aes  <- setdiff(c("x", "y"), aes)
  size     <- switch(aes, x = "width",  y = "height")
  alt_size <- switch(aes, x = "height", y = "width")

  # Draw axis line
  if (is_second) {
    line <- exec(element_grob, line_elem,
                 !!aes := unit(c(0, 1), "npc"),
                 !!alt_aes := unit(c(0, 0), "npc"))
  } else {
    line <- exec(element_grob, line_elem,
                 !!aes := unit(c(0, 1), "npc"),
                 !!alt_aes := unit(c(1, 1), "npc"))
  }
  if (nrow(key) == 0) {
    return(
      ggplot2:::absoluteGrob(
        gList(line), width = grobWidth(line), height = grobHeight(line)
      )
    )
  }

  # Digest labels
  labels <- strsplit(key$.label, guide$sep)
  levels <- guide$levels
  n_row <- length(levels); n_col <- length(labels)
  alt_breaks <- rescale(seq_along(levels), from = c(0.5, n_row + 0.5))

  # Build combination matrix
  combs <- vapply(labels, function(x) {levels %in% x}, logical(n_row))
  pos <- as.vector(combs)

  # Build points
  pts_main <- key[[aes]][as.vector(col(combs))]
  pts_alt  <- alt_breaks[as.vector(row(combs))]
  points <- list(
    exec(pointsGrob, !!aes := pts_main[pos],  !!alt_aes := pts_alt[pos],
         pch = guide$pos_shape, gp = guide$pos_gp),
    exec(pointsGrob, !!aes := pts_main[!pos], !!alt_aes := pts_alt[!pos],
         pch = guide$neg_shape, gp = guide$neg_gp)
  )

  # Connector lines
  index <- ifelse(combs, seq_len(length(combs)), NA)
  keep  <- colSums(index, na.rm = TRUE) > 1
  index <- apply(index[, keep, drop = FALSE], 2, range, na.rm = TRUE)
  index <- index[, index[1,] != index[2,]]
  connector <- exec(
    element_grob, connector,
    !!aes     := pts_main[as.vector(index)],
    !!alt_aes := pts_alt[as.vector(index)],
    id.lengths = rep(nrow(index), ncol(index))
  )

  # Draw text
  titles <- ggplot2:::draw_axis_labels(
    break_positions = alt_breaks,
    break_labels  = levels,
    label_element = text_elem,
    is_vertical   = !is_vertical,
    check.overlap = guide$check.overlap
  )

  # Adjust text viewports
  if (is_vertical) {
    title_height <- ggplot2:::height_cm(titles)
    titles$vp$parent$y <- unit(-0.5 * title_height, "cm") - tick_len
    titles$vp$parent$height <- unit(title_height, "cm")
  } else {
    title_width  <- ggplot2:::width_cm(titles)
    titles$vp$parent$x <- unit(-0.5 * title_width, "cm") - tick_len
    titles$vp$parent$width <- unit(title_width, "cm")
  }

  # Draw zebra stripes
  zebra_size <- rescale(1, from = c(0, n_row))
  even <- alt_breaks[seq_along(levels) %% 2 == 0]
  odd  <- alt_breaks[seq_along(levels) %% 2 != 0]
  zebra_even <- exec(element_grob, zebra_even,
                     !!alt_aes := even, !!alt_size := zebra_size)
  zebra_odd  <- exec(element_grob, zebra_odd,
                     !!alt_aes := even, !!alt_size := zebra_size)

  # Assemble output
  if (is_vertical) {
    axis_size <- unit(text_elem$size, "pt") +
      text_elem$margin[2] + text_elem$margin[4]
  } else {
    axis_size <- unit(text_elem$size, "pt") +
      text_elem$margin[1] + text_elem$margin[3]
  }


  gt <- exec(init_gtable, "axis", list(line),
             !!size := unit.c(unit(1, "npc")),
             !!alt_size := axis_size * 1.2 * (n_row + 1))
  gt <- gtable_add_grob(
    gt,
    list(zebra_even, zebra_odd, points[[1]], points[[2]], connector, titles),
    t = 1, l = 1, b = 1, r = 1, clip = c(rep("on", 5), "off"),
    name = c("zebra_even", "zebra_odd", "pos_points",
             "neg_points", "line", "labels"),
    z = c(1, 1, 4, 2, 3, 5)
  )
  gTree(
    children = gList(gt),
    width  = gtable_width(gt),
    height = gtable_height(gt),
    cl = "absoluteGrob"
  )
}

# Helpers -----------------------------------------------------------------

# Based on ggplot2:::combine_elements
inherit_element <- function(child, parent) {

  if (is.null(parent) || inherits(child, "element_blank")) {
    return(child)
  }

  if (is.null(child)) {
    return(parent)
  }

  if (!inherits(child, "element") && !inherits(parent, "element")) {
    return(child)
  }

  if (inherits(parent, "element_blank")) {
    if (child$inherit.blank) {
      return(parent)
    } else {
      return(child)
    }
  }

  n <- names(child)[vapply(child, is.null, logical(1))]
  child[n] <- parent[n]

  if (inherits(child$size, "rel")) {
    child$size <- parent$size * unclass(child$size)
  }
  return(child)
}