R/enhanced.R

In eilslabs/ComplexHeatmap: Making Complex Heatmaps

# == title
# Enhanced version of basic barplot and boxplot
#
# == param
# -data a matrix, a list or a simple numeric vector. If your data is a data frame
#       please convert it to a matrix in the first place.
# -... pass to `Heatmap`
# -ylim ranges on y axis
# -ylab label on y axis
# -title title of the plot
# -title_gp graphic parameters for the title
# -type type of the plot
# -width relative width of the bar or box
# -gp graphic parameters for hte bar or box
# -pch shape of outlier points in the boxplot
# -size size of hte outlier points in the boxplot
# -axis_gp graphic parameters for the axis
# -padding padding of the plot
# -heatmap_legend_list a list of `grid::grob` which contains legend. It can be generated by `color_mapping_legend,ColorMapping-method`.
#
# == details
# This function adds annotations to the barplot or boxplot.
#
# This function is still quite experimental.
# 
# == value
# No value is returned
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
# == example
# mat = matrix(runif(100), 10)
# enhanced_basicplot(mat)
# ha = HeatmapAnnotation(char = sample(letters[1:2], 10, replace = TRUE),
#                        num = runif(10))
# enhanced_basicplot(mat, top_annotation = ha)
# enhanced_basicplot(mat, type = "barplot", top_annotation = ha)
#
enhanced_basicplot = function(data, ..., ylim = NULL,
    ylab = deparse(substitute(data)), title = NULL, title_gp = gpar(fontsize = 14),
    type = c("boxplot", "barplot"), width = 0.8, gp = gpar(), 
    pch = 1, size = unit(2, "mm"), axis_gp = gpar(fontsize = 8),
    padding = unit(c(2, 18, 2, 2), "mm"),
    heatmap_legend_list = list()) {

    if(!(inherits(data, c("matrix")) || inherits(data, "list") || is.atomic(data))) {
        stop("`data` should be a matrix or a list.")
    }

    type = match.arg(type)[1]

    if(is.matrix(data)) {
        n = ncol(data)
    } else if(is.list(data)) {
        n = length(data)
        if(type == "barplot") {
            stop("You can only make boxplot with a list.")
        }
    } else if(is.atomic(data)) {
        n = length(data)
        if(type == "boxplot") {
            stop("You can only make barplot with a vector.")
        }
    }

    mat_foo = matrix("foo", nrow = 2, ncol = n)
    if(is.matrix(data)) {
        colnames(mat_foo) = colnames(data)
    } else if(is.list(data)) {
        colnames(mat_foo) = names(data)
    } else if(is.atomic(data)) {
        colnames(mat_foo) = names(data)
    }
        
    ht = Heatmap(mat_foo, col = c("foo" = "white"), name = "main", cluster_rows = FALSE, cluster_columns = FALSE,
                rect_gp = gpar(type = "none"), show_heatmap_legend = FALSE, ...)

    draw(ht, padding = padding, column_title = title, column_title_gp = title_gp,
        heatmap_legend_list = heatmap_legend_list)

    if(is.matrix(data)) {
        if(type == "barplot") {
            ymin = 0
            ymax = max(colSums(data))
        } else {
            ymin = min(data)
            ymax = max(data)
        }
    } else if(is.list(data)) {
        ymin = min(unlist(data))
        ymax = max(unlist(data))
    } else if(is.atomic(data)) {
        ymin = 0
        ymax = max(data)
    }

    if(!is.null(ylim)) {
        ymin = ylim[1]
        ymax = ylim[2]
    }

    single_barplot = function(x, h, gp) {
        grid.rect(x, h, width = width, height = h, just = "top", default.units = "native", gp = gp)
    }

    single_stacked_bar = function(x, value, gp) {
        y = cumsum(value)
        grid.rect(x, y, width = width, height = value, just = "top",
            default.units = "native", gp = gp)
    }

    single_boxplot = function(x, value, gp) {
        boxplot_stats = as.vector(boxplot(value, plot = FALSE)$stats)
        grid.rect(x, boxplot_stats[2], width = width, height = boxplot_stats[4] - boxplot_stats[2], just = "bottom",
            default.units = "native", gp = gp)
        grid.segments(c(x - width/3, x, x, x - width/3, x - width/2),
                      c(boxplot_stats[1], boxplot_stats[1], boxplot_stats[4], boxplot_stats[5], boxplot_stats[3]),
                      c(x + width/3, x, x, x + width/3, x + width/2),
                      c(boxplot_stats[1], boxplot_stats[2], boxplot_stats[5], boxplot_stats[5], boxplot_stats[3]),
                      default.units = "native")
        l1 = value > boxplot_stats[5]
        if(sum(l1)) grid.points(rep(x, sum(l1)), value[l1], pch = pch, size = size)
        l2 = value < boxplot_stats[1]
        if(sum(l2)) grid.points(rep(x, sum(l2)), value[l2], pch = pch, size = size)
    }

    if(!(is.matrix(data) && type == "barplot")) {
        gp = recycle_gp(gp, n)
    }

    yrange = ymax - ymin
    if(type == "barplot") {
        ymin = 0
        yrange = 0
    }

    decorate_heatmap_body("main", {
        pushViewport(viewport(xscale = c(0, n),
                              yscale = c(ymin - 0.05*yrange, ymax + 0.05*yrange),
                              name = "main_plotting_region"))
        grid.rect()
        axis = yaxisGrob(gp = axis_gp)
        grid.draw(axis)
        grid.text(ylab, x = unit(0, "npc") - unit(10, "mm"), y = unit(0.5, "npc"), just = "bottom", rot = 90)
        for(i in seq_len(n)) {
            x = i - 0.5
            if(is.matrix(data)) {
                if(type == "barplot") {
                    single_stacked_bar(x, data[, i], gp = gp)
                } else {
                    single_boxplot(x, data[, i], gp = subset_gp(gp, i))
                }
            } else if(is.list(data)) {
                single_boxplot(x, data[[i]], gp = subset_gp(gp, i))
            } else if(is.atomic(data)) {
                single_barplot(x, data[i], gp = subset_gp(gp, i))
            }
        }
        
        upViewport()
    }) 
}