ggplot2中如何修改draw_key传入数据实现双向箭头自定义图例

问题描述 #

• 最初思路是将direction（箭头方向）作为自定义美学映射传入，调整segmentsGrob绘制方向，让图例字形的箭头方向自动和数据匹配。
• 调试发现draw_key系列函数接收到的单行数据框仅包含colour、size、linetype、alpha四个默认字段，没有传入自定义的direction字段，调试打印的data内容如下：

## browser()
# Browse[1]> data
# colour size linetype alpha
# 1 #F8766D  0.5        1    NA

• 尝试通过以下代码将direction新增为GeomArrow的必需美学属性未生效，运行代码会提示Ignoring unknown aesthetics: direction警告：

GeomArrow <- ggproto(NULL, GeomSegment)
GeomArrow$required_aes <- c("x", "y", "xend", "yend", "direction")

• 自定义的draw_key_segment_custom函数因为拿不到direction取值，无法实现根据方向切换segmentsGrob起止坐标的逻辑。

可复现问题代码 #

library(ggplot2)
foo <- structure(list(direction = c("backward", "forward"), 
                      x = c(0, 0), xend = c(1, 1), y = c(1, 1.2), 
                      yend = c(1, 1.2)), row.names = 1:2, class = "data.frame")

StatArrow <- ggproto(NULL, StatIdentity)
StatArrow$compute_layer <- function (self, data, params, layout) 
{
    swap <- data$direction == "backward"
    v1 <- data$x
    v2 <- data$xend
    data$x[swap] <- v2[swap]
    data$xend[swap] <- v1[swap]
    data
}

# draw_key_segment_custom <- function(data, params = list(direction), size) {
## ... 
## 核心思路：根据自定义美学取值切换segmentsGrob中的x0、x1坐标
## 该逻辑无法生效，因为传入的data中不存在direction字段
## 逻辑示意如下：
  # if(data$direction == "backward") {
  #   x0 = 0.9
  #   x1 = 0.1
  # } else {
  #   x0 = 0.1
  #   x1 = 0.9
  # }
  ## 后续调用网格绘制函数
  # grid::segmentsGrob(x0, 0.5, x1, 0.5,
                     ##...)
## 调试查看传入的data仅包含四个字段：
## browser()
# Browse[1]> data
# colour size linetype alpha
# 1 #F8766D  0.5        1    NA

# }

## 图层层面direction可以正常生效，但图例无法匹配
ggplot() +
  geom_segment(data = foo,
               stat = "arrow",
               aes(x, y, xend = xend, yend = yend, col = direction, 
                   direction = direction),
               arrow = arrow(length = unit(0.3, "cm"), type = "closed")
             ## 因无法获取direction取值，自定义key_glyph无法正常工作
             # key_glyph = "segment_custom"
)
#> Warning: Ignoring unknown aesthetics: direction

问题代码运行效果：

^{本示例于2022-06-28通过reprex包（v2.0.1）创建生成。}

解决方案 #

问题根源是ggplot2构建图例数据时，会自动过滤掉两类字段：不在官方内置美学列表的字段、没有在对应Geom的default_aes中声明的自定义字段。仅修改required_aes不会让自定义字段透传到draw_key的入参中。
完整修复步骤如下：

• 自定义继承GeomSegment的Geom类，在default_aes中注册direction字段，设置默认值
• 在自定义Geom中重写draw_panel方法，处理反向箭头的坐标交换逻辑
• 编写自定义draw_key函数，根据入参data中的direction值调整箭头grob的起止坐标，绑定到自定义Geom上
• 封装对外的geom调用函数，直接使用即可

完整可运行代码：

library(ggplot2)
library(grid)

# 自定义图例箭头绘制逻辑
draw_key_arrow_dir <- function(data, params, size) {
  # 根据方向调整箭头起止点
  if(data$direction == "backward") {
    x0 <- 0.9
    x1 <- 0.1
  } else {
    x0 <- 0.1
    x1 <- 0.9
  }
  
  segmentsGrob(
    x0 = x0, y0 = 0.5, x1 = x1, y1 = 0.5,
    gp = gpar(
      col = alpha(data$colour, data$alpha),
      lwd = data$size * .pt,
      lty = data$linetype,
      lineend = "butt"
    ),
    arrow = params$arrow
  )
}

# 自定义箭头Geom类
GeomArrowDir <- ggproto(
  "GeomArrowDir", GeomSegment,
  # 注册自定义direction美学，设置默认值
  default_aes = aes(
    colour = "black", size = 0.5, linetype = 1, alpha = NA,
    direction = "forward"
  ),
  # 处理坐标交换逻辑
  draw_panel = function(self, data, panel_params, coord, arrow = NULL,
                        arrow.fill = NULL, lineend = "butt", linejoin = "round",
                        na.rm = FALSE) {
    swap_idx <- which(data$direction == "backward")
    if(length(swap_idx) > 0) {
      # 交换x轴起止
      tmp_x <- data$x[swap_idx]
      tmp_xend <- data$xend[swap_idx]
      data$x[swap_idx] <- tmp_xend
      data$xend[swap_idx] <- tmp_x
      # 交换y轴起止（适配垂直方向箭头场景）
      tmp_y <- data$y[swap_idx]
      tmp_yend <- data$yend[swap_idx]
      data$y[swap_idx] <- tmp_yend
      data$yend[swap_idx] <- tmp_y
    }
    # 调用父类GeomSegment的绘制逻辑
    ggproto_parent(GeomSegment, self)$draw_panel(
      data, panel_params, coord, arrow = arrow, arrow.fill = arrow.fill,
      lineend = lineend, linejoin = linejoin, na.rm = na.rm
    )
  },
  # 绑定自定义图例绘制函数
  draw_key = draw_key_arrow_dir
)

# 封装对外调用的geom函数
geom_arrow_dir <- 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 = GeomArrowDir,
    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, ...
    )
  )
}

# 测试绘图
foo <- structure(list(direction = c("backward", "forward"), 
                      x = c(0, 0), xend = c(1, 1), y = c(1, 1.2), 
                      yend = c(1, 1.2)), row.names = 1:2, class = "data.frame")

ggplot(foo, aes(x = x, y = y, xend = xend, yend = yend, color = direction)) +
  geom_arrow_dir(arrow = arrow(length = unit(0.3, "cm"), type = "closed"))

运行后不会再出现未知美学警告，图中图例的箭头方向会和数据完全匹配：backward对应朝左箭头，forward对应朝右箭头。

内容的提问来源于stack exchange，提问作者tjebo