From 5a17b9f62de7890ff142fff7fca6b084f9bd38cf Mon Sep 17 00:00:00 2001
From: Ned Cauley <ned.cauley@nih.gov>
Date: Wed, 21 Feb 2024 16:33:09 -0500
Subject: [PATCH] feat: custom colors for heatmap annotations

---
 R/heatmap.R                           | 31 +++++++++++++++++++++++++--
 R/qc_preprocessing.R                  |  2 ++
 vignettes/Integration_Test_Kidney.Rmd |  9 +++++++-
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/R/heatmap.R b/R/heatmap.R
index cef7105..3e6c158 100644
--- a/R/heatmap.R
+++ b/R/heatmap.R
@@ -49,6 +49,8 @@
 #' @param heatmap.color Colors of heatmap to match breaks above 
 #'                      (Default: colorRampPalette(
 #'                       c("blue", "white", "red"))(120))
+#' @param anno.colors A list of named vectors that assigns the levels of each 
+#'                    annotation to a specific color
 #'
 #' @importFrom NanoStringNCTools assayDataApply
 #' @importFrom Biobase assayDataElement
@@ -74,7 +76,8 @@ heatMap <- function(
   clustering.distance.cols = "correlation",
   annotation.row = NA,
   breaks.by.values = seq(-3, 3, 0.05),
-  heatmap.color = colorRampPalette(c("blue", "white", "red"))(120)) {
+  heatmap.color = colorRampPalette(c("blue", "white", "red"))(120), 
+  annotation.colors = NULL) {
 
   # norm.method must be either quant or neg
   if((norm.method != "quant") && (norm.method != "neg")){
@@ -119,6 +122,9 @@ heatMap <- function(
   anno.col <- 
     pData(object)[, annotation.col]
   
+  # annotation colors
+  anno.colors <- annotation.colors
+  
   # make plot
   p <- pheatmap(
     plot.genes[1:ngenes, ],
@@ -134,9 +140,30 @@ heatMap <- function(
     clustering_distance_cols = clustering.distance.cols,
     breaks = breaks.by.values,
     color = col.palette,
-    annotation_col = anno.col
+    annotation_col = anno.col, 
+    annotation_colors = anno.colors
   )
   
+  p <- pheatmap(
+    plot.genes[1:ngenes, ],
+    main = "Clustering high CV genes",
+    scale = scale.by.row.or.col,
+    show_rownames = show.rownames,
+    show_colnames = show.colnames,
+    border_color = NA,
+    clustering_method = clustering.method,
+    cluster_rows = cluster.rows,
+    cluster_cols = cluster.cols,
+    clustering_distance_rows = clustering.distance.rows,
+    clustering_distance_cols = clustering.distance.cols,
+    breaks = breaks.by.values,
+    color = col.palette,
+    annotation_col = anno.col, 
+    annotation_colors = anno.colors
+  )
+  
+  
+  
   ## gene.df converts to data frame
   gene.df <- as.data.frame(plot.genes)
   
diff --git a/R/qc_preprocessing.R b/R/qc_preprocessing.R
index 6475c59..0821cc2 100644
--- a/R/qc_preprocessing.R
+++ b/R/qc_preprocessing.R
@@ -107,6 +107,8 @@ qcProc <-  function(object,
   ## settings ####
   ## shift counts (useDALogic=TRUE adds 1 only to 0s)
   object <- shiftCountsOne(object, useDALogic = shift.counts.zero)
+  
+  
   ## list of user-defined segment QC params for annotations and variables
   ## expected to be always present in the input object
   qc.params <-
diff --git a/vignettes/Integration_Test_Kidney.Rmd b/vignettes/Integration_Test_Kidney.Rmd
index 43fc012..cbdb7a0 100644
--- a/vignettes/Integration_Test_Kidney.Rmd
+++ b/vignettes/Integration_Test_Kidney.Rmd
@@ -195,6 +195,12 @@ qc.output <-  qcProc(object = sdesign.list$object,
  
 ```{r Clustering high CV Genes, echo=TRUE}
     
+    anno.colors = list(
+       class = c(DKD = "yellow", normal = "orange"), 
+       region = c(glomerulus = "purple", tubule = "red"), 
+       segment = c(GeometricSegment = "blue", neg = "coral", PanCK = "green")
+    )
+    
     heatmap.output <- heatMap(object = unsupervised.output$object, 
                         ngenes = 200, 
                         scale.by.row.or.col = "row", 
@@ -209,7 +215,8 @@ qc.output <-  qcProc(object = sdesign.list$object,
                         annotation.col = c("class", "segment", "region"), 
                         breaks.by.values = seq(-3, 3, 0.05), 
                         heatmap.color = colorRampPalette(c("blue", "white", "red"))(120), 
-                        norm.method = "quant")
+                        norm.method = "quant", 
+                        annotation.colors = anno.colors)
 
 
     print(heatmap.output$plot)