Efficient way to map data to legend text color in ggplot2

Question

I'm wondering if there's an efficient way to map data onto legend text color in ggplot2, just like we can do with axis text. Reproducible example follows.

First, let's make a plot:

library(ggplot2)
library(dplyr)

drv_counts <- mutate(mpg,
                     drv = case_when(drv == "r" ~ "rear wheel drive",
                                     drv == "4" ~ "4 wheel drive",
                                     drv == "f" ~ "front wheel drive"),
                     model_drv = interaction(model, drv)) %>%
  group_by(model_drv) %>%
  summarize(model = model[1], drv = drv[1], count = n()) %>%
  arrange(drv, count) %>%
  mutate(model = factor(model, levels = model))

p <- ggplot(drv_counts, aes(x=model, y=count, fill=drv)) +
  geom_col() + coord_flip() + guides(fill = guide_legend(reverse=T)) +
  theme_minimal()
p

Now let's color the axis labels by drive train. This is very easy:

# ggplot2 colors
cols <- c("4 wheel drive" = "#F8766D", "front wheel drive" = "#00BA38", "rear wheel drive" = "#619CFF")

p2 <- p + theme(axis.text.y = element_text(color = cols[drv_counts$drv]))
p2

Now let's try the same trick on the legend. It doesn't work:

p2 + theme(legend.text = element_text(color = cols))

The reason this doesn't work for legend text but does work for axis text is that all the axis labels are drawn in one grob, and hence we can give that grob a vector of colors, but the legend labels are drawn in separate grobs.

We can go in and color all the grobs manually, but that's super ugly and cumbersome:

g <- ggplotGrob(p2)
g$grobs[[15]]$grobs[[1]]$grobs[[9]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[9]]$children[[1]]$label]
g$grobs[[15]]$grobs[[1]]$grobs[[10]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[10]]$children[[1]]$label]
g$grobs[[15]]$grobs[[1]]$grobs[[11]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[11]]$children[[1]]$label]
grid::grid.newpage()
grid::grid.draw(g)

My question is: Can somebody think of a way of getting this effect without having to dig down into the grob tree? I'm Ok with a patch to ggplot2 if it's only a few modified lines. Alternatively, can the digging down into the grob tree be automated so I don't have to access child grobs by manually setting list indices that will change the moment I make a minor change to the figure?

Update: A related question can be found here. To make my question distinct, let's add the requirement that colors aren't copied over from the symbols but rather can be set to any arbitrary values. This added requirement has real-world relevance because I usually use a darker color for text than for symbols.

Mesothorium

Following your example I can't reproduce your final visual with the legend text color matching the boxes. R version 3.4.3 64bit windows, ggplot2_2.2.1 – Supersedure

The labels are still the default black text color – Supersedure

a workaround / patch would be to replace the current lapply loop over labels to mapply with both labels and colours, creating a new 'theme' for each grob rather than one for all – Banerjee

that being said, the guides code is quite obscure (to everyone, as I understand it), and I believe a long-term approach would involve its complete refactoring. I personally like the idea of treating guides as simplified graphics, which could be generated with a grammar also. – Banerjee

@Supersedure I'm running the development branch of ggplot2 (v. 2.2.1.9000). It's quite possible that the exact layout of the grobtree is different between the versions. That's exactly my point here: the manual grob digging is not reliable. – Mesothorium

@Banerjee Thanks for the mapply suggestion, I'll look into it. And I agree with your vision for a long-term approach for guides. – Mesothorium

Answer 1

Here's a pretty mediocre method of hacking grobs together to make a legend. I setup a palette based on the unique values of the drv variable (so it can be scaled to larger datasets or more colors). Then I mapped over the values of the palette to make each legend item: a rectGrob and a textGrob, both with the corresponding color from the palette. These could definitely be tweaked to look better. All of these get arranged into a new grob and stuck alongside the plot with cowplot. It isn't gorgeous but it might be a start.

library(tidyverse)
library(grid)
library(gridExtra)

pal <- colorspace::qualitative_hcl(n = length(unique(drv_counts$drv)), l = 60, c = 70) %>%
  setNames(unique(drv_counts$drv))

p2 <- ggplot(drv_counts, aes(x=model, y=count, fill=drv)) +
  geom_col() + 
  coord_flip() +
  theme_minimal() +
  scale_fill_manual(values = pal, guide = F) +
  theme(axis.text.y = element_text(color = pal[drv_counts$drv]))

legend <- pal %>%
  imap(function(col, grp) {
    rect <- rectGrob(x = 0, width = unit(0.5, "line"), height = unit(0.5, "line"), gp = gpar(col = col, fill = col), hjust = 0)
    label <- textGrob(label = grp, gp = gpar(col = colorspace::darken(col, 0.4), fontsize = 10), x = 0, hjust = 0)
    cowplot::plot_grid(rect, label, nrow = 1, rel_widths = c(0.12, 1))
  }) %>%
  arrangeGrob(grobs = rev(.), padding = unit(0.1, "line"), heights = rep(unit(1.1, "line"), 3))

cowplot::plot_grid(p2, legend, rel_widths = c(1, 0.45))

Created on 2018-05-26 by the reprex package (v0.2.0).

Monafo