Updated geom_segment_plus

geom_segment_plus.R

# an updated version of https://pastebin.com/0BRwUzpu by SO user mo-seph (https://stackoverflow.com/a/14692588/4560765)
# The version above no longer works

#' # The example:
#' points <- data.frame( x=runif(10), y=runif(10),class=1:10, size=runif(10,min=1000,max=100000) )
#' trans <- data.frame( from=rep(1:10,times=10), to=rep(1:10,each=10), amount=runif(100)^3 )
#' trans <- merge( trans, points, by.x="from", by.y="class" )
#' trans <- merge( trans, points, by.x="to", by.y="class", suffixes=c(".to",".from") )
#' ggplot( points, aes( x=x, y=y ) ) + geom_point(aes(size=size),color="red",shape=1) + 
#'   scale_size_continuous(range=c(4,20)) +
#'   geom_segment( data=trans[trans$amount>0.6,], aes( x=x.from, y=y.from, xend=x.to, yend=y.to ),lineend="round",arrow=arrow(),alpha=0.5, size=0.3)
#' 
#' ggplot( points, aes( x=x, y=y ) ) + geom_point(aes(size=size),color="red",shape=1) +
#'   scale_size_continuous(range=c(4,20)) + 
#'   geom_segment_plus( data=trans[trans$amount>0.3,], 
#'                      aes( x=x.from, y=y.from, xend=x.to, yend=y.to ),
#'                      lineend="round",arrow=arrow(length=unit(0.15, "inches")),
#'                      alpha=0.5, size=1.3, 
#'                      offset=0.01, shorten.start=0.03, shorten.end=0.03)




geom_segment_plus <- function (mapping = NULL, data = NULL, stat = "identity",
                               position = "identity", ..., 
                               arrow = NULL, arrow.fill = NULL, lineend = "butt", 
                               linejoin = "round", 
                               na.rm = FALSE, show.legend = NA, inherit.aes = TRUE) {
  layer(data = data, mapping = mapping, stat = stat, geom = GeomSegmentPlus, 
        position = position, show.legend = show.legend, inherit.aes = inherit.aes, 
        params = list(arrow = arrow, arrow.fill = arrow.fill, 
                      lineend = lineend, linejoin = linejoin, na.rm = na.rm, 
                      ...))
}

GeomSegmentPlus <- ggplot2::ggproto(
  "GeomSegmentPlus",
  ggplot2::GeomSegment,
  
  default_aes = aes(colour="black", size=0.5, linetype=1, alpha = NA,shorten.start=0,shorten.end=0,offset=0),
  # draw_key    = ggplot2::draw_key_blank,
  # required_aes = GeomText$required_aes,
  required_aes = c("x", "y", "xend", "yend"),
  
  draw_panel = function (data, panel_params, coord, arrow = NULL, arrow.fill = NULL, 
                         lineend = "butt", linejoin = "round", na.rm = FALSE) 
  {
    data <- remove_missing(data, na.rm = na.rm, 
                           c("x", "y", "xend", "yend", "linetype", "size", "shape", "shorten.start","shorten.end","offset"),
                           name = "geom_segment_plus")
    if (ggplot2:::empty(data)) 
      return(zeroGrob())
    if (coord$is_linear()) {
      data <- coord$transform(data, panel_params)
      arrow.fill <- arrow.fill %||% data$colour
      
      for (i in 1:dim(data)[1]) {
        match = data$xend == data$x[i] & data$x == data$xend[i] & data$yend == data$y[i] & data$y == data$yend[i]
        if( sum( match ) == 0 ) data$offset[i] <- 0
      }
      
      data$dx = data$xend - data$x
      data$dy = data$yend - data$y
      data$dist = sqrt( data$dx^2 + data$dy^2 )
      data$px = data$dx/data$dist
      data$py = data$dy/data$dist
      
      data$x = data$x + data$px * data$shorten.start
      data$y = data$y + data$py * data$shorten.start
      data$xend = data$xend - data$px * data$shorten.end
      data$yend = data$yend - data$py * data$shorten.end
      data$x = data$x - data$py * data$offset
      data$xend = data$xend - data$py * data$offset
      data$y = data$y + data$px * data$offset
      data$yend = data$yend + data$px * data$offset
      
      return(with(data, 
                  segmentsGrob(x, y, xend, yend, default.units="native",
                               gp = gpar(col=alpha(colour, alpha), 
                                         fill = alpha(colour, alpha),
                                         lwd = size * ggplot2::.pt, 
                                         lty = linetype, 
                                         lineend = lineend, 
                                         linejoin = linejoin),
                               arrow = arrow)
      ))

    }
    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)
  }
)