Accessing ggplot's inherited data object inside custom layer

[holer]

I am trying to create geom_observation, which annotates the plot with the number of observations in each group, but I am struggling with the positioning of the annotation.

I would like the position to be based on the "global" data-set, and not the grouped/faceted data-set. Further, I would like to achieve this without assigning the plot (i.e. p <- ggplot(...)).

Lets have a look at the code:

library(tidyverse)

StatObservation <- ggproto("StatObservation", Stat,
  compute_group = function(data, scales) {
    data %>% 
      summarize(cnt=n(), x = max(data$x), y = max(data$y)) %>%
      mutate(label = paste("n =", cnt), .keep="unused")
  },
  required_aes = c("x", "y")
)

GeomObservation <- ggproto("GeomObservation", GeomText,
  default_aes = aes(
    colour = "red", size = 4, angle = 0, hjust = "right", vjust = "top", 
    alpha = NA, family = "", fontface = 1, lineheight = 1.2
  )
)

geom_observation <- function(mapping = NULL, data = NULL, position = "identity", 
                             na.rm = FALSE, show.legend = NA, inherit.aes = TRUE, ...) {
  layer(
    stat = StatObservation, geom = GeomObservation, data = data, mapping = mapping,
    position = position, show.legend = show.legend, inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, ...)
  )
}

And the plots (I know I am assigning the plot here, but that is only to make the code more easily understandable):

# standard plot
p <- mpg %>%
#  filter(manufacturer == "audi") %>%
  ggplot(aes(x = cty, y = hwy)) +
  geom_point() +
  facet_grid(drv ~ factor(year))
p  

# plot with poorly positioned annotations
p + geom_observation()

# plot with correctly positioned annotations
p + geom_observation(
  aes(x = max(p$data[[rlang::get_expr(p$mapping$x)]]),
      y = max(p$data[[rlang::get_expr(p$mapping$y)]]))
)

How can I achieve the result above without assigning the plot?

BONUS QUESTION (in case we solve the main question): How to handle with colour aesthetics?

1. How to align the annotation above/below each other?
2. Is it possible to re-group the colors, i.e. as if aes(colour) was not assigned?

p + aes(colour = class) +
  geom_observation(
    aes(x = max(p$data[[rlang::get_expr(p$mapping$x)]]),
        y = max(p$data[[rlang::get_expr(p$mapping$y)]]))
  )

[holer]

I was able to solve this by replacing copy-pasting the source-code for GeomText, and using the ViewScales class, which is passed to each panel as panel_params in the draw_panel function, to get the limits of the original data.

And while this solves my current problem I am still curious if there exists a more elegant solution!

This solution/work-around didn't get access to the un-grouped data, but it instead leveraged the fact that the limits were passed as an argument to each panel. So I will leave this as unsolved in hopes of proper solution.

This code should reproduce a working example:

library(tidyverse)
library(grid)

StatObservation <- ggproto("StatObservation", Stat,
  compute_group = function(data, scales) {
    data %>% 
      summarize(cnt=n(), x = max(data$x), y = max(data$y)) %>%
      mutate(label = paste("n =", cnt), .keep="unused")
  }
)

GeomObservation <- ggproto("GeomObservation", GeomText,
  default_aes = aes(
    colour = "red", size = 4, angle = 0, hjust = "right", vjust = "top", 
    alpha = NA, family = "", fontface = 1, lineheight = 1.2, force_position = TRUE
  ),
  draw_panel = function(data, panel_params, coord, parse = FALSE,
                        na.rm = FALSE, check_overlap = FALSE) {
    # logic to find the fixed position (from vjust and hjust)
    if (data$force_position) {
      
      # replace the x-coordinate
      if (data$hjust == "left") {
        data$x <- panel_params$x$limits[[1]]
      } else if (data$hjust == "middle") {
        data$x <- panel_params$x$limits %>% mean() # NOTE: This may fail with "categorical" data
      } else if (data$hjust == "right") {
        data$x <- panel_params$x$limits[[2]]
      }
      
      # replace the y-coordinate
      if (data$vjust == "top") {
        data$y <- panel_params$y$limits[[2]]
      } else if (data$vjust == "center") {
        data$y <- panel_params$y$limits %>% mean()
      } else if (data$vjust == "bottom") {
        data$y <- panel_params$y$limits[[1]]
      } 
    }
    
    lab <- data$label
    if (parse) {
      lab <- parse_safe(as.character(lab))
    }
    
    data <- coord$transform(data, panel_params)
    
    if (is.character(data$vjust)) {
      data$vjust <- compute_just(data$vjust, data$y, data$x, data$angle)
    }
    if (is.character(data$hjust)) {
      data$hjust <- compute_just(data$hjust, data$x, data$y, data$angle)
    }
    
    textGrob(
      lab,
      data$x, data$y, default.units = "native",
      hjust = data$hjust, vjust = data$vjust,
      rot = data$angle,
      gp = gpar(
        col = alpha(data$colour, data$alpha),
        fontsize = data$size * .pt,
        fontfamily = data$family,
        fontface = data$fontface,
        lineheight = data$lineheight
      ),
      check.overlap = check_overlap
    )
  }
)

geom_observation <- function(mapping = NULL, data = NULL, position = "identity", 
                             na.rm = FALSE, show.legend = NA, inherit.aes = TRUE, ...) {
  layer(
    stat = StatObservation, geom = GeomObservation, data = data, mapping = mapping,
    position = position, show.legend = show.legend, inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, ...)
  )
}

# plots
p <- mpg %>%
  ggplot(aes(x = cty, y = hwy)) +
  geom_point() +
  facet_grid(drv ~ factor(year))

p
p + geom_observation()
p + geom_observation(vjust="top", hjust="left")



### HELPER FUNCTIONS (from geom-text.R) ###
compute_just <- function(just, a, b = a, angle = 0) {
  #  As justification direction is relative to the text, not the plotting area
  #  we need to swap x and y if text direction is rotated so that hjust is
  #  applied along y and vjust along x.
  if (any(grepl("outward|inward", just))) {
    # ensure all angles are in -360...+360
    angle <- angle %% 360
    # ensure correct behaviour for angles in -360...+360
    angle <- ifelse(angle > 180, angle - 360, angle)
    angle <- ifelse(angle < -180, angle + 360, angle)
    rotated_forward <-
      grepl("outward|inward", just) & (angle > 45 & angle < 135)
    rotated_backwards <-
      grepl("outward|inward", just) & (angle < -45 & angle > -135)
    
    ab <- ifelse(rotated_forward | rotated_backwards, b, a)
    just_swap <- rotated_backwards | abs(angle) > 135
    inward <-
      (just == "inward" & !just_swap | just == "outward" & just_swap)
    just[inward] <- c("left", "middle", "right")[just_dir(ab[inward])]
    outward <-
      (just == "outward" & !just_swap) | (just == "inward" & just_swap)
    just[outward] <- c("right", "middle", "left")[just_dir(ab[outward])]
    
  }
  
  unname(c(left = 0, center = 0.5, right = 1,
           bottom = 0, middle = 0.5, top = 1)[just])
}

just_dir <- function(x, tol = 0.001) {
  out <- rep(2L, length(x))
  out[x < 0.5 - tol] <- 1L
  out[x > 0.5 + tol] <- 3L
  out
}