Fix ordering

1 files changed, 23 insertions, 22 deletions

diff --git a/R/heatmap.R b/R/heatmap.R
index 3351045..cf34d5f 100644
--- a/R/heatmap.R
+++ b/R/heatmap.R
@@ -25,6 +25,10 @@ heatmap <- function(R, title = NULL, labels = FALSE, order = NULL){
 			Rm <- as.matrix(R)
 		}
 
+		if(nrow(Rm) != ncol(Rm)){
+			stop('Input must be coercible to a square matrix')
+		}
+
 		cnames <- colnames(Rm)
 		rnames <- rownames(Rm)
 		if(is.null(cnames)){
@@ -34,41 +38,38 @@ heatmap <- function(R, title = NULL, labels = FALSE, order = NULL){
 			rnames <- as.character(seq_len(nrow(Rm)))
 		}
 
-		if(nrow(Rm) == ncol(Rm)){
-			if(!is.null(order)){
-				if(length(order) != nrow(Rm)){
-					stop('Provided order is not the same length as the width of R')
-				}else{
-					rownames(Rm) <- rnames
-					colnames(Rm) <- cnames
-					Rm <- Rm[order, order]
-				}
+			
+		if(!is.null(order)){
+			if(length(order) != nrow(Rm)){
+				stop('Provided order is not the same length as the width of R')
+			}else{
+				rownames(Rm) <- rnames
+				colnames(Rm) <- cnames
+				Rm <- Rm[order, order, drop = FALSE]
 			}
-			list(
-				Rm = Rm[upper.tri(Rm)],
-				cnames = cnames,
-				rnames = rnames,
-				p = ncol(Rm)
-			)
-		}else{
-			stop('Input must be coercible to a square matrix')
 		}
+		list(
+			cnames = colnames(Rm),
+			rnames = rownames(Rm),
+			R = Rm[upper.tri(Rm)],
+			p = ncol(Rm)
+		)
 
 	}, error = function(e){
 		stop(paste('Invalid input for heatmap:', e$message))
 	})
 
 	# Now get a df from the upper triangle
-	m <- length(R_list$R)
-	n <- as.integer((1 + sqrt(1 + 8 * m)) / 2)
+	n <- R_list$p
 	M <- matrix(NA, n, n)
 	M[upper.tri(M)] <- R_list$R
 	colnames(M) <- R_list$cnames
 	rownames(M) <- R_list$rnames
 	df <- reshape2::melt(M, na.rm = TRUE)
-	colnames(df) <- c('i', 'j', 'value')
-	df$i <- as.numeric(df$i)
-	df$j <- as.numeric(df$j)
+	names(df) <- c('i', 'j', 'value')
+
+	df$i <- match(df$i, rownames(M))
+	df$j <- match(df$j, colnames(M))
 
 	# And plot it
 	p <- ggplot2::ggplot(