heatmap of table

gradient_heatmap.R

library(RColorBrewer)
library(fields) #to use designer.colors
library(reshape2)
library(dplyr)


gradient_heatmap <- 
  function(
    df, 
    num_percentiles = 10, 
    reorder_col = T,
    reorder_row = T,
    base_size = 9,
    grouping = c("column", "row"),
    transformation = c("none", "scale", "residual"),
    scale_high = "Frequent", 
    scale_low = "Uncommon",
    continuous = F,
    ...
  )
  {



    if(length(grouping) > 1) grouping <- "row"
    if(length(transformation) > 1) transformation <- "residual"

    extra_args <- list(...)

    if(!"pallete" %in% names(extra_args)) 
      pallete <- brewer.pal(min(9, num_percentiles-1), "Spectral")
    else
      pallete <- extra_args$pallete

    col_ramp <- rev(designer.colors(n=num_percentiles + 1, col=pallete))

    df_matrix <- data.matrix(df)

    if(transformation == "scale")
    {
      print("warning: transformation = scale does column-wise scaling")
      
      df_matrix <- df_matrix %>% scale

      df_q <- 
        data.matrix(df) %>% 
        scale %>% 
        as.data.frame %>% 
        mutate(row = rownames(.)) %>% 
        melt(id = "row", .)
    }

    if(transformation == "none")
    {
      df_q <- df %>%
        mutate(row = rownames(.)) %>% 
        melt(id = "row", .)
    }

    if(transformation == "residual")
    {
      df_q <- chisq.test(df)  
      df_q <- df_q$residuals %>%
        as.data.frame %>% 
        mutate(row = rownames(.)) %>% 
        melt(id = "row", .)
    }

    if(grouping == "column")
    {
      df_q <- 
        df_q %>% 
        group_by(variable) %>% 
        mutate(
          quantile = cut(
            value, 
            unique(
              quantile(
                value, 
                probs =0:num_percentiles/num_percentiles)), 
            include.lowest = T
          ) %>% as.integer %>% as.factor,
          value = value - mean(value)
        )
    }

    if(grouping == "row")
    {
      df_q <- 
        df_q %>% 
        group_by(row) %>% 
        mutate(
          quantile = cut(
            value, 
            unique(
              quantile(
                value, 
                probs =0:num_percentiles/num_percentiles)), 
            include.lowest = T
          ) %>% as.integer %>% as.factor,
          value = value - mean(value)
      )
    }
  
    if(reorder_row == T)
    {
      row_order <- 
        df_matrix %>% 
        dist %>% 
        hclust %>% 
        as.dendrogram %>% 
        order.dendrogram
      
      df_q <- 
        df_q %>% 
        ungroup %>% 
        mutate(
          row = ordered(
            row, 
            levels = rownames(df)[row_order]
          )
        )
    } 
    
    if(reorder_col == T)
    {
      col_order <- 
        df_matrix %>% 
        t %>% 
        dist %>% 
        hclust %>% 
        as.dendrogram %>% 
        order.dendrogram
      
      df_q <- 
        df_q %>% 
        ungroup %>% 
        mutate(
          variable = ordered(
            variable, 
            levels = colnames(df)[col_order]
          )
        )
    }
    
    if(!"angle" %in% names(extra_args)) 
      angle <- 330
    else
      angle <- extra_args$angle
    
    p <- 
      ggplot(df_q, aes(variable, row)) + 
      geom_tile(
        aes(fill = quantile),
        colour = "white"
      ) + 
      scale_fill_manual(
        values = col_ramp,
        breaks = c(num_percentiles, 1),
        labels = c(scale_high, scale_low),
        name = ""
      )
    
    if(continuous == T)
    {
      p <- 
        ggplot(df_q, aes(variable, row)) + 
        geom_tile(
          aes(fill = value),
          colour = "white"
        ) + 
        scale_fill_gradientn(
          colors = col_ramp,
          breaks = c(max(df_q$value), min(df_q$value)),
          labels = c(scale_high, scale_low),
          name = ""
        )
    }
    
    p <- p + 
      theme_grey(base_size = base_size) + 
      labs(x = "", y = "") + 
      scale_x_discrete(expand = c(0, 0)) +
      scale_y_discrete(expand = c(0, 0)) + 
      theme(
        # legend.position = "none",
        axis.ticks = element_blank(), 
        axis.text.x = 
          element_text(
            family = "Consolas",
            size = base_size * 0.8, 
            angle = angle, 
            hjust = 0, 
            colour = "grey50"
          )
        )
    
    if("plot_margin" %in% names(extra_args)) 
      p <- p + theme(plot.margin = extra_args$plot_margin)
    
    if(exists("flip_coord") && flip_coord == T) 
      p <- p + coord_flip()
    
    
    return(p)
  }