Implement filled 2d density contours (#3864)

* initial implementation idea; still tons of stuff to do

* geom_density2d_filled, doesn't work yet

* fix geom_density2d_filled()

* fix unit test

* update usage examples

* improve documentation

* add news item, allow manual setting of breaks, add nlevel column to contour output

* compute density scaled to samle size, just like in geom_density().

* Make `level` numeric for filled contours. Fixes #3875.

* Document computed variables.

* properly contour in a two-step process

* cleanup, better docs

* clean-up code, docs

* unit tests

* add a proper `stat_density_2d_filled()`. Closes #3846.

* update news

* fix unit tests

* keep track of the type of contour needed rather than the stat that does the work.

Author: clauswilke

NAMESPACE

@@ -165,11 +165,13 @@ export(GeomBlank)
 export(GeomBoxplot)
 export(GeomCol)
 export(GeomContour)
+export(GeomContourFilled)
 export(GeomCrossbar)
 export(GeomCurve)
 export(GeomCustomAnn)
 export(GeomDensity)
 export(GeomDensity2d)
+export(GeomDensity2dFilled)
 export(GeomDotplot)
 export(GeomErrorbar)
 export(GeomErrorbarh)
@@ -231,6 +233,7 @@ export(StatContourFilled)
 export(StatCount)
 export(StatDensity)
 export(StatDensity2d)
+export(StatDensity2dFilled)
 export(StatEcdf)
 export(StatEllipse)
 export(StatFunction)
@@ -341,7 +344,9 @@ export(geom_crossbar)
 export(geom_curve)
 export(geom_density)
 export(geom_density2d)
+export(geom_density2d_filled)
 export(geom_density_2d)
+export(geom_density_2d_filled)
 export(geom_dotplot)
 export(geom_errorbar)
 export(geom_errorbarh)
@@ -591,7 +596,9 @@ export(stat_contour_filled)
 export(stat_count)
 export(stat_density)
 export(stat_density2d)
+export(stat_density2d_filled)
 export(stat_density_2d)
+export(stat_density_2d_filled)
 export(stat_ecdf)
 export(stat_ellipse)
 export(stat_function)

NEWS.md

@@ -3,6 +3,10 @@
 * `annotation_raster()` adds support for native rasters. For large rasters,
   native rasters render significantly faster than arrays (@kent37, #3388)
 
+* A newly added geom `geom_density_2d_filled()` and associated stat 
+  `stat_density_2d_filled()` can draw filled density contours
+  (@clauswilke, #3846).
+  
 * Support graphics devices that use the `file` argument instead of `fileneame` 
   in `ggsave()` (@bwiernik, #3810)
 

R/geom-contour.r

@@ -1,14 +1,15 @@
-#' 2d contours of a 3d surface
+#' 2D contours of a 3D surface
 #'
-#' ggplot2 can not draw true 3d surfaces, but you can use `geom_contour`
-#' and [geom_tile()] to visualise 3d surfaces in 2d. To be a valid
-#' surface, the data must contain only a single row for each unique combination
-#' of the variables mapped to the `x` and `y` aesthetics. Contouring
+#' ggplot2 can not draw true 3D surfaces, but you can use `geom_contour()`,
+#' `geom_contour_filled()`, and [geom_tile()] to visualise 3D surfaces in 2D.
+#' To specify a valid surface, the data must contain `x`, `y`, and `z` coordinates,
+#' and each unique combination of `x` and `y` can appear exactly once. Contouring
 #' tends to work best when `x` and `y` form a (roughly) evenly
 #' spaced grid. If your data is not evenly spaced, you may want to interpolate
-#' to a grid before visualising.
+#' to a grid before visualising, see [geom_density_2d()].
 #'
 #' @eval rd_aesthetics("geom", "contour")
+#' @eval rd_aesthetics("geom", "contour_filled")
 #' @inheritParams layer
 #' @inheritParams geom_point
 #' @inheritParams geom_path
@@ -20,7 +21,7 @@
 #' @seealso [geom_density_2d()]: 2d density contours
 #' @export
 #' @examples
-#' #' # Basic plot
+#' # Basic plot
 #' v <- ggplot(faithfuld, aes(waiting, eruptions, z = density))
 #' v + geom_contour()
 #'
@@ -33,7 +34,7 @@
 #' v + geom_contour_filled()
 #'
 #' # Setting bins creates evenly spaced contours in the range of the data
-#' v + geom_contour(bins = 2)
+#' v + geom_contour(bins = 5)
 #' v + geom_contour(bins = 10)
 #'
 #' # Setting binwidth does the same thing, parameterised by the distance
@@ -95,7 +96,7 @@ geom_contour_filled <- function(mapping = NULL, data = NULL,
     data = data,
     mapping = mapping,
     stat = stat,
-    geom = GeomPolygon,
+    geom = GeomContourFilled,
     position = position,
     show.legend = show.legend,
     inherit.aes = inherit.aes,
@@ -123,3 +124,11 @@ GeomContour <- ggproto("GeomContour", GeomPath,
     alpha = NA
   )
 )
+
+#' @rdname ggplot2-ggproto
+#' @format NULL
+#' @usage NULL
+#' @export
+#' @include geom-polygon.r
+GeomContourFilled <- ggproto("GeomContourFilled", GeomPolygon)
+

R/geom-density2d.r

@@ -1,27 +1,41 @@
-#' Contours of a 2d density estimate
+#' Contours of a 2D density estimate
 #'
 #' Perform a 2D kernel density estimation using [MASS::kde2d()] and
 #' display the results with contours. This can be useful for dealing with
-#' overplotting. This is a 2d version of [geom_density()].
+#' overplotting. This is a 2D version of [geom_density()]. `geom_density_2d()`
+#' draws contour lines, and `geom_density_2d_filled()` draws filled contour
+#' bands.
 #'
 #' @eval rd_aesthetics("geom", "density_2d")
-#' @seealso [geom_contour()] for information about how contours
-#'  are drawn; [geom_bin2d()] for another way of dealing with
+#' @eval rd_aesthetics("geom", "density_2d_filled")
+#' @seealso [geom_contour()], [geom_contour_filled()] for information about
+#'  how contours are drawn; [geom_bin2d()] for another way of dealing with
 #'  overplotting.
 #' @param geom,stat Use to override the default connection between
 #'   `geom_density_2d` and `stat_density_2d`.
 #' @inheritParams layer
 #' @inheritParams geom_point
 #' @inheritParams geom_path
+#' @param contour_var Character string identifying the variable to contour
+#'   by. Can be one of `"density"`, `"ndensity"`, or `"count"`. See the section
+#'   on computed variables for details.
 #' @export
 #' @examples
 #' m <- ggplot(faithful, aes(x = eruptions, y = waiting)) +
 #'  geom_point() +
 #'  xlim(0.5, 6) +
 #'  ylim(40, 110)
+#'
+#' # contour lines
 #' m + geom_density_2d()
+#'
 #' \donttest{
-#' m + stat_density_2d(aes(fill = after_stat(level)), geom = "polygon")
+#' # contour bands
+#' m + geom_density_2d_filled(alpha = 0.5)
+#'
+#' # contour bands and contour lines
+#' m + geom_density_2d_filled(alpha = 0.5) +
+#'   geom_density_2d(size = 0.25, colour = "black")
 #'
 #' set.seed(4393)
 #' dsmall <- diamonds[sample(nrow(diamonds), 1000), ]
@@ -30,23 +44,29 @@
 #' # set of contours for each value of that variable
 #' d + geom_density_2d(aes(colour = cut))
 #'
-#' # Similarly, if you apply faceting to the plot, contours will be
-#' # drawn for each facet, but the levels will calculated across all facets
-#' d + stat_density_2d(aes(fill = after_stat(level)), geom = "polygon") +
-#'   facet_grid(. ~ cut) + scale_fill_viridis_c()
-#' # To override this behavior (for instace, to better visualize the density
-#' # within each facet), use after_stat(nlevel)
-#' d + stat_density_2d(aes(fill = after_stat(nlevel)), geom = "polygon") +
-#'   facet_grid(. ~ cut) + scale_fill_viridis_c()
+#' # If you draw filled contours across multiple facets, the same bins are
+#' # used across all facets
+#' d + geom_density_2d_filled() + facet_wrap(vars(cut))
+#' # If you want to make sure the peak intensity is the same in each facet,
+#' # use `contour_var = "ndensity"`.
+#' d + geom_density_2d_filled(contour_var = "ndensity") + facet_wrap(vars(cut))
+#' # If you want to scale intensity by the number of observations in each group,
+#' # use `contour_var = "count"`.
+#' d + geom_density_2d_filled(contour_var = "count") + facet_wrap(vars(cut))
 #'
-#' # If we turn contouring off, we can use use geoms like tiles:
-#' d + stat_density_2d(geom = "raster", aes(fill = after_stat(density)), contour = FALSE)
+#' # If we turn contouring off, we can use other geoms, such as tiles:
+#' d + stat_density_2d(
+#'   geom = "raster",
+#'   aes(fill = after_stat(density)),
+#'   contour = FALSE
+#' ) + scale_fill_viridis_c()
 #' # Or points:
 #' d + stat_density_2d(geom = "point", aes(size = after_stat(density)), n = 20, contour = FALSE)
 #' }
 geom_density_2d <- function(mapping = NULL, data = NULL,
-                            stat = "density2d", position = "identity",
+                            stat = "density_2d", position = "identity",
                             ...,
+                            contour_var = "density",
                             lineend = "butt",
                             linejoin = "round",
                             linemitre = 10,
@@ -65,6 +85,8 @@ geom_density_2d <- function(mapping = NULL, data = NULL,
       lineend = lineend,
       linejoin = linejoin,
       linemitre = linemitre,
+      contour = TRUE,
+      contour_var = contour_var,
       na.rm = na.rm,
       ...
     )
@@ -84,3 +106,42 @@ geom_density2d <- geom_density_2d
 GeomDensity2d <- ggproto("GeomDensity2d", GeomPath,
   default_aes = aes(colour = "#3366FF", size = 0.5, linetype = 1, alpha = NA)
 )
+
+#' @export
+#' @rdname geom_density_2d
+geom_density_2d_filled <- function(mapping = NULL, data = NULL,
+                                   stat = "density_2d_filled", position = "identity",
+                                   ...,
+                                   contour_var = "density",
+                                   na.rm = FALSE,
+                                   show.legend = NA,
+                                   inherit.aes = TRUE) {
+  layer(
+    data = data,
+    mapping = mapping,
+    stat = stat,
+    geom = GeomDensity2dFilled,
+    position = position,
+    show.legend = show.legend,
+    inherit.aes = inherit.aes,
+    params = list(
+      na.rm = na.rm,
+      contour = TRUE,
+      contour_var = contour_var,
+      ...
+    )
+  )
+}
+
+#' @export
+#' @rdname geom_density_2d
+#' @usage NULL
+geom_density2d_filled <- geom_density_2d_filled
+
+#' @rdname ggplot2-ggproto
+#' @format NULL
+#' @usage NULL
+#' @export
+#' @include geom-polygon.r
+GeomDensity2dFilled <- ggproto("GeomDensity2dFilled", GeomPolygon)
+

R/stat-contour.r

@@ -2,11 +2,21 @@
 #' @inheritParams geom_contour
 #' @export
 #' @eval rd_aesthetics("stat", "contour")
+#' @eval rd_aesthetics("stat", "contour_filled")
 #' @section Computed variables:
+#' The computed variables differ somewhat for contour lines (computed by
+#' `stat_contour()`) and contour bands (filled contours, computed by `stat_contour_filled()`).
+#' The variables `nlevel` and `piece` are available for both, whereas `level_low`, `level_high`,
+#' and `level_mid` are only available for bands. The variable `level` is a numeric or a factor
+#' depending on whether lines or bands are calculated.
 #' \describe{
-#'  \item{level}{height of contour}
-#'  \item{nlevel}{height of contour, scaled to maximum of 1}
-#'  \item{piece}{contour piece (an integer)}
+#'  \item{`level`}{Height of contour. For contour lines, this is numeric vector that
+#'    represents bin boundaries. For contour bands, this is an ordered factor that
+#'    represents bin ranges.}
+#'  \item{`level_low`, `level_high`, `level_mid`}{(contour bands only) Lower and upper
+#'    bin boundaries for each band, as well the mid point between the boundaries.}
+#'  \item{`nlevel`}{Height of contour, scaled to maximum of 1.}
+#'  \item{`piece`}{Contour piece (an integer).}
 #' }
 #' @rdname geom_contour
 stat_contour <- function(mapping = NULL, data = NULL,
@@ -39,7 +49,7 @@ stat_contour <- function(mapping = NULL, data = NULL,
 #' @rdname geom_contour
 #' @export
 stat_contour_filled <- function(mapping = NULL, data = NULL,
-                                geom = "polygon", position = "identity",
+                                geom = "contour_filled", position = "identity",
                                 ...,
                                 bins = NULL,
                                 binwidth = NULL,
@@ -74,11 +84,15 @@ StatContour <- ggproto("StatContour", Stat,
   required_aes = c("x", "y", "z"),
   default_aes = aes(order = after_stat(level)),
 
-  compute_group = function(data, scales, bins = NULL, binwidth = NULL,
+  setup_params = function(data, params) {
+    params$z.range <- range(data$z, na.rm = TRUE, finite = TRUE)
+    params
+  },
+
+  compute_group = function(data, scales, z.range, bins = NULL, binwidth = NULL,
                            breaks = NULL, na.rm = FALSE) {
 
-    z_range <- range(data$z, na.rm = TRUE, finite = TRUE)
-    breaks <- contour_breaks(z_range, bins, binwidth, breaks)
+    breaks <- contour_breaks(z.range, bins, binwidth, breaks)
 
     isolines <- xyz_to_isolines(data, breaks)
     path_df <- iso_to_path(isolines, data$group[1])
@@ -99,16 +113,23 @@ StatContourFilled <- ggproto("StatContourFilled", Stat,
   required_aes = c("x", "y", "z"),
   default_aes = aes(order = after_stat(level), fill = after_stat(level)),
 
-  compute_group = function(data, scales, bins = NULL, binwidth = NULL, breaks = NULL, na.rm = FALSE) {
+  setup_params = function(data, params) {
+    params$z.range <- range(data$z, na.rm = TRUE, finite = TRUE)
+    params
+  },
 
-    z_range <- range(data$z, na.rm = TRUE, finite = TRUE)
-    breaks <- contour_breaks(z_range, bins, binwidth, breaks)
+  compute_group = function(data, scales, z.range, bins = NULL, binwidth = NULL, breaks = NULL, na.rm = FALSE) {
+    breaks <- contour_breaks(z.range, bins, binwidth, breaks)
 
     isobands <- xyz_to_isobands(data, breaks)
     names(isobands) <- pretty_isoband_levels(names(isobands))
     path_df <- iso_to_polygon(isobands, data$group[1])
 
     path_df$level <- ordered(path_df$level, levels = names(isobands))
+    path_df$level_low <- breaks[as.numeric(path_df$level)]
+    path_df$level_high <- breaks[as.numeric(path_df$level) + 1]
+    path_df$level_mid <- 0.5*(path_df$level_low + path_df$level_high)
+    path_df$nlevel <- rescale_max(path_df$level_high)
 
     path_df
   }