Add order option

1 files changed, 19 insertions, 31 deletions

diff --git a/R/heatmap.R b/R/heatmap.R
index c40946f..cb6c8f8 100644
--- a/R/heatmap.R
+++ b/R/heatmap.R
@@ -5,6 +5,7 @@
 #' @param R A square numeric correlation matrix
 #' @param title Optional plot title
 #' @param labels Logical to plot column names or not (default FALSE)
+#' @param order Optional character vector of the order you want the plot sorted in
 #'
 #' @return A ggplot object
 #'
@@ -16,45 +17,32 @@
 #' @importFrom reshape2 melt
 #' @import ggplot2
 #' @export
-heatmap <- function(R, title = NULL, labels = FALSE){
+heatmap <- function(R, title = NULL, labels = FALSE, order = NULL){
 	R_list <- tryCatch({
 		if(is.matrix(R)){
-			cnames <- colnames(R)
-			rnames <- rownames(R)
-			if(nrow(R) == ncol(R)){
-				list(
-					R = R[upper.tri(R)],
-					cnames = cnames,
-					rnames = rnames,
-					p = ncol(R)
-				)
-			}
-		}else if(is.vector(R)){
-			m <- length(R)
-			n <- (1 + sqrt(1 + 8 * m)) / 2
-			if(n == floor(n)){
-				list(
-					R = R,
-					cnames = NULL,
-					rnames = NULL,
-					p = n
-				)
-			}else{
-				stop('Input length not compatible with upper triangle')
-			}
+			Rm <- R
+			cnames <- colnames(Rm)
+			rnames <- rownames(Rm)
 		}else{
 			Rm <- as.matrix(R)
 			cnames <- colnames(Rm)
 			rnames <- rownames(Rm)
-			if(nrow(Rm) == ncol(Rm)){
-				list(
-					R = Rm[upper.tri(Rm)],
-					cnames = cnames,
-					rnames = rnames,
-					p = ncol(Rm)
-				)
+		}
+		if(nrow(Rm) == ncol(Rm)){
+			if(!is.null(order)){
+				order <- get_order(Rm)
+				M <- Rm[order, order]
 			}
+			list(
+				M = M[upper.tri(M)],
+				cnames = cnames,
+				rnames = rnames,
+				p = ncol(M)
+			)
+		}else{
+			stop('Input must be coercible to a square matrix')
 		}
+
 	}, error = function(e){
 		stop(paste('Invalid input for heatmap:', e$message))
 	})