R/beta_heatmap.R

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

#' @title Create heatmaps of the supplied dissimilarity matrices
#' @name beta_heatmap
#' @description Create heatmaps of the supplied dissimilarity matrices ordered by supplied grouping variables
#' @param beta_div A dissimilarity matrix calculated by \code{beta_div}
#' @param micro_set A tidy_micro data set
#' @param ... Variables for ordering
#' @param low_grad Colors for the corelation magnitude. Will be fed into scale_fill_gradient
#' @param high_grad Colors for the corelation magnitude. Will be fed into scale_fill_gradient
#' @param main Plot title
#' @param xlab x-axis label
#' @param ylab y-axis label
#' @param subtitle Plot label
#' @param natural_order Keep order of axes in the conventional order for dissimilarity matrices
#' @param legend_title Title for the legend
#' @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
#'
#' ## Bray-Curtis beta diversity
#' bray <- set %>% beta_div(table = "Family")
#'
#' bray %>% beta_heatmap(set, bpd1)
#' @export
beta_heatmap <- function(beta_div, micro_set, ..., low_grad, high_grad,
                         main = NULL, xlab = NULL, ylab = NULL, subtitle = NULL,
                         natural_order = TRUE,
                         legend_title = "Dissimilarity"){

  if(!missing(low_grad) & missing(high_grad)) stop("Must specify both low_grad and high_grad.")
  if(missing(low_grad) & !missing(high_grad)) stop("Must specify both low_grad and high_grad.")

  micro_set %<>%
    dplyr::distinct(Lib, .keep_all = TRUE) %>% ## unique subjects from micro_set
    dplyr::select(Lib, !!!rlang::quos(...))

  if(ncol(micro_set) > 2) stop("Must use one factor variable.")
  if(class(micro_set[,2]) %nin% c("character", "factor")) stop("Must use one factor variable.")

  ## Ensuring colnames are the Lib names for gather step
  colnames(beta_div) <- micro_set$Lib

  ## For reording leves
  Var <- micro_set[, 2]

  CC <- beta_div %>%
    as.data.frame %>% ## Getting beta diversity file into format for ggplot
    dplyr::mutate(Lib = rownames(.)) %>%
    dplyr::full_join(micro_set, by = "Lib") %>%
    tidyr::unite(II, !!!rlang::quos(...), Lib, sep = ":_:") %>%
    ## Supposed to be a 'unique' sep so we can locate it later
    ## Will break if ':_:' is in someones Libs or variable name
    tidyr::gather(Lib2, val, -II) %>%
    dplyr::full_join(micro_set %>% dplyr::rename(Lib2 = Lib), by = "Lib2") %>%
    tidyr::unite(III, !!!rlang::quos(...), Lib2, sep = ":_:")

  if(natural_order){
    CC %<>%
      dplyr::mutate(II = factor(II),
                    II = factor(II, levels = rev(levels(II)))
      ) ## gets the same subject beta divs on the main diagonal

    Var_y <- factor(Var, levels = rev(levels(Var)))

    gg <- CC %>%
      ggplot2::ggplot(aes(II, III, fill = val)) +
      ggplot2::geom_tile() +
      ggplot2::labs(x = xlab, y = ylab, subtitle = subtitle, title = main, fill = legend_title) +
      ggplot2::theme(axis.text.x = element_text(angle = 45)) +
      ggplot2::scale_y_discrete(labels = rep(levels(Var_y), table(Var_y))) +
      ggplot2::scale_x_discrete(labels = rep(levels(Var), table(Var)))
  } else {
    gg <- CC %>%
      ggplot2::ggplot(aes(II, III, fill = val)) +
      ggplot2::geom_tile() +
      ggplot2::labs(x = xlab, y = ylab, subtitle = subtitle, title = main, fill = legend_title) +
      ggplot2::theme(axis.text.x = element_text(angle = 45)) +
      ggplot2::scale_y_discrete(labels = rep(levels(Var), table(Var))) +
      ggplot2::scale_x_discrete(labels = rep(levels(Var), table(Var)))
  }

  if(!missing(low_grad) & !missing(high_grad)){
    gg <- gg + ggplot2::scale_fill_gradient(low = low_grad, high = high_grad)
  }

  gg
}