R/areaplot.R

In plotthis: High-Level Plotting Built Upon 'ggplot2' and Other Plotting Packages

#' Atomic area plot (internal)
#'
#' @description
#' Core implementation for drawing a single stacked area plot.  This is the
#' workhorse behind the exported \code{\link{AreaPlot}} function — it takes a
#' **single** data frame (no \code{split_by} support) and returns a
#' \code{ggplot} object.  The plot shows how one or more groups' numeric values
#' (or counts) accumulate across a discrete x-axis, with each group rendered
#' as a filled area stacked from baseline.
#'
#' @section Architecture:
#' \enumerate{
#'   \item \strong{Column resolution} — \code{x}, \code{y}, \code{group_by},
#'         and \code{facet_by} are validated and transformed via
#'         \code{\link{check_columns}}.  Multi-column inputs for \code{x} and
#'         \code{group_by} are concatenated into single columns using their
#'         respective separators (\code{x_sep}, \code{group_by_sep}).
#'   \item \strong{NA / empty-level handling} — \code{\link{process_keep_na_empty}()}
#'         applies \code{keep_na} and \code{keep_empty} policies.  Per-column
#'         \code{keep_empty} settings are extracted for \code{x},
#'         \code{group_by}, and \code{facet_by} independently.  The facet
#'         columns must have identical \code{keep_empty} values.
#'   \item \strong{Count aggregation} — when \code{y = NULL}, the count of
#'         observations in each unique (\code{x}, \code{group_by},
#'         \code{facet_by}) combination is computed as a new \code{.count}
#'         column.  Factor levels are preserved after aggregation.
#'   \item \strong{Proportion scaling} — when \code{scale_y = TRUE}, the
#'         y-values are divided by the sum within each (\code{x},
#'         \code{facet_by}) group, producing a proportion (0–1).  Percent
#'         labels are used automatically on the y-axis.
#'   \item \strong{Empty-fill guard} — when \code{group_by = NULL} (no
#'         grouping), a dummy \code{.fill} factor is created so the single
#'         area still draws with the first palette colour.  The legend is
#'         suppressed (\code{legend.position = "none"}).
#'   \item \strong{Colour mapping} — \code{\link{palette_this}()} assigns
#'         colours to all \code{group_by} levels, including \code{NA}
#'         (defaulting to \code{"grey80"}).
#'   \item \strong{Data completion} — \code{\link[tidyr]{complete}()} pads
#'         all \code{x} × \code{group_by} (× \code{facet_by}) combinations
#'         with \code{y = 0}.  This prevents \code{\link[ggplot2]{geom_area}()}
#'         from interpolating across missing groups, which would otherwise
#'         cause stacked areas to exceed the correct total.
#'   \item \strong{x-axis numeric mapping} — the discrete x variable is
#'         converted to a numeric position column (\code{.x_numeric}) so
#'         \code{geom_area()} can draw continuous area fills between x
#'         positions.  \code{NA} levels are placed at position
#'         \code{n_levels + 1}.
#'   \item \strong{Plot assembly} — the \code{ggplot} object is built with
#'         \code{geom_area(position = position_stack(vjust = 0.5))},
#'         \code{scale_x_discrete()} (breaks from factor levels),
#'         \code{scale_y_continuous()} (percent labels when scaled), and
#'         \code{scale_fill_manual()} with per-group colours.  The fill
#'         scale \code{drop} argument is controlled by
#'         \code{keep_empty_group}.
#'   \item \strong{Dimension calculation} — \code{\link{calculate_plot_dimensions}()}
#'         computes plot height and width from the x-axis category count,
#'         \code{aspect.ratio}, and legend metrics.  The resulting
#'         \code{height} / \code{width} attributes are stored on the
#'         \code{ggplot} object.
#'   \item \strong{Faceting} — \code{\link{facet_plot}()} wraps the plot
#'         with \code{facet_wrap} / \code{facet_grid} if \code{facet_by} is
#'         provided, respecting the \code{keep_empty} setting for facet
#'         variables.
#' }
#'
#' @inheritParams common_args
#' @param x A character string specifying the column name to plot on the
#'  x-axis.  Must be character or factor.  Multiple columns can be provided;
#'  they are concatenated with \code{x_sep} as the separator.
#' @param x_sep A character string used to join multiple \code{x} columns.
#'  Default \code{"_"}.  Ignored when \code{x} is a single column.
#' @param y A character string specifying the numeric column for the y-axis.
#'  When \code{NULL}, the count of observations in each (\code{x},
#'  \code{group_by}, \code{facet_by}) combination is used.
#' @param scale_y A logical value.  When \code{TRUE}, y-values are scaled to
#'  proportions within each (\code{x}, \code{facet_by}) group so that each
#'  x position stacks to 1.0.  The y-axis labels switch from numeric to
#'  percent format automatically.
#' @param group_by A character vector of column names to fill the areas by.
#'  Each unique combination becomes a separate stacked area.  Multiple
#'  columns are concatenated with \code{group_by_sep}.  When \code{NULL},
#'  a single filled area is drawn (no grouping) and the legend is hidden.
#' @param group_by_sep A character string to separate concatenated
#'  \code{group_by} columns.  Default \code{"_"}.
#' @param group_name A character string used as the fill legend title.
#'  When \code{NULL}, the \code{group_by} column name is used.
#' @keywords internal
#' @return A \code{ggplot} object with \code{height} and \code{width}
#'  attributes (in inches) attached.
#' @importFrom rlang syms :=
#' @importFrom dplyr summarise n %>%
#' @importFrom ggplot2 geom_area scale_x_discrete scale_y_continuous scale_fill_manual
#' @importFrom ggplot2 labs theme element_line element_text position_stack waiver
AreaPlotAtomic <- function(
    data,
    x,
    y = NULL,
    x_sep = "_",
    group_by = NULL,
    group_by_sep = "_",
    group_name = NULL,
    scale_y = FALSE,
    theme = "theme_this",
    theme_args = list(),
    palette = "Paired",
    palcolor = NULL,
    palreverse = FALSE,
    alpha = 1,
    facet_by = NULL,
    facet_scales = "fixed",
    facet_ncol = NULL,
    facet_nrow = NULL,
    facet_byrow = TRUE,
    x_text_angle = 0,
    aspect.ratio = 1,
    legend.position = waiver(),
    legend.direction = "vertical",
    title = NULL,
    subtitle = NULL,
    xlab = NULL,
    ylab = NULL,
    keep_na = FALSE,
    keep_empty = FALSE,
    ...
) {
    ggplot <- if (getOption("plotthis.gglogger.enabled", FALSE)) {
        gglogger::ggplot
    } else {
        ggplot2::ggplot
    }
    x <- check_columns(
        data,
        x,
        force_factor = TRUE,
        allow_multi = TRUE,
        concat_multi = TRUE,
        concat_sep = x_sep
    )
    y <- check_columns(data, y)
    group_by <- check_columns(
        data,
        group_by,
        force_factor = TRUE,
        allow_multi = TRUE,
        concat_multi = TRUE,
        concat_sep = group_by_sep
    )
    facet_by <- check_columns(
        data,
        facet_by,
        force_factor = TRUE,
        allow_multi = TRUE
    )
    data <- process_keep_na_empty(data, keep_na, keep_empty)
    # TRUE: unused levels will be kept on X axis
    # FALSE/level: unused levels will be dropped
    keep_empty_x <- keep_empty[[x]]
    # TRUE: unused levels will be kept in group_by
    # FALSE: unused levels will be dropped
    # level: unused levels will be dropped, but the group colors will be identified using all levels
    keep_empty_group <- if (!is.null(group_by)) keep_empty[[group_by]] else NULL
    # TRUE: unused levels will be kept in facet_by
    # FALSE/level: unused levels will be dropped
    # 2-column facet_by is not supported yet
    keep_empty_facet <- if (!is.null(facet_by)) {
        keep_empty[[facet_by[1]]]
    } else {
        NULL
    }
    if (length(facet_by) > 1) {
        stopifnot(
            "[AreaPlot] `keep_empty` for `facet_by` variables must be identical." = identical(
                keep_empty_facet,
                keep_empty[[facet_by[2]]]
            )
        )
    }

    orig_data <- data
    if (is.null(y)) {
        y <- ".count"
        data <- data %>%
            dplyr::group_by(!!!syms(unique(c(x, group_by, facet_by)))) %>%
            summarise(.count = n(), .groups = "drop")
        # keep the levels
        for (col in unique(c(x, group_by, facet_by))) {
            data[[col]] <- factor(
                data[[col]],
                levels = levels(orig_data[[col]])
            )
        }
    }

    if (isTRUE(scale_y)) {
        data <- data %>%
            dplyr::group_by(!!!syms(unique(c(x, facet_by)))) %>%
            mutate(!!sym(y) := !!sym(y) / sum(!!sym(y))) %>%
            ungroup()
        # keep the levels
        for (col in unique(c(x, group_by, facet_by))) {
            data[[col]] <- factor(
                data[[col]],
                levels = levels(orig_data[[col]])
            )
        }
    }
    rm(orig_data)

    if (is.null(group_by)) {
        data$.fill <- factor("")
        group_by <- ".fill"
        legend.position <- ifelse(
            inherits(legend.position, "waiver"),
            "none",
            "right"
        )
    } else {
        legend.position <- ifelse(
            inherits(legend.position, "waiver"),
            "right",
            legend.position
        )
    }
    if (!isTRUE(keep_empty_x)) {
        data[[x]] <- droplevels(data[[x]])
    }

    x_vals <- levels(data[[x]])
    if (anyNA(data[[x]])) {
        x_vals <- c(x_vals, NA)
    }

    # group_by_vals <- if (!isTRUE(keep_empty_group) || !anyNA(data[[group_by]])) {
    #     levels(data[[group_by]])
    # } else {
    #     c(levels(data[[group_by]]), NA)
    # }
    group_by_vals <- levels(data[[group_by]])
    if (anyNA(data[[group_by]])) {
        group_by_vals <- c(group_by_vals, NA)
    }
    group_colors <- palette_this(
        group_by_vals,
        palette = palette,
        palcolor = palcolor,
        NA_keep = TRUE,
        reverse = palreverse
    )

    just <- calc_just(x_text_angle)

    # Complete all x * group_by (and facet_by) combinations so that geom_area
    # doesn't interpolate across missing groups, which would cause stacked
    # areas to exceed the correct total.
    complete_vars <- unique(c(x, group_by, facet_by))
    complete_fill <- setNames(list(0), y)
    data_complete <- data %>%
        tidyr::complete(!!!syms(complete_vars), fill = complete_fill)
    # Restore factor levels that complete() may have altered
    for (col in complete_vars) {
        if (is.factor(data[[col]])) {
            data_complete[[col]] <- factor(
                data_complete[[col]],
                levels = levels(data[[col]])
            )
        }
    }

    # Convert x to numeric, handling NA values by assigning them the next position
    data_complete$.x_numeric <- as.numeric(data_complete[[x]])
    if (anyNA(data_complete[[x]])) {
        data_complete$.x_numeric[is.na(data_complete[[
            x
        ]])] <- length(levels(data_complete[[x]])) + 1
    }

    p <- ggplot(
        data_complete,
        aes(x = !!sym(".x_numeric"), y = !!sym(y), fill = !!sym(group_by))
    ) +
        geom_area(
            alpha = alpha,
            color = "grey50",
            position = position_stack(vjust = 0.5),
            show.legend = TRUE
        ) +
        scale_x_discrete(
            expand = c(0, 0),
            breaks = x_vals,
            limits = x_vals,
            drop = isFALSE(keep_empty_x)
        ) +
        scale_y_continuous(
            expand = c(0, 0),
            labels = if (isFALSE(scale_y)) scales::number else scales::percent
        ) +
        labs(
            title = title,
            subtitle = subtitle,
            x = xlab %||% x,
            y = ylab %||% y
        ) +
        do_call(theme, theme_args) +
        ggplot2::theme(
            aspect.ratio = aspect.ratio,
            legend.position = legend.position,
            legend.direction = legend.direction,
            panel.grid.major = element_line(colour = "grey80", linetype = 2),
            axis.text.x = element_text(
                angle = x_text_angle,
                hjust = just$h,
                vjust = just$v
            )
        )

    if (isTRUE(keep_empty_group)) {
        p <- p +
            scale_fill_manual(
                name = group_name %||% group_by,
                values = group_colors,
                na.value = group_colors["NA"] %||% "grey80",
                breaks = group_by_vals,
                limits = group_by_vals,
                drop = FALSE
            )
    } else {
        p <- p +
            scale_fill_manual(
                name = group_name %||% group_by,
                values = group_colors,
                na.value = group_colors["NA"] %||% "grey80"
            )
    }

    # Calculate plot dimensions with aspect ratio consideration
    dims <- calculate_plot_dimensions(
        base_height = 4.5,
        aspect.ratio = aspect.ratio,
        n_x = length(x_vals),
        x_scale_factor = 0.5,
        legend.position = legend.position,
        legend.direction = legend.direction,
        legend_n = length(group_by_vals),
        legend_nchar = max(nchar(as.character(group_by_vals)), na.rm = TRUE),
        flip = FALSE
    )

    attr(p, "height") <- dims$height
    attr(p, "width") <- dims$width

    facet_plot(
        p,
        facet_by,
        facet_scales,
        facet_nrow,
        facet_ncol,
        facet_byrow,
        legend.position = legend.position,
        legend.direction = legend.direction,
        drop = !isTRUE(keep_empty_facet)
    )
}

#' Area plot
#'
#' @description
#' Draws a stacked area plot showing how one or more groups' numeric values
#' (or counts) accumulate across the progression of a discrete x-axis
#' variable.  Each group is rendered as a filled area stacked from baseline,
#' making it easy to compare both individual magnitudes and the total across
#' categories.
#'
#' The function supports \strong{count aggregation} (omit \code{y} to plot
#' observation counts per x-category), \strong{proportion scaling}
#' (\code{scale_y = TRUE} normalises each x position to 100\%), per-group
#' colour control, faceting, and splitting into separate sub-plots via
#' \code{split_by}.
#'
#' @section split_by workflow:
#' When \code{split_by} is provided:
#' \enumerate{
#'   \item \code{\link{check_keep_na}()} and \code{\link{check_keep_empty}()}
#'         normalise the \code{keep_na} / \code{keep_empty} arguments for all
#'         columns (\code{x}, \code{split_by}, \code{group_by}, \code{facet_by}).
#'   \item The \code{split_by} column is validated and its NA / empty levels
#'         are processed via \code{\link{process_keep_na_empty}()}.  It is
#'         then removed from the per-column \code{keep_na} / \code{keep_empty}
#'         lists.
#'   \item The data frame is split by \code{split_by} (preserving level
#'         order).  If \code{split_by} is \code{NULL}, the data is wrapped in
#'         a single-element list with name \code{"..."}.
#'   \item Per-split \code{palette}, \code{palcolor},
#'         \code{legend.position}, and \code{legend.direction} are resolved
#'         via \code{\link{check_palette}()}, \code{\link{check_palcolor}()},
#'         and \code{\link{check_legend}()}.
#'   \item \code{\link{AreaPlotAtomic}()} is called for each split.  If
#'         \code{title} is a function, it receives the split level name and
#'         can generate dynamic titles.
#'   \item Results are combined via \code{\link{combine_plots}()} (when
#'         \code{combine = TRUE}) or returned as a named list.
#' }
#'
#' @inheritParams common_args
#' @inheritParams AreaPlotAtomic
#' @param split_by The column(s) to split the data by and produce separate
#'  sub-plots.  Multiple columns are concatenated with \code{split_by_sep}.
#' @param split_by_sep A character string to separate concatenated
#'  \code{split_by} columns.  Default \code{"_"}.
#' @param seed A numeric seed for reproducibility.  Passed to
#'  \code{\link{validate_common_args}()}.
#' @param combine Logical; when \code{TRUE} (default), returns a combined
#'  \code{patchwork} object.  When \code{FALSE}, returns a named list of
#'  individual \code{ggplot} objects.
#' @param ncol,nrow Integer number of columns / rows for the combined layout
#'  (passed to \code{\link[patchwork]{wrap_plots}}).
#' @param byrow Logical; fill the combined layout by row.  Default \code{TRUE}
#'  (passed to \code{\link[patchwork]{wrap_plots}}).
#' @param axes A character string specifying how axes should be treated across
#'  the combined layout (passed to \code{\link[patchwork]{wrap_plots}}).
#' @param axis_titles A character string specifying how axis titles should be
#'  treated across the combined layout.  Defaults to \code{axes}.
#' @param guides A character string specifying how guides (legends) should be
#'  collected across panels.  Default \code{"collect"} (passed to
#'  \code{\link{combine_plots}()}).
#' @param design A custom layout design for the combined plot (passed to
#'  \code{\link{combine_plots}()}).
#' @return A \code{ggplot} object, a \code{patchwork} object, or a named list
#'  of \code{ggplot} objects (when \code{combine = FALSE}), each with
#'  \code{height} and \code{width} attributes in inches.
#' @importFrom ggplot2 waiver
#' @export
#' @examples
#' \donttest{
#' set.seed(8525)
#' data <- data.frame(
#'     x = rep(c("A", "B", "C", "D"), 2),
#'     y = c(1, 3, 6, 4, 2, 5, 7, 8),
#'     group = rep(c("F1", "F2"), each = 4),
#'     split = rep(c("X", "Y"), 4)
#' )
#' # Basic stacked area
#' AreaPlot(data, x = "x", y = "y", group_by = "group")
#'
#' # Scaled to proportions
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          scale_y = TRUE)
#'
#' # Split into sub-plots (no group_by — single-colour fill)
#' AreaPlot(data, x = "x", y = "y", split_by = "group")
#'
#' # Per-split palettes
#' AreaPlot(data, x = "x", y = "y", split_by = "group",
#'          palette = c(F1 = "Blues", F2 = "Reds"))
#'
#' # Per-split legend positioning
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          split_by = "split",
#'          legend.direction = c(X = "horizontal", Y = "vertical"),
#'          legend.position = c(X = "top", Y = "right"))
#'
#' # How keep_na and keep_empty work
#' data <- data.frame(
#'     x = factor(rep(c("A", NA, "C", "D"), 3),
#'                levels = c("A", "B", "C", "D")),
#'     y = c(1, 3, 6, 4, 2, 5, 7, 8, 4, 2, 3, 5),
#'     group = factor(sample(rep(c("F1", NA, "F3"), each = 4)),
#'                    levels = c("F1", "F2", "F3")),
#'     split = factor(sample(rep(c("X", "Y", NA), 4)),
#'                    levels = c("X", "Y", "Z")),
#'     facet = factor(sample(rep(c("M", "N", NA), 4)),
#'                    levels = c("M", "N", "O"))
#' )
#'
#' # Default: NA and empty levels dropped
#' AreaPlot(data, x = "x", y = "y", group_by = "group")
#'
#' # Keep NA and empty levels
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          keep_na = TRUE, keep_empty = TRUE)
#'
#' # Keep NA, assign empty levels colours but don't show them
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          keep_na = TRUE, keep_empty = "level")
#'
#' # Drop NA, keep empty levels
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          keep_na = FALSE, keep_empty = TRUE)
#'
#' # Per-column keep_na / keep_empty via named lists
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          keep_na = list(x = TRUE, group = FALSE),
#'          keep_empty = list(x = FALSE, group = TRUE))
#' AreaPlot(data, x = "x", y = "y", group_by = "group",
#'          keep_na = list(x = FALSE, group = TRUE),
#'          keep_empty = list(x = TRUE, group = FALSE))
#' }
AreaPlot <- function(
    data,
    x,
    y = NULL,
    x_sep = "_",
    split_by = NULL,
    split_by_sep = "_",
    group_by = NULL,
    group_by_sep = "_",
    group_name = NULL,
    scale_y = FALSE,
    theme = "theme_this",
    theme_args = list(),
    palette = "Paired",
    palcolor = NULL,
    palreverse = FALSE,
    alpha = 1,
    facet_by = NULL,
    facet_scales = "fixed",
    facet_ncol = NULL,
    facet_nrow = NULL,
    facet_byrow = TRUE,
    x_text_angle = 0,
    aspect.ratio = 1,
    legend.position = waiver(),
    legend.direction = "vertical",
    title = NULL,
    subtitle = NULL,
    xlab = NULL,
    ylab = NULL,
    keep_na = FALSE,
    keep_empty = FALSE,
    seed = 8525,
    combine = TRUE,
    nrow = NULL,
    ncol = NULL,
    byrow = TRUE,
    axes = NULL,
    axis_titles = axes,
    guides = NULL,
    design = NULL,
    ...
) {
    validate_common_args(seed, facet_by = facet_by)
    keep_na <- check_keep_na(keep_na, c(x, split_by, group_by, facet_by))
    keep_empty <- check_keep_empty(
        keep_empty,
        c(x, split_by, group_by, facet_by)
    )
    theme <- process_theme(theme)
    split_by <- check_columns(
        data,
        split_by,
        force_factor = TRUE,
        allow_multi = TRUE,
        concat_multi = TRUE,
        concat_sep = split_by_sep
    )

    if (!is.null(split_by)) {
        data <- process_keep_na_empty(data, keep_na, keep_empty, col = split_by)
        keep_na[[split_by]] <- NULL
        keep_empty[[split_by]] <- NULL
        datas <- split(data, data[[split_by]])
        # keep the order of levels
        datas <- datas[levels(data[[split_by]])]
    } else {
        datas <- list(data)
        split_by <- names(datas) <- "..."
    }
    palette <- check_palette(palette, names(datas))
    palcolor <- check_palcolor(palcolor, names(datas))
    legend.direction <- check_legend(
        legend.direction,
        names(datas),
        "legend.direction"
    )
    legend.position <- check_legend(
        legend.position,
        names(datas),
        "legend.position"
    )

    plots <- lapply(
        names(datas),
        function(nm) {
            default_title <- if (length(datas) == 1 && identical(nm, "...")) {
                NULL
            } else {
                nm
            }
            if (is.function(title)) {
                title <- title(default_title)
            } else {
                title <- title %||% default_title
            }
            AreaPlotAtomic(
                datas[[nm]],
                x = x,
                y = y,
                x_sep = x_sep,
                group_by = group_by,
                group_by_sep = group_by_sep,
                group_name = group_name,
                scale_y = scale_y,
                theme = theme,
                theme_args = theme_args,
                palette = palette[[nm]],
                palcolor = palcolor[[nm]],
                palreverse = palreverse,
                alpha = alpha,
                facet_by = facet_by,
                facet_scales = facet_scales,
                facet_ncol = facet_ncol,
                facet_nrow = facet_nrow,
                facet_byrow = facet_byrow,
                x_text_angle = x_text_angle,
                aspect.ratio = aspect.ratio,
                legend.position = legend.position[[nm]],
                legend.direction = legend.direction[[nm]],
                title = title,
                subtitle = subtitle,
                xlab = xlab,
                ylab = ylab,
                keep_na = keep_na,
                keep_empty = keep_empty,
                ...
            )
        }
    )

    names(plots) <- names(datas)

    combine_plots(
        plots,
        combine = combine,
        split_by = split_by,
        nrow = nrow,
        ncol = ncol,
        byrow = byrow,
        axes = axes,
        axis_titles = axis_titles,
        guides = guides,
        design = design
    )
}