1 files changed, 73 insertions, 8 deletions
diff --git a/R/corplot.R b/R/corplot.R
index 08221a4..e849a92 100644
--- a/R/corplot.R
+++ b/R/corplot.R
@@ -16,43 +16,100 @@
 #' @import ggplot2
 #' @export
 corplot <- function(R, title = NULL){
-	R <- tryCatch({
+	R_list <- tryCatch({
 		if(is.matrix(R)){
+			cnames <- colnames(R)
+			rnames <- rownames(R)
 			if(nrow(R) == ncol(R)){
-				R[upper.tri(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)){
-				R
+				list(
+					R = R,
+					cnames = NULL,
+					rnames = NULL,
+					p = n
+				)
 			}else{
 				stop('Input length not compatible with upper triangle')
 			}
 		}else{
 			Rm <- as.matrix(R)
+			cnames <- colnames(Rm)
+			rnames <- rownames(Rm)
 			if(nrow(Rm) == ncol(Rm)){
-				Rm[upper.tri(Rm)]
+				list(
+					R = Rm[upper.tri(Rm)],
+					cnames = cnames,
+					rnames = rnames,
+					p = ncol(Rm)
+				)
 			}
 		}
 	}, error = function(e){
 		stop(paste('Invalid input for corplot:', e$message))
 	})
 
+	if(is.null(R_list$cnames)){
+		R_list$cnames <- as.character(seq_len(R_list$p))
+	}
+	if(is.null(R_list$rnames)){
+		R_list$rnames <- as.character(seq_len(R_list$p))
+	}
+
 	# Now get a df from the upper triangle
-	m <- length(R)
+	m <- length(R_list$R)
 	n <- as.integer((1 + sqrt(1 + 8 * m)) / 2)
 	M <- matrix(NA, n, n)
-	M[upper.tri(M)] <- R
+	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')
 
+	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',
@@ -63,8 +120,16 @@ corplot <- function(R, title = NULL){
 		guide = 'none'
 	) +
 	ggplot2::scale_y_reverse() +
-	ggplot2::coord_fixed() +
-	ggplot2::theme_void()
+	ggplot2::coord_fixed(clip = 'off') +
+	ggplot2::theme_void() +
+	ggplot2::theme(
+		plot.margin = ggplot2::margin(
+			t = top_margin,
+			r = 5,
+			b = 5,
+			l = 5
+		)
+	)
 
 	if(!is.null(title)){
 		p <- p +