R/geom-arrowsegment.R

In ggarchery: Flexible Segment Geoms with Arrows for 'ggplot2'

#' @include legend-draw-ggarchery.R
NULL

#' Line segments with flexible arrows
#'
#' @description The basic `geom_arrowsegment()` is equivalent to `geom_segment(arrow = arrow())`.
#' (It is assumed that the user wants some kind of arrow.) The extended functionality
#' is to allow free placement of the arrowhead anywhere along the segment, and also
#' multiple arrowheads, and to allow a fill aesthetic (which will only be visible for
#' closed arrowheads).
#'
#' The function works by dividing the line up into 1 or more segment grobs, each of
#' which is generated by [`grid::arrow()`] except potentially the last (the one closest
#' to the point (`xend`, `yend`)). The vector `arrow_positions`, whose entries must
#' lie between 0 and 1, defines where each arrow segment ends, as a proportional
#' position along the line. If the last entry of `arrow_positions` is 1, then the last
#' grob has an arrow; otherwise it does not.
#'
#' The function is designed with the expectation that arrows point from (`x`, `y`) to
#' (`xend`, `yend`) but the `arrows` argument will happily accept `arrow(ends = "first")`
#' or `arrow(ends = "both")` if you prefer. Just remember that the final segment is
#' only an arrow at all if the last entry of `arrow_positions` is 1.
#'
#' @inheritParams ggplot2::layer
#' @inheritParams ggplot2::geom_segment
#' @param arrow_positions A vector of distinct points on the unit interval. 0 is not
#' permitted but arbitrarily small values are; 1 is permitted. The default behaviour is that
#' arrowheads will be placed proportionally along the line connecting (`x`, `y`)
#' to (`xend`,`yend`) at these points.
#' In more detail: The first arrow segment begins at (`x`, `y`) and ends a proportional
#' distance along the straight line joining (`x`, `y`) and (`xend`, `yend`) equal to the first
#' entry of this vector. The second bridges the first two entries, and so on. If the final
#' entry is 1 then the last segment is an arrow (and hence usually an arrowhead will be
#' placed at the end of the line). If it is not, then  the last segment is simply a line.
#' These will be sorted into order from 0 to 1 if they are not already.
#' @param arrows Either an arrow generated by [`grid::arrow()`] of a list of such arrows. In
#' the former case or if the list has length 1, the arrowhead so defined is used every
#' time; otherwise the list is expected to have the same length as
#' `arrow_positions` and each segment defined by that argument is ended by the respective
#' element of this one. The default is [`grid::arrow()`] with default parameters.
#' @param arrow_fills A vector of fill colours for the arrowheads, behaves as
#' the `arrow_fill` option in [`geom_segment`]. This will overrule a fill aesthetic in
#' the same way that specifying a single `fill` outside `aes` specification will.
#' @import ggplot2 tidyr dplyr purrr
#' @importFrom magrittr %>%
#' @export
#' @return A ggproto object
#' @examples
#'
#'  library(ggplot2)
#'  library(magrittr)
#'  library(tidyr)
#'
#'  # Generate some dummy data
#'
#'  ten.points <- data.frame(line.no = rep(1:5, each = 2), x = runif(10), y = runif(10),
#'                           position = rep(c("start", "end"), 5))
#'  five.segments <- ten.points %>% pivot_wider(names_from = position, values_from = c(x,y))
#'
#'  # Default behaviour
#'
#'  ggplot(five.segments) +
#'     geom_point(data = ten.points, aes(x = x, y = y)) +
#'     geom_arrowsegment(aes(x = x_start, xend = x_end, y = y_start, yend = y_end))
#'
#'  # Midpoint arrowheads
#'
#'  ggplot(five.segments) +
#'     geom_point(data = ten.points, aes(x = x, y = y)) +
#'     geom_arrowsegment(aes(x = x_start, xend = x_end, y = y_start, yend = y_end),
#'                       arrow_positions = 0.5)
#'
#'  # Double arrows
#'
#'  ggplot(five.segments) +
#'     geom_point(data = ten.points, aes(x = x, y = y)) +
#'     geom_arrowsegment(aes(x = x_start, xend = x_end, y = y_start, yend = y_end),
#'                       arrow_positions = c(0.25, 0.75))
#'
#'  # Double arrows, last arrowhead at the end point
#'
#'  ggplot(five.segments) +
#'     geom_point(data = ten.points, aes(x = x, y = y)) +
#'     geom_arrowsegment(aes(x = x_start, xend = x_end, y = y_start, yend = y_end),
#'                       arrow_positions = c(0.25, 1))
#'
#'  # Double arrowheads of varying appearance and position
#'
#'  ggplot(five.segments) +
#'      geom_point(data = ten.points, aes(x = x, y = y)) +
#'      geom_arrowsegment(aes(x = x_start, xend = x_end, y = y_start, yend = y_end),
#'                        arrow_positions = c(0.25, 0.75),
#'                        arrows = list(arrow(angle = 45, type = "closed"),
#'                                      arrow(angle = 25, ends = "both")),
#'                        arrow_fills = "indianred")
#'
geom_arrowsegment <- function(mapping = NULL, data = NULL,
                              stat = "identity", position = "identity",
                              ...,
                              arrows = list(arrow()),
                              arrow_fills = NULL,
                              arrow_positions = 1,
                              lineend = "butt",
                              linejoin = "round",
                              na.rm = FALSE,
                              show.legend = NA,
                              inherit.aes = TRUE) {

  layer(
    data = data,
    mapping = mapping,
    stat = stat,
    geom = GeomArrowsegment,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      arrows = arrows,
      arrow_fills = arrow_fills,
      arrow_positions = arrow_positions,
      lineend = lineend,
      linejoin = linejoin,
      na.rm = na.rm,
      ...
    )
  )
}


#' Find where to place arrows
#' @param x,xend,y,yend The coordinates of the start and end of the segment
#' @param splits A vector of points between 0 and 1 determining where on the line between start and end to split the
#' segment into sub-segments, which will each receive an arrowhead
#' @return A data frame describing the start and end coordinates of the sub-segments. The final segment, ending in 1, is not given (and not given an arrowhead) unless the final element of `splits` is 1.
#' @keywords internal
#' @importFrom purrr map_dbl
#' @export
split_arrows <- function(x, xend, y, yend, splits){

  # the final segment is assumed

  if(splits[length(splits)] == 1){
    if(length(splits) == 1){
      return(data.frame(segment = 1, x = x, xend = xend, y = y, yend = yend))
    }
    splits <- splits[1:(length(splits)-1)]
  }

  x.starts <- c(x, map_dbl(splits, function(sp){
    x + (xend-x)*sp
  }))
  x.ends <- c(x.starts[2:length(x.starts)], xend)

  y.starts <- c(y, map_dbl(splits, function(sp){
    y + (yend-y)*sp
  }))
  y.ends <- c(y.starts[2:length(y.starts)], yend)

  data.frame(segment = 1:length(x.starts), x = x.starts, xend = x.ends, y = y.starts, yend = y.ends)
}

#' @rdname ggarchery-ggproto
#' @format NULL
#' @usage NULL
#' @import ggplot2 grid
#' @export
GeomArrowsegment <- ggproto("GeomArrowsegment", GeomSegment,
                            default_aes = aes(colour = "black", fill = "black", linewidth = 0.5, linetype = 1, alpha = NA),
                            draw_panel = function(self, data, panel_params, coord, arrows = list(arrow()), arrow_fills = NULL, arrow_positions = 1,
                                                  lineend = "butt", linejoin = "round", na.rm = FALSE) {

                              # if the arrows argument is not a list of arrows, make it one

                              if(!is.null(attr(arrows,'class'))){
                                if(attr(arrows,'class') == "arrow"){
                                  arrows <- list(arrows)
                                }
                              }

                              if(arrow_positions[1] == 0){
                                # The arrow can't orient itself properly if the first segment is of length zero
                                abort("First arrowhead position cannot be 0; try a very small positive value instead")
                              }
                              if(any(arrow_positions < 0) | any(arrow_positions > 1)){
                                abort("Arrowhead positions must lie between 0 and 1")
                              }

                              if((length(arrows) != length(arrow_positions)) & (length(arrows) != 1 )){
                                abort("Number of arrows and arrowhead positions do not match")
                              }

                              if((length(arrow_fills) != length(arrow_positions)) & (length(arrow_fills) != 1) & !is.null(arrow_fills)){
                                abort("Number of arrow fills and arrowhead positions do not match")
                              }

                              if(any(duplicated(arrow_positions))){
                                abort("Arrowhead positions must be distinct")
                              }

                              # Some users will inevitably place these out of order. Well, they can't.

                              arrow_positions <- sort(arrow_positions)

                              data <- remove_missing(data, na.rm = na.rm,
                                                     c("x", "y", "xend", "yend", "linetype", "shape"),
                                                     name = "geom_segment")

                              if (nrow(data) == 0) return(zeroGrob())

                              if (coord$is_linear()) {
                                coord <- coord$transform(data, panel_params)


                                newcoord <- coord %>%
                                  mutate(new_locations = pmap(list(x,xend,y,yend), split_arrows, splits = arrow_positions)) %>%
                                  select(-x, -xend, -y, -yend) %>%
                                  unnest(new_locations)

                                # if only one arrow specification is given we need to repeat it for each segment

                                if(length(arrows) == 1 & length(arrow_positions) > 1){
                                  arrows <- rep(arrows, length(arrow_positions))
                                }

                                if(!is.null(arrow_fills) &  length(arrow_fills) == 1 & length(arrow_positions) > 1){
                                  arrow_fills <- rep(arrow_fills, length(arrow_positions))
                                }


                                if(!is.null(arrow_fills) & arrow_positions[length(arrow_positions)] != 1){
                                  # arrow.positions = list(0.5) has one arrow at 0.5 but the line continues with no arrow
                                  # arrow.positions = as.list(c(0.5, 1)) puts two arrowheads at 0.5 and 1
                                  # hence the last segment needs a NULL arrow in the former case
                                  arrow_fills <- c(arrow_fills, NA_character_)
                                }

                                out <- map(1:max(newcoord$segment), function(sg){

                                  bundle <- newcoord %>% filter(segment == sg)

                                  if(sg <= length(arrows)){
                                    current.arrow <- arrows[[sg]]
                                  } else {
                                    current.arrow <- NULL
                                  }

                                  if(is.null(arrow_fills)){
                                    fill.value <- alpha(bundle$fill, bundle$alpha)
                                  } else {
                                    fill.value <- alpha(arrow_fills[sg], bundle$alpha)
                                  }

                                  segmentsGrob(bundle$x, bundle$y, bundle$xend, bundle$yend,
                                               default.units = "native",
                                               gp = gpar(
                                                 col = alpha(bundle$colour, bundle$alpha),
                                                 fill = fill.value,
                                                 lwd = bundle$linewidth * .pt,
                                                 lty = bundle$linetype,
                                                 lineend = lineend,
                                                 linejoin = linejoin
                                               ),
                                               arrow = current.arrow
                                  )
                                })

                                return(do.call("gList", out))
                              }

                              abort("geom_arrowsegment() supports linear coordinate systems only at present")

                              # data$group <- 1:nrow(data)
                              # starts <- subset(data, select = c(-xend, -yend))
                              # ends <- rename(subset(data, select = c(-x, -y)), c("xend" = "x", "yend" = "y"))
                              #
                              # pieces <- rbind(starts, ends)
                              # pieces <- pieces[order(pieces$group),]
                              #
                              # GeomPath$draw_panel(pieces, panel_params, coord, arrow = arrow,
                              #                     lineend = lineend)
                            },
                            draw_key = draw_key_arrowpath,
                            rename_size = TRUE

)