geom_raster() of a matrix: Performance analysis and improvements

[zeehio]

Summary

• We see how a plot of a 4k x 3k matrix can be made around 10 times
  faster than when using geom_raster() (45 seconds to 5.6 seconds).

• The differences in performance tell us that geom_raster() may not
  be the best choice to rasterize a matrix.

• We can bring the timing further down to 1.5 seconds if we omit
  handling of missing values and reduce the palette of colours, but
  these are shortcuts ggplot2 can’t make.

• In this issue we will see an efficient way of rasterizing a matrix. We'll see which are the main ggplot2 bottlenecks affecting the performance of that efficient approach and we'll get to some pull requests to address those issues.

• This issue is structured in several messages:

  • This message explains the problem and shows that ggplot is slower than the expected. It shows what performance we can aim to achieve.
  • The second message shows how performance is improved with a custom geom, available as a ggplot extension, going down from 45 seconds to 18 seconds without any ggplot2 changes.
  • The third message profiles the custom geom, and points to two pull requests.
  • Afterwards we'll discuss the scale mapping system proposing improvements to it

• All code is included just for reproducibility, but it is not expected that the you linger in the details.

• If I happen to call your attention, I'm looking for opportunities ideally starting around summer 2023. Happy to do remote work from Barcelona (Spain) or Mexico and open to relocation if needed.

Introduction: A small example

On my field of work it is a common case to have a matrix that we have to
plot with something similar to filled.contour or geom_raster. The
matrix has two associated axes, one for rows and one for columns. We
often need a scale transformation.

Here is a small example of the data:

library(bench)
library(ggplot2)
library(cowplot)
library(reshape2) # for melt
library(forcats)

x_small <- c(3,5,7,9)
y_small <- c(100,200,300)
# The matrix values are usually "random numbers" acquired from an instrument
# Here I choose them so they will have a nice scale when they are log2 transformed:
mat_small <- matrix(
  2^seq(from = 1, to = 64, length.out = length(x_small)*length(y_small)),
  nrow = length(x_small),
  ncol = length(y_small)
)
dimnames(mat_small) <- list(x = x_small, y = y_small)

mat_small

##    y
## x           100          200          300
##   3      2.0000 1.575265e+07 1.240729e+14
##   5    105.9524 8.345155e+08 6.572914e+15
##   7   5612.9576 4.420947e+10 3.482081e+17
##   9 297353.2218 2.342050e+12 1.844674e+19

We can use geom_raster() to plot it. Let’s call this the naïve
strategy, because we just use ggplot in a naïve way. This approach is
great and it works, but we’ll see that it does not scale very well…

naive_strategy <- function(mat) {
    df <- melt(mat, value.name = "intensity")
    gplt <- ggplot(df) +
      geom_raster(aes(x = x, y = y, fill = intensity)) +
      scale_fill_gradient(trans = "log2")
    gplt
}

gplt_naive_small <- naive_strategy(mat_small)
gplt_naive_small

Scaling with the naïve strategy:

The same problem, with a 4000 x 3000 matrix:

x_big <- seq(from = 3, by = 2, length.out = 3000)
y_big <- seq(from = 100, by = 100, length.out = 4000)
# The matrix values are usually "random numbers" acquired from an instrument
mat_big <- matrix(2^seq(1, 64, length.out = length(x_big)*length(y_big)), nrow = length(x_big), ncol = length(y_big))
dimnames(mat_big) <- list(x = x_big, y = y_big)

mat_big[1:5,1:5]

##     y
## x         100      200      300      400      500
##   3  2.000000 2.021954 2.044148 2.066587 2.089272
##   5  2.000007 2.021961 2.044156 2.066594 2.089279
##   7  2.000015 2.021968 2.044163 2.066602 2.089287
##   9  2.000022 2.021976 2.044171 2.066609 2.089294
##   11 2.000029 2.021983 2.044178 2.066617 2.089302

This is how it looks like:

naive_strategy(mat_big)

bm_baseline <- bench::mark(
  naive_strategy = {
    gplt <- naive_strategy(mat_big)
    benchplot(gplt)
  },
  iterations = 1L
)

## Warning: Some expressions had a GC in every iteration; so filtering is disabled.

bm_baseline

## # A tibble: 1 × 6
##   expression          min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 naive_strategy    45.4s    45.4s    0.0220    28.6GB    0.683

Cutting corners strategy:

This is “as fast as I can make it”. It is useful as a reference for what
can be done, but it is not realistic to expect ggplot2 to be this
fast, due to these shortcuts being taken:

• Ignore handling missing values
• Limit to 256 colours in the image

cut_corners_strategy <- function(mat, low = "#132B43", high = "#56B1F7", num_levels = 256L) {
  # Helper function:
  mat_to_nativeRaster <- function(x, colormap, rangex)  {
    cols_to_ints <- farver::encode_native(colormap)
    breaks <- seq(from = rangex[1], to = rangex[2], length.out = length(colormap))
    xdim <- dim(x)
    rev_cols <- seq.int(ncol(x), 1L, by = -1L)
    x <- x[, rev_cols]
    x <- findInterval(x, breaks, rightmost.closed = TRUE)
    x <- cols_to_ints[x]
    x <- matrix(x, nrow = xdim[1], ncol = xdim[2], byrow = FALSE)
    
    structure(
      x,
      dim = c(xdim[2], xdim[1]),
      class = "nativeRaster",
      channels = 4L
    )
  }
  
  minmax <- range(mat)
  mat_trans <- log2(mat)
  palette <- scales::seq_gradient_pal(low = low, high = high)
  colormap <- palette(seq(from=0, to = 1, length.out = num_levels))
  nr <- mat_to_nativeRaster(mat_trans, colormap, log2(minmax))
  xmin <- as.numeric(rownames(mat)[1L])
  xmax <- as.numeric(rownames(mat)[nrow(mat)])
  ymin <- as.numeric(colnames(mat)[1L])
  ymax <- as.numeric(colnames(mat)[ncol(mat)])

  # The geom_rect is fake and it is only used to force the fill legend to appear
  # The geom_rect  limits are used to help set the plot limits
  gplt <- ggplot() +
    geom_rect(
      data = data.frame(
        x = NA_real_
      ),
      mapping = aes(fill = x),
      xmin = xmin, xmax = xmax,
      ymin = ymin, ymax = ymax
    ) +
    annotation_raster(
      nr,
      xmin = xmin, xmax = xmax,
      ymin = ymin, ymax = ymax
    ) +
    scale_fill_gradient( # This has to match with the colormap above
      low = low, high = high,
      limits = minmax,
      na.value = "#00000000",
      trans = "log2"
    ) +
    labs(fill="Intensity") +
    lims(
      x = c(xmin, xmax),
      y = c(ymin, ymax)
    )
  gplt
}

Let’s apply this to the small matrix:

cowplot::plot_grid(
  naive_strategy(mat_small) + labs(title = "naive"),
  cut_corners_strategy(mat_small) + labs(title = "cut corners"),
  ncol = 2
)

bm_cut_corners <- bench::mark(
  cut_corners = {
    gplt <- cut_corners_strategy(mat_big)
    benchplot(gplt)
  },
  iterations = 1L
)

bm_cut_corners

## # A tibble: 1 × 6
##   expression       min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr>  <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 cut_corners    1.43s    1.43s     0.700     687MB        0

Fast and fair strategy

If we avoid cutting corners, we can still get quite decent performance.
Here we take care of missing values (if there was any) and we don’t
limit the palette to 256 colours.

fast_fair_strategy <- function(mat, low = "#132B43", high = "#56B1F7", na_colour = "grey30") {
  # Helper function:
  mat_to_nativeRaster <- function(x, palette, rangex)  {
    rev_cols <- seq.int(ncol(x), 1L, by = -1L)
    x <- scales::rescale(x, to = c(0, 1), from = rangex)
    colours <- palette(x)
    colours <- farver::encode_native(colours)
    colours <- ifelse(is.na(colours), na_colour, colours)
    xdim <- dim(x)
    x <- matrix(colours, nrow = xdim[1], ncol = xdim[2], byrow = FALSE)
    x <- x[, rev_cols]
    
    structure(
      x,
      dim = c(xdim[2], xdim[1]),
      class = "nativeRaster",
      channels = 4L
    )
  }
  
  minmax <- range(mat)
  mat_trans <- log2(mat)
  if (any(is.na(mat_trans) & !is.na(mat))) {
    warning("NA introduced by log transform")
  }
  palette <- scales::seq_gradient_pal(low = low, high = high)
  nr <- mat_to_nativeRaster(mat_trans, palette, log2(minmax))
  xmin <- as.numeric(rownames(mat)[1L])
  xmax <- as.numeric(rownames(mat)[nrow(mat)])
  ymin <- as.numeric(colnames(mat)[1L])
  ymax <- as.numeric(colnames(mat)[ncol(mat)])

  # The geom_rect is fake and it is only used to force the fill legend to appear
  # The geom_rect  limits are used to help set the plot limits
  gplt <- ggplot() +
    geom_rect(
      data = data.frame(
        x = NA_real_
      ),
      mapping = aes(fill = x),
      xmin = xmin, xmax = xmax,
      ymin = ymin, ymax = ymax
    ) +
    annotation_raster(
      nr,
      xmin = xmin, xmax = xmax,
      ymin = ymin, ymax = ymax
    ) +
    scale_fill_gradient( # This has to match with the colormap above
      low = low, high = high,
      limits = minmax,
      na.value = "#00000000",
      trans = "log2"
    ) +
    labs(fill="Intensity") +
    lims(
      x = c(xmin, xmax),
      y = c(ymin, ymax)
    )
  gplt
}

Let’s apply this to the small matrix to assess correctness visually:

cowplot::plot_grid(
  naive_strategy(mat_small) + labs(title = "naive"),
  fast_fair_strategy(mat_small) + labs(title = "fast_fair"),
  ncol = 2
)

bm_fast_fair <- bench::mark(
  fast_fair = {
    gplt <- fast_fair_strategy(mat_big)
    benchplot(gplt)
  },
  iterations = 1L
)

## Warning: Some expressions had a GC in every iteration; so filtering is disabled.

bm_fast_fair

## # A tibble: 1 × 6
##   expression      min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 fast_fair      5.9s     5.9s     0.169    1.79GB    0.678

Comparison of all strategies:

We can see how ggplot2 creates a lot of extra copies, it allocates and deallocates 28GB of RAM. This hints for room for improvement.

strategies <- rbind(
  bm_baseline,
  bm_cut_corners,
  bm_fast_fair
)
strategies

## # A tibble: 3 × 6
##   expression          min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 naive_strategy   45.41s   45.41s    0.0220   28.59GB    0.683
## 2 cut_corners       1.43s    1.43s    0.700   687.35MB    0    
## 3 fast_fair          5.9s     5.9s    0.169     1.79GB    0.678

ggplot(strategies) +
  geom_col(aes(
    x = forcats::fct_reorder(as.character(expression), as.numeric(min)),
    y = as.numeric(min), 
    fill=as.character(expression))
  ) +
  coord_flip() +
  labs(x = "Strategy", y = "CPU time (s)") + 
  guides(fill = "none")

ggplot2 is taking more than 40 seconds, while other approaches need
close to 5 seconds.

There is clearly room for improvement and that’s my goal to address in this issue.

[zeehio]

A custom geom to improve efficiency

library(bench)
library(ggplot2)
library(cowplot)
library(reshape2) # for melt
library(forcats)
library(grid)
library(rlang)

Summary

• We profile the geom_raster() naïve strategy. We see most of the
  time is spent training positional scales. We know that on a matrix
  we just care about the corners.

• We define a custom geom_matrix_raster() that simplifies the
  handling of the positional scales.

• The time spent goes from 45 seconds to 18 seconds. Still far from
  the 5.5 seconds we are targeting.

• On the next message, we will profile the custom geom to understand where is room for optimization

Profiling:

pv <- profvis::profvis({
    gplt <- naive_strategy(mat_big)
    benchplot(gplt)
})
pv

I don’t know how to summarize and represent the output of profvis in a
plot easily, so I’ll just provide here the interpretation of my result:

Total time: 38 seconds

• 28s ggplot_build
  • 11s map_position -> scale_apply
    • 7s lapply
      • 3s self$oob
      • 1s ifelse/which
    • 1.1s unlist
    • 1.1s order…
  • 11s self$map (fill)
    • 3s ifelse
    • 2.8s palette(x)
    • 2.7s pal[match(x, uniq)]
    • 1.4s unique.default
  • 5s train_position
    • 3.8s scale_apply
  • 1.2s by_layer / compute_geom1 / unique.default
• 10s ggplot_gtable
  • 6s draw_geom
    • 1.8s coord$transform
    • 1.7s alpha()
    • 1.2s resolution
  • 2.5s split.data.frame
  • 1s remove_missing

Our naive strategy is spending a huge amount of time in mapping the
position, but our matrix should only care for the corners. Let’s write
first a better strategy, with our own geom.

geom_matrix_raster:

If you want to try this geom, you can try installing my ggmatrix
package:

remotes::install_github("zeehio/ggmatrix")

# Copyright 2022 Sergio Oller Moreno <sergioller@gmail.com>
# This file is part of the ggmatrix package and it is distributed under the MIT license terms.
# Check the ggmatrix package license information for further details.

#' Raster a matrix as a rectangle, efficiently
#'
#'
#' @param matrix The matrix we want to render in the plot
#' @param xmin,xmax,ymin,ymax Coordinates where the corners of the matrix will
#'  be centered By default they are taken from rownames (x) and colnames (y) respectively.
#' @param interpolate If `TRUE`, interpolate linearly, if `FALSE` (the default) don't interpolate.
#' @param flip_cols,flip_rows Flip the rows and columns of the matrix. By default we flip the columns.
#' @inheritParams ggplot2::geom_raster
#'
#' @export
geom_matrix_raster <- function(matrix, xmin = NULL, xmax = NULL, ymin = NULL, ymax = NULL,
                               interpolate = FALSE,
                               flip_cols = TRUE,
                               flip_rows = FALSE,
                               show.legend = NA,
                               inherit.aes = TRUE)
{
  data <- data.frame(values = c(matrix))
  mapping <- aes(fill = .data$values)

  if (is.null(xmin)) {
    xmin <- as.numeric(rownames(matrix)[1L])
  }
  if (is.null(xmax)) {
    xmax <- as.numeric(rownames(matrix)[nrow(matrix)])
  }
  if (is.null(ymin)) {
    ymin <- as.numeric(colnames(matrix)[1L])
  }
  if (is.null(ymax)) {
    ymax <- as.numeric(colnames(matrix)[ncol(matrix)])
  }

  if (nrow(matrix) > 1L) {
    x_step <- (xmax - xmin)/(nrow(matrix) - 1L)
  } else {
    x_step <- 1
  }
  if (ncol(matrix) > 1L) {
    y_step <- (ymax - ymin)/(ncol(matrix) - 1L)
  } else {
    y_step <- 1
  }


  # we return two layers, one blank to create the axes and handle limits, another
  # rastering the matrix.
  corners <- data.frame(
    x = c(xmin - x_step/2, xmax + x_step/2),
    y = c(ymin - y_step/2, ymax + y_step/2)
  )
  corners_xy <- corners
  x_y_names <- names(dimnames(matrix))
  if (is.null(x_y_names)) {
    x_y_names <- c("rows", "columns")
  }
  colnames(corners) <- x_y_names
  x_name <- rlang::sym(x_y_names[1L])
  y_name <- rlang::sym(x_y_names[2L])

  list(
    layer(
      data = corners, mapping = aes(x=!!x_name, y=!!y_name), stat = StatIdentity, geom = GeomBlank,
      position = PositionIdentity, show.legend = show.legend, inherit.aes = inherit.aes,
      params = list(), check.aes = FALSE
    ),
    layer(
      data = data,
      mapping = mapping,
      stat = StatIdentity,
      geom = GeomMatrixRaster,
      position = PositionIdentity,
      show.legend = show.legend,
      inherit.aes = inherit.aes,
      params = list2(
        mat = matrix,
        matrix_nrows = nrow(matrix),
        matrix_ncols = ncol(matrix),
        corners = corners_xy,
        flip_cols = flip_cols,
        flip_rows = flip_rows,
        interpolate = interpolate
      )
    )
  )
}

GeomMatrixRaster <- ggproto(
  "GeomMatrixRaster", Geom,
  non_missing_aes = c("fill"),
  required_aes = c("fill"),
  default_aes = aes(fill = "grey35"),

  draw_panel = function(self, data, panel_params, coord, mat, matrix_nrows, matrix_ncols,
                        corners, flip_cols, flip_rows, interpolate) {
    if (!inherits(coord, "CoordCartesian")) {
      rlang::abort(c(
        "GeomMatrixRaster only works with coord_cartesian"
      ))
    }
    corners <- coord$transform(corners, panel_params)
    if (inherits(coord, "CoordFlip")) {
      byrow <- TRUE
      mat_nr <- matrix_ncols
      mat_nc <- matrix_nrows
      nr_dim <- c(matrix_nrows, matrix_ncols)
    } else {
      byrow <- FALSE
      mat_nr <- matrix_nrows
      mat_nc <- matrix_ncols
      nr_dim <- c(matrix_ncols, matrix_nrows)
    }
    x_rng <- range(corners$x, na.rm = TRUE)
    y_rng <- range(corners$y, na.rm = TRUE)

    mat <- matrix(
      farver::encode_native(data$fill),
      nrow = mat_nr,
      ncol = mat_nc,
      byrow = byrow
    )

    if (flip_cols) {
      rev_cols <- seq.int(mat_nc, 1L, by = -1L)
      mat <- mat[, rev_cols, drop = FALSE]
    }
    if (flip_rows) {
      rev_rows <- seq.int(mat_nr, 1L, by = -1L)
      mat <- mat[rev_rows, drop = FALSE]
    }

    nr <- structure(
      mat,
      dim = nr_dim,
      class = "nativeRaster",
      channels = 4L
    )

    rasterGrob(nr, x_rng[1], y_rng[1],
               diff(x_rng), diff(y_rng), default.units = "native",
               just = c("left","bottom"), interpolate = interpolate)
  },
  draw_key = draw_key_rect
)

efficient_strategy <- function(mat) {
    gplt <- ggplot() +
      geom_matrix_raster(matrix = mat) +
      scale_fill_gradient(trans = "log2")
    gplt
}

cowplot::plot_grid(
  naive_strategy(mat_small) + labs(title = "naive"),
  efficient_strategy(mat_small) + labs(title = "efficient"),
  ncol = 2
)

Same results, how about performance?

bm_efficient <- bench::mark(
  efficient = {
    gplt <- efficient_strategy(mat_big)
    benchplot(gplt)
  },
  iterations = 1L
)

## Warning: Some expressions had a GC in every iteration; so filtering is disabled.

bm_efficient

## # A tibble: 1 × 6
##   expression      min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 efficient     19.4s    19.4s    0.0516    5.61GB    0.826

Comparison of all strategies:

Our new strategy is much better than geom_raster(), but we are still far from our target.

strategies <- rbind(
  bm_baseline,
  bm_fast_fair,
  bm_efficient
)
strategies

## # A tibble: 3 × 6
##   expression          min   median `itr/sec` mem_alloc `gc/sec`
##   <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
## 1 naive_strategy   49.34s   49.34s    0.0203   28.59GB    0.669
## 2 fast_fair         6.35s    6.35s    0.157     1.79GB    0.630
## 3 efficient        19.36s   19.36s    0.0516    5.61GB    0.826

ggplot(strategies) +
  geom_col(aes(
    x = forcats::fct_reorder(as.character(expression), as.numeric(min)),
    y = as.numeric(min), 
    fill=as.character(expression))
  ) +
  coord_flip() +
  labs(x = "Strategy", y = "CPU time (s)") + 
  guides(fill = "none")

Our efficient strategy is twice as fast compared to geom_raster(), but
according to our fast but fair implementation it still could be three
times faster.

We will next profile and target the slowest parts of the efficient
strategy, to improve its performance.

[zeehio]

Profiling the efficient strategy

library(bench)
library(ggplot2)
library(cowplot)
library(reshape2) # for melt
library(forcats)
library(grid)
library(rlang)

Summary

We profile the efficient strategy described before. The main bottlenecks
in the code are:

• The algorithm mapping the matrix values to colours (10/18 seconds)
• Mapping positions (2.5 seconds)
• Drawing the geom (1.8 seconds split.data.frame + 0.4 sec converting
  colours to native)

Profiling the efficient strategy

pv <- profvis::profvis({
    gplt <- efficient_strategy(mat_big)
    benchplot(gplt)
})
pv

I don’t know how to summarize and represent the output of profvis in a
plot easily, so I’ll just provide here the interpretation of my result:

Total time: 18 seconds

• 2.5 seconds training and mapping positions, Spent in
  match_id <- match(layer_data$PANEL, layout$PANEL),

• 10 seconds in self$map, mapping the fill to colours. This
  includes

  • 1.4 seconds finding unique values
  • 3 seconds converting the unique values to colours (2.3s spent in
    the farver package)
  • 2.5 seconds matching the unique values back to the original
    vector
  • 3 seconds replacing any remaining missing value (ifelse) with
    self$na.value.

• 2.8 seconds drawing the geom

  • 1.7s in split.data.frame
  • 0.4s in the farver package, encoding colours to native format.

Proposed changes to ggplot2

• 2.5s: Fix the mapping position code. Pull request: Performance improvement on layers without positional scales #4990
• 1.7s: Fix split() call for drawing the geom. Pull request: Performance: Avoid split.data.frame if there's only one panel #4991

The next message will cover improvements in scale mapping.