R/plotting_ident.R

In ncchung/SeuratAddon: Unsupervised evaluation of cell identities in Seurat

#' Dimensional reduction plot, with probabilities for cell identities
#'
#' Graphs the output of a dimensional reduction technique (PCA by default).
#' Based on \code{Seurat::DimPlot}.
#' Cells are colored by their identity class. This identity class \code{object@ident}
#' should be set prior to running this code. See \code{ClusterCellsKmeans} and \code{EvaluateIdentKmeans}.
#'
#' @param object Seurat object
#' @param reduction.use Which dimensionality reduction to use. Default is
#' "pca", can also be "tsne", or "ica", assuming these are precomputed.
#' @param ident.threshold A probability threshold for cell identities,
#' using \code{object@meta.data["ident_prob"]}.
#' @param dim.1 Dimension for x-axis (default 1)
#' @param dim.2 Dimension for y-axis (default 2)
#' @param cells.use Vector of cells to plot (default is all cells)
#' @param pt.size Adjust point size for plotting
#' @param do.return Return a ggplot2 object (default : FALSE)
#' @param do.bare Do only minimal formatting (default : FALSE)
#' @param cols.use Vector of colors, each color corresponds to an identity
#' class. By default, ggplot assigns colors.
#' @param group.by Group (color) cells in different ways (for example, orig.ident)
#' @param pt.shape If NULL, all points are circles (default). You can specify any
#' cell attribute (that can be pulled with FetchData) allowing for both
#' different colors and different shapes on cells.
#' @param do.hover Enable hovering over points to view information
#' @param data.hover Data to add to the hover, pass a character vector of
#' features to add. Defaults to cell name and ident. Pass 'NULL' to clear extra
#' information.
#' @param do.identify Opens a locator session to identify clusters of cells.
#' @param do.label Whether to label the clusters
#' @param label.size Sets size of labels
#' @param no.legend Setting to TRUE will remove the legend
#' @param coord.fixed Use a fixed scale coordinate system (for spatial coordinates). Default is FALSE.
#' @param no.axes Setting to TRUE will remove the axes
#' @param dark.theme Use a dark theme for the plot
#' @param plot.order Specify the order of plotting for the idents. This can be
#' useful for crowded plots if points of interest are being buried. Provide
#' either a full list of valid idents or a subset to be plotted last (on top).
#' @param cells.highlight A list of character or numeric vectors of cells to
#' highlight. If only one group of cells desired, can simply
#' pass a vector instead of a list. If set, colors selected cells to the color(s)
#' in \code{cols.highlight} and other cells black (white if dark.theme = TRUE);
#'  will also resize to the size(s) passed to \code{sizes.highlight}
#' @param cols.highlight A vector of colors to highlight the cells as; will
#' repeat to the length groups in cells.highlight
#' @param sizes.highlight Size of highlighted cells; will repeat to the length
#' groups in cells.highlight
#' @param plot.title Title for plot
#' @param vector.friendly FALSE by default. If TRUE, points are flattened into
#' a PNG, while axes/labels retain full vector resolution. Useful for producing
#' AI-friendly plots with large numbers of cells.
#' @param png.file Used only if vector.friendly is TRUE. Location for temporary
#' PNG file.
#' @param png.arguments Used only if vector.friendly is TRUE. Vector of three
#' elements (PNG width, PNG height, PNG DPI) to be used for temporary PNG.
#' Default is c(10,10,100)
#' @param na.value Color value for NA points when using custom scale.
#' @param ... Extra parameters to FeatureLocator for do.identify = TRUE
#'
#' @return If do.return==TRUE, returns a ggplot2 object. Otherwise, only
#' graphical output.
#'
#' @import SDMTools
#' @importFrom stats median
#' @importFrom dplyr summarize group_by
#' @importFrom png readPNG
#'
#' @export DimPlot2
DimPlot2 <- function(
  object,
  reduction.use = "pca",
  ident.threshold = 0.9,
  dim.1 = 1,
  dim.2 = 2,
  cells.use = NULL,
  pt.size = 1,
  do.return = FALSE,
  do.bare = FALSE,
  cols.use = NULL,
  group.by = "ident",
  pt.shape = NULL,
  do.hover = FALSE,
  data.hover = 'ident',
  do.identify = FALSE,
  do.label = FALSE,
  label.size = 4,
  no.legend = FALSE,
  coord.fixed = FALSE,
  no.axes = FALSE,
  dark.theme = FALSE,
  plot.order = NULL,
  cells.highlight = NULL,
  cols.highlight = 'red',
  sizes.highlight = 1,
  plot.title = NULL,
  vector.friendly = FALSE,
  png.file = NULL,
  png.arguments = c(10,10, 100),
  na.value = 'grey50',
  ...
) {
  #first, consider vector friendly case
  if (vector.friendly) {
    previous_call <- blank_call <- png_call <-  match.call()
    blank_call$pt.size <- -1
    blank_call$do.return <- TRUE
    blank_call$vector.friendly <- FALSE
    png_call$no.axes <- TRUE
    png_call$no.legend <- TRUE
    png_call$do.return <- TRUE
    png_call$vector.friendly <- FALSE
    png_call$plot.title <- NULL
    blank_plot <- eval(blank_call, sys.frame(sys.parent()))
    png_plot <- eval(png_call, sys.frame(sys.parent()))
    png.file <- Seurat:::SetIfNull(x = png.file, default = paste0(tempfile(), ".png"))
    ggsave(
      filename = png.file,
      plot = png_plot,
      width = png.arguments[1],
      height = png.arguments[2],
      dpi = png.arguments[3]
    )
    to_return <- AugmentPlot(plot1 = blank_plot, imgFile = png.file)
    file.remove(png.file)
    if (do.return) {
      return(to_return)
    } else {
      print(to_return)
    }
  }
  embeddings.use <- GetDimReduction(
    object = object,
    reduction.type = reduction.use,
    slot = "cell.embeddings"
  )
  if (length(x = embeddings.use) == 0) {
    stop(paste(reduction.use, "has not been run for this object yet."))
  }
  cells.use <- Seurat:::SetIfNull(x = cells.use, default = colnames(x = object@data))
  dim.code <- GetDimReduction(
    object = object,
    reduction.type = reduction.use,
    slot = "key"
  )
  dim.codes <- paste0(dim.code, c(dim.1, dim.2))
  data.plot <- as.data.frame(x = embeddings.use)
  # data.plot <- as.data.frame(GetDimReduction(object, reduction.type = reduction.use, slot = ""))
  cells.use <- intersect(x = cells.use, y = rownames(x = data.plot))
  data.plot <- data.plot[cells.use, dim.codes]
  ident.use <- as.factor(x = object@ident[cells.use])
  if (group.by != "ident") {
    ident.use <- as.factor(x = FetchData(
      object = object,
      vars.all = group.by
    )[cells.use, 1])
  }
  data.plot$ident <- ident.use
  data.plot$x <- data.plot[, dim.codes[1]]
  data.plot$y <- data.plot[, dim.codes[2]]
  data.plot$pt.size <- pt.size
  # Add p-values and PIPs for cell identity from meta.data
  data.plot$ident_prob <- object@meta.data["ident_prob"][[1]]
  # ##
  if (!is.null(x = cells.highlight)) {
    # Ensure that cells.highlight are in our data.frame
    if (is.character(x = cells.highlight)) {
      cells.highlight <- list(cells.highlight)
    } else if (is.data.frame(x = cells.highlight) || !is.list(x = cells.highlight)) {
      cells.highlight <- as.list(x = cells.highlight)
    }
    cells.highlight <- lapply(
      X = cells.highlight,
      FUN = function(cells) {
        cells.return <- if (is.character(x = cells)) {
          cells[cells %in% rownames(x = data.plot)]
        } else {
          cells <- as.numeric(x = cells)
          cells <- cells[cells <= nrow(x = data.plot)]
          rownames(x = data.plot)[cells]
        }
        return(cells.return)
      }
    )
    # Remove groups that had no cells in our dataframe
    cells.highlight <- Filter(f = length, x = cells.highlight)
    if (length(x = cells.highlight) > 0) {
      if (!no.legend) {
        no.legend <- is.null(x = names(x = cells.highlight))
      }
      names.highlight <- if (is.null(x = names(x = cells.highlight))) {
        paste0('Group_', 1L:length(x = cells.highlight))
      } else {
        names(x = cells.highlight)
      }
      sizes.highlight <- rep_len(
        x = sizes.highlight,
        length.out = length(x = cells.highlight)
      )
      cols.highlight <- rep_len(
        x = cols.highlight,
        length.out = length(x = cells.highlight)
      )
      highlight <- rep_len(x = NA_character_, length.out = nrow(x = data.plot))
      if (is.null(x = cols.use)) {
        cols.use <- 'black'
      }
      cols.use <- c(cols.use[1], cols.highlight)
      size <- rep_len(x = pt.size, length.out = nrow(x = data.plot))
      for (i in 1:length(x = cells.highlight)) {
        cells.check <- cells.highlight[[i]]
        index.check <- match(x = cells.check, rownames(x = data.plot))
        highlight[index.check] <- names.highlight[i]
        size[index.check] <- sizes.highlight[i]
      }
      plot.order <- sort(x = unique(x = highlight), na.last = TRUE)
      plot.order[is.na(x = plot.order)] <- 'Unselected'
      highlight[is.na(x = highlight)] <- 'Unselected'
      highlight <- as.factor(x = highlight)
      data.plot$ident <- highlight
      data.plot$pt.size <- size
      if (dark.theme) {
        cols.use[1] <- 'white'
      }
    }
  }
  if (!is.null(x = plot.order)) {
    if (any(!plot.order %in% data.plot$ident)) {
      stop("invalid ident in plot.order")
    }
    plot.order <- rev(x = c(
      plot.order,
      setdiff(x = unique(x = data.plot$ident), y = plot.order)
    ))
    data.plot$ident <- factor(x = data.plot$ident, levels = plot.order)
    data.plot <- data.plot[order(data.plot$ident), ]
  }
  # Add ident_prob on ggplot
  if ("soft" %in% ident.threshold) {
    p <- ggplot(data = data.plot, mapping = aes(x = x, y = y)) +
      geom_point(mapping = aes(colour = factor(x = ident), alpha = ident_prob, size = pt.size))
  } else if (is.numeric(ident.threshold)) {
    if(ident.threshold >= 1 | ident.threshold <= 0) {
      stop("invalid ident.threshold: outside of a range (0,1)")
    }
    p <- ggplot(data = data.plot[data.plot$ident_prob > ident.threshold,], mapping = aes(x = x, y = y)) +
      geom_point(mapping = aes(colour = factor(x = ident), size = pt.size))
  } else {
    stop("invalid ident.threshold")
  }
  if (!is.null(x = pt.shape)) {
    shape.val <- FetchData(object = object, vars.all = pt.shape)[cells.use, 1]
    if (is.numeric(shape.val)) {
      shape.val <- cut(x = shape.val, breaks = 5)
    }
    data.plot[, "pt.shape"] <- shape.val
    p <- ggplot(data = data.plot, mapping = aes(x = x, y = y)) +
      geom_point(mapping = aes(
        colour = factor(x = ident),
        shape = factor(x = pt.shape),
        size = pt.size
      ))

    # Add ident_prob on ggplot
    if ("soft" %in% ident.threshold) {
      p <- ggplot(data = data.plot, mapping = aes(x = x, y = y)) +
        geom_point(mapping = aes(
          colour = factor(x = ident),
          shape = factor(x = pt.shape),
          size = pt.size,
          alpha = ident_prob
        ))
    } else if (is.numeric(ident.threshold)) {
      if(ident.threshold >= 1 | ident.threshold <= 0) {
        stop("invalid ident.threshold: outside of a range (0,1)")
      }
      p <- ggplot(data = data.plot[data.plot$ident_prob > ident.threshold,], mapping = aes(x = x, y = y)) +
        geom_point(mapping = aes(
          colour = factor(x = ident),
          shape = factor(x = pt.shape),
          size = pt.size
        ))
    } else {
      stop("invalid ident.threshold")
    }
  }
  if (!is.null(x = cols.use)) {
    p <- p + scale_colour_manual(values = cols.use, na.value=na.value)
  }
  if(coord.fixed){
    p <- p + coord_fixed()
  }
  p <- p + guides(size = FALSE)
  p2 <- p +
    xlab(label = dim.codes[[1]]) +
    ylab(label = dim.codes[[2]]) +
    scale_size(range = c(min(data.plot$pt.size), max(data.plot$pt.size)))
  p3 <- p2 +
    Seurat:::SetXAxisGG() +
    Seurat:::SetYAxisGG() +
    Seurat:::SetLegendPointsGG(x = 6) +
    Seurat:::SetLegendTextGG(x = 12) +
    Seurat:::no.legend.title +
    theme_bw() +
    Seurat:::NoGrid()
  if (dark.theme) {
    p <- p + Seurat::DarkTheme()
    p3 <- p3 + Seurat::DarkTheme()
  }
  p3 <- p3 + theme(legend.title = element_blank())
  if (!is.null(plot.title)) {
    p3 <- p3 + ggtitle(plot.title) + theme(plot.title = element_text(hjust = 0.5))
  }
  if (do.label) {
    data.plot %>%
      dplyr::group_by(ident) %>%
      summarize(x = median(x = x), y = median(x = y)) -> centers
    p3 <- p3 +
      geom_point(data = centers, mapping = aes(x = x, y = y), size = 0, alpha = 0) +
      geom_text(data = centers, mapping = aes(label = ident), size = label.size)
  }
  if (no.legend) {
    p3 <- p3 + theme(legend.position = "none")
  }
  if (no.axes) {
    p3 <- p3 + theme(
      axis.line = element_blank(),
      axis.text.x = element_blank(),
      axis.text.y = element_blank(),
      axis.ticks = element_blank(),
      axis.title.x = element_blank(),
      axis.title.y = element_blank(),
      panel.background = element_blank(),
      panel.border = element_blank(),
      panel.grid.major = element_blank(),
      panel.grid.minor = element_blank(),
      plot.background = element_blank()
    )
  }
  if (do.identify || do.hover) {
    if (do.bare) {
      plot.use <- p
    } else {
      plot.use <- p3
    }
    if (do.hover) {
      if (is.null(x = data.hover)) {
        features.info <- NULL
      } else {
        features.info <- Seurat::FetchData(object = object, vars.all = data.hover)
      }
      return(HoverLocator(
        plot = plot.use,
        data.plot = data.plot,
        features.info = features.info,
        dark.theme = dark.theme
      ))
    } else if (do.identify) {
      return(FeatureLocator(
        plot = plot.use,
        data.plot = data.plot,
        dark.theme = dark.theme,
        ...
      ))
    }
  }
  if (do.return) {
    if (do.bare) {
      return(p)
    } else {
      return(p3)
    }
  }
  if (do.bare) {
    print(p)
  } else {
    print(p3)
  }
}

#' Plot PCA map, with probabilities for computational cell identities
#'
#' Graphs the output of a PCA analysis
#' Cells are colored by their identity class.
#' Based on \code{Seurat::PCAPlot}.
#'
#' This function is a wrapper for DimPlot. See ?DimPlot for a full list of possible
#' arguments which can be passed in here.
#'
#' @param object Seurat object
#' @param ident.threshold A probability threshold for cell identities,
#' using \code{object@meta.data["ident_prob"]}.
#' @param \dots Additional parameters to DimPlot, for example, which dimensions to plot.
#'
#' @export PCAPlot2
PCAPlot2 <- function(object, ident.threshold = .9, ...) {
  return(DimPlot2(object = object,
    reduction.use = "pca",
    ident.threshold = ident.threshold,
    label.size = 4, ...))
}

#' Plot tSNE map, with probabilities for computational cell identities
#'
#' Graphs the output of a tSNE analysis
#' Cells are colored by their identity class.
#' Based on \code{Seurat::TSNEPlot}.
#'
#' This function is a wrapper for DimPlot2. See ?DimPlot2 for a full list of possible
#' arguments which can be passed in here.
#'
#' @param object Seurat object
#' @param ident.threshold A probability threshold for cell identities,
#' using \code{object@meta.data["ident_prob"]}.
#' @param do.label FALSE by default. If TRUE, plots an alternate view where the center of each
#' cluster is labeled
#' @param pt.size Set the point size
#' @param label.size Set the size of the text labels
#' @param cells.use Vector of cell names to use in the plot.
#' @param colors.use Manually set the color palette to use for the points
#' @param \dots Additional parameters to DimPlot, for example, which dimensions to plot.
#'
#' @seealso DimPlot2
#'
#' @export TSNEPlot2
TSNEPlot2 <- function(
  object,
  ident.threshold = 0.9,
  do.label = FALSE,
  pt.size = 1,
  label.size = 4,
  cells.use = NULL,
  colors.use = NULL,
  ...
) {
  return(DimPlot2(
    object = object,
    reduction.use = "tsne",
    cells.use = cells.use,
    pt.size = pt.size,
    do.label = do.label,
    label.size = label.size,
    cols.use = colors.use,
    ident.threshold = ident.threshold,
    ...
  ))
}