R/stacked_dodged.R

In pik-piam/quitte: Bits and pieces of code to use with quitte-style data frames

#' Stacked and dodged bar plots
#'
#' Utility functions for plotting stacked (on top of each other) _and_ dodged
#' (next to each other) bars in the same figure.
#'
#' `add_stacked_dodged_xpos()` adds x-axis positions to a data frame for
#' plotting two categorical variables within a bar plot.
#' `calc_stacked_dodged_xlabels()` calculates matching label positions on
#' the x-axis.
#' `ggplot_bar_stacked_dodged()` uses both functions to generate a plot.
#'
#' @md
#' @param data A data frame
#' @param ... A selection of two columns.  Both will be combined to form x-axis
#'   coordinates.  The first will form the outer iteration (groups), the second
#'   the inner iteration (bars within a group).  If unnamed, the column with
#'   calculated positions will be called `xpos`.
#' @param mapping An aesthetic mapping generated by [ggplot2:aes][aes()],
#'   containing the aesthetics `x`, `y`, `fill`, and `dodge`.  The aesthetic
#'   `x` will form groups on the x-axis, while `dodge` will form individual
#'   bars within the groups.
#' @param gap The width of the gap between bars, relative to the width of the
#'   bars themselves (default: `1`).
#'
#' @return `add_stacked_dodged_xpos()` returns the input data frame with an
#'   additional column.  Row and column order are preserved.
#'   `calc_stacked_dodged_xlabels()` returns a named character vector for use
#'   with [ggplot2:scale_x_continuous][scale_x_continuous()].
#'   `ggplot_bar_stacked_dodged()` returns a [ggplot2:ggplot][ggplot()] object.
#'
#' @author Michaja Pehl
#'
#' @examples
#' require(tidyverse)
#'
#' set.seed(0)
#' (data <- crossing(a = factorise(c('left', 'center', 'right')),
#'                   b = factorise(c('top', 'middle', 'bottom')),
#'                   c = letters[1:4],
#'                   d = LETTERS[25:26]) %>%
#'         mutate(value = abs(rnorm(n())) + 0.2))
#'
#' (plot.data <- add_stacked_dodged_xpos(data, c('c', 'a')))
#'
#' (xlabels <- calc_stacked_dodged_xlabels(data, c('c', 'a')))
#'
#' ggplot(data = plot.data) +
#'     scale_x_continuous(breaks = xlabels) +
#'     facet_wrap(~ d, ncol = 1, scales = 'free_x')
#'
#' ggplot_bar_stacked_dodged(data, aes(x = a, y = value, fill = b, dodge = c),
#'                           gap = 1/3) +
#'     facet_wrap(~ d, ncol = 1, scales = 'free_x')

#' @importFrom dplyr mutate filter left_join
#' @importFrom ggplot2 aes geom_col ggplot scale_x_continuous
#' @importFrom tidyr crossing
#' @importFrom rlang get_expr sym

#' @rdname stacked_dodged
#' @export
ggplot_bar_stacked_dodged <- function(data, mapping, gap = 1) {
    if (!is.data.frame(data))
        stop('requires a data frame')

    x     <- get_expr(mapping$x)
    y     <- get_expr(mapping$y)
    fill  <- get_expr(mapping$fill)
    dodge <- get_expr(mapping$dodge)

    if (any(missing.mappings <- c('x' = is.null(x),
                                  'y' = is.null(y),
                                  'fill' = is.null(fill),
                                  'dodge' = is.null(dodge))))
        stop('ggplot_bar_stacked_dodged requires the following missing ',
             'aesthetics: ', paste(names(which(missing.mappings)),
                                   collapse = ', '))

    if (length(missing.columns <- setdiff(as.character(c(x, y, fill, dodge)),
                                          colnames(data))))
        stop('missing column', ifelse(1 == length(missing.columns), ' ', 's '),
             paste(missing.columns, collapse = ', '))

    xpos <- setNames(list(as.character(c(dodge, x))),
                     ifelse('xpos' %in% colnames(data), paste0(colnames(data)),
                            'xpos'))

    return(
        ggplot() +
            geom_col(
                data = add_stacked_dodged_xpos(data, xpos, gap = gap),
                mapping = aes(x = !!sym(names(xpos)),
                              y = !!sym(as.character(y)),
                              fill = !!sym(as.character(fill)))) +
            scale_x_continuous(
                breaks = calc_stacked_dodged_xlabels(data, xpos[[1]],
                                                     gap = gap)))
}

#' @rdname stacked_dodged
#' @export
add_stacked_dodged_xpos <- function(data, ..., gap = 1) {

    if (!is.data.frame(data))
        stop('requires a data frame')

    .dots <- list(...)
    if (is.list(.dots[[1]]))
        .dots <- .dots[[1]]

    # ignore all but the first unspecified argument
    xpos <- .dots[1]

    xpos.column <- ifelse(is.null(names(xpos)), 'xpos', names(xpos))

    if (xpos.column %in% colnames(data))
        stop(ifelse(is.null(names(xpos)), 'Default c', 'C'),
             'olumn name "', xpos.column, '" already in use.')

    if (anyNA(xpos[[1]][1:2]))
        stop('requires two columns to be combined')

    if (0 != length(missing.columns <- setdiff(xpos[[1]][1:2], colnames(data))))
        stop('column', ifelse(1 < length(missing.columns), 's ', ' '),
             paste(missing.columns, collapse = ', '), ' not found')

    group <- xpos[[1]][1]
    bar   <- xpos[[1]][2]

    # use the union of all elements within the <bar> column for the gap dummies
    # this way, there can't be any name clashes
    if (!is.factor(data[[bar]])) {
        bar.gap <- factor(paste(unique(data[[bar]]), collapse = ''),
                          levels = c(unique(data[[bar]]),
                                     paste(unique(data[[bar]]), collapse = '')))
    } else {
        # if the <bar> column is a factor, make sure the gap element is the
        # last level in the factor
        bar.gap <- paste(levels(data[[bar]]), collapse = '')
        bar.gap <- factor(bar.gap, levels = c(levels(data[[bar]]), bar.gap))
    }

    # generate all combinations of the <group> and <bar> columns, including
    # dummies for the gaps
    tmp <- crossing(!!sym(group) := factorise(unique_or_levels(data[[group]])),
                    !!sym(bar)   := factorise(levels(bar.gap))) %>%
        mutate(
            !!sym(xpos.column) :=
                # bars have width 1, gaps have width gap
                ifelse(bar.gap != !!sym(bar), 1, gap) %>%
                # sum everything up (first bar and shift to position 1 cancel)
                cumsum()) %>%
        # remove gap, they were only needed to count out positions
        filter(!!sym(bar) != bar.gap) %>%
        droplevels()

    if (!is.factor(data[[group]])) {
        tmp[[group]] <- as.character(tmp[[group]])
        if (is.numeric(data[[group]]))
            tmp[[group]] <- as.numeric(tmp[[group]])
    }
    if (!is.factor(data[[bar]])) {
        tmp[[bar]] <- as.character(tmp[[bar]])
        if (is.numeric(data[[bar]]))
            tmp[[bar]] <- as.numeric(tmp[[bar]])
    }

    left_join(data, tmp, c(group, bar)) %>%
        return()
}

#' @rdname stacked_dodged
#' @export
calc_stacked_dodged_xlabels <- function(data, ..., gap = 1) {
    if (!is.data.frame(data))
        stop('requires a data frame')

    # capture unspecified arguments
    xpos <- list(..1)
    if (is.list(xpos[[1]]))
        xpos <- xpos[[1]]

    if (anyNA(xpos[[1]][1:2]))
        stop('requires two columns to be combined')

    n.groups <- length(unique_or_levels(data[[xpos[[1]][1]]]))
    n.bars   <- length(unique_or_levels(data[[xpos[[1]][2]]]))

    return(setNames(
        seq(from = (n.bars + 1) / 2,
            length.out = n.groups,
            by = n.bars + gap),
        unique_or_levels(data[[xpos[[1]][1]]])))
}