R/cluster_dimplot.R

In zdebruine/scNMF: Fast Non-negative matrix factorization toolkit for single cell data

#' Plot UMAP dimensional reduction with cluster centers as bubbles
#'
#' @description This visualization mode is useful for qualitative assessment of how well a given clustering samples transcriptional space of a given dataset. Note that any ggplot grammar can be applied to the returned plot object.
#'
#' @param seuratObj A seurat object with UMAP dimensional reduction and cluster idents in the active slot (unless providing centers)
#' @param assay Assay to use for cluster center averaging (default is the active assay)
#' @param slot Slot in the assay to use for center averaging (default is counts)
#' @param verbose Boolean
#' @param reduction.name If other than "umap", the name of the reduction containing two-dimensional cell embeddings
#' @param centers Optionally provide a matrix of averaged centers to be projected onto the UMAP coordinates (returned from average.expression with detail = TRUE)
#' @param cells.color Color code for single cells (default #C0C0C0)
#' @param bubble.color Color code for bubbles as cluster centers (default #F26815)
#' @param bubble.alpha Transparency for bubbles (1 = no transparency, 0 = fully transparent, default = 0.8)
#' @param scale.breaks Legend for bubble size based on number of cells in the cluster center represented by that bubble. Breaks specify what bubble sizes will be shown. Default is c(10,50,100) corresponding to 10, 50, and 100 cells.
#' @param min.scale.size Minimum number of cells a bubble should be scaled to (default is 1). Note that bubbles with fewer than this number of cells will not be displayed.
#' @param max.scale.size Maximum number of cells a bubble should be scaled to (default is 1000). Note that bubbles with more than this number of cells will not be displayed.
#' @param bubble.scale Scale to apply to cluster sizes (i.e. "sqrt", "log2", "log", or default which is NULL for linear scaling)
#' @return ggplot object
cluster.dimplot <- function(seuratObj,
                          centers = NULL,
                          assay = "active.assay",
                          slot = "counts",
                          ident = "active.ident",
                          verbose = TRUE,
                          reduction.name = "umap",
                          cells.color = "#C0C0C0",
                          bubble.color = "#F26815",
                          bubble.alpha = 0.8,
                          scale.breaks = c(10, 50, 100),
                          min.scale.size = 1,
                          max.scale.size = 1000,
                          bubble.scale = NULL) {
  require(ggplot2)

  umap.coords <- data.frame(seuratObj@reductions[[reduction.name]]@cell.embeddings)
  if (ncol(umap.coords) != 2) stop("Dimensional reduction object is not of valid dimensions")
  colnames(umap.coords) <- c("x", "y")
  umap.coords[, "size"] <- rep(1, nrow(umap.coords))
  umap.coords[, "group"] <- rep("input", nrow(umap.coords))

  # calculate averages for cluster centers
  if (is.null(centers)) centers <- average.expression(seuratObj, ident = ident, assay, slot, verbose, details = TRUE)

  # get sizes of each cluster
  if (is.null(names(centers))) stop("Could not find centers$cells.per.center. Most likely reason is that an average.expression(..., details = TRUE) object is not being submitted")
  cluster_sizes <- centers$cells.per.center

  # for each cluster center matrix, calculate UMAP coordinates as the average of the UMAP coordinates of all cells in that cluster
  center.umap <- list()
  for (j in 1:length(centers$cells.in.center)) {
    umap.coords.subset <- umap.coords[centers$cells.in.center[[j]],]
    center.umap[[j]] <- c(mean(umap.coords.subset$x), mean(umap.coords.subset$y), cluster_sizes[j], "center")
  }
  center.umap <- data.frame(do.call(rbind, center.umap))
  colnames(center.umap) <- c("x", "y", "size", "group")

  dat <- rbind(umap.coords, center.umap)

  p <- ggplot(dat = dat, aes(x = as.numeric(x), y = as.numeric(y), color = as.factor(group))) +
  geom_point(shape = 16, stroke = 0, aes(size = as.numeric(size)), alpha = bubble.alpha) +
  scale_color_manual(values = c(bubble.color, cells.color), guide = "none") +
  theme_classic() +
  xlab("umap 1") +
  ylab("umap 2") +
  theme(plot.title = element_text(hjust = 0.5, size = 12))

  ifelse(!is.null(bubble.scale),
  p <- p + scale_size_continuous(limits = c(min.scale.size, max.scale.size), range = c(.1, 10), name = "Cells/cluster", breaks = scale.breaks, labels = paste0(scale.breaks, " cells"), trans = bubble.scale),
  p <- p + scale_size_continuous(limits = c(min.scale.size, max.scale.size), range = c(.1, 10), name = "Cells/cluster", breaks = scale.breaks, labels = paste0(scale.breaks, " cells"))
  )

  return(p)
}