Backtransform data before mapping statistics (#4194)

Author: yutannihilation

.github/workflows/R-CMD-check.yaml

@@ -59,6 +59,7 @@ jobs:
 
       - uses: r-lib/actions/setup-r-dependencies@v2
         with:
+          cache-version: 2
           extra-packages: >
             any::rcmdcheck,
             maps=?ignore-before-r=3.5.0,

NEWS.md

@@ -1,5 +1,10 @@
 # ggplot2 (development version)
 
+* `stage()` now properly refers to the values without scale transformations for
+  the stage of `after_stat`. If your code requires the scaled version of the
+  values for some reason, you have to apply the same transformation by yourself,
+  e.g. `sqrt()` for `scale_{x,y}_sqrt()` (@yutannihilation and @teunbrand, #4155).
+
 * A `linewidth` aesthetic has been introduced and supersedes the `size` 
   aesthetic for scaling the width of lines in line based geoms. `size` will 
   remain functioning but deprecated for these geoms and it is recommended to 

R/layer.r

@@ -349,10 +349,13 @@ Layer <- ggproto("Layer", NULL,
     new <- strip_dots(aesthetics[is_calculated_aes(aesthetics) | is_staged_aes(aesthetics)])
     if (length(new) == 0) return(data)
 
+    # data needs to be non-scaled
+    data_orig <- scales_backtransform_df(plot$scales, data)
+
     # Add map stat output to aesthetics
     env <- child_env(baseenv(), stat = stat, after_stat = after_stat)
     stage_mask <- child_env(emptyenv(), stage = stage_calculated)
-    mask <- new_data_mask(as_environment(data, stage_mask), stage_mask)
+    mask <- new_data_mask(as_environment(data_orig, stage_mask), stage_mask)
     mask$.data <- as_data_pronoun(mask)
 
     new <- substitute_aes(new)

R/scales-.r

@@ -80,13 +80,48 @@ scales_map_df <- function(scales, df) {
 
 # Transform values to cardinal representation
 scales_transform_df <- function(scales, df) {
-  if (empty(df) || length(scales$scales) == 0) return(df)
+  if (empty(df)) return(df)
+
+  # if the scale contains no trans or the trans is of identity, it doesn't need
+  # to be transformed.
+  idx_skip <- vapply(scales$scales, function(x) {
+    is.null(x$trans) ||
+      identical(x$trans$transform, identity)
+  }, logical(1L))
+  scale_list <- scales$scales[!idx_skip]
+
+  if (length(scale_list) == 0L) return(df)
 
-  transformed <- unlist(lapply(scales$scales, function(s) s$transform_df(df = df)),
+  transformed <- unlist(lapply(scale_list, function(s) s$transform_df(df = df)),
     recursive = FALSE)
   new_data_frame(c(transformed, df[setdiff(names(df), names(transformed))]))
 }
 
+scales_backtransform_df <- function(scales, df) {
+  # NOTE: no need to check empty(data) because it should be already checked
+  # before this function is called.
+
+  # if the scale contains no trans or the trans is of identity, it doesn't need
+  # to be backtransformed.
+  idx_skip <- vapply(scales$scales, function(x) {
+    is.null(x$trans) ||
+      identical(x$trans$inverse, identity)
+  }, logical(1L))
+  scale_list <- scales$scales[!idx_skip]
+
+  if (length(scale_list) == 0L) return(df)
+
+  backtransformed <- unlist(lapply(scale_list, function(scale) {
+    aesthetics <- intersect(scale$aesthetics, names(df))
+
+    if (length(aesthetics) == 0) return()
+
+    lapply(df[aesthetics], scale$trans$inverse)
+  }), recursive = FALSE)
+
+  new_data_frame(c(backtransformed, df[setdiff(names(df), names(backtransformed))]))
+}
+
 # @param aesthetics A list of aesthetic-variable mappings. The name of each
 #   item is the aesthetic, and the value of each item is the variable in data.
 scales_add_defaults <- function(scales, data, aesthetics, env) {

tests/testthat/test-scales.r

@@ -459,3 +459,12 @@ test_that("breaks and labels are correctly checked", {
   p <- ggplot(mtcars) + geom_bar(aes(mpg)) + scale_x_binned(labels = function(x) 1:2)
   expect_snapshot_error(ggplotGrob(p))
 })
+
+test_that("staged aesthetics are backtransformed properly (#4155)", {
+  p <- ggplot(data.frame(value = 16)) +
+    geom_point(aes(stage(value, after_stat = x / 2), 0)) +
+    scale_x_sqrt(limits = c(0, 16), breaks = c(2, 4, 8))
+
+  # x / 2 should be 16 / 2 = 8, thus the result should be sqrt(8) on scale_x_sqrt()
+  expect_equal(layer_data(p)$x, sqrt(8))
+})