R/signature_heatmap.R

In cellGeometry: Geometric Single Cell Deconvolution

#' Gene signature heatmap
#' 
#' Produces a heatmap of genes signatures for each cell subclass using
#' ComplexHeatmap.
#' 
#' @param x Either a gene signature matrix with genes in rows and cell
#'   subclasses in columns, an object of S3 class 'cellMarkers' generated by
#'   [cellMarkers()], or an object of class 'deconv' generated by
#'   [deconvolute()].
#' @param type Either "subclass" or "group" specifying whether to show the cell
#'   subclass or cell group signature from a 'cellMarkers' or 'deconv' object.
#'   "groupsplit" shows the distribution of mean gene expression for the group
#'   signature across subclasses.
#' @param top Specifies the number of genes per subclass/group to be displayed.
#' @param use_filter Logical whether to show denoised gene signature.
#' @param arith_mean Logical whether to show log2(arithmetic mean), if
#'   calculated, instead of usual mean(log2(counts +1)).
#' @param rank Either "max" or "angle" controlling whether genes (rows) are
#'   ordered in the heatmap by max expression (the default) or lowest angle
#'   (a measure of specificity of the gene as a cell marker).
#' @param scale Character value controlling scaling of genes: "none" for no scaling, "max" to equalise the maximum mean
#'   expression between genes, "sphere" to scale genes to the unit hypersphere
#'   where cell subclasses or groups are dimensions.
#' @param col Vector of colours passed to [ComplexHeatmap::Heatmap()].
#' @param text Logical whether to show values of the maximum cell in each row.
#' @param fontsize Numeric value for font size for cell values when
#'   `text = TRUE`.
#' @param outlines Logical whether to outline boxes with maximum values in each
#'   row. This supercedes `text`.
#' @param outline_col Colour for the outline boxes when `outlines = TRUE`.
#' @param subset Character vector of groups to be subsetted.
#' @param add_genes Character vector of gene names to be added to the heatmap.
#' @param ... Optional arguments passed to [ComplexHeatmap::Heatmap()].
#' @returns A 'Heatmap' class object.
#' @importFrom grDevices hcl.colors
#' @importFrom grid gpar grid.rect
#' @importFrom ComplexHeatmap Heatmap pindex
#' @export

signature_heatmap <- function(x,
                              type = c("subclass", "group", "groupsplit"),
                              top = Inf,
                              use_filter = NULL,
                              arith_mean = FALSE,
                              rank = c("max", "angle"),
                              scale = c("none", "max", "sphere"),
                              col = rev(hcl.colors(10, "Greens3")),
                              text = TRUE,
                              fontsize = 6.5,
                              outlines = FALSE,
                              outline_col = "black",
                              subset = NULL,
                              add_genes = NULL,
                              ...) {
  type <- match.arg(type)
  rank <- match.arg(rank)
  scale <- match.arg(scale)
  cell_table <- NULL
  if (inherits(x, "deconv")) {
    x <- x$mk
    if (is.null(use_filter)) use_filter <- x$opt$use_filter
  }
  if (inherits(x, "cellMarkers")) {
    if (is.null(use_filter)) use_filter <- TRUE
    gset <- if (type == "subclass") x$geneset else x$group_geneset
    if (is.finite(top)) {
      best_angle <- if (type == "subclass") x$best_angle else x$group_angle
      gset <- lapply(seq_along(best_angle), function(i) {
        rownames(best_angle[[i]])[seq_len(top)]
      })
      gset <- unique(unlist(gset))
    }
    if (!is.null(add_genes)) {
      ok <- add_genes %in% rownames(x$genemeans)
      if (!all(ok)) {
        message("Genes not found: ", paste(add_genes[!ok], collapse = ", "))
      }
      add_genes <- add_genes[ok]
      gset <- unique(c(gset, add_genes))
    }
    if (type != "group") cell_table <- x$cell_table
    if (arith_mean) {
      if (type == "group") stop("arithmetic mean not available for group means")
      gmat <- if (use_filter) x$genemeans_filtered_ar else x$genemeans_ar
      if (is.null(gmat)) stop("arithmetic mean not available")
    } else {
      gmat <- if (type == "group") {
        if (use_filter) x$groupmeans_filtered else x$groupmeans
      } else {
        if (use_filter) x$genemeans_filtered else x$genemeans
      }
    }
    gene_signature <- gmat[gset, ]
  } else {
    gene_signature <- x
  }
  if (!is.null(subset)) {
    if (type != "subclass") stop("subset can only be used with subclass heatmaps")
    s <- which(x$cell_table %in% subset)
    if (length(s) == 0) stop("no such subgroup")
    nsub <- min(c(x$opt$nsubclass, top), na.rm = TRUE)
    genes <- lapply(x$best_angle[s], function(i) rownames(i)[1:nsub])
    genes <- unique(c(unlist(genes), add_genes))
    gs <- gene_signature[genes, s, drop = FALSE]
    return(signature_heatmap(x = gs, rank = rank, scale = scale, col = col,
                             text = text, fontsize = fontsize,
                             outlines = outlines, outline_col = outline_col,
                             ...))
  }
  
  whmax <- max.col(gene_signature)
  rmax <- rowMaxs(gene_signature)
  if (rank == "max") {
    ord <- order(whmax, -rmax)
  } else {
    ord <- seq_len(nrow(gene_signature))
  }
  rs <- cell_table[whmax]
  title <- "mean\nexpr"
  if (scale == "max") {
    gene_signature <- gene_signature / rmax
    title <- "max\nscaled\nexpr"
  } else if (scale == "sphere") {
    gene_signature <- scaleSphere(gene_signature)
    title <- "sphere\nscaled\nexpr"
  }
  
  layer_fun <- NULL
  if (text) {
    layer_fun <- function(j, i, x, y, width, height, fill) {
      v <- pindex(gene_signature, i, j)
      ind <- which(v == rowMaxs(gene_signature[i, , drop = FALSE]))
      if (length(ind) > 0) {
        grid.text(sprintf("%.1f", v[ind]), x[ind], y[ind],
                  gp = gpar(fontsize = fontsize))
      }
    }
  }
  
  if (outlines) {
    if (rank == "max") {
      layer_fun <- function(j, i, x, y, width, height, fill) {
        for (jj in unique(j)) {
          ind1 <- j == jj
          ind <- which(whmax[i[ind1]] == jj)
          if (length(ind) > 0) {
            dy <- y[min(ind)] - y[max(ind)] + height
            yy <- (y[min(ind)] + y[max(ind)]) / 2
            grid.rect(x[min(which(ind1))], yy, width, dy,
                      gp = gpar(col = outline_col, lwd = 0.5, fill = NA))
          }
        }
      }
    } else {
      layer_fun <- function(j, i, x, y, width, height, fill) {
        ind <- which(whmax[i] == j)
        if (length(ind) > 0) {
          grid.rect(x[ind], y[ind], width, height,
                    gp = gpar(col = outline_col, lwd = 0.5, fill = NA))
        }
      }
    }
  }
  
  dots <- list(...)
  args <- list(gene_signature,
          cluster_rows = FALSE,
          row_order = ord, row_split = rs,
          cluster_columns = FALSE, column_split = cell_table,
          cluster_column_slices = FALSE,
          column_title_gp = gpar(fontsize = 6),
          row_names_gp = gpar(fontsize = 6),
          column_names_rot = 75, column_names_gp = gpar(fontsize = 6),
          row_title_gp = gpar(fontsize = 6),
          col = col,
          layer_fun = layer_fun,
          heatmap_legend_param = list(title = title))
  if (length(dots)) args[names(dots)] <- dots
  do.call(Heatmap, args)
}