R/PomaHeatmap.R

In pcastellanoescuder/POMA: Tools for Omics Data Analysis

#' Heatmap Plot
#'
#' @description `PomaHeatmap` generates a heatmap. 
#' 
#' @param data A `SummarizedExperiment` object.
#' @param covs Character vector. Indicates the names of `colData` columns to be included as covariates. Default is NULL (no covariates).
#' @param sample_names Logical. Indicates if sample names should be displayed or not. Default is TRUE.
#' @param feature_names Logical. Indicates if feature names should be displayed or not. Default is FALSE.
#' @param show_legend Logical. Indicates if legend should be displayed or not. Default is TRUE.
#' 
#' @export
#'
#' @return A `ggplot` object.
#' @author Pol Castellano-Escuder
#' 
#' @examples 
#' data <- POMA::st000284 %>% # Example SummarizedExperiment object included in POMA
#'   PomaNorm()
#' 
#' # Basic heatmap
#' data %>% 
#'   PomaHeatmap()
#' 
#' # Heatmap with one covariate  
#' data %>% 
#'   PomaHeatmap(covs = "factors")
#' 
#' # Heatmap with two covariates
#' data %>% 
#'   PomaHeatmap(covs = c("factors", "smoking_condition"))
PomaHeatmap <- function(data, 
                        covs = NULL,
                        sample_names = TRUE,
                        feature_names = FALSE,
                        show_legend = TRUE) {
  
  if(!is(data, "SummarizedExperiment")){
    stop("data is not a SummarizedExperiment object. \nSee POMA::PomaCreateObject or SummarizedExperiment::SummarizedExperiment")
  }
  
  plot_data <- SummarizedExperiment::assay(data)
  
  if (ncol(SummarizedExperiment::colData(data)) > 0 & !is.null(covs)) {
    metadata <- SummarizedExperiment::colData(data) %>%
      as.data.frame() %>%
      dplyr::select(dplyr::all_of(covs))
    
    heatmap_annotations <- ComplexHeatmap::HeatmapAnnotation(df = metadata, show_legend = show_legend)
    
    poma_heatmap <- suppressMessages(
      ComplexHeatmap::Heatmap(plot_data, 
                              name = "Value", 
                              top_annotation = heatmap_annotations,
                              show_row_names = feature_names, 
                              show_column_names = sample_names,
                              show_heatmap_legend = show_legend)
    )
  } else {
    poma_heatmap <- suppressMessages(
      ComplexHeatmap::Heatmap(plot_data, 
                              name = "Value", 
                              show_row_names = feature_names, 
                              show_column_names = sample_names,
                              show_heatmap_legend = show_legend)
    )
  }
  
  # Convert to ggplot
  hm_grob <- grid::grid.grabExpr(ComplexHeatmap::draw(poma_heatmap))
  poma_heatmap <- ggplot2::ggplot() +
    ggplot2::annotation_custom(hm_grob, xmin = -Inf, xmax = Inf, ymin = -Inf, ymax = Inf) +
    ggplot2::theme_void()
  
  return(poma_heatmap)
}