geom_grid.R

require(ggplot2)
require(grid)
require(dplyr)

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

geom_grid <- function(mapping = NULL, data = NULL,
                         position = "identity",
                         ...,
                         # bins = NULL,
                         binwidth = NULL,
                         binaxis = "y",
                         method = "dotdensity",
                         binpositions = "bygroup",
                         stackdir = "up",
                         stackratio = 1,
                         dotsize = 1,
                         stackgroups = FALSE,
                         origin = NULL,
                         right = TRUE,
                         width = 0.9,
                         drop = FALSE,
                         na.rm = FALSE,
                         show.legend = NA,
                         inherit.aes = TRUE) {
  # If identical(position, "stack") or position is position_stack(), tell them
  # to use stackgroups=TRUE instead. Need to use identical() instead of ==,
  # because == will fail if object is position_stack() or position_dodge()
  if (!is.null(position) &&
      (identical(position, "stack") || (inherits(position, "PositionStack"))))
    message("position=\"stack\" doesn't work properly with geom_dotplot. Use stackgroups=TRUE instead.")
  
  if (stackgroups && method == "dotdensity" && binpositions == "bygroup")
    message('geom_dotplot called with stackgroups=TRUE and method="dotdensity". You probably want to set binpositions="all"')
  
  layer(
    data = data,
    mapping = mapping,
    stat = StatBindot,
    geom = GeomGrid,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    # Need to make sure that the binaxis goes to both the stat and the geom
    params = list(
      binaxis = binaxis,
      # bins = bins,
      binwidth = binwidth,
      binpositions = binpositions,
      method = method,
      origin = origin,
      right = right,
      width = width,
      drop = drop,
      stackdir = stackdir,
      stackratio = stackratio,
      dotsize = dotsize,
      stackgroups = stackgroups,
      na.rm = na.rm,
      ...
    )
  )
}

GeomGrid <- ggproto("GeomGrid", Geom,
                       required_aes = c("x", "y"),
                       non_missing_aes = c("size", "shape"),
                       
                       default_aes = aes(colour = "black", fill = "white", alpha = NA, stroke = 1, linetype = "solid"),
                       
                       setup_data = function(data, params) {
                         data$width <- data$width %||%
                           params$width %||% (resolution(data$x, FALSE) * 0.9)
                         
                         # Set up the stacking function and range
                         if (is.null(params$stackdir) || params$stackdir == "up") {
                           stackdots <- function(a)  a - .5
                           stackaxismin <- 0
                           stackaxismax <- 1
                         } else if (params$stackdir == "down") {
                           stackdots <- function(a) -a + .5
                           stackaxismin <- -1
                           stackaxismax <- 0
                         } else if (params$stackdir == "center") {
                           stackdots <- function(a)  a - 1 - max(a - 1) / 2
                           stackaxismin <- -.5
                           stackaxismax <- .5
                         } else if (params$stackdir == "centerwhole") {
                           stackdots <- function(a)  a - 1 - floor(max(a - 1) / 2)
                           stackaxismin <- -.5
                           stackaxismax <- .5
                         }
                         
                         
                         # Fill the bins: at a given x (or y), if count=3, make 3 entries at that x
                         data <- data[rep(1:nrow(data), data$count), ]
                         
                         # Next part will set the position of each dot within each stack
                         # If stackgroups=TRUE, split only on x (or y) and panel; if not stacking, also split by group
                         # plyvars <- params$binaxis %||% "x"
                         # plyvars <- params$binaxis %||% "x"
                         # plyvars <- c(plyvars, "PANEL")
                         # if (is.null(params$stackgroups) || !params$stackgroups)
                         #   plyvars <- c(plyvars, "group")
                         
                         # Within each x, or x+group, set countidx=1,2,3, and set stackpos according to stack function
                         # data <- dapply(data, plyvars, function(xx) {
                         #   xx$countidx <- 1:nrow(xx)
                         #   xx$stackpos <- stackdots(xx$countidx)
                         #   xx
                         # })
                         data <-  data %>%
                           dplyr::group_by_("x","y","PANEL","group") %>%
                           dplyr::mutate(countidx = row_number()) %>%
                           dplyr::mutate(stackpos = stackdots(countidx))
                         
                         # Set the bounding boxes for the dots
                         if (is.null(params$binaxis) || params$binaxis == "x") {
                           # ymin, ymax, xmin, and xmax define the bounding rectangle for each stack
                           # Can't do bounding box per dot, because y position isn't real.
                           # After position code is rewritten, each dot should have its own bounding box.
                           data$xmin <- data$x - data$binwidth / 2
                           data$xmax <- data$x + data$binwidth / 2
                           data$ymin <- stackaxismin
                           data$ymax <- stackaxismax
                           data$y    <- 0
                           
                         } else if (params$binaxis == "y") {
                           # ymin, ymax, xmin, and xmax define the bounding rectangle for each stack
                           # Can't do bounding box per dot, because x position isn't real.
                           # xmin and xmax aren't really the x bounds, because of the odd way the grob
                           # works. They're just set to the standard x +- width/2 so that dot clusters
                           # can be dodged like other geoms.
                           # After position code is rewritten, each dot should have its own bounding box.
                           data <- dplyr::mutate(dplyr::group_by(data, group, PANEL),
                                          ymin = min(y) - binwidth / 2,
                                          ymax = max(y) + binwidth / 2)
                           
                           data$xmin <- data$x + data$width * stackaxismin
                           data$xmax <- data$x + data$width * stackaxismax
                           # Unlike with y above, don't change x because it will cause problems with dodging
                         }
                         data
                       },
                       
                       
                       draw_group = function(data, panel_params, coord, na.rm = FALSE,
                                             binaxis = "x", stackdir = "up", stackratio = 1,
                                             dotsize = 0.9, stackgroups = FALSE) {
                         if (!coord$is_linear()) {
                           warning("geom_grid does not work properly with non-linear coordinates.")
                         }
                         
                         tdata <- coord$transform(data, panel_params)
                         
                         # Swap axes if using coord_flip
                         if (inherits(coord, "CoordFlip"))
                           binaxis <- ifelse(binaxis == "x", "y", "x")
                         
                         if (binaxis == "x") {
                           stackaxis = "y"
                           dotdianpc <- dotsize * tdata$binwidth[1] / (max(panel_params$x.range) - min(panel_params$x.range))
                           
                         } else if (binaxis == "y") {
                           stackaxis = "x"
                           dotdianpc <- dotsize * tdata$binwidth[1] / (max(panel_params$y.range) - min(panel_params$y.range))
                         }
                         
                         ggplot2:::ggname("geom_grid",
                                rectstackGrob(stackaxis = stackaxis, x = tdata$x, y = tdata$y, dotdia = dotdianpc,
                                             stackposition = tdata$stackpos, stackratio = stackratio,
                                             default.units = "npc",
                                             gp = gpar(col = alpha(tdata$colour, tdata$alpha),
                                                       fill = alpha(tdata$fill, tdata$alpha),
                                                       lwd = tdata$stroke, lty = tdata$linetype))
                         )
                       },
                       
                       draw_key = draw_key_dotplot
)

rectstackGrob <- function(
  x = unit(0.5, "npc"),     # x pos of the dotstack's origin
  y = unit(0.5, "npc"),     # y pos of the dotstack's origin
  stackaxis = "y",
  dotdia = unit(1, "npc"),  # Dot diameter in the non-stack axis, should be in npc
  stackposition = 0,        # Position of each dot in the stack, relative to origin
  stackratio = 1,           # Stacking height of dots (.75 means 25% dot overlap)
  default.units = "npc", name = NULL, gp = gpar(), vp = NULL)
{
  if (!grid::is.unit(x))
    x <- unit(x, default.units)
  if (!grid::is.unit(y))
    y <- unit(y, default.units)
  if (!grid::is.unit(dotdia))
    dotdia <- unit(dotdia, default.units)
  if (attr(dotdia,"unit") != "npc")
    warning("Unit type of dotdia should be 'npc'")
  
  grob(x = x, y = y, stackaxis = stackaxis, dotdia = dotdia,
       stackposition = stackposition, stackratio = stackratio,
       name = name, gp = gp, vp = vp, cl = "rectstackGrob")
}


makeContext.rectstackGrob <- function(x, recording = TRUE) {
  # Need absolute coordinates because when using npc coords with circleGrob,
  # the radius is in the _smaller_ of the two axes. We need the radius
  # to instead be defined in terms of the non-stack axis.
  xmm <- convertX(x$x, "mm", valueOnly = TRUE)
  ymm <- convertY(x$y, "mm", valueOnly = TRUE)
  
  if (x$stackaxis == "x") {
    dotdiamm <- convertY(x$dotdia, "mm", valueOnly = TRUE)
    xpos <- xmm + dotdiamm * (x$stackposition * x$stackratio + (1 - x$stackratio) / 2)
    ypos <- ymm
  } else if (x$stackaxis == "y") {
    dotdiamm <- convertX(x$dotdia, "mm", valueOnly = TRUE)
    xpos <- xmm
    ypos <- ymm + dotdiamm * (x$stackposition * x$stackratio + (1 - x$stackratio) / 2)
  }
  
  grid::rectGrob(
    x = xpos, y = ypos, width = dotdiamm , height = dotdiamm ,
    default.units = "mm",
    name = x$name, gp = x$gp, vp = x$vp
  )
}