Update docs, improve params, add X annotations to heatmap

Author: lahuuki

R/layer_stat_cor_plot_complex.R

@@ -1,13 +1,34 @@
-
-
 #' Visualize the correlation of layer modeling t-statistics with ComplexHeatmap
-#' @param query_colors named vector of colors for query row annotations
-#' @param reference_colors named vector of colors for reference column annotations
+#' 
+#' This function updates [layer_stat_cor_plot()], using ComplexHeatmap to plot
+#' the correlation matrix between a reference and query modeling statistics 
+#' from [layer_stat_cor()]. Includes functionality to add color annotations, 
+#' (helpful to match to colors in Visium spot plots), and annotations from 
+#' [annotate_registered_clusters()].
+#' 
+#' @param color_max A `numeric(1)` specifying the highest correlation value for 
+#' the color scale (should be between 0 and 1).
+#' @param color_min A `numeric(1)` specifying the lowest correlation value for 
+#' the color scale (should be between 0 and -1).
+#' @param color_scale A `character` vector specifying the color scale for the 
+#' fill of the heatmap, defaults to classic purple -> green
+#' @param query_colors named `character` vector of colors, Adds colors to query 
+#' row annotations
+#' @param reference_colors named `character` vector of colors, Adds colors to 
+#' reference column annotations
+#' @param annotation annotation data.frame output of [annotate_registered_clusters()],
+#' adds 'X' for good confidence annotations, '*' for poor confidence
 #'
 #' @inheritParams layer_stat_cor_plot
 #'
-#' @return ComplexHeatmap plot of t-stat correlations
+#' @return ([Heatmap-class][ComplexHeatmap::Heatmap-class]) plot of t-stat correlations
 #' @export
+#' @author Louise Huuki-Myers
+#' @family Layer correlation functions
+#'
+#' @importFrom RColorBrewer brewer.pal
+#' @importFrom grDevices colorRampPalette
+#' @importFrom ComplexHeatmap columnAnnotation rowAnnotation Heatmap
 #'
 #' @examples
 #' ## Obtain the necessary data
@@ -48,16 +69,33 @@
 #'                             reference_colors = libd_layer_colors)
 #' 
 #' ## Apply additional ComplexHeatmap param
-#' layer_stat_cor_plot_complex(cor_stats_layer, cluster_rows = FALSE)
+#' layer_stat_cor_plot_complex(cor_stats_layer, cluster_rows = FALSE, cluster_columns = FALSE)
 #' 
+#' ## Add annotation
+#' annotation_df <- annotate_registered_clusters(cor_stats_layer, confidence_threshold = .55)
+#' layer_stat_cor_plot_complex(cor_stats_layer, annotation = annotation_df)
+#' 
+#' ## All together
+#' layer_stat_cor_plot_complex(cor_stats_layer, 
+#'                             query_colors = cluster_colors,
+#'                             reference_colors = libd_layer_colors, 
+#'                             annotation = annotation_df,
+#'                             rect_gp = gpar(col = "black", lwd = 1))
+#'                             
 layer_stat_cor_plot_complex <- function(cor_stats_layer,
-                                        theSeq = seq(min(cor_stats_layer), max(cor_stats_layer), by = 0.01),
-                                        my.col = grDevices::colorRampPalette(RColorBrewer::brewer.pal(7, "PRGn"))(length(theSeq)),
+                                        color_max = max(cor_stats_layer),
+                                        color_min = min(cor_stats_layer),
+                                        color_scale = RColorBrewer::brewer.pal(7, "PRGn"),
                                         query_colors = NULL,
                                         reference_colors = NULL,
+                                        annotation = NULL,
                                         ...
                                         ){
 
+  ## define color pallet
+  theSeq = seq(color_min, color_max, by = 0.01)
+  my.col = grDevices::colorRampPalette(color_scale)(length(theSeq))
+  
   # ## query annotations on row
   if(!is.null(query_colors)){
 
@@ -84,16 +122,55 @@ layer_stat_cor_plot_complex <- function(cor_stats_layer,
     )
   } else ref_col_annotation <- NULL
 
-  
+  ## add annotation
+  if(!is.null(annotation)){
+    anno_matrix <- create_annotation_matrix(annotation, cor_stats_layer)
+    
+    ## plot heatmap
+    return(
+      ComplexHeatmap::Heatmap(
+        matrix = cor_stats_layer,
+        col = my.col,
+        bottom_annotation = ref_col_annotation,
+        right_annotation = query_row_annotation,
+        cell_fun = function(j, i, x, y, width, height, fill) {
+          grid.text(anno_matrix[i, j], x, y, gp = gpar(fontsize = 10))
+        },
+        ...
+      ))
+  }
 
   ## plot heatmap
+  return(
   ComplexHeatmap::Heatmap(
     matrix = cor_stats_layer,
     col = my.col,
     bottom_annotation = ref_col_annotation,
     right_annotation = query_row_annotation,
     ...
-  )
+  ))
 
 
 }
+
+create_annotation_matrix <- function(annotation_df, cor_stats_layer){
+  
+  anno_list <- lapply(rownames(cor_stats_layer), 
+                      function(cluster){
+    # look up confidence                    
+    confidence <- annotation_df[match(cluster, annotation_df$cluster),"layer_confidence"]
+    sym <- ifelse(confidence=="good", "X","*")
+    # match annotations
+    anno <- annotation_df[match(cluster, annotation_df$cluster),"layer_label"]
+    return(ifelse(unlist(lapply(colnames(cor_stats_layer), grepl, anno)),sym,""))
+  })
+  
+  anno_matrix <- t(data.frame(anno_list))
+  rownames(anno_matrix) <- rownames(cor_stats_layer)
+  colnames(anno_matrix) <- colnames(cor_stats_layer)
+  
+  return(anno_matrix)
+}
+
+
+