R/results_heatmap.R

In NathanSkene/MAGMA_Celltyping: Find Causal Cell-Types Underlying Complex Trait Genetics

#' Plot heatmap of enrichment results
#' 
#' Plot of GWAS x cell-type enrichment results from 
#' \link[MAGMA.Celltyping]{celltype_associations_pipeline}.
#' 
#' @param merged_results Enrichment results generated by 
#' \link[MAGMA.Celltyping]{celltype_associations_pipeline} and merged by 
#' \link[MAGMA.Celltyping]{merge_results}.
#' @param title Plot title.
#' @param x_lab Plot x-axis label.
#' @param fdr_thresh FDR filtering threshold.
#' @param facet_formula Facet formula for plot,
#'  passed to \link[ggplot2]{facet_grid}.
#' @param x_var x-axis variable.
#' @param y_var y-axis variable.
#' @param fill_var Fill variable.
#' @param show_plot Whether to print the plot.
#' @param save_path Path to save plot to.
#' @param height Plot height.
#' @param width Plot width. 
#' @inheritParams ggplot2::facet_grid
#' @inheritParams ggplot2::ggsave
#' 
#' @return ggplot object.
#' 
#' @examples 
#' MAGMA_results <- MAGMA.Celltyping::enrichment_results
#' merged_results <- MAGMA.Celltyping::merge_results(MAGMA_results)
#' heat <- MAGMA.Celltyping::results_heatmap(
#'     merged_results = merged_results,
#'     fdr_thresh = 1)  
#' @export
#' @importFrom stats as.formula
results_heatmap <- function(merged_results,
                            title = NULL,
                            x_lab = NULL,
                            fdr_thresh = .05,
                            facet_formula = "EnrichmentMode ~ .",
                            x_var = "Celltype",
                            y_var = "GWAS",
                            fill_var = "-log1p(FDR)",
                            scales = "free_y",
                            space = "fixed",
                            show_plot = TRUE,
                            height = 5, 
                            width = 7,
                            dpi = 300,
                            save_path = file.path(
                                tempdir(),
                                "MAGMA_Celltyping.heatmap.jpg")
                            ) { 
    #### Avoid confusing checks ####
    FDR <- NULL;
    #### Check args #####
    if (!"GWAS" %in% colnames(merged_results)) {
        merged_results$GWAS <- merged_results$dataset
    }
    if (is.null(fdr_thresh)) fdr_thresh <- 1
    subtitle <- if (!is.null(fdr_thresh)) paste0("FDR < ", fdr_thresh) else NULL
    #### Filter by FDR ####
    plot_dat <- subset(merged_results, FDR < fdr_thresh)
    messager(formatC(nrow(plot_dat), big.mark = ","),
             "results @ FDR <",fdr_thresh)
    if(nrow(plot_dat)==0) stop("Filtered data must contain >0 rows.")
    #### Plot ####
    heat <- ggplot2::ggplot(
        data = plot_dat,
        ggplot2::aes_string(x = x_var, y = y_var, fill = fill_var)
    ) +
        ggplot2::geom_tile(color = "white") +
        ggplot2::facet_grid(
            facets = stats::as.formula(facet_formula),
            scales = scales,
            space = space
        ) +
        ggplot2::labs(
            title = title,
            subtitle = subtitle,
            x = x_lab
        ) +
        ggplot2::scale_fill_gradient(low = "blue", high = "red", 
                                     na.value = "white") +
        ggplot2::theme_bw() +
        ggplot2::theme(
            axis.text.x = ggplot2::element_text(
                angle = 45, hjust = 1, size = 7),
            strip.background = ggplot2::element_rect(fill = "grey20"),
            strip.text = ggplot2::element_text(color = "white"),
            panel.grid.minor = ggplot2::element_blank()
        )
    #### Print ####
    if (show_plot) print(heat)
    #### Save ####
    if (!is.null(save_path)){ 
        dir.create(dirname(save_path), showWarnings = FALSE, recursive = TRUE)
        ggplot2::ggsave(filename = save_path, 
                        plot = heat, 
                        dpi = dpi,
                        height = height, width = width)
    }
    return(heat)
}