R/plot_decontx.R

In campbio/celda: CEllular Latent Dirichlet Allocation

#' @title Plots contamination on UMAP coordinates
#' @description A scatter plot of the UMAP dimensions generated by DecontX with
#' cells colored by the estimated percentation of contamation.
#' @param x Either a \linkS4class{SingleCellExperiment} with \code{decontX}
#' results stored in \code{metadata(x)$decontX} or the result from running
#' decontX on a count matrix.
#' @param batch Character. Batch of cells to plot. If \code{NULL}, then
#' the first batch in the list will be selected. Default \code{NULL}.
#' @param colorScale Character vector. Contains the color spectrum to be passed
#' to \code{scale_colour_gradientn} from package 'ggplot2'. Default
#' c("blue","green","yellow","orange","red").
#' @param size Numeric. Size of points in the scatterplot. Default 1.
#' @return Returns a \code{ggplot} object.
#' @author Shiyi Yang, Joshua Campbell
#' @seealso See \code{\link{decontX}} for a full example of how to estimate
#' and plot contamination.
#' @export
plotDecontXContamination <- function(x,
                                     batch = NULL,
                                     colorScale = c(
                                       "blue",
                                       "green",
                                       "yellow",
                                       "orange",
                                       "red"
                                     ),
                                     size = 1) {
  if (inherits(x, "SingleCellExperiment")) {
    estimates <- S4Vectors::metadata(x)$decontX$estimates
  } else {
    estimates <- x$estimates
  }
  if (is.null(estimates)) {
    stop("decontX estimates not found. Estimates will be found in
          'metadata(x)$decontX$estimates' if 'x' is a
          SingleCellExperiment or 'x$estimates' if decontX was run
          on a count matrix. Are you sure 'x' is output from decontX?")
  }
  batches <- names(estimates)

  if (is.null(batch)) {
    i <- batches[1]
  } else {
    if (!(batch %in% batches)) {
      stop(
        "'", batch, "' is not one of the batches in 'x'. Batches available",
        " for plotting: '", paste(batches, collapse = ","), "'"
      )
    }
    i <- batch
  }

  contamin <- estimates[[i]]$contamination
  umap <- estimates[[i]]$UMAP

  ## Create data.frame
  df <- data.frame(umap, "Contamination" = contamin)
  naIx <- is.na(umap[, 1]) | is.na(umap[, 2])
  df <- df[!naIx, ]

  ## Generate ggplot scatterplot
  gg <- ggplot2::ggplot(
    df,
    ggplot2::aes_string(
      x = colnames(umap)[1],
      y = colnames(umap)[2]
    )
  ) +
    ggplot2::geom_point(
      stat = "identity",
      size = size,
      ggplot2::aes_string(color = "Contamination")
    ) +
    ggplot2::theme_bw() +
    ggplot2::scale_colour_gradientn(
      colors = colorScale,
      name = "Contamination",
      limits = c(0, 1)
    ) +
    ggplot2::theme(
      panel.grid.major = ggplot2::element_blank(),
      panel.grid.minor = ggplot2::element_blank(),
      axis.text = ggplot2::element_text(size = 15),
      axis.title = ggplot2::element_text(size = 15)
    )
  return(gg)
}


#' @title Plots percentage of cells cell types expressing markers
#' @description Generates a barplot that shows the percentage of
#' cells within clusters or cell types that have detectable levels
#' of given marker genes. Can be used to view the expression of
#' marker genes in different cell types before and after
#' decontamination with \code{\link{decontX}}.
#' @param x Either a \linkS4class{SingleCellExperiment} or
#' a matrix-like object of counts.
#' @param markers List. A named list indicating the marker genes
#' for each cell type of
#' interest. Multiple markers can be supplied for each cell type. For example,
#' \code{list(Tcell_Markers=c("CD3E", "CD3D"),
#' Bcell_Markers=c("CD79A", "CD79B", "MS4A1")}
#' would specify markers for human T-cells and B-cells.
#' A cell will be considered
#' "positive" for a cell type if it has a count greater than \code{threshold}
#' for at least one of the marker genes in the list.
#' @param groupClusters List. A named list that allows
#' cell clusters labels coded in
#' \code{z} to be regrouped and renamed on the fly. For example,
#' \code{list(Tcells=c(1, 2), Bcells=7)} would recode
#' clusters 1 and 2 to "Tcells"
#' and cluster 7 to "Bcells". Note that if this is
#' used, clusters in \code{z} not found
#' in \code{groupClusters} will be excluded from the barplot.
#' Default \code{NULL}.
#' @param assayName Character vector. Name(s) of the assay(s) to
#' plot if \code{x} is a
#' \linkS4class{SingleCellExperiment}. If more than one assay
#' is listed, then side-by-side barplots will be generated.
#' Default \code{c("counts", "decontXcounts")}.
#' @param z Character, Integer, or Vector. Indicates the cluster labels
#' for each cell.
#' If \code{x} is a \linkS4class{SingleCellExperiment} and \code{z = NULL},
#' then the cluster labels from \code{\link{decontX}} will be retived from the
#' \code{colData} of \code{x} (i.e. \code{colData(x)$decontX_clusters}).
#' If \code{z} is a single character or integer,
#' then that column will be retrived
#' from \code{colData} of \code{x}. (i.e. \code{colData(x)[,z]}). If \code{x}
#' is a counts matrix, then \code{z} will need
#' to be a vector the same length as
#' the number of columns in \code{x} that indicate
#' the cluster to which each cell
#' belongs. Default \code{NULL}.
#' @param threshold Numeric. Markers greater than or equal to this value will
#' be considered detected in a cell. Default 1.
#' @param exactMatch Boolean. Whether to only identify exact matches
#' for the markers or to identify partial matches using \code{\link{grep}}. See
#' \code{\link{retrieveFeatureIndex}} for more details. Default \code{TRUE}.
#' @param by Character. Where to search for the markers if \code{x} is a
#' \linkS4class{SingleCellExperiment}. See \code{\link{retrieveFeatureIndex}}
#' for more details. If \code{x} is a matrix,
#' then this must be set to \code{"rownames"}.Default \code{"rownames"}.
#' @param ncol Integer. Number of columns to make in the plot.
#' Default \code{round(sqrt(length(markers))}.
#' @param labelBars Boolean. Whether to display percentages above each bar
#' Default \code{TRUE}.
#' @param labelSize Numeric. Size of the percentage labels in the barplot.
#' Default 3.
#' @return Returns a \code{ggplot} object.
#' @author Shiyi Yang, Joshua Campbell
#' @seealso See \code{\link{decontX}} for a full example of how to estimate
#' and plot contamination.
#' @export
plotDecontXMarkerPercentage <- function(x, markers, groupClusters = NULL,
                                        assayName = c(
                                          "counts",
                                          "decontXcounts"
                                        ),
                                        z = NULL, threshold = 1,
                                        exactMatch = TRUE, by = "rownames",
                                        ncol = round(sqrt(length(markers))),
                                        labelBars = TRUE, labelSize = 3) {
  cellTypeLabels <- percent <- NULL # fix check note

  legend <- "none"
  # Check that list arguments are named
  if (!is(markers, "list") || is.null(names(markers))) {
    stop("'markers' needs to be a named list.")
  }

  temp <- .processPlotDecontXMarkerInupt(
    x = x,
    z = z,
    markers = markers,
    groupClusters = groupClusters,
    by = by,
    exactMatch = exactMatch
  )
  x <- temp$x
  z <- temp$z
  geneMarkerIndex <- temp$geneMarkerIndex
  geneMarkerCellTypeIndex <- temp$geneMarkerCellTypeIndex
  groupClusters <- temp$groupClusters
  xlab <- temp$xlab

  if (inherits(x, "SingleCellExperiment")) {
    # If 'x' is SingleCellExperiment, then get percentage
    # for each matrix in 'assayName'
    df.list <- list()
    for (i in seq_along(assayName)) {
      counts <- SummarizedExperiment::assay(
        x[geneMarkerIndex, ],
        assayName[i]
      )
      df <- .calculateDecontXBarplotPercent(
        counts,
        z,
        geneMarkerCellTypeIndex,
        threshold
      )
      df.list[[i]] <- cbind(df, assay = assayName[i])
    }
    df <- do.call(rbind, df.list)
    assay <- as.factor(df$assay)
    if (length(assayName) > 1) {
      legend <- "right"
    }
  } else {
    ## If 'x' is matrix, then calculate percentages directly
    counts <- x[geneMarkerIndex, ]
    df <- .calculateDecontXBarplotPercent(
      counts,
      z,
      geneMarkerCellTypeIndex,
      threshold
    )
    assay <- "red3"
    legend <- "none"
  }

  # Build data.frame for ggplots
  df <- cbind(df, cellTypeLabels = names(groupClusters)[df$cellType])
  df$cellTypeLabels <- factor(df$cellTypeLabels,
    levels = names(groupClusters)
  )
  df <- cbind(df, markerLabels = names(markers)[df$markers])
  df$markerLabels <- factor(df$markerLabels, levels = names(markers))

  plt <- ggplot2::ggplot(df, ggplot2::aes_string(
    x = "cellTypeLabels",
    y = "percent", fill = "assay"
  )) +
    ggplot2::geom_bar(
      stat = "identity",
      position = ggplot2::position_dodge2(width = 0.9, preserve = "single")
    ) +
    ggplot2::xlab(xlab) +
    ggplot2::ylab(paste0("Percentage of cells expressing markers")) +
    ggplot2::facet_wrap(. ~ df$markerLabels, ncol = ncol) +
    ggplot2::theme(
      panel.background = ggplot2::element_rect(
        fill = "white",
        color = "grey"
      ),
      panel.grid = ggplot2::element_line("grey"),
      legend.position = legend,
      legend.key = ggplot2::element_rect(
        fill = "white",
        color = "white"
      ),
      panel.grid.minor = ggplot2::element_blank(),
      panel.grid.major = ggplot2::element_blank(),
      text = ggplot2::element_text(size = 10),
      axis.text.x = ggplot2::element_text(
        size = 8, angle = 45,
        hjust = 1
      ),
      axis.text.y = ggplot2::element_text(size = 9),
      legend.key.size = grid::unit(8, "mm"),
      legend.text = ggplot2::element_text(size = 10),
      strip.text.x = ggplot2::element_text(size = 10)
    )

  if (isTRUE(labelBars)) {
    plt <- plt + ggplot2::geom_text(ggplot2::aes(
      x = cellTypeLabels,
      y = percent + 2.5,
      label = percent
    ),
    position = ggplot2::position_dodge2(width = 0.9, preserve = "single"),
    size = labelSize
    )
  }
  return(plt)
}


#' @title Plots expression of marker genes before and after decontamination
#' @description Generates a violin plot that shows the counts of marker
#' genes in cells across specific clusters or cell types. Can be used to view
#' the expression of marker genes in different cell types before and after
#' decontamination with \code{\link{decontX}}.
#' @param x Either a \linkS4class{SingleCellExperiment}
#' or a matrix-like object of counts.
#' @param markers Character Vector or List. A character vector
#' or list of character vectors
#' with the names of the marker genes of interest.
#' @param groupClusters List. A named list that allows
#' cell clusters labels coded in
#' \code{z} to be regrouped and renamed on the fly. For example,
#' \code{list(Tcells=c(1, 2), Bcells=7)} would recode clusters
#' 1 and 2 to "Tcells"
#' and cluster 7 to "Bcells". Note that if this is used, clusters
#' in \code{z} not found
#' in \code{groupClusters} will be excluded. Default \code{NULL}.
#' @param assayName Character vector. Name(s) of the assay(s) to
#' plot if \code{x} is a
#' \linkS4class{SingleCellExperiment}. If more than one assay is listed, then
#' side-by-side violin plots will be generated.
#' Default \code{c("counts", "decontXcounts")}.
#' @param z Character, Integer, or Vector.
#' Indicates the cluster labels for each cell.
#' If \code{x} is a \linkS4class{SingleCellExperiment} and \code{z = NULL},
#' then the cluster labels from \code{\link{decontX}} will be retreived from the
#' \code{colData} of \code{x} (i.e. \code{colData(x)$decontX_clusters}).
#' If \code{z} is a single character or integer, then that column will be
#' retrived from \code{colData} of \code{x}. (i.e. \code{colData(x)[,z]}).
#' If \code{x} is a counts matrix, then \code{z} will need to be a vector
#' the same length as the number of columns in \code{x} that indicate
#' the cluster to which each cell belongs. Default \code{NULL}.
#' @param exactMatch Boolean. Whether to only identify exact matches
#' for the markers or to identify partial matches using \code{\link{grep}}.
#' See \code{\link{retrieveFeatureIndex}} for more details.
#' Default \code{TRUE}.
#' @param by Character. Where to search for the markers if \code{x} is a
#' \linkS4class{SingleCellExperiment}. See \code{\link{retrieveFeatureIndex}}
#' for more details. If \code{x} is a matrix, then this must be set to
#' \code{"rownames"}. Default \code{"rownames"}.
#' @param log1p Boolean. Whether to apply the function \code{log1p} to the data
#' before plotting. This function will add a pseudocount of 1 and then log
#' transform the expression values. Default \code{FALSE}.
#' @param ncol Integer. Number of columns to make in the plot.
#' Default \code{NULL}.
#' @param plotDots Boolean. If \code{TRUE}, the
#'  expression of features will be plotted as points in addition to the violin
#'  curve. Default \code{FALSE}.
#' @param dotSize Numeric. Size of points if \code{plotDots = TRUE}.
#' Default \code{0.1}.
#' @return Returns a \code{ggplot} object.
#' @author Shiyi Yang, Joshua Campbell
#' @seealso See \code{\link{decontX}} for a full example of how to estimate
#' and plot contamination.
#' @export
plotDecontXMarkerExpression <- function(x, markers, groupClusters = NULL,
                                        assayName = c(
                                          "counts",
                                          "decontXcounts"
                                        ),
                                        z = NULL, exactMatch = TRUE,
                                        by = "rownames", log1p = FALSE,
                                        ncol = NULL,
                                        plotDots = FALSE, dotSize = 0.1) {
  legend <- "none"
  temp <- .processPlotDecontXMarkerInupt(
    x = x,
    z = z,
    markers = markers,
    groupClusters = groupClusters,
    by = by,
    exactMatch = exactMatch
  )
  x <- temp$x
  z <- temp$z
  geneMarkerIndex <- temp$geneMarkerIndex
  groupClusters <- temp$groupClusters
  xlab <- temp$xlab

  if (inherits(x, "SingleCellExperiment")) {
    # If 'x' is SingleCellExperiment, then get percentage
    # for each matrix in 'assayName'
    df.list <- list()
    for (i in seq_along(assayName)) {
      counts <- SummarizedExperiment::assay(
        x[geneMarkerIndex, ],
        assayName[i]
      )
      df <- reshape2::melt(as.matrix(counts),
        varnames = c("Marker", "Cell"),
        value.name = "Expression"
      )
      df.list[[i]] <- cbind(df, assay = assayName[i])
    }
    df <- do.call(rbind, df.list)
    assay <- factor(df$assay, levels = assayName)
    if (length(assayName) > 1) {
      legend <- "right"
    }
  } else {
    ## If 'x' is matrix, then calculate percentages directly
    counts <- x[geneMarkerIndex, ]
    df <- reshape2::melt(counts,
      varnames = c("Marker", "Cell"),
      value.name = "Expression"
    )
    assay <- "red3"
    legend <- "none"
  }

  # Create data.frame and add cell type groups back in
  names(z) <- colnames(x)
  df <- cbind(df, Cluster = z[df$Cell])

  ylab <- "Expression"
  if (isTRUE(log1p)) {
    df$Expression <- log1p(df$Expression)
    ylab <- "Expression (log1p)"
  }
  Expression <- df$Expression
  Marker <- df$Marker
  Assay <- factor(df$assay, levels = assayName)
  Cluster <- df$Cluster
  if (!is.null(groupClusters)) {
    df <- cbind(df, Cell_Type = names(groupClusters)[Cluster])
    Cell_Type <- factor(df$Cell_Type, levels = names(groupClusters))
    plt <- ggplot2::ggplot(df, ggplot2::aes(
      x = Cell_Type,
      y = Expression,
      fill = Assay
    )) +
      ggplot2::facet_wrap(~ Cell_Type + Marker,
        scales = "free",
        labeller = ggplot2::label_context,
        ncol = ncol
      )
  } else {
    plt <- ggplot2::ggplot(df, ggplot2::aes(
      x = Cluster,
      y = Expression,
      fill = Assay
    )) +
      ggplot2::facet_wrap(~ Cluster + Marker,
        scales = "free",
        labeller = ggplot2::label_context,
        ncol = ncol
      )
  }
  plt <- plt + ggplot2::geom_violin(
    trim = TRUE,
    scale = "width"
  ) +
    ggplot2::theme_bw() + ggplot2::theme(
      axis.text.x = ggplot2::element_blank(),
      axis.ticks.x = ggplot2::element_blank(),
      axis.title.x = ggplot2::element_blank(),
      strip.text = ggplot2::element_text(size = 8),
      panel.grid = ggplot2::element_blank(),
      legend.position = legend
    ) + ggplot2::ylab(ylab)


  if (isTRUE(plotDots)) {
    plt <- plt + ggplot2::geom_jitter(height = 0, size = dotSize)
  }

  return(plt)
}



.processPlotDecontXMarkerInupt <- function(x, z, markers, groupClusters,
                                           by, exactMatch) {

  # Process z and convert to a factor
  if (is.null(z) & inherits(x, "SingleCellExperiment")) {
    cn <- colnames(SummarizedExperiment::colData(x))
    if (!("decontX_clusters" %in% cn)) {
      stop("'decontX_clusters' not found in 'colData(x)'. Make sure you have
           run 'decontX' or supply 'z' directly.")
    }
    z <- SummarizedExperiment::colData(x)$decontX_clusters
  } else if (length(z) == 1 & inherits(x, "SingleCellExperiment")) {
    if (!(z %in% colnames(SummarizedExperiment::colData(x)))) {
      stop("'", z, "' not found in 'colData(x)'.")
    }
    z <- SummarizedExperiment::colData(x)[, z]
  } else if (length(z) != ncol(x)) {
    stop("If 'x' is a SingleCellExperiment, then 'z' needs to be",
          " a single character or integer specifying the column in",
          " 'colData(x)'. Alternatively to specify the cell cluster",
          " labels directly as a vector, the length of 'z' needs to",
          " be the same as the number of columns in 'x'. This is",
          " required if 'x' is a matrix.")
  }
  if (!is.factor(z)) {
    z <- as.factor(z)
  }

  if (!is.null(groupClusters)) {
    if (!is(groupClusters, "list") || is.null(names(groupClusters))) {
      stop("'groupClusters' needs to be a named list.")
    }

    # Check that groupClusters are found in 'z'
    cellMappings <- unlist(groupClusters)
    if (any(!(cellMappings %in% z))) {
      missing <- cellMappings[!(cellMappings %in% z)]
      stop(
        "'groupClusters' not found in 'z': ",
        paste(missing, collapse = ",")
      )
    }

    # Check for duplicate groupClusters
    ta <- table(unlist(groupClusters))
    if (any(ta > 1)) {
      dup <- names(ta)[ta > 1]
      stop("'groupClusters' had duplicate values for the following clusters: ",
           paste(dup, collapse = ","), ". Clusters need be assigned to a",
           "single group.")
    }

    labels <- rep(NA, ncol(x))
    for (i in seq_along(groupClusters)) {
      labels[z %in% groupClusters[[i]]] <- names(groupClusters)[i]
    }
    na.ix <- is.na(labels)
    labels <- labels[!na.ix]
    x <- x[, !na.ix]
    z <- as.integer(factor(labels, levels = names(groupClusters)))
    xlab <- "Cell types"
  } else {
    labels <- as.factor(z)
    groupClusters <- levels(labels)
    names(groupClusters) <- levels(labels)
    xlab <- "Clusters"
  }

  # Find index of each feature in 'x'
  geneMarkerCellTypeIndex <- rep(
    seq(length(markers)),
    lapply(markers, length)
  )
  geneMarkerIndex <- retrieveFeatureIndex(unlist(markers),
    x,
    by = by,
    removeNA = FALSE,
    exactMatch = exactMatch
  )

  # Remove genes that did not match
  na.ix <- is.na(geneMarkerIndex)
  geneMarkerCellTypeIndex <- geneMarkerCellTypeIndex[!na.ix]
  geneMarkerIndex <- geneMarkerIndex[!na.ix]

  return(list(
    x = x,
    z = z,
    geneMarkerIndex = geneMarkerIndex,
    geneMarkerCellTypeIndex = geneMarkerCellTypeIndex,
    groupClusters = groupClusters,
    xlab = xlab
  ))
}


.calculateDecontXBarplotPercent <- function(counts,
                                            z,
                                            geneMarkerCellTypeIndex,
                                            threshold) {

  # Get counts matrix and convert to DelayedMatrix
  counts <- DelayedArray::DelayedArray(counts)

  # Convert to boolean matrix and sum markers in same cell type
  # The "+ 0" is to convert boolean to numeric
  counts <- counts >= threshold
  countsByMarker <- DelayedArray::rowsum(counts + 0, geneMarkerCellTypeIndex)
  countsByCellType <- DelayedArray::colsum((countsByMarker > 0) + 0, z)

  # Calculate percentages within each cell cluster
  zTotals <- tabulate(z)
  percentByCellType <- round(sweep(countsByCellType, 2, zTotals, "/") * 100)
  df <- reshape2::melt(percentByCellType,
    varnames = c("markers", "cellType"),
    value.name = "percent"
  )

  return(df)
}