R/plot_heatmap.R

In PoisonAlien/peakseason: Functions for faster ChIP/ATAC-seq analysis and visualization.

#' Draw heatmap
#' @details This function takes output from extract_matrix and draws a heatmap
#' @param mat_list Input matrix list generated by \code{\link{extract_matrix}}
#' @param sortBy Sort matrix by.., Can be mean, median. Default mean.
#' @param rcb_pal Color palettte to use. RColorBrewer::display.brewer.all() to see available palettes
#' @param sample_bames Manually specifcy sample names.
#' @param title_size size of title. Default 0.8
#' @param top_profile Boolean. Whether to draw top profile plot.
#' @param zmins Manually specify min scores to include in heatmap
#' @param zmaxs Manually specify max scores to include in heatmap
#' @param scale Whether to row scale the matrix. Default FALSE
#' @param file_name Default NULL. If provided saves plot as a tiff.
#' @param hm_width Width of the plot
#' @param hm_height Height of the plot
#' @return If matrix was generated with `k`, returns BED file with cluster informartion
#' @export
plot_heatmap = function(mat_list, sortBy = "mean", rcb_pal = "Blues", sample_names = NULL,
                        title_size = 0.8, top_profile = FALSE, zmins = NULL, zmaxs = NULL,
                        scale = FALSE, file_name = NULL, hm_width = NULL, hm_height = 12, ...){


  if(!sortBy %in% c("mean", "median")){
    stop("sortBy can only be mean, median")
  }

  hmcols = suppressWarnings(RColorBrewer::brewer.pal(11, rcb_pal))

  if(!is.na(mat_list$param["k_num"])){
    mat_list = clustered_hm(mat_list, rcb_pal = rcb_pal, sample_names = sample_names,
                            title_size = title_size, top_profile = top_profile, zmins = zmins, zmaxs = zmaxs,
                            scale = scale, file_name = file_name, hm_width = hm_width, hm_height = hm_height)
    return(mat_list)
  }else{
    mat_list = order_matrix(mats = mat_list, sortBy = sortBy)
  }

  if(top_profile){
    profile_dat = summarizeMats(mats = mat_list, summarizeBy = "mean")
    yl = c(min(unlist(x = profile_dat[1:(length(profile_dat)-1)]), na.rm = TRUE),
           max(unlist(x = profile_dat[1:(length(profile_dat)-1)]), na.rm = TRUE))
    yl = round(x = yl, digits = 2)

  }

  param = mat_list$param
  cdata = mat_list$cdata
  size = param["size"]
  mat_list = mat_list[1:(length(mat_list)-2)]

  nsamps = length(mat_list)

  if(!is.null(zmins)){
    if(length(zmins) != length(mat_list)){
      warning("zmins are recycled")
      zmins = rep(zmins, length(mat_list))
    }
  }

  if(!is.null(zmaxs)){
    if(length(zmaxs) != length(mat_list)){
      warning("zmaxs are recycled")
      zmaxs = rep(zmaxs, length(mat_list))
    }
  }

  if(!is.null(sample_names)){
    if(length(sample_names) != length(mat_list)){
      stop("Number of sample names should be equal to number of samples in the matrix")
    }else{
      sample_names = cdata$sample_names
      names(mat_list) = sample_names
    }
  }

  xlabs = c(sapply(strsplit(x = as.character(size), split = ":", fixed = TRUE), "[[", 1), 0,
            sapply(strsplit(x = as.character(size), split = ":", fixed = TRUE), "[[", 2))
  xticks = c(0,
             1/(sum(as.integer(xlabs[1]), as.integer(xlabs[3]))/as.integer(xlabs[1])),
             1)

  if(top_profile){
    lo_mat = matrix(data = 1:(nsamps*2), nrow = 2, ncol = nsamps)
    lo = layout(mat = lo_mat, heights = c(1, 5))
  }else{
    lo_mat = matrix(data = 1:nsamps, nrow = 1)
    lo = layout(mat = lo_mat)
  }

  for(i in 1:nsamps){

    if(top_profile){
      plot_profile_mini(plot_dat = profile_dat, index = i, ylims = yl)
    }

    par(mar = c(3,2,1.5,1), las=1, tcl=-.25, font.main=4)
    hm.dat = mat_list[[i]]
    data.table::setDF(x = hm.dat)

    if(is.null(zmins)){
      #colMin = apply(hm.dat, 2, min, na.rm = TRUE)
      #zmin = floor(min(hm.dat, na.rm = TRUE))
      zmin = 0
    }else{
      zmin = zmins[i]
    }

    if(scale){
      hm.dat = scale(x = t(hm.dat), scale = TRUE)
    }else{
      hm.dat =  t(apply(hm.dat, 2, rev))
    }

    if(is.null(zmaxs)){
      colMax = apply(hm.dat, 1, max, na.rm = TRUE)
      colMax = which(x = colMax == max(colMax), arr.ind = TRUE)[1]
      zmax = ceiling(max(boxplot.stats(unlist(hm.dat[colMax,]))$stats))
    }else{
      zmax = zmaxs[i]
    }

    hmcols = grDevices::colorRampPalette(hmcols)(255)
    hm.dat[hm.dat >= zmax] = zmax
    #return(hm.dat)
    image(hm.dat, axes = FALSE, col = hmcols, useRaster = TRUE,
          zlim = c(zmin, zmax), xlim = c(-0.2, 1))

    #Add legend
    image(x = c(-0.1, -0.05), y = seq(0, 1, length.out = length(hmcols)-1),
          z = matrix(data = 1:(length(hmcols)-1), nrow = 1), add = TRUE, col = hmcols)
    axis(side = 2, at = seq(0, 1, length.out = 5),
         labels = NA,
         line = -0.6, font.axis = 2, yaxp  = c(1.1, 1.2, 3), lwd = 1.5)
    mtext(text = round(seq(zmin, zmax, length.out = 5), digits = 2), side = 2,
          line = -0.1, at = seq(0, 1, length.out = 5), font = 2)

    title(main = names(mat_list)[i], cex.main = title_size, font.main = 2)
    axis(side = 1, at = xticks,
         labels = NA, lty = 1, lwd = 1.5,
         font.axis = 2, cex.axis = 1, line = 0.7)
    mtext(text = c(paste0("-", xlabs[1]), xlabs[2], xlabs[3]), side = 1, line = 1, at = xticks, font = 2)

    rect(xleft = 0, ybottom = 0, xright = 1, ytop = 1, border = "black", lwd = 2)
  }
}