Set continuous scales to be identical across ggplot2 plots automatically

set_scales_constant.R

# Check https://www.zachburchill.ml/constant_scales/ for a walkthrough.  
#  Note that this code is slightly "better" than the code in the post, since it checks 
#   to see if the scales being supplied and used are all continuous.

# Examples included after the function definitions.

# Given a bunch of plots and a scale you want to apply to them, this returns that scale,
#   but with limits that encompass the union of all the ranges of values in those plots
# Meant to be for users
get_shared_scale <- function(..., scale) {
  if (!inherits(scale$range, "RangeContinuous"))
    stop("Supplied scale needs to be continuous")
  plots <- list(...)
  ranges <- purrr::map(plots, ~simple_range_extracter(., scale))
  scale$limits <- range(unlist(ranges))
  scale
}
# Given a plot and a scale, get the range of the values used in the scale for that plot
#   Not really intended for users
simple_range_extracter <- function(p, scale) {
  d <- ggplot2::ggplot_build(p)
  p_range <- d$plot$scales$get_scales(scale$aesthetics)$range
  if (!inherits(p_range, "RangeContinuous"))
    stop("Plot's scale for `", scale$aesthetics,
         "` is not continuous")
  p_range$range
}

# Given the unquoted variable names of a bunch of plots that have been 
#   saved to the current environment, and a scale that you want to apply
#   to them, this will call `get_shared_scale()` using those plots and scale
#   add that scale to the plots, and then assign these new plots to those 
#   variable names, essentially editing the plots 'in place'
set_scale_union <- function(..., scale) {
  exprs <- rlang::enexprs(...)
  scale <- get_shared_scale(..., scale = scale)
  var_nms <- purrr::map_chr(exprs, rlang::as_name)
  edit_plots_in_place(var_nms, env = parent.frame(),
                      scale = scale)
  # Invisibly return the scale, in case you need it later
  invisible(scale)
}
# A sub-function. *Definitely* not intended for users
edit_plots_in_place <- function(names, env, scale) {
  vars <- rlang::env_has(env = env, nms = names)
  if (!all(vars))
    stop("Environment does not have variables for ",
         paste(names(vars[!vars]), collapse=", "))
  
  purrr:::walk(names, function(nm) {
    og_plot <- rlang::env_get(env, nm = nm)
    message("Changing plot `", nm, "`")
    # Muffles messages about already having scales
    withCallingHandlers(
      assign(x = nm, envir = env,
             value = og_plot + scale),
      message = function(err) {
        if (grepl("already present", err$message))
          invokeRestart("muffleMessage")
      })
  })
}

# Examples -------------------------------------------------

# Make the plots
p1 <- data.frame(x = runif(400),
                 y = runif(400)) %>%
  ggplot(aes(x, y)) +
  stat_density_2d(geom = "raster",
                  aes(fill = after_stat(density)),
                  contour = FALSE) + 
  scale_fill_viridis_c("z") +
  theme(legend.position = "bottom")
p2 <- data.frame(x = c(runif(200), runif(200, 0.25, 0.75)),
                 y = c(runif(200), runif(200, 0.25, 0.75))) %>%
  ggplot(aes(x, y)) +
  stat_density_2d(geom = "raster",
                  aes(fill = after_stat(density)),
                  contour = FALSE) + 
  scale_fill_viridis_c("z") +
  theme(legend.position = "bottom")

# The color scales don't match:
p1
p2

# I'm setting the scale I want to eventually use as a variable
#   so I don't have to keep retyping it for these examples
my_viridis_scale <- scale_fill_viridis_c(name = "z")

# Test the range:
simple_range_extracter(p1, my_viridis_scale)

# Get the shared scale
shared_scale <- get_shared_scale(p1, p2, scale=my_viridis_scale)
print(shared_scale)

# Test it out
p1 + shared_scale
p2 + shared_scale

# Now edit the plots in place
set_scale_union(p1, p2, scale=my_viridis_scale)

# Now they match:
p1
p2