Fix stat_function() for transformed scales (#3908)

* Attempt to fix stat_function(). Doesn't work yet.

* fix stat_function to work with and without explicit y mapping

Author: clauswilke

NEWS.md

@@ -8,6 +8,9 @@
 
 * Default continuous color scales (i.e., the `options()` `ggplot2.continuous.colour` and `ggplot2.continuous.fill`, which inform the `type` argument of `scale_fill_continuous()` and `scale_colour_continuous()`) now accept a function, which allows more control over these default `continuous_scale()`s (@cpsievert, #3827)
 
+* `stat_function()` now works with transformed y axes, e.g. `scale_y_log10()`
+  (@clauswilke, #3905).
+  
 * A newly added geom `geom_density_2d_filled()` and associated stat 
   `stat_density_2d_filled()` can draw filled density contours
   (@clauswilke, #3846).

R/stat-function.r

@@ -97,7 +97,7 @@ stat_function <- function(mapping = NULL, data = NULL,
 #' @usage NULL
 #' @export
 StatFunction <- ggproto("StatFunction", Stat,
-  default_aes = aes(y = after_stat(y)),
+  default_aes = aes(y = after_scale(y)),
 
   compute_group = function(data, scales, fun, xlim = NULL, n = 101, args = list()) {
     range <- xlim %||% scales$x$dimension()
@@ -113,9 +113,15 @@ StatFunction <- ggproto("StatFunction", Stat,
 
     if (is.formula(fun)) fun <- as_function(fun)
 
+    y_out <- do.call(fun, c(list(quote(x_trans)), args))
+    if (!is.null(scales$y) && !scales$y$is_discrete()) {
+      # For continuous scales, need to apply transform
+      y_out <- scales$y$trans$transform(y_out)
+    }
+
     new_data_frame(list(
       x = xseq,
-      y = do.call(fun, c(list(quote(x_trans)), args))
+      y = y_out
     ))
   }
 )

tests/testthat/test-stats-function.r

@@ -35,6 +35,51 @@ test_that("works with discrete x", {
   expect_equal(ret$y, 1:2)
 })
 
+test_that("works with transformed scales", {
+  dat <- data_frame(x = 1:10, y = (1:10)^2)
+
+  # first without explicit mapping of y
+  base <- ggplot(dat, aes(x, group = 1)) +
+    stat_function(fun = ~ .x^2, n = 5)
+
+  ret <- layer_data(base)
+  expect_equal(nrow(ret), 5)
+  expect_equal(ret$y, ret$x^2)
+
+  ret <- layer_data(base + scale_x_log10())
+  expect_equal(nrow(ret), 5)
+  expect_equal(ret$y, (10^ret$x)^2)
+
+  ret <- layer_data(base + scale_y_log10())
+  expect_equal(nrow(ret), 5)
+  expect_equal(10^ret$y, ret$x^2)
+
+  ret <- layer_data(base + scale_x_log10() + scale_y_log10())
+  expect_equal(nrow(ret), 5)
+  expect_equal(10^ret$y, (10^ret$x)^2)
+
+  # now with explicit mapping of y
+  base <- ggplot(dat, aes(x, y)) + geom_point() +
+    stat_function(fun = ~ .x^2, n = 5)
+
+  ret <- layer_data(base, 2)
+  expect_equal(nrow(ret), 5)
+  expect_equal(ret$y, ret$x^2)
+
+  ret <- layer_data(base + scale_x_log10(), 2)
+  expect_equal(nrow(ret), 5)
+  expect_equal(ret$y, (10^ret$x)^2)
+
+  ret <- layer_data(base + scale_y_log10(), 2)
+  expect_equal(nrow(ret), 5)
+  expect_equal(10^ret$y, ret$x^2)
+
+  ret <- layer_data(base + scale_x_log10() + scale_y_log10(), 2)
+  expect_equal(nrow(ret), 5)
+  expect_equal(10^ret$y, (10^ret$x)^2)
+})
+
+
 test_that("works with formula syntax", {
   dat <- data_frame(x = 1:10)