R/micro_heatmap.R

In CharlieCarpenter/tidy.micro: A Pipeline For Microbiome Analysis And Visualization

#' @title Create heatmaps of estiamted coefficients from negative binomial models
#' @name micro_heatmap
#' @description A function to create heatmaps of estimated beta coeffients from each model fit by nb_mods
#' @param modsum The output from nb_mods
#' @param low_grad The low gradient colors for the coefficient magnitude. Will be fed into scale_fill_gradient
#' @param high_grad The high gradient colors for the coefficient magnitude. Will be fed into scale_fill_gradient
#' @param mid_grad The medium gradient colors for the coefficient magnitude. Will be fed into scale_fill_gradient
#' @param midpoint Midpoint for coefficient magnitude in legend
#' @param top_taxa Only plot X taxa with the largest magnitude beta coefficients
#' @param low_lim Lower limits of the fill gradient. Will default to the largest magnitude effect size
#' @param high_lim Upper limits of the fill gradient. Will default to the largest magnitude effect size
#' @param mute_cols Mute the colors of the fill gradients
#' @param alpha Mark beta coefficient cells with p-values below this cutoff
#' @param dot_size size of marker in cells
#' @param shape shape of marker in cells
#' @param main Plot title
#' @param xlab x-axis label
#' @param ylab y-axis label
#' @param subtitle Plot label
#' @param xaxis Labels for the x-axis ticks
#' @param legend_title Title for the legend
#' @details The output will give gray columns if there are missing values in the supplied continuous variable
#' @return Returns a ggplot that you can add geoms to if you'd like
#' @examples
#' data(phy); data(cla); data(ord); data(fam); data(met)
#'
#' otu_tabs = list(Phylum = phy, Class = cla, Order = ord, Family = fam)
#' set <- tidy_micro(otu_tabs = otu_tabs, meta = met) \%>\%
#' filter(day == 7) ## Only including the first week
#'
#' ## Creating negative binomial models on filtered tidy_micro set
#' nb_fam <- set \%>\%
#' mutate(bpd1 = factor(bpd1)) \%>\% ## making bpd1 a factor
#' otu_filter(ra_cutoff = 0.1, exclude_taxa = c("Unclassified", "Bacteria")) \%>\%
#' nb_mods(table = "Family", bpd1)
#'
#' nb_fam \%>\% micro_heatmap
#' @export
micro_heatmap <- function(modsum, low_grad, high_grad, mid_grad, midpoint = 0, top_taxa = 10,
                       low_lim, high_lim, mute_cols = T, alpha = 0.05, dot_size = 2,
                       dot_shape = 8, main = NULL, xlab = NULL, ylab = NULL, subtitle = NULL,
                       xaxis = NULL,  caption = latex2exp::TeX("'*' denotes significance at $\\alpha$ = 0.05")){

  if(is.null(xaxis)) xaxis <- modsum$Model_Coef$Coef %>% unique %>% .[!grepl("(Intercept)",.)]

  if(missing(low_grad)) low_grad <- "blue"; if(missing(high_grad)) high_grad <- "red"
  if(missing(mid_grad)) mid_grad <- "beige"

  ## selecting the Taxa, model Coefs, and RR, and arranging them by RR within Taxa
  CC <- modsum$Convergent_Summary %>%
    dplyr::select(Taxa, Coef, Beta, FDR_Pval) %>%
    dplyr::filter(!grepl("(Intercept)", Coef))

  if(missing(low_lim)) low_lim <- -max(abs(max(CC$Beta)),
                                       abs(min(CC$Beta)))
  if(missing(high_lim)) high_lim <- max(abs(max(CC$Beta)),
                                        abs(min(CC$Beta)))

  ## Organizing by strongest and most consistent effect size
  ar <- CC %>%
    dplyr::group_by(Taxa) %>%
    dplyr::summarise(max = max(Beta), avg = mean(Beta)) %>%
    dplyr::left_join(CC, by = "Taxa") %>%
    unique %>%
    dplyr::arrange(desc(max), desc(avg))

  ## Pulling out top taxa by effect size and plotting
  CC %>% dplyr::filter(Taxa %in% unique(ar$Taxa)[1:top_taxa]) %>%

    ggplot2::ggplot(aes(Coef, Taxa)) +
    ggplot2::geom_tile(aes(fill = Beta)) +
    ggplot2::theme_bw() +
    ggplot2::scale_fill_gradient2(limits = c(low_lim, high_lim),
                         low = low_grad,
                         high = high_grad,
                         mid = mid_grad,
                         midpoint = midpoint) +
    # ggplot2::labs(title = main, x = xlab, y = ylab,
    #               subtitle = subtitle, caption = caption,
    #               fill = latex2exp::TeX("$\\hat{\\beta}$")) +
    ggplot2::scale_x_discrete(labels = xaxis) +
    ggplot2::geom_point(aes(size=ifelse(FDR_Pval < alpha, "dot", "no_dot")), shape = 8) +
    ggplot2::scale_size_manual(values=c(dot=2, no_dot=NA), guide="none")
}