geom_flat_violin.R

# Source: https://gist.github.com/dgrtwo/eb7750e74997891d7c20
# Source: https://github.com/hadley/ggplot2/blob/master/R/geom-violin.r
# Last update: Wednesday, July 29, 2020 | ggplot2 3.3.2

library(ggplot2)
library(dplyr)
library(grid)


"%||%" <- function(a, b) {
  if (!is.null(a))
    a
  else
    b
}

geom_flat_violin <- function(mapping = NULL,
                             data = NULL,
                             stat = "ydensity",
                             position = "dodge",
                             ...,
                             draw_quantiles = NULL,
                             trim = TRUE,
                             scale = "area",
                             na.rm = FALSE,
                             orientation = NA,
                             show.legend = NA,
                             inherit.aes = TRUE) {
  layer(
    data = data,
    mapping = mapping,
    stat = stat,
    geom = GeomFlatViolin,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      trim = trim,
      scale = scale,
      draw_quantiles = draw_quantiles,
      na.rm = na.rm,
      orientation = orientation,
      ...
    )
  )
}

#' @rdname ggplot2-ggproto
#' @format NULL
#' @usage NULL
#' @export
GeomFlatViolin <- ggproto(
  "GeomFlatViolin",
  Geom,
  setup_params = function(data, params) {
    params$flipped_aes <-
      has_flipped_aes(data, params, ambiguous = TRUE)
    params
  },
  
  extra_params = c("na.rm", "orientation"),
  
  setup_data = function(data, params) {
    data$flipped_aes <- params$flipped_aes
    data <-
      flip_data(data, params$flipped_aes)
    data$width <- data$width %||%
      params$width %||% (resolution(data$x, FALSE) * 0.9)
    
    # ymin, ymax, xmin, and xmax define the bounding rectangle for each group
    data <-
      ggplot2:::dapply(data,
                       "group",
                       transform,
                       xmin = x,
                       # Original: x - width / 2
                       xmax = x + width / 2)
    flip_data(data, params$flipped_aes)
  },
  
  draw_group = function(self,
                        data,
                        ...,
                        draw_quantiles = NULL,
                        flipped_aes = FALSE) {
    data <- flip_data(data, flipped_aes)
    # Find the points for the line to go all the way around
    data <- transform(data,
                      xminv = x,
                      # Original: x - violinwidth * (x - xmin)
                      xmaxv = x + violinwidth * (xmax - x))
    
    # Make sure it's sorted properly to draw the outline
    newdata <- rbind(transform(data, x = xminv)[order(data$y),],
                     transform(data, x = xmaxv)[order(data$y, decreasing = TRUE),])
    
    # Close the polygon: set first and last point the same
    # Needed for coord_polar and such
    newdata <- rbind(newdata, newdata[1, ])
    newdata <-
      flip_data(newdata, flipped_aes)
    
    # Draw quantiles if requested, so long as there is non-zero y range
    if (length(draw_quantiles) > 0 &
        !scales::zero_range(range(data$y))) {
      if (!(all(draw_quantiles >= 0) && all(draw_quantiles <= 1))) {
        abort("`draw_quantiles must be between 0 and 1")
      }
      
      # Compute the quantile segments and combine with existing aesthetics
      quantiles <-
        create_quantile_segment_frame(data, draw_quantiles)
      aesthetics <- data[rep(1, nrow(quantiles)),
                         setdiff(names(data), c("x", "y", "group")),
                         drop = FALSE]
      aesthetics$alpha <-
        rep(1, nrow(quantiles))
      both <- cbind(quantiles, aesthetics)
      both <-
        both[!is.na(both$group), , drop = FALSE]
      both <- flip_data(both, flipped_aes)
      quantile_grob <-
        if (nrow(both) == 0) {
          zeroGrob()
        } else {
          GeomPath$draw_panel(both, ...)
        }
      
      ggplot2:::ggname("geom_flat_violin",
                       grobTree(GeomPolygon$draw_panel(newdata, ...),
                                quantile_grob))
    } else {
      ggplot2:::ggname("geom_flat_violin", GeomPolygon$draw_panel(newdata, ...))
    }
  },
  
  draw_key = draw_key_polygon,
  
  default_aes = aes(
    weight = 1,
    colour = "grey20",
    fill = "white",
    size = 0.5,
    alpha = NA,
    linetype = "solid"
  ),
  
  required_aes = c("x", "y")
)

# Returns a data.frame with info needed to draw quantile segments.
create_quantile_segment_frame <- function(data, draw_quantiles) {
  dens <- cumsum(data$density) / sum(data$density)
  ecdf <- stats::approxfun(dens, data$y)
  ys <-
    ecdf(draw_quantiles) # these are all the y-values for quantiles
  
  # Get the violin bounds for the requested quantiles.
  violin.xminvs <- (stats::approxfun(data$y, data$xminv))(ys)
  violin.xmaxvs <- (stats::approxfun(data$y, data$xmaxv))(ys)
  
  # We have two rows per segment drawn. Each segment gets its own group.
  ggplot2:::new_data_frame(list(
    x = ggplot2:::interleave(violin.xminvs, violin.xmaxvs),
    y = rep(ys, each = 2),
    group = rep(ys, each = 2)
  ))
}