plot a segment connecting points across facets in ggplot

across_facet_plot.R

library(tidyverse)
library(grid)
library(gtable)


plot_segment_across_facets <- function(p, from=1, to=2, 
                                       from_point_id=1, 
                                       to_point_id=1, 
                                       plotout = F,
                                       gp=gpar(lty=1, alpha=0.5)){
  if (TRUE %in% grepl("ggplot", class(p))) {
    g <- ggplot_gtable(ggplot_build(p))
  } else {
    g <- p
  }
  
  # if one of them contains NA, return g 
  if (NA %in% c(from, to, from_point_id, to_point_id)) return(g)
  
  # collect panel viewport names and index numbers in the grob
  panel_vps <- c()
  id_n <- c()
  for (i in 1:length(g$grobs)) {
    if (str_detect(g$layout[i, "name"], "panel") & g$grobs[[i]]$name != "NULL") {
      p_name <- g$layout[i, "name"]
      panel_vps <- c(panel_vps, p_name)
      id_n <- c(id_n, i)
    }
  }
  
  panel_vps %>%
    str_replace("panel-", "") %>%
    str_split("[\\-\\.]") %>%
    map_chr(1) -> ind_col
  ind_col <- as.numeric(ind_col)

  panel_vps %>%
    str_replace("panel-", "") %>%
    str_split("[\\-\\.]") %>%
    map_chr(2) -> ind_row
  ind_row <- as.numeric(ind_row)

  my_dim <- c(max(ind_row), max(ind_col))
  x <- 1:length(id_n)
  L <- length(x)
  x[(L+1):(my_dim[1]*my_dim[2])] <- NA
  m1 <- as.vector(matrix(x, nrow=my_dim[1], byrow=T))

  x2 <- 1:L
  xx <- as.vector(!is.na(m1))
  xx[xx] <- x2
  xx[!xx] <- NA
  m2 <- as.vector(matrix(xx, nrow=my_dim[1]))

  from <- m2[m1==from]
  from <- from[complete.cases(from)]
  to <- m2[m1==to]
  to <- to[complete.cases(to)]
  
  # select points to be connected
  pnames1 <- names(g$grobs[[id_n[from]]]$children)
  pnames2 <- names(g$grobs[[id_n[to]]]$children)
  
  pname1 <- pnames1[str_detect(pnames1, "geom_point.points")]
  pname2 <- pnames2[str_detect(pnames2, "geom_point.points")]
  
  p1 <- g$grobs[[id_n[from]]]$children[[pname1[1]]]
  p2 <- g$grobs[[id_n[to]]]$children[[pname2[1]]]
  
  g <- with(g$layout[id_n[from],], 
            gtable_add_grob(g, 
                            moveToGrob(p1$x[from_point_id], 
                                       p1$y[from_point_id]), 
                            t=t, l=l))
  g <- with(g$layout[id_n[to],],
            gtable_add_grob(g,
                            lineToGrob(p2$x[to_point_id],
                                       p2$y[to_point_id], gp=gp),
                            t=t, l=l))
  
  # tried curve, but it seems no function for curve across viewports, this is within viewport plot
  # g <- with(g$layout[id_n[to],],
  #           gtable_add_grob(g,
  #                           curveGrob(p1$x[from_point_id],
  #                                     p1$y[from_point_id],
  #                                     p2$x[to_point_id],
  #                                     p2$y[to_point_id], gp=gp),
  #                           t=t, l=l))
  
  
  g$layout$clip <- "off"
  if (plotout==TRUE) grid.draw(g) 
  return(g)
}

# # sample plot
# d <- data.frame(x = rnorm(10), y = rnorm(10), id = c("a", "a", "b", "b", "c",
#                                                      "d", "e", "f", "g", "g"))
# 
# p <- ggplot(d, aes(x, y)) +
#   geom_point() +
#   facet_wrap(~id) +
#   theme_bw() 
# 
# g <- plot_segment_across_facets(p, 1, 2)
# ps <- c(2,3,3,4,4,5,5,6,6,7)
# while (length(ps) > 0){
#   g <- plot_segment_across_facets(g, from=ps[1], to=ps[2])
#   ps <- ps[-c(1:2)]
# }
# 
# grid.draw(g)
# # 
# # 
# # 
# # # save plot
# # # ggsave(grid.draw(g), "sample.pdf")

http://rpubs.com/katzkagaya/410976