R/geom_edge_density.R

In ggraph: An Implementation of Grammar of Graphics for Graphs and Networks

#' Show edges as a density map
#'
#' This geom makes it possible to add a layer showing edge presence as a density
#' map. Each edge is converted to `n` points along the line and a jitter is
#' applied. Based on this dataset a two-dimensional kernel density estimation is
#' applied and plotted as a raster image. The density is mapped to the alpha
#' level, making it possible to map a variable to the fill.
#'
#' @inheritSection geom_edge_link Edge aesthetic name expansion
#'
#' @section Aesthetics:
#' `geom_edge_density` understand the following aesthetics. Bold aesthetics are
#' automatically set, but can be overwritten.
#'
#' **x**
#' **y**
#' **xend**
#' **yend**
#' edge_fill
#' filter
#'
#'
#' @section Computed variables:
#'
#' \describe{
#'  \item{x, y}{The coordinates for each pixel in the raster}
#'  \item{density}{The density associated with the pixel}
#' }
#'
#' @inheritParams geom_edge_link
#' @inheritParams ggplot2::geom_raster
#'
#' @param n The number of points to estimate in the x and y direction, i.e. the
#' resolution of the raster.
#'
#' @author Thomas Lin Pedersen
#'
#' @family geom_edge_*
#'
#' @examples
#' require(tidygraph)
#' gr <- create_notable('bull') %>%
#'   activate(edges) %>%
#'   mutate(class = sample(letters[1:3], n(), replace = TRUE))
#'
#' ggraph(gr, 'stress') +
#'   geom_edge_density(aes(fill = class)) +
#'   geom_edge_link() + geom_node_point()
#' @rdname geom_edge_density
#' @name geom_edge_density
#'
NULL

#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom ggforce StatLink
#' @importFrom MASS bandwidth.nrd kde2d
#' @export
StatEdgeDensity <- ggproto('StatEdgeDensity', Stat,
  compute_group = function(data, scales, na.rm = FALSE, h = NULL,
                             n = 100, bins = NULL, binwidth = NULL) {
    group <- data$group[1]
    x_range <- diff(range(c(data$x, data$xend)))
    y_range <- diff(range(c(data$y, data$yend)))
    x_extend <- x_range / 10
    y_extend <- y_range / 10
    data <- StatLink$compute_panel(data, n = 50)
    data <- PositionJitter$compute_layer(data, list(
      width = x_extend,
      height = y_extend
    ))

    if (is.null(h)) {
      h <- c(MASS::bandwidth.nrd(data$x), MASS::bandwidth.nrd(data$y))
    }
    dens <- MASS::kde2d(data$x, data$y,
      h = h, n = n,
      lims = c(
        scales$x$dimension(),
        scales$y$dimension()
      ) + c(
        -x_extend,
        x_extend,
        -y_extend,
        y_extend
      )
    )
    df <- expand.grid(x = dens$x, y = dens$y)
    df$z <- as.vector(dens$z)
    df$group <- group
    names(df) <- c('x', 'y', 'density', 'group')
    df
  },
  setup_data = function(data, params) {
    StatFilter$setup_data(data, params)
  },
  default_aes = aes(filter = TRUE),
  required_aes = c('x', 'y', 'xend', 'yend')
)
#' @rdname ggraph-extensions
#' @format NULL
#' @usage NULL
#' @importFrom grid gTree gList grobName rasterGrob rectGrob gpar
GeomEdgeDensity <- ggproto('GeomEdgeDensity', GeomRaster,
  draw_panel = function(self, data, panel_scales, coord, ...) {
    groups <- split(data, factor(data$group))
    max_density <- max(rowSums(inject(cbind(!!!lapply(groups, `[[`, i = 'density')))))
    grobs <- lapply(groups, function(group) {
      self$draw_group(group, panel_scales, coord,
        max.alpha = max_density,
        ...
      )
    })
    grobs <- gTree(children = inject(gList(!!!grobs)))
    grobs$name <- grobName(grobs, 'geom_edge_density')
    grobs
  },
  draw_group = function(self, data, panel_scales, coord, max.alpha) {
    if (!inherits(coord, 'CoordCartesian')) {
      cli::cli_abort('{.fn {snake_class(self)}} only works with {.fn coord_cartesian')
    }
    data <- coord$transform(data, panel_scales)
    x_pos <- as.integer((data$x - min(data$x)) / resolution(
      data$x,
      FALSE
    ))
    y_pos <- as.integer((data$y - min(data$y)) / resolution(
      data$y,
      FALSE
    ))
    nrow <- max(y_pos) + 1
    ncol <- max(x_pos) + 1
    raster <- matrix(NA_character_, nrow = nrow, ncol = ncol)
    raster[cbind(nrow - y_pos, x_pos + 1)] <- alpha(
      data$edge_fill,
      data$density / max.alpha
    )
    x_rng <- c(min(data$xmin, na.rm = TRUE), max(data$xmax, na.rm = TRUE))
    y_rng <- c(min(data$ymin, na.rm = TRUE), max(data$ymax, na.rm = TRUE))
    rasterGrob(raster,
      x = mean(x_rng), y = mean(y_rng),
      width = diff(x_rng), height = diff(y_rng),
      default.units = 'native', interpolate = TRUE
    )
  },
  draw_key = function(data, params, size) {
    rectGrob(gp = gpar(
      col = NA,
      fill = alpha(data$edge_fill %||% data$fill %||% "darkgrey"),
      lty = 0
    ))
  },
  required_aes = c('x', 'y'),
  default_aes = aes(edge_fill = 'darkgrey')
)
#' @rdname geom_edge_density
#'
#' @importFrom ggforce StatLink
#' @export
geom_edge_density <- function(mapping = NULL, data = get_edges('short'),
                              position = 'identity', show.legend = NA,
                              n = 100, ...) {
  mapping <- complete_edge_aes(mapping)
  mapping <- aes_intersect(mapping, aes(x = x, y = y,
                                        xend = xend, yend = yend))
  layer(
    data = data, mapping = mapping, stat = StatEdgeDensity,
    geom = GeomEdgeDensity, position = position,
    show.legend = show.legend, inherit.aes = FALSE,
    params = expand_edge_aes(
      list2(n = n, ...)
    )
  )
}