velociraptor: Toolkit for Single-Cell Velocity

Documented in plotVelocityStream

#' Velocity stream plot in low-dimensional space
#'
#' Plot velocities embedded into low-dimensional space as a stream plot. Stream
#' lines are lines that follow the gradient in the velocity field and illustrate
#' paths that cells could follow based on observed RNA velocities.
#'
#' @param sce A \linkS4class{SingleCellExperiment} object containing
#'   low-dimensional coordinates, e.g., after t-SNE, in its
#'   \code{\link{reducedDims}}.
#' @param embedded A low-dimensional projection of the velocity vectors into the
#'   embedding of \code{sce}. This should be of the same dimensions as \code{sce}
#'   and is typically produced by \code{\link{embedVelocity}}.
#' @param use.dimred String or integer scalar specifying the reduced dimensions
#'   to retrieve from \code{sce}.
#' @param color_by A character scalar specifying a column in \code{colData(sce)}
#'   to color cells in the phase graph. Alternatively, \code{color_by} can be
#'   set to a valid R color to be used to color cells.
#' @param color.alpha An integer scalar giving the transparency of colored
#'   cells. Possible values are between 0 (fully transparent) and 1.0 (opaque).
#' @param grid.resolution Integer scalar specifying the resolution of the grid,
#'   in terms of the number of grid intervals along each axis.
#' @param scale Logical scalar indicating whether the averaged vectors should be
#'   scaled by the grid resolution.
#' @param stream.L Integer scalar giving the typical length of a streamline
#'   low-dimensional space units.
#' @param stream.min.L A numeric scalar with the minimum length of segments to be shown.
#' @param stream.res Numeric scalar specifying the resolution of estimated
#'   streamlines (higher numbers increase smoothness of lines but also the time
#'   for computation).
#' @param stream.width A numeric scalar controlling the width of streamlines.
#' @param color.streamlines Logical scalar. If \code{TRUE} streamlines will
#'   be colored by local velocity. Arrows cannot be shown in that case.
#' @param color.streamlines.map A character vector specifying the
#'   color range used for mapping local velocities to streamline colors. The
#'   default is \code{viridisLite::viridis(11)}.
#' @param arrow.angle,arrow.length Numeric scalars giving the \code{angle} and
#'   \code{length} of arrowheads.
#'   
#' @details \code{grid.resolution} and \code{scale} are passed to
#'   \code{\link{gridVectors}}, which is used to summarized the velocity vectors
#'   into an initial grid. A full regular grid is computed from that and used
#'   in \code{\link[metR]{geom_streamline}} to calculate streamlines. The
#'   following arguments are passed to the arguments given in parenthesis of
#'   \code{\link[metR]{geom_streamline}}:
#'   \code{stream.L} (\code{L}), \code{stream.res} (\code{res}),
#'   \code{stream.min.L} (\code{min.L}), \code{arrow.angle} (\code{arrow.angle})
#'   and \code{arrow.length} (\code{arrow.length}).
#'   Streamlines are computed by simple integration with a forward Euler method,
#'   and \code{stream.L} and \code{stream.res} are used to compute the number of
#'   steps and the time interval between steps for the integration.
#'   \code{stream.width} is multiplied with \code{..step..} estimated by
#'   \code{\link[metR]{geom_streamline}} to control the width of streamlines.
#'   
#' @return A \code{ggplot2} object with the streamline plot.
#' 
#' @author Michael Stadler
#' 
#' @examples
#' library(scuttle)
#' set.seed(42)
#' sce1 <- mockSCE(ncells = 100, ngenes = 500)
#' sce2 <- mockSCE(ncells = 100, ngenes = 500)
#'
#' datlist <- list(X=counts(sce1), spliced=counts(sce1), unspliced=counts(sce2))
#'
#' out <- scvelo(datlist, mode = "dynamical")
#' 
#' em <- embedVelocity(reducedDim(out, 1), out)[,1:2]
#' 
#' # https://github.com/kevinrue/velociraptor/issues/90
#' 
#' if (!basilisk.utils::isMacOSX()){
#' plotVelocityStream(out, em)
#' plotVelocityStream(out, em, color.streamlines = TRUE)
#' }
#' 
#' @seealso \code{\link{gridVectors}} used to summarize velocity vectors into
#'   a grid (velocity field), the \pkg{ggplot2} package used for plotting,
#'   \code{\link[metR]{geom_streamline}} in package \pkg{metR} used to
#'   calculate and add streamlines from the RNA velocity field to the plot,
#'   \code{\link[viridisLite]{viridis}} for creation of color palettes.
#' 
#' @export
#' @importFrom S4Vectors DataFrame
plotVelocityStream <- function(
    sce, embedded, use.dimred = 1,
  color_by = "#444444", color.alpha = 0.2,
  grid.resolution = 60, scale = TRUE,
  stream.L = 10, stream.min.L = 0, stream.res = 4,
  stream.width = 8,
  color.streamlines = FALSE,
  color.streamlines.map = c("#440154", "#482576", "#414487",
    "#35608D", "#2A788E", "#21908C",
    "#22A884", "#43BF71", "#7AD151",
    "#BBDF27", "#FDE725"),
  arrow.angle = 8, arrow.length = 0.8) {
  if (basilisk.utils::isMacOSX()){ # both Intel and Arm
    warning(
      "This function has been observed to run indefinitely on macOS. ",
      "The issue is related to the metR::geom_streamline() layer but unsolved yet. ",
      "Use with caution and consider terminating R if rendering the plot freezes your session.",
      immediate. = TRUE
    )
  }
  
  if (!identical(ncol(sce), nrow(embedded))) {
    stop("'sce' and 'embedded' do not have consistent dimensions.")
  }
  if (is.numeric(use.dimred)) {
    stopifnot(exprs = {
      identical(length(use.dimred), 1L)
      use.dimred <= length(reducedDims(sce))
    })
    use.dimred <- reducedDimNames(sce)[use.dimred]
  }
  else if (is.character(use.dimred)) {
    stopifnot(exprs = {
      length(use.dimred) == 1L
      use.dimred %in% reducedDimNames(sce)
    })
  }
  else {
    stop("'use.dimred' is not a valid value for use in reducedDim(sce, use.dimred)")
  }
  if (!requireNamespace("ggplot2")) {
    stop("'plotVelocityStream' requires the package 'ggplot2'.")
  }
  
  # get coordinates in reduced dimensional space
  xy <- reducedDim(sce, use.dimred)[, 1:2]
  
  # summarize velocities in a grid
  gr <- gridVectors(x = xy, embedded = embedded,
    resolution = grid.resolution, scale = scale,
    as.data.frame = FALSE,
    return.intermediates = TRUE)
  
  # now make it a regular grid needed for metR::geom_streamline
  xbreaks <- seq(gr$limits[1,1], gr$limits[2,1], by = gr$delta[1])
  ybreaks <- seq(gr$limits[1,2], gr$limits[2,2], by = gr$delta[2])
  plotdat2 <- expand.grid(x = xbreaks + gr$delta[1] / 2,
    y = ybreaks + gr$delta[2] / 2,
    dx = 0, dy = 0)
  allcategories <- DataFrame(expand.grid(V1 = seq(0, grid.resolution),
    V2 = seq(0, grid.resolution)))
  ivec <- match(gr$categories[sort(unique(gr$grp)), ], allcategories)
  plotdat2[ivec, c("dx", "dy")] <- gr$vec
  
  
  # plot it using ggplot2 and metR::geom_streamline
  plotdat1 <- data.frame(xy)
  colnames(plotdat1) <- c("x", "y")
  if (is.character(color_by) && length(color_by) == 1L && color_by %in% colnames(colData(sce))) {
    plotdat1 <- cbind(plotdat1, col = colData(sce)[, color_by])
    colByFeat <- TRUE
  } else {
    colByFeat <- FALSE
  }
  p <- ggplot2::ggplot(plotdat1, ggplot2::aes(x = !!ggplot2::sym("x"), y = !!ggplot2::sym("y"))) +
    ggplot2::labs(x = paste(use.dimred, "1"), y = paste(use.dimred, "2"))
  if (!colByFeat) {
    colMatrix <- grDevices::col2rgb(col = color_by, alpha = TRUE)
    if (any(colMatrix[4, ] != 255)) {
      warning("ignoring 'color.alpha' as 'color_by' already specifies alpha channels")
      color.alpha <- colMatrix[4, ] / 255
    }
    p <- p + ggplot2::geom_point(color = color_by, alpha = color.alpha)
  } else {
    p <- p + ggplot2::geom_point(ggplot2::aes(color = !!ggplot2::sym("col")), alpha = color.alpha) +
      ggplot2::labs(color = color_by)
  }
  if (color.streamlines) {
    # remark: when coloring streamlines, we currently cannot have any arrows
    # remark: ..dx.., ..dy.. and ..step.. are calculated by metR::geom_streamline
    p <- p +
      metR::geom_streamline(
        mapping = ggplot2::aes(
          x = !!ggplot2::sym("x"),
          y = !!ggplot2::sym("y"),
          dx = !!ggplot2::sym("dx"),
          dy = !!ggplot2::sym("dy"),
          size = stream.width * !!ggplot2::sym("..step.."),
          alpha = !!ggplot2::sym("..step.."),
          color = ggplot2::stat(sqrt((!!ggplot2::sym("..dx.."))^2 +
              (!!ggplot2::sym("..dy.."))^2))),
        arrow = NULL, lineend = "round",
        data = plotdat2, size = 0.6, jitter = 2,
        L = stream.L, min.L = stream.min.L,
        res = stream.res, inherit.aes = FALSE) +
      ggplot2::scale_color_gradientn(colors = color.streamlines.map,
        guide = "none") +
      ggplot2::scale_alpha_continuous(guide = "none") +
      ggplot2::theme_minimal() +
      ggplot2::theme(axis.text = ggplot2::element_blank(),
        panel.grid.major = ggplot2::element_blank(),
        panel.grid.minor = ggplot2::element_blank())
  } else {
    p <- p +
      metR::geom_streamline(
        mapping = ggplot2::aes(
          x = !!ggplot2::sym("x"),
          y = !!ggplot2::sym("y"),
          dx = !!ggplot2::sym("dx"),
          dy = !!ggplot2::sym("dy"),
          size = stream.width * !!ggplot2::sym("..step..")),
        data = plotdat2, size = 0.3, jitter = 2,
        L = stream.L, min.L = stream.min.L,
        res = stream.res, arrow.angle = arrow.angle,
        arrow.length = arrow.length, inherit.aes = FALSE) +
      ggplot2::theme_minimal() +
      ggplot2::theme(axis.text = ggplot2::element_blank(),
        panel.grid.major = ggplot2::element_blank(),
        panel.grid.minor = ggplot2::element_blank())
  }
  
  return(p)
}