ComplexHeatmap: Making Complex Heatmaps

#####################################
# class and methods to map values to colors
#

# == title
# Class to map values to colors
#
# == details
# The `ColorMapping-class` handles color mapping with both discrete values and continuous values.
# Discrete values are mapped by setting a vector of colors and continuous values are mapped by setting
# a color mapping function.
#
# == methods
# The `ColorMapping-class` provides following methods:
#
# - `ColorMapping`: contructor methods.
# - `map_to_colors,ColorMapping-method`: mapping values to colors.
# - `color_mapping_legend,ColorMapping-method`: draw legend or get legend as a `grid::grob` object.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
ColorMapping = setClass("ColorMapping",
	slots = list(
		colors  = "character", # a list of colors
		levels  = "ANY", # levels which colors correspond to
		col_fun = "function", # function to map values to colors
		type    = "character",  # continuous or discrete
		name    = "character",  # used to map to the dataset and taken as the title of the legend
		na_col  = "character"	
	)
)

# == title
# Constructor methods for ColorMapping class
#
# == param
# -name name for this color mapping. The name is automatically generated if it is not specified.
# -colors discrete colors.
# -levels levels that correspond to ``colors``. If ``colors`` is name indexed, 
#         ``levels`` can be ignored.
# -col_fun color mapping function that maps continuous values to colors.
# -breaks breaks for the continuous color mapping. If ``col_fun`` is
#         generated by `circlize::colorRamp2`, ``breaks`` can be ignored.
# -na_col colors for ``NA`` values.
#
# == detail
# ``colors`` and ``levels`` are used for discrete color mapping, ``col_fun`` and 
# ``breaks`` are used for continuous color mapping.
#
# == value
# A `ColorMapping-class` object.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
ColorMapping = function(name, colors = NULL, levels = NULL, 
	col_fun = NULL, breaks = NULL, na_col = "#FFFFFF") {

	.Object = new("ColorMapping")

	if(missing(name)) {
		increase_color_mapping_index()
		name = paste0("color_mapping_", get_color_mapping_index())
	}
	if(!is.null(colors)) {
		if(is.null(levels)) {
			if(is.null(names(colors))) {
				stop("either provide `levels` or provide named `colors`.\n")
			}
			levels = names(colors)
		}
		if(length(colors) != length(levels)) {
			stop("length of colors and length of levels should be the same.\n")
		}
		colors = t(col2rgb(colors, alpha = TRUE))
		colors = rgb(colors[, 1:3, drop = FALSE], alpha = colors[, 4], maxColorValue = 255)
		.Object@colors = colors
		if(is.numeric(levels)) {
			.Object@levels = as.character(levels)
			#attr(.Object@levels, "numeric") = TRUE
		} else {
			.Object@levels = levels
		}
		names(.Object@colors) = levels
		.Object@type = "discrete"
	} else if(!is.null(col_fun)) {
		if(is.null(breaks)) {
			breaks = attr(col_fun, "breaks")
			if(is.null(breaks)) {
				stop("You should provide breaks.\n")
			}
		}
		
		le1 = grid.pretty(range(breaks))
		le2 = pretty(breaks, n = 3)
		if(abs(length(le1) - 5) < abs(length(le2) - 5)) {
			le = le1
		} else {
			le = le2
		}
		
		.Object@colors = col_fun(le)
		.Object@levels = le
		.Object@col_fun = col_fun
		.Object@type = "continuous"
	} else {
		stop("initialization failed. Either specify `colors` + `levels` or `col_fun` + `breaks`\n")
	}

	.Object@name = name
	na_col = t(col2rgb(na_col, alpha = TRUE))
	na_col = rgb(na_col[, 1:3, drop = FALSE], alpha = na_col[, 4], maxColorValue = 255)
	.Object@na_col = na_col[1]

	return(.Object)
}

# == title
# Print ColorMapping object
#
# == param
# -object a `ColorMapping-class` object.
#
# == value
# This function returns no value.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
setMethod(f = "show",
	signature = "ColorMapping",
	definition = function(object) {
	if(object@type == "discrete") {
		cat("Discrete color mapping:\n")
		cat("name:", object@name, "\n")
		cat("levels:\n")
		print(object@levels)
		cat("\n")
		cat("colors:\n")
		col = object@colors; names(col) = NULL
		print(col)
		cat("\n")
	} else if(object@type == "continuous") {
		cat("Continuous color mapping:\n")
		cat("name:", object@name, "\n")
		cat("default breaks:\n")
		print(object@levels)
		cat("\n")
		cat("colors:\n")
		col = object@colors; names(col) = NULL
		print(col)
		cat("\n")
	}
})

# == title
# Map values to colors
#
# == param
# -object a `ColorMapping-class` object.
# -x input values.
#
# == details
# It maps a vector of values to a vector of colors.
#
# == value
# A vector of colors.
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
setMethod(f = "map_to_colors",
	signature = "ColorMapping",
	definition = function(object, x) {
	
	if(is.factor(x)) x = as.vector(x)
	original_attr = attributes(x)
	x2 = vector(length = length(x))

	if(object@type == "discrete") {
		x[grepl("^\\s*$", x)] = NA
		lna = is.na(x)

		if(is.numeric(x)) x = as.character(x)
		if(any(! x[!lna] %in% object@levels)) {
			msg = paste0(object@name, ": cannot map colors to some of the levels:\n", paste(setdiff(x[!lna], object@levels), sep = ", ", collapse = ", "))
			stop(msg)
		}
		
		x2[lna] = object@na_col
		x2[!lna] = object@colors[ x[!lna] ]
	} else {
		lna = is.na(x)
		x2[lna] = object@na_col
		x2[!lna] = object@col_fun(x[!lna])
	}

	# keep original attributes, such as dimension
	attributes(x2) = original_attr
	return(x2)
})


# == title
# Draw legend based on color mapping
#
# == param
# -object a `ColorMapping-class` object.
# -plot whether to plot or just return the size of the legend viewport.
# -title title of the legend, by default it is the name of the legend
# -title_gp graphical parameters for legend title
# -title_position position of the title
# -color_bar a string of "continous" or "discrete". If the mapping is continuous, whether show the legend as discrete color bar or continuous color bar
# -grid_height height of each legend grid.
# -grid_width width of each legend grid.
# -border color for legend grid borders.
# -at break values of the legend
# -labels labels corresponding to break values
# -labels_gp graphcial parameters for legend labels
# -nrow if there are too many legend grids, they can be put as an array, this controls number of rows
# -ncol if there are too many legend grids, they can be put as an array, this controls number of columns
# -legend_height height of the legend, only works when ``color_bar`` is ``continuous`` and ``direction`` is ``vertical``
# -legend_width width of the legend, only works when ``color_bar`` is ``continuous`` and ``direction`` is ``horizontal``
# -legend_direction when ``color_bar`` is ``continuous``, should the legend be vertical or horizontal?
# -param will be parsed if the parameters are specified as a list
# -... pass to `grid::viewport`.
#
# == details
# A viewport is created which contains a legend title, legend grids and corresponding labels.
#
# This function will be improved in the future to support more types of legends.
#
# == value
# A `grid::grob` object which contains the legend
#
# == author
# Zuguang Gu <z.gu@dkfz.de>
#
setMethod(f = "color_mapping_legend",
	signature = "ColorMapping",
	definition = function(object, ...,
	plot = TRUE, 
	title = object@name,
	title_gp = gpar(fontsize = 10, fontface = "bold"),
	title_position = c("topleft", "topcenter", "leftcenter", "lefttop"),
	color_bar = object@type,
	grid_height = unit(4, "mm"),
	grid_width = unit(4, "mm"),
	border = NULL,
	at = object@levels,
	labels = at,
	labels_gp = gpar(fontsize = 10),
	nrow = NULL,
	ncol = 1,
	legend_height = NULL, legend_width = NULL,
	legend_direction = c("vertical", "horizontal"),
	param = NULL) {

	e = environment()
	if(!is.null(param)) {
		for(nm in names(param)) {
			assign(nm, param[[nm]], envir = e)
		}
	}

	title_gp = check_gp(title_gp)
	labels_gp = check_gp(labels_gp)

	# color_bar = match.arg(color_bar)

	if(object@type == "discrete" && color_bar == "continuous") {
		stop("'color_bar' can only be set to 'discrete' only if the color mapping is discrete")
	}

	# get labels
	if(length(at) != length(labels)) {
		stop("Length of 'at' should be same as length of 'labels'.")
	}
	# if it is character color mapping, remove items in `at` which are not in the available optinos
	if(color_bar == "discrete" && is.character(at)) {
		l = which(at %in% object@levels)
		at = at[l]
		labels = labels[l]
	}

	if(color_bar == "discrete") {
		if(object@type == "continuous") {
			at = rev(at)
			labels = rev(labels)
		}
		gf = Legend(at = at, labels = labels, title = title, title_gp = title_gp, grid_height = grid_height,
			grid_width = grid_width, border = border, labels_gp = labels_gp, nrow = nrow, ncol = ncol,
			legend_gp = gpar(fill = map_to_colors(object, at)), title_position = title_position)

	} else {

		gf = Legend(at = at, labels = labels, col_fun = object@col_fun, title = title, title_gp = title_gp, grid_height = grid_height,
			grid_width = grid_width, border = border, labels_gp = labels_gp, direction = legend_direction,
			legend_width = legend_width, legend_height = legend_height, title_position = title_position)

	}
	
	if(plot) {
		pushViewport(viewport(..., width = grobWidth(gf), height = grobHeight(gf), name = paste0("legend_", object@name)))
		grid.draw(gf)
		upViewport()
	}

	#size = unit.c(vp_width, vp_height)
	return(invisible(gf))
})