R/heatmaps.R

In demuellae/muRtools: Mueller's R tools

################################################################################
# Utilities and wrappers for plotting heatmaps
################################################################################
#' plotCorPhm
#' 
#' Plot a correlation matrix as heatmap. Wraps around \code{pheatmap}. Note that no clustering will be performed if not supplied with an appropriate clustering dendrogram
#' @param cc        a correlation matrix (as returned by \code{cor()})
#' @param clustDend a clustering dendrogram to be used. Set to \code{NULL} to disable clustering dendrogram
#' @param sampleAnnot a data.frame containing sample information to color by (corrsponds to \code{annotation_row} and \code{annotation_col} parameters of \code{pheatmap()})
#' @param color     see \code{?pheatmap} for details
#' @param breaks    see \code{?pheatmap} for details. In this wrapper, the default value corresponds to splitting color across the full range of correlation coefficients [-1,1]
#' @param border_color     see \code{?pheatmap} for details
#' @param ...       parameters passed on to \code{pheatmap()}
#' @return invisibly the result of a call to \code{pheatmap()}
#' @author Fabian Mueller
#' @export 
plotCorPhm <- function(
		cc,
		clustDend=NULL,
		sampleAnnot=NA,
		color=colorRampPalette(rev(colpal.cont(n=11, name="cb.RdBu")))(100),
		breaks=seq(-1-1e-6, 1+1e-6, length.out=length(color)+1),
		border_color=NA,
		...){
	require(pheatmap)
	clustr <- FALSE
	if (!is.null(clustDend)){
		if (!is.element("dendrogram", class(clustDend))){
			stop("clustDend must be a dendrogram")
		}
		clustr <- as.hclust(clustDend)
	}

	pheatmap(cc, color=color, breaks=breaks, border_color=border_color, annotation_row=sampleAnnot, annotation_col=sampleAnnot, cluster_rows=clustr, cluster_cols=clustr, ...)
}
# todo: check out display_numbers parameter for pheatmap (https://github.com/raivokolde/pheatmap/issues/24)

#' getColorFun
#' 
#' Retrieve a color function mapping values to colors. Uses and modeled after \code{circlize::colorRamp2}.
#' @param x   vector or matrix including potential values
#' @param colPal a color palette to be used. Should be a character vector specifying colors. can be named if specific
#'            colors should be used for specific values
#' @return a function mapping a value to a corresponding color character string
#' @author Fabian Mueller
#' @export
#' @examples
#' randomLetters <- sample(letters[1:6], 50, replace=TRUE)
#' cf_cat <- getColorFun(randomLetters)
#' cf_cat("f")
#' cf_cat("x") # NA
#' cf_cat2 <- getColorFun(randomLetters, colPal=c(a="#009FE3", b="#DE7E00", c="#8EC041", d="#FFCC00", e="#951B81", f="#BE1716"))
#' cf_num <- getColorFun(runif(50))
#' cf_num2 <- getColorFun(runif(50), colpal.cont(9, "cb.YlOrRd")))
#' cf_num2(0.5)
getColorFun <- function(x, colPal=NULL){
	colFun <- NULL
	isCat <- is.character(x) || is.factor(x)

	getMapFun <- function(vals, colors){
		vals <- vals[!is.na(vals)]
		mapVec <- colors
		if (!all(vals %in% names(mapVec))){
			if (length(vals) > length(mapVec)){
				mapVec <- rep(mapVec, length.out=length(vals))
			}
			names(mapVec) <- vals
		}
		mapVec <- mapVec[vals]
		mapVec <- c(mapVec, "#d9d9d9") # add NA color
		names(mapVec)[length(mapVec)] <- NA
		return(function(x){mapVec[x]})
	}

	if (isCat){
		lvls <- sort(as.character(unique(as.vector(x))))
		cvals <- colpal.mu.cat
		if (!is.null(colPal)) cvals <- colPal
		colFun <- getMapFun(lvls, cvals)
	} else {
		if (is.null(colPal)){
			colFun <- circlize::colorRamp2(seq(min(x, na.rm=TRUE), max(x, na.rm=TRUE), length.out=9), colpal.cont(9, "cb.YlGnBu"))
		} else if (is.character(colPal) && length(colPal)>1){
			colFun <- circlize::colorRamp2(seq(min(x, na.rm=TRUE), max(x, na.rm=TRUE), length.out=length(colPal)), colPal)
		}
	}
	return(colFun)
}
#' diagDivHeatmap
#' 
#' Plot a diagonally divided heatmap u
#' @param ml        the first value matrix (the left-lower diagonal matrix will be in the result)
#' @param mr        the first value matrix (the right-upper diagonal matrix will be in the result)
#' @param col.l     color scheme for the left-lower diagonal matrix. Should be generated by \code{circlize::colorRamp2}.
#' 					Alternatively can be a character vector specifying color levels/breaks.
#'                  If \code{NULL} a default color scheme will be used.
#' @param col.r     color scheme for the upper-right diagonal matrix. Should be generated by \code{circlize::colorRamp2}.
#'                  Alternatively can be a character vector specifying color levels/breaks.
#'                  If \code{NULL} a default color scheme will be used.
#' @param name.l    Name for the lower-left submatrix
#' @param name.r    Name for the upper-right submatrix
#' @param cluster   logical or clustering object. will be passed to the \code{cluster_rows} and \code{cluster_columns} arguments of \code{Heatmap}
#' @param cell.val.text  logical indicating whether the cells value should be added as text
#' @param cell.val.text.round if the cell value is numeric, the number of digits to which the cell text is rounded
#' @param ...       parameters passed on to \code{ComplexHeatmap::Heatmap}
#' @return a \code{ComplexHeatmap::Heatmap} object containing the heatmap
#' @author Fabian Mueller
#' @export
#' @examples
#' m1 <- matrix(rnorm(100, mean=0), ncol=10)
#' m2 <- matrix(rnorm(100, mean=2), ncol=10)
#' rownames(m1) <- rownames(m2) <- colnames(m1) <- colnames(m2) <- paste0("Idx", 1:10)
#' diagDivHeatmap(m1, m2)
#' diagDivHeatmap(m1, m2, cell.val.text=TRUE, cell.val.text.round=3)
#' cres <- as.hclust(muRtools::getClusteringDendrogram(m1, distMethod="euclidean", linkMethod="ward.D", corMethod="pearson"))
#' diagDivHeatmap(m1, m2, cluster=cres, cell.val.text=TRUE, cell.val.text.round=2)
diagDivHeatmap <- function(ml, mr, col.l=NULL, col.r=NULL, name.l="Lower left", name.r="Upper right", cluster=FALSE, cell.val.text=FALSE, cell.val.text.round=2, ...){
	require(ComplexHeatmap)
	if (nrow(ml)!=nrow(mr)) stop("Numbers of rows of the two matrices must match")
	if (ncol(ml)!=ncol(mr)) stop("Numbers of columns of the two matrices must match")
	if (nrow(ml)!=ncol(ml)) stop("Matrices must be quadratic")
	if (!is.element(class(cluster), c("hclust", "logical"))) stop("parameter 'cluster' must be hclust or logical")
	if (is.logical(cluster) && cluster) stop("parameter 'cluster' currently can not be set to TRUE (would mess up the ordering)")

	if (!is.function(col.l)) col.l <- getColorFun(ml, col.l)
	if (!is.function(col.r)) col.r <- getColorFun(mr, col.r)
	colPalLegend_l <- col.l
	if (is.character(ml) || is.factor(ml)){
		colPalLegend_l <- environment(col.l)[["mapVec"]]
	}
	colPalLegend_r <- col.r
	if (is.character(mr) || is.factor(mr)){
		colPalLegend_r <- environment(col.r)[["mapVec"]]
	}

	mc <- ml
	mc[upper.tri(mc)] <- mr[upper.tri(mr)]

	cell.text.fun <- identity
	if (is.numeric(mc)) cell.text.fun <- function(x){sprintf(paste0("%.", cell.val.text.round, "f"), x)}

	idxMap <- 1:nrow(ml)
	if (class(cluster)=="hclust") idxMap <- order(cluster$order)
	
	dotArgs <- list(...)
	res <- Heatmap(ml,
		col=colPalLegend_l,
		cluster_rows=cluster, cluster_columns=cluster,
		rect_gp=gpar(type="none"),
		cell_fun = function(j, i, x, y, width, height, fill) {
			grid.rect(x=x, y=y, width=width, height=height, gp=gpar(col=NA, fill="grey"))
			if (idxMap[i] == idxMap[j]) {
				grid.text(rownames(ml)[i], x=x, y=y)
			} else if(idxMap[i] > idxMap[j]) {
				if (!is.na(ml[i, j])){
					grid.rect(x=x, y=y, width=width, height=height, gp=gpar(col=NA, fill=col.l(ml[i, j])))
					if (cell.val.text){
						grid.text(cell.text.fun(ml[i, j]), x, y)
					}
				}
			} else {
				if (!is.na(mr[i, j])) {
					grid.rect(x=x, y=y, width=width, height=height, gp=gpar(col=NA, fill=col.r(mr[i, j])))
					if (cell.val.text){
						grid.text(cell.text.fun(mr[i, j]), x, y)
					}
				}
			}
		},
		name=name.l,
		...
	)
	# dummy heatmap for color legend
	dummyM <- matrix(rep(NA, length.out=nrow(mr)), nrow=nrow(mr), ncol=1)
	rownames(dummyM) <- rownames(mr)
	dummyHm <- Heatmap(dummyM, col=colPalLegend_r, cluster_rows=cluster, cluster_columns=FALSE, width=unit(0, "mm"), name=name.r)
	return(res + dummyHm)
}

#' diagDivCellHeatmap
#' 
#' Plot a heatmap in which each cell is subdivided into a lower-left and upper-right triangle representing
#' the values of two matrices
#' @param ml        the first value matrix (will be the left-lower diagonal in the cells of the resulting heatmap)
#' @param mr        the second value matrix (will be the right-upper diagonal in the cells of the resulting heatmap)
#' @param col.l     color scheme for the left-lower diagonal matrix. Should be generated by \code{circlize::colorRamp2}.
#'                  Alternatively can be a character vector specifying color levels/breaks.
#'                  If \code{NULL} a default color scheme will be used.
#' @param col.r     color scheme for the upper-right diagonal matrix. Should be generated by \code{circlize::colorRamp2}.
#'                  Alternatively can be a character vector specifying color levels/breaks.
#'                  If \code{NULL} a default color scheme will be used.
#' @param name.l    Name for the lower-left submatrix
#' @param name.r    Name for the upper-right submatrix
#' @param ...       parameters passed on to \code{ComplexHeatmap::Heatmap}
#' @return a \code{ComplexHeatmap::Heatmap} object containing the heatmap
#' @author Fabian Mueller
#' @export
#' @examples
#' m1 <- matrix(rnorm(100, mean=0), ncol=10)
#' m2 <- matrix(rnorm(100, mean=2), ncol=10)
#' rownames(m1) <- rownames(m2) <- colnames(m1) <- colnames(m2) <- paste0("Idx", 1:10)
#' diagDivCellHeatmap(m1, m2, cluster_rows=FALSE, cluster_columns=FALSE)
#' cres <- as.hclust(muRtools::getClusteringDendrogram(m1, distMethod="euclidean", linkMethod="ward.D", corMethod="pearson"))
#' diagDivCellHeatmap(m1, m2, cluster_rows=cres, cluster_columns=cres)
diagDivCellHeatmap <- function(ml, mr, col.l=NULL, col.r=NULL, name.l="Lower left", name.r="Upper right", ...){
	require(ComplexHeatmap)
	if (nrow(ml)!=nrow(mr)) stop("Numbers of rows of the two matrices must match")
	if (ncol(ml)!=ncol(mr)) stop("Numbers of columns of the two matrices must match")

	if (!is.function(col.l)) col.l <- getColorFun(ml, col.l)
	if (!is.function(col.r)) col.r <- getColorFun(mr, col.r)
	colPalLegend_l <- col.l
	if (is.character(ml) || is.factor(ml)){
		colPalLegend_l <- environment(col.l)[["mapVec"]]
	}
	colPalLegend_r <- col.r
	if (is.character(mr) || is.factor(mr)){
		colPalLegend_r <- environment(col.r)[["mapVec"]]
	}

	dummyColor <- circlize::colorRamp2(seq(0, 1, length.out=2), rep("grey", 2))
	res <- Heatmap(ml,
		col=colPalLegend_l, rect_gp=gpar(type="none"),
		cell_fun = function(j, i, x, y, width, height, fill) {
			# grid.rect(x=x, y=y, width=width, height=height, gp=gpar(col="grey", fill="grey"))
			grid.polygon(
				x=unit.c(x-0.5*width, x+0.5*width, x-0.5*width),
				y=unit.c(y-0.5*width, y-0.5*width, y+0.5*width),
				gp=gpar(col=NA, fill=col.l(ml[i, j]))
			)
			grid.polygon(
				x=unit.c(x+0.5*width, x-0.5*width, x+0.5*width),
				y=unit.c(y+0.5*width, y+0.5*width, y-0.5*width),
				gp=gpar(col=NA, fill=col.r(mr[i, j]))
			)
		},
		name=name.l,
		...
	)
	# dummy heatmap for color legend
	dummyM <- matrix(rep(NA, length.out=nrow(mr)), nrow=nrow(ml), ncol=1)
	rownames(dummyM) <- rownames(ml)
	dummyHm <- Heatmap(dummyM, col=colPalLegend_r, width=unit(0, "mm"), name=name.r)

	return(res + dummyHm)
}

#' randomGroupedHeatmap
#' 
#' generate a random grouped heatmap using \code{ComplexHeatmap}.
#' @param n.row     number of rows
#' @param n.col     number of columns
#' @param ngrps.row number of groups to group rows into
#' @param ngrps.col number of groups to group columns into
#' @param cols      color scheme. Should be a color character vector.
#'                  If \code{NULL} a default color scheme will be used.
#' @param ...       parameters passed on to \code{ComplexHeatmap::Heatmap}
#' @return a \code{ComplexHeatmap::Heatmap} object containing the heatmap
#' @author Fabian Mueller
#' @export
#' @examples
#' randomGroupedHeatmap(n.row=10, n.col=3, ngrps.row=2, ngrps.col=3, cols=colpal.cont(n=9, name="viridis"))
#' randomGroupedHeatmap(n.row=100, n.col=18, ngrps.row=3, ngrps.col=3, cols=colpal.cont(n=9, name="cb.BrBG"))
#' pdftemp()
#' draw(randomGroupedHeatmap(n.row=100, n.col=6, ngrps.row=3, ngrps.col=3, cols=colpal.PhFr.a))
#' dev.off()
randomGroupedHeatmap <- function(n.row=20, n.col=6, ngrps.row=2, ngrps.col=3, cols=NULL, ...){
	require(ComplexHeatmap)
	npergroup_r <- ceiling(n.row/ngrps.row)
	ridx <- split(1:n.row, ceiling(1:n.row/npergroup_r))
	npergroup_c <- ceiling(n.col/ngrps.col)
	cidx <- split(1:n.col, ceiling(1:n.col/npergroup_c))
	norm_means <- matrix(rnorm(ngrps.row*ngrps.col), nrow=ngrps.row, ncol=ngrps.col)
	mm <- do.call(rbind, lapply(1:nrow(norm_means), FUN=function(i){
		do.call(cbind, lapply(1:ncol(norm_means), FUN=function(j){
			ng_r <- length(ridx[[i]])
			ng_c <- length(cidx[[j]])
			matrix(rnorm(ng_r*ng_c, mean=norm_means[i,j]), nrow=ng_r, ncol=ng_c)
		}))
	}))
	if (is.null(cols)) cols <- colpal.cont(n=9, name="solarextra")
	hm <- Heatmap(mm, name="Random",
		col = circlize::colorRamp2(seq(min(mm), max(mm), length.out=length(cols)), cols),
		cluster_columns=TRUE, show_column_dend=FALSE,
		cluster_rows=TRUE, show_row_dend=FALSE,
		show_row_names = FALSE, show_column_names = FALSE,
		...
	)
	return(hm)
}