CitFuns/R/ggheatmap.R

ggheatmap<-function(df, x=NULL, y=NULL, value=NULL, grouping="mean", exclude_group=NULL, scale="none",color="#FFFFFF00"){
  if (is.null(x)){x=colnames(df)[1]}
  if (is.null(y)){y=colnames(df)[2]}
  if (is.null(value)){value=colnames(df)[3]}

  df<-rename(df, "VarX"=all_of(x), "VarY"=all_of(y), "Value"=all_of(value))

  if (is.null(exclude_group)){
    df<-df %>% group_by(VarX,VarY)
  }else{
    df<-df %>% group_by_("VarX","VarY",exclude_group) #%>% rename(exclude_group="all_of(exclude_group)")
  }
  if (grouping == "mean"){
    df<-df %>% summarise(Value=mean(Value)) %>% ungroup
  }
  if (grouping == "median"){
    df<-df %>% summarise(Value=median(Value)) %>% ungroup
  }

  if (length(unique(df$VarX)) > 1 & length(unique(df$VarY)) >
      1) {
    order <- clustsort(df %>% spread(VarY, Value) %>% select(!all_of(exclude_group)) %>%
                         as.data.frame)
  }
  else {
    order <- list()
    if (length(unique(df$VarX)) > 1) {
      xhclust <- df %>% spread(VarY, Value) %>%
        select(!all_of(exclude_group)) %>% as.data.frame
      order[["x"]] <- pull(xhclust, 1)[hclust(dist(xhclust %>%
                                                     select(-1)))$order]
    }
    else {
      order[["x"]] <- df %>% pull(VarX) %>% unique
    }
    if (length(unique(df$VarY)) > 1) {
      yhclust <- clustsort(df %>% spread(VarY, Value) %>%
                             select(!all_of(exclude_group)) %>% as.data.frame)
      order[["y"]] <- colnames(yhclust)[2:ncol(yhclust)][hclust(dist(t(yhclust %>%
                                                                         select(-1))))$order]
    }
    else {
      order[["y"]] <- df %>% pull(VarY) %>% unique
    }
  }

  if (scale != "none"){
    if (scale == "rows"){
      cols<-unique(df$VarY)
      sca.df<-spread(df, VarY, Value)
      for (i in cols){sca.df[,i]<-scale(sca.df[,i])}
      df<-gather(sca.df, VarY, Value, all_of(cols))
    }
    if (scale == "cols"){
      cols<-unique(df$VarX)
      sca.df<-spread(df, VarX, Value)
      for (i in cols){sca.df[,i]<-scale(sca.df[,i])}
      df<-gather(sca.df, VarX, Value, all_of(cols))
    }
  }

  df$VarX<-factor(df$VarX, levels=order$x)
  df$VarY<-factor(df$VarY, levels=order$y)

  df %>%
    ggplot(aes(VarX, VarY, fill=Value))+
    labs(x=x, y=y)+
    geom_tile(aes(fill=Value), color=color)+
    scale_fill_gradientn(colors=col2(200))+
    theme(axis.text.x = element_text(angle = 90, hjust = 1, vjust=0.5),
          panel.background = element_blank(),
          axis.ticks = element_blank())
}