Add flags for axis labels

1 files changed, 46 insertions, 40 deletions

diff --git a/R/heatmap.R b/R/heatmap.R
index 569f0ec..77721ce 100644
--- a/R/heatmap.R
+++ b/R/heatmap.R
@@ -4,6 +4,7 @@
 #' in the top-right corner with no axis labels.
 #' @param R A square numeric correlation matrix
 #' @param title Optional plot title
+#' @param labels Logical to plot column names or not (default FALSE)
 #'
 #' @return A ggplot object
 #'
@@ -15,7 +16,7 @@
 #' @importFrom reshape2 melt
 #' @import ggplot2
 #' @export
-heatmap <- function(R, title = NULL){
+heatmap <- function(R, title = NULL, labels = FALSE){
 	R_list <- tryCatch({
 		if(is.matrix(R)){
 			cnames <- colnames(R)
@@ -55,7 +56,7 @@ heatmap <- function(R, title = NULL){
 			}
 		}
 	}, error = function(e){
-		stop(paste('Invalid input for corplot:', e$message))
+		stop(paste('Invalid input for heatmap:', e$message))
 	})
 
 	if(is.null(R_list$cnames)){
@@ -75,41 +76,12 @@ heatmap <- function(R, title = NULL){
 	df <- reshape2::melt(M, na.rm = TRUE)
 	colnames(df) <- c('i', 'j', 'value')
 
-	col_label_df <- data.frame(
-		j = seq_len(n)[-1],
-		i = 0,
-		label = R_list$cnames[-1]
-	)
-	row_label_df <- data.frame(
-		j = seq_len(n)[-n],
-		i = seq_len(n)[-n],
-		label = R_list$cnames[-n]
-	)
-
-	max_nchar <- max(nchar(R_list$cnames))
-	top_margin <- 5 + max_nchar * 2
-
 	# And plot it
 	p <- ggplot2::ggplot(
 		data = df,
 		ggplot2::aes(x = .data$j, y = .data$i, fill = .data$value)
 	) +
 	ggplot2::geom_tile() +
-	ggplot2::geom_text(
-		data = col_label_df,
-		ggplot2::aes(x = .data$j, y = .data$i, label = .data$label),
-		inherit.aes = FALSE,
-		angle = 90,
-		hjust = 0,
-		vjust = 0.5
-	) +
-	ggplot2::geom_text(
-		data = row_label_df,
-		ggplot2::aes(x = .data$j, y = .data$i, label = .data$label),
-		inherit.aes = FALSE,
-		hjust = 1,
-		vjust = 0.5
-	) +
 	ggplot2::scale_fill_gradient2(
 		low = 'red',
 		mid = 'white',
@@ -121,15 +93,7 @@ heatmap <- function(R, title = NULL){
 	) +
 	ggplot2::scale_y_reverse() +
 	ggplot2::coord_fixed(clip = 'off') +
-	ggplot2::theme_void() +
-	ggplot2::theme(
-		plot.margin = ggplot2::margin(
-			t = top_margin,
-			r = 5,
-			b = 5,
-			l = 5
-		)
-	)
+	ggplot2::theme_void()
 
 	if(!is.null(title)){
 		p <- p +
@@ -138,5 +102,47 @@ heatmap <- function(R, title = NULL){
 			plot.title = ggplot2::element_text(hjust = 0.5)
 		)
 	}
+
+	if(labels){
+		col_label_df <- data.frame(
+			j = seq_len(n)[-1],
+			i = 0,
+			label = R_list$cnames[-1]
+		)
+		row_label_df <- data.frame(
+			j = seq_len(n)[-n],
+			i = seq_len(n)[-n],
+			label = R_list$cnames[-n]
+		)
+
+		max_nchar <- max(nchar(R_list$cnames))
+		top_margin <- 5 + max_nchar * 2
+
+		p <- p +
+		ggplot2::geom_text(
+			data = col_label_df,
+			ggplot2::aes(x = .data$j, y = .data$i, label = .data$label),
+			inherit.aes = FALSE,
+			angle = 90,
+			hjust = 0,
+			vjust = 0.5
+		) +
+		ggplot2::geom_text(
+			data = row_label_df,
+			ggplot2::aes(x = .data$j, y = .data$i, label = .data$label),
+			inherit.aes = FALSE,
+			hjust = 1,
+			vjust = 0.5
+		) +
+		ggplot2::theme(
+			plot.margin = ggplot2::margin(
+				t = top_margin,
+				r = 5,
+				b = 5,
+				l = 5
+			)
+		)
+	}
+
 	return(p)
 }