average_expression_heatmaps.R

# Include docstring package to render documentation
if (!requireNamespace("docstring", quietly = TRUE)) {
  install.packages("docstring")
}

library(docstring)

average_expression_heatmaps <- function(object,
                                        assay = NULL,
                                        columns_to_keep = NULL,
                                        scale_limits = c(-2, 2),
                                        heatmap_pal = NULL,
                                        colnames_pal = NULL,
                                        annotation_label = "",
                                        extra_heatmap_args = list(),
                                        ...) {
  #' Create three average expression heatmaps
  #'
  #' This function is a wrapper around [Seurat:.AverageExpresssion()] and
  #' [ComplexHeatmap::Heatmap()] that allows more styling flexibility.
  #'
  #' @param object Seurat object.
  #' @param assay Assay to use to calculate [Seurat::AverageExpression()].
  #'   Default is equal to `Seurat::DefaultAssay(object = object),
  #'  for example `RNA` or `ADT`.
  #' @param columns_to_keep A vector with the names of columns to keep after the
  #'   calculation of [Seurat::AverageExpression()], if creating mini heatmaps.
  #' @param scale_limits Output range (numeric vector of length two).
  #' @param heatmap_pal A vector of colors to be interpolated.
  #' @param colnames_pal Named list specifying which color correspond to each column.
  #' @param annotation_label Label to accompany the column annotation.
  #' @param extra_heatmap_args Extra arguments passed to [ComplexHeatmap::Heatmap()],
  #'   except matrix, col, top_annotation, and name.
  #' @inheritDotParams Seurat::AverageExpression -object -assays -return.seurat -slot -add.ident
  #'
  #' @return A named list with three [ComplexHeatmap::Heatmap()] objects:
  #'   1. normalized_heatmap: Uses `data` slot.
  #'   2. scaled_heatmap: Uses `scale.data` slot and then [scales::rescale()].
  #'   3. scaled_markerwise_heatmap: Uses `scale.data` slot and then apply
  #'      [scales::rescale()] per each column.
  #'
  #' @examples
  #' library(Seurat)
  #'
  #' average_expression_heatmaps(
  #'   object = pbmc_small,
  #'   features = Seurat::VariableFeatures(pbmc_small),
  #'   columns_to_keep = c("0", "2"),
  #'   annotation_label = "Cluster",
  #'   extra_heatmap_args = list(
  #'     rect_gp = grid::gpar(col = "white"),
  #'     row_title = "Marker",
  #'     row_title_side = "left",
  #'     row_title_gp = grid::gpar(fontsize = 10),
  #'     column_title = "Cluster",
  #'     column_title_side = "bottom",
  #'     column_title_gp = grid::gpar(fontsize = 10),
  #'     row_dend_side = "right",
  #'     column_dend_side = "bottom",
  #'     row_names_side = "left",
  #'     row_names_gp = grid::gpar(fontsize = 5),
  #'     column_names_side = "top",
  #'     column_names_gp = grid::gpar(fontsize = 7),
  #'     column_names_rot = 0
  #'   )
  #' )
  #' @md

  ## Get default assay ----
  if (is.null(assay)) {
    assay <- Seurat::DefaultAssay(object = object)
  }

  ## Get average expression ----
  object_average <- Seurat::AverageExpression(
    object = object,
    assays = assay,
    return.seurat = TRUE,
    ...
  )

  ## Get data and scaled matrices ----
  data_matrix <- Seurat::GetAssayData(
    object = object_average,
    slot = "data"
  )

  scaled_data_matrix <- Seurat::GetAssayData(
    object = object_average,
    slot = "scale.data"
  )

  column_names <- colnames(data_matrix) |>
    purrr::set_names()

  ## Set plotting defaults ----
  if (is.null(heatmap_pal)) {
    heatmap_pal <- rev(RColorBrewer::brewer.pal(11, "RdYlBu"))
  }

  if (is.null(colnames_pal)) {
    colnames_pal <- CATALYST:::.cluster_cols[
      seq_len(length.out = length(column_names))
    ] |>
      purrr::set_names(nm = column_names)
  }

  ## If `columns_to_keep` is set, then subset matrices ----
  ## now that we set the color palette
  if (is.null(columns_to_keep)) {
    columns_to_keep <- column_names
  }

  ## Columns annotation ----
  if ("top_annotation" %in% names(extra_heatmap_args)) {
    column_annotation <- purrr::pluck(extra_heatmap_args, "top_annotation")
    purrr::pluck(extra_heatmap_args, "top_annotation") <- NULL
  } else {
    column_annotation <- ComplexHeatmap::columnAnnotation(
      annotation = ComplexHeatmap::anno_simple(
        x = column_names[columns_to_keep],
        col = colnames_pal
      ),
      annotation_label = annotation_label
    )
  }

  ## Continious color palette creation ----
  custom_color_palette <- function(colors, matrix) {
    circlize::colorRamp2(
      breaks = seq(
        from = min(matrix, na.rm = TRUE),
        to = max(matrix, na.rm = TRUE),
        length.out = length(colors)
      ),
      colors = colors,
      space = "LAB"
    )
  }

  ## Normalized heatmap ----
  normalized_heatmap <- data_matrix[, columns_to_keep] |>
    (\(x) {
      purrr::exec(
        .fn = ComplexHeatmap::Heatmap,
        matrix = x,
        name = "Normalized\nexpression",
        top_annotation = column_annotation,
        col = custom_color_palette(heatmap_pal, x),
        !!!extra_heatmap_args
      )
    })()

  ## Scaled heatmap ----
  scaled_heatmap <- scaled_data_matrix |>
    t() |>
    scales::rescale(to = scale_limits) |>
    t() |>
    (\(x) {
      x[, columns_to_keep]
    })() |>
    (\(x) {
      purrr::exec(
        .fn = ComplexHeatmap::Heatmap,
        matrix = x,
        name = "Scaled\nexpression",
        top_annotation = column_annotation,
        col = custom_color_palette(heatmap_pal, x),
        !!!extra_heatmap_args
      )
    })()

  ## Scaled marker-wise ----
  scaled_marker_wise <- scaled_data_matrix |>
    t() |>
    asplit(MARGIN = 2) |>
    lapply(function(x) scales::rescale(x = x, to = scale_limits)) |>
    as.data.frame() |>
    as.matrix() |>
    t() |>
    (\(x) {
      x[, columns_to_keep]
    })() |>
    (\(x) {
      purrr::exec(
        .fn = ComplexHeatmap::Heatmap,
        matrix = x,
        name = "Scaled\nexpression\nby marker",
        top_annotation = column_annotation,
        col = custom_color_palette(heatmap_pal, x),
        !!!extra_heatmap_args
      )
    })()

  # Return all heatmaps
  list(
    normalized_heatmap = normalized_heatmap,
    scaled_heatmap = scaled_heatmap,
    scaled_marker_wise = scaled_marker_wise
  )
}