Add rasterization support for simple heatmap annotations

R/HeatmapAnnotation-class.R

@@ -71,6 +71,14 @@ HeatmapAnnotation = setClass("HeatmapAnnotation",
 # -width Width of the whole heatmap annotations.
 # -simple_anno_size Size of the simple annotation.
 # -simple_anno_size_adjust Whether also adjust the size of simple annotations when adjusting the whole heatmap annotation.
+# -use_raster Whether render the annotations as raster images. It helps to reduce file size when there are a huge number
+#      of columns (or rows for row annotations). The value is passed to all single annotations.
+# -raster_device Graphic device which is used to generate the raster image.
+# -raster_quality A value larger than 1.
+# -raster_device_param A list of further parameters for the selected graphic device.
+# -raster_by_magick Whether to use `magick::image_resize` to scale the image.
+# -raster_magick_filter Pass to ``filter`` argument of `magick::image_resize`. A character scalar and all possible values
+#      are in `magick::filter_types`. The default is ``"Lanczos"``.
 #
 # == details
 # For arguments ``show_legend``, ``border``, ``annotation_name_offset``, ``annotation_name_side``, ``annotation_name_rot``,
@@ -116,7 +124,14 @@ HeatmapAnnotation = function(...,
 	height = NULL,
 	width = NULL,
 	simple_anno_size = ht_opt$simple_anno_size,
-	simple_anno_size_adjust = FALSE
+	simple_anno_size_adjust = FALSE,
+
+	use_raster = FALSE,
+	raster_device = NULL,
+	raster_quality = 1,
+	raster_device_param = list(),
+	raster_by_magick = requireNamespace("magick", quietly = TRUE),
+	raster_magick_filter = NULL
 	) {
 
 	dev.null()
@@ -136,6 +151,13 @@ HeatmapAnnotation = function(...,
 
 	fun_args = names(as.list(environment()))
 
+	# annotations built inside wrapper functions (oncoPrint(), UpSet(), pheatmap())
+	# never get `use_raster` passed to them, so the global option is the only way
+	# to rasterize those
+	if(missing(use_raster) && !is.null(ht_opt$annotation_use_raster)) {
+		use_raster = ht_opt$annotation_use_raster
+	}
+
 	verbose = ht_opt$verbose
 
 	.Object = new("HeatmapAnnotation")
@@ -331,7 +353,13 @@ HeatmapAnnotation = function(...,
 	        	name_offset = annotation_name_offset[[i_anno]], 
 	        	name_side = annotation_name_side[i_anno], 
 	        	name_rot = annotation_name_rot[[i_anno]],
-	        	border = border[i_anno])
+	        	border = border[i_anno],
+	        	use_raster = use_raster,
+	        	raster_device = raster_device,
+	        	raster_quality = raster_quality,
+	        	raster_device_param = raster_device_param,
+	        	raster_by_magick = raster_by_magick,
+	        	raster_magick_filter = raster_magick_filter)
 
 		if(inherits(anno_value_list[[ag]], c("function", "AnnotationFunction"))) {
 			arg_list$fun = anno_value_list[[ag]]

R/SingleAnnotation-class.R

@@ -36,7 +36,8 @@ SingleAnnotation = setClass("SingleAnnotation",
         width = "ANY",
         height = "ANY",
         extended = "ANY",
-        subsettable = "logical"
+        subsettable = "logical",
+        raster_param = "list"
 	),
 	prototype = list(
 		color_mapping = NULL,
@@ -45,7 +46,8 @@ SingleAnnotation = setClass("SingleAnnotation",
         color_is_random = FALSE,
 		name_to_data_vp = FALSE,
         extended = unit(c(0, 0, 0, 0), "mm"),
-        subsettable = FALSE
+        subsettable = FALSE,
+        raster_param = list(use_raster = FALSE)
 	)
 )
 
@@ -80,6 +82,14 @@ SingleAnnotation = setClass("SingleAnnotation",
 #        the width must be an absolute unit.
 # -height The height of the plotting region (the viewport) that the annotation is drawn. If it is a column annotation,
 #        the width must be an absolute unit.
+# -use_raster Whether render the annotation as a raster image. It helps to reduce file size when there are a
+#      huge number of columns (or rows for a row annotation).
+# -raster_device Graphic device which is used to generate the raster image.
+# -raster_quality A value larger than 1.
+# -raster_device_param A list of further parameters for the selected graphic device.
+# -raster_by_magick Whether to use `magick::image_resize` to scale the image.
+# -raster_magick_filter Pass to ``filter`` argument of `magick::image_resize`. A character scalar and all possible values
+#      are in `magick::filter_types`. The default is ``"Lanczos"``.
 #
 # == details
 # A single annotation is a basic unit of complex heatmap annotations where the heamtap annotations
@@ -156,7 +166,13 @@ SingleAnnotation = function(name, value, col, fun,
 	name_side = ifelse(which == "column", "right", "bottom"),
     name_rot = NULL,
     simple_anno_size = ht_opt$simple_anno_size,
-    width = NULL, height = NULL) {
+    width = NULL, height = NULL,
+    use_raster = FALSE,
+    raster_device = NULL,
+    raster_quality = 1,
+    raster_device_param = list(),
+    raster_by_magick = requireNamespace("magick", quietly = TRUE),
+    raster_magick_filter = NULL) {
 
     .ENV$current_annotation_which = NULL
 	which = match.arg(which)[1]
@@ -582,6 +598,18 @@ SingleAnnotation = function(name, value, col, fun,
         .Object@subsettable = .Object@fun@subsettable
     }
 
+    if(is.null(raster_device)) {
+        raster_device = if(requireNamespace("Cairo", quietly = TRUE)) "CairoPNG" else "png"
+    }
+    .Object@raster_param = list(
+        use_raster = use_raster,
+        raster_device = raster_device,
+        raster_quality = raster_quality,
+        raster_device_param = raster_device_param,
+        raster_by_magick = raster_by_magick,
+        raster_magick_filter = raster_magick_filter
+    )
+
     return(.Object)
 }
 
@@ -666,8 +694,20 @@ setMethod(f = "draw",
         xscale = data_scale$x, yscale = data_scale$y))
 
     if(verbose) qqcat("execute annotation function\n")
-    draw(object@fun, index = index, k = k, n = n)
-
+    rp = object@raster_param
+    if(isTRUE(rp$use_raster)) {
+        rasterize_in_viewport(
+            draw_fun = function() draw(object@fun, index = index, k = k, n = n),
+            raster_device = rp$raster_device,
+            raster_quality = rp$raster_quality,
+            raster_device_param = rp$raster_device_param,
+            raster_by_magick = rp$raster_by_magick,
+            raster_magick_filter = rp$raster_magick_filter
+        )
+    } else {
+        draw(object@fun, index = index, k = k, n = n)
+    }
+
 	# add annotation name
     draw_name = object@name_param$show
 	if(object@name_param$show && n > 1) {

R/global.R

@@ -209,6 +209,8 @@ ht_opt = setGlobalOptions(
 		.value = FALSE
 	),
 	"validate_names" = TRUE,
+	annotation_use_raster = list(
+		.value = NULL),
 	raster_temp_image_max_width = 30000,
 	raster_temp_image_max_height = 30000,
 	COLOR = c("blue", "#EEEEEE", "red")

R/utils.R

@@ -1137,6 +1137,95 @@ setAs("list", "HeatmapList", function(from) {
 })
 
 
+rasterize_in_viewport = function(draw_fun,
+    raster_device = "png",
+    raster_quality = 1,
+    raster_device_param = list(),
+    raster_by_magick = FALSE,
+    raster_magick_filter = NULL) {
+
+    # calculate current viewport size in pixels
+    vp_width_pt = max(1, ceiling(convertWidth(unit(1, "npc"), "bigpts", valueOnly = TRUE)))
+    vp_height_pt = max(1, ceiling(convertHeight(unit(1, "npc"), "bigpts", valueOnly = TRUE)))
+
+    if(raster_quality < 1) raster_quality = 1
+    vp_width_pt = ceiling(vp_width_pt * raster_quality)
+    vp_height_pt = ceiling(vp_height_pt * raster_quality)
+
+    device_info = switch(raster_device,
+        png = c("grDevices", "png", "readPNG"),
+        jpeg = c("grDevices", "jpeg", "readJPEG"),
+        tiff = c("grDevices", "tiff", "readTIFF"),
+        CairoPNG = c("Cairo", "png", "readPNG"),
+        CairoJPEG = c("Cairo", "jpeg", "readJPEG"),
+        CairoTIFF = c("Cairo", "tiff", "readTIFF"),
+        agg_png = c("ragg", "png", "readPNG")
+    )
+
+    if(!requireNamespace(device_info[1], quietly = TRUE)) {
+        stop_wrap(paste0("Need ", device_info[1], " package to write image."))
+    }
+    if(!requireNamespace(device_info[2], quietly = TRUE)) {
+        stop_wrap(paste0("Need ", device_info[2], " package to read image."))
+    }
+
+    if(raster_device %in% c("png", "jpeg", "tiff")) {
+        if(!"type" %in% names(raster_device_param)) {
+            if(capabilities("cairo")) {
+                raster_device_param$type = "cairo"
+            }
+        }
+    }
+
+    temp_image_width = as.integer(vp_width_pt)
+    temp_image_height = as.integer(vp_height_pt)
+
+    if(!is.na(ht_opt$raster_temp_image_max_width)) {
+        temp_image_width = min(temp_image_width, ht_opt$raster_temp_image_max_width)
+    }
+    if(!is.na(ht_opt$raster_temp_image_max_height)) {
+        temp_image_height = min(temp_image_height, ht_opt$raster_temp_image_max_height)
+    }
+
+    temp_image = tempfile(pattern = ".annotation_raster_", tmpdir = tempdir(),
+        fileext = paste0(".", device_info[2]))
+    device_fun = getFromNamespace(raster_device, ns = device_info[1])
+
+    # Scale res with raster_quality so the device's physical size (in inches)
+    # matches the original viewport. Without this, annotation draw functions
+    # that push viewports with absolute units (e.g. unit(5, "mm")) would only
+    # fill a fraction of the enlarged temp device.
+    if(!"res" %in% names(raster_device_param)) {
+        raster_device_param$res = as.integer(72 * raster_quality)
+    }
+
+    oe = try(do.call(device_fun, c(list(filename = temp_image,
+        width = temp_image_width, height = temp_image_height), raster_device_param)))
+    if(inherits(oe, "try-error")) {
+        stop_wrap(qq("The temporary image size for annotation rasterization is too large (@{temp_image_width} x @{temp_image_height} px)."))
+    }
+
+    draw_fun()
+    dev.off2()
+
+    if(raster_by_magick) {
+        image = magick::image_read(temp_image)
+        image = magick::image_resize(image,
+            paste0(vp_width_pt, "x", vp_height_pt, "!"),
+            filter = raster_magick_filter)
+        image = as.raster(image)
+    } else {
+        image = getFromNamespace(device_info[3], ns = device_info[2])(temp_image)
+    }
+
+    grid.raster(image, width = unit(1, "npc"), height = unit(1, "npc"), interpolate = FALSE)
+
+    file.remove(temp_image)
+
+    invisible(NULL)
+}
+
+
 draw_heatmap_in_jupyter = function(ht, ...) {
     width = getOption("repr.plot.width")
     height = getOption("repr.plot.height")

tests/testthat/testthat-annotation-raster.R

@@ -0,0 +1,56 @@
+
+context("Test annotation rasterization")
+
+test_that("raster parameters are passed down to single annotations", {
+	ha = HeatmapAnnotation(foo = 1:10, use_raster = TRUE, raster_quality = 3)
+	rp = ha@anno_list[["foo"]]@raster_param
+	expect_true(rp$use_raster)
+	expect_equal(rp$raster_quality, 3)
+
+	ha = HeatmapAnnotation(foo = 1:10)
+	expect_false(ha@anno_list[["foo"]]@raster_param$use_raster)
+})
+
+# annotations built inside UpSet(), oncoPrint() and pheatmap() are never given a
+# `use_raster` argument, so `ht_opt$annotation_use_raster` is the only way to
+# rasterize them
+test_that("ht_opt$annotation_use_raster reaches internally built annotations", {
+	on.exit(ht_opt$annotation_use_raster <- NULL)
+
+	m = make_comb_mat(list(a = 1:5, b = 3:8, c = 4:10))
+	expect_false(UpSet(m)@top_annotation@anno_list[[1]]@raster_param$use_raster)
+
+	ht_opt$annotation_use_raster = TRUE
+	expect_true(UpSet(m)@top_annotation@anno_list[[1]]@raster_param$use_raster)
+
+	# an explicit argument still wins over the global
+	expect_false(HeatmapAnnotation(foo = 1:10, use_raster = FALSE)@anno_list[["foo"]]@raster_param$use_raster)
+})
+
+# the graphics drawn by `draw_fun` are in absolute units, so the rasterized
+# result must occupy the same area no matter how large the temporary image is
+test_that("rasterized annotation does not shrink when raster_quality increases", {
+	skip_if_not_installed("png")
+
+	filled_fraction = function(quality) {
+		f = tempfile(fileext = ".png")
+		png(f, width = 200, height = 200)
+		pushViewport(viewport(width = unit(1, "npc"), height = unit(1, "npc")))
+		rasterize_in_viewport(
+			function() grid.rect(width = unit(10, "mm"), height = unit(10, "mm"),
+				gp = gpar(fill = "red", col = NA)),
+			raster_device = "png", raster_quality = quality)
+		popViewport()
+		dev.off()
+
+		img = png::readPNG(f)
+		file.remove(f)
+		mean(img[, , 1] > 0.5 & img[, , 2] < 0.5)
+	}
+
+	q1 = filled_fraction(1)
+	q3 = filled_fraction(3)
+
+	expect_gt(q1, 0.01)
+	expect_lt(abs(q1 - q3)/q1, 0.1)
+})