R/plot_umap.R

In niaid/HDStIM: High Dimensional Stimulation Immune Mapping ('HDStIM')

#' @title Diagnostic UMAP plots showing the partitioning of cells into responding and non-responding groups by HDStIM
#'
#' @param mapped_data      Returned list from the \code{\link{HDStIM}} function.
#' @param path             Path to the folder to save figures generated by this function.
#' @param verbose          Logical. To make function more verbose. Default is FALSE.
#' @import ggplot2
#' @return                 A list of ggplot objects. If the path is not NULL, PNG files of the plots are saved in the specified folder.
#' @export
#' @examples
#' \donttest{
#' mapped_data <- HDStIM(chi11$expr_data, chi11$state_markers,
#'                       chi11$cluster_col, chi11$stim_label,
#'                       chi11$unstim_label, seed_val = 123, umap = TRUE, 
#'                       umap_cells = 50, verbose = FALSE)
#'
#' pu <- plot_umap(mapped_data, path = NULL, verbose = FALSE)
#' }
plot_umap <- function(mapped_data, path = NULL, verbose = FALSE){
  # Check if the list contains data from UMAPSs.
  if(is.null(mapped_data$umap_plot_data)){
    stop("No UMAP data in the list.")
  }

  # Check if path exists; if not then create it.
  if(!is.null(path)){
    if(!dir.exists(path)){
      if(verbose){message(paste("Creating %s folder", path))}
      dir.create(path, recursive = TRUE)
    } else {
      if(verbose){message(paste(path, "folder already exists. Output will be over written."))}
    }
  }

  # Bind global variables.
  UMAP1 <- UMAP2 <- cell_type <- cell_population  <- stim_type <- response_status <- NULL

  # Group data according to clusters and stimulation type.
  group_data <- group_by(mapped_data$umap_plot_data, cell_population, stim_type)

  # Split groups into a list of individual tables.
  split_groups <- group_split(group_data)

  # Define colors.
  cbPalette <- c("#009E73", "#0072B2", "#E69F00", "#CC79A7")

  # Create ggplot for a single group.
  allplots <- list()
  for(i in 1:length(split_groups)){
    clust <- unique(split_groups[[i]]$cell_population)
    stim <- unique(split_groups[[i]]$stim_type)
    if(verbose){message(paste("Plotting for the cell population", clust, "and stimulation", stim))}

    plot_dat <- split_groups[[i]]
    plot_dat$response_status <- factor(plot_dat$response_status, levels=c("Resp. Stim.", "Resp. Unstim.",
                                                              "Non-resp. Stim.", "Non-resp. Unstim."))
    umap_plot <- ggplot(plot_dat, aes(x = UMAP1, y = UMAP2, color = response_status)) +
    geom_point(size = 1) +
    scale_colour_manual(values=cbPalette) +
    labs(color = "Response Status", title = paste0("Cell Pop.: ", clust,
                                             "; Stim.: ", stim,
                                             "\nTotal Cells: ", unique(split_groups[[i]]$tot_of_cells),
                                             "; Cells Plotted: ", unique(split_groups[[i]]$no_of_cells))) +
    scale_x_continuous("UMAP1") +
    scale_y_continuous("UMAP2")

    allplots[[i]] <- umap_plot

    if(!is.null(path)){
      ggsave(paste0("umap_",clust,"_", stim, ".png"),
             plot = umap_plot, device = "png", path = file.path(path), width = 7, height = 5, dpi = 300)
    }
  }
  return(allplots)
}