Add groupings to the network graph

1 files changed, 92 insertions, 11 deletions

diff --git a/R/networkplot.R b/R/networkplot.R
index cb6aad8..1055b08 100644
--- a/R/networkplot.R
+++ b/R/networkplot.R
@@ -3,10 +3,12 @@
 #' Draw an undirected graph from a square matrix
 #' where each variable is a node and the pairwise
 #' entries are the strength of the relationship
-#' (ex: correlation, partial correlation)
+#' (ex: correlation, partial correlation).
+#' Based on https://r-graph-gallery.com/310-custom-hierarchical-edge-bundling.html
 #'
 #' @param R A square numeric matrix with names columns and rows
 #' @param title Optional plot title
+#' @param groups Optional data.frame mapping nodes to groups
 #'
 #' @return A ggraph object
 #'
@@ -14,12 +16,18 @@
 #' X <- matrix(data = rnorm(100), nrow = 10)
 #' R <- cor(X)
 #' networkplot(R, title = 'Example Correlation Plot')
+#' colnames(R) <- as.character(seq_len(ncol(R)))
+#' g <- data.frame(
+#'	node = colnames(R),
+#'	group = rep(c('A', 'B'), length.out = ncol(R))
+#' )
+#' networkplot(R, groups = g, title = 'Example Correlation Plot')
 #'
 #' @import ggraph
 #' @import igraph
 #'
 #' @export
-networkplot <- function(R, title = NULL){
+networkplot <- function(R, groups = NULL, title = NULL){
 	Rm <- tryCatch({
 		if(is.matrix(R)){
 			if(nrow(R) != ncol(R)) stop('R must be square')
@@ -51,26 +59,86 @@ networkplot <- function(R, title = NULL){
 		stop(paste('Invalid input for networkplot:', e$message))
 	})
 
+	# Rm is p x p matrix
+	# upper.tri(Rm) returns a logical p x p matrix with upper right TRUE
+	# which() gives the coordinate pairs for unique edges
 	idx <- which(upper.tri(Rm), arr.ind = TRUE)
+	
+	# edges is a data.frame of k rows with all unique
+	# combinations of pairwise entries and value r
 	edges <- data.frame(
 		from = rownames(Rm)[idx[,1]],
 		to = colnames(Rm)[idx[,2]],
 		r = Rm[idx]
 	)
+
+	# groups must be a data.frame with 2 columns 'node' and 'group'
+	if(!is.null(groups)){
+		if(!all(c('node','group') %in% colnames(groups))){
+			stop("groups must have columns 'node' and 'group'")
+		}
+		if(!all(rownames(Rm) %in% groups$node)){
+			stop('All nodes must appear in groups$node')
+		}
+
+		# map hierarchy connecting root -> each group
+		d1 <- data.frame(from = 'root', to = unique(groups$group))
+		# map hierarchy of each group -> each node
+		d2 <- data.frame(from = groups$group, to = groups$node)
+		# concatenate them (for 2 levels of the hierarhcy)
+		hierarchy <- rbind(d1, d2)
+	}else{
+		# fallback to single group
+		# map root -> each node (no grouping)
+		hierarchy <- data.frame(from = 'root', to = rownames(Rm))
+	}
+
+	# this is a data.frame with one column listing all vertices
+	# this includes the hierarchy vertices (root, groups)
+	# this is required by igraph as a node ID table
+	vertices <- data.frame(name = unique(c(hierarchy$from, hierarchy$to)))
+
+	# If groups were passed, add a column to vertices called group
+	if(!is.null(groups)){
+		# maps each vertex to its group, returns NA for root and the
+		# group vertices themselves
+		# we will use this to also color the nodes
+		vertices$group <- groups$group[match(vertices$name, groups$node)]
+	}else{
+		vertices$group <- 'all'
+	}
+
+	# hierarchy is a data.frame listing all the edges
+	# vertices is a data.frame listing all the nodes with their attributes
+	# returns an igraph object representing the tree structure
 	g <- igraph::graph_from_data_frame(
-		edges,
-		directed = FALSE
+		hierarchy,
+		vertices = vertices
 	)
 
-	p <- ggraph::ggraph(graph = g, layout = 'circle') +
-	ggraph::geom_edge_arc(
+	# converts the node names to vertex indeces
+	# required by ggraph::get_con()
+	from <- match(edges$from, vertices$name)
+	to <- match(edges$to, vertices$name)
+
+	# takes the integer vectors from and to and returns the paths along tree
+	# transforms node-node edges into paths through hierarchy
+	con <- ggraph::get_con(from = from, to = to, value = edges$r)
+
+	# g is the igraph tree
+	# dendrogram defines tree embedding
+	# circlular arranges the nodes in a circle
+	# g contains vertices
+	p <- ggraph::ggraph(graph = g, layout = 'dendrogram', circular = TRUE) +
+	ggraph::geom_conn_bundle(
+		data = con,
 		ggplot2::aes(
-			edge_alpha = abs(.data$r),
-			edge_color = .data$r
+			edge_alpha = abs(.data$value),
+			edge_colour = .data$value
 		),
-		strength = 0.5
+		tension = 0.5
 	) +
-	ggraph::scale_edge_color_gradient2(
+	ggraph::scale_edge_colour_gradient2(
 		low = 'red',
 		mid = 'white',
 		high = 'blue',
@@ -78,9 +146,22 @@ networkplot <- function(R, title = NULL){
 		limits = c(-1, 1),
 		guide = 'none'
 	) +
-	ggraph::geom_node_point(size = 3) +
+	# operates on node data from the layout
+	# leaf is a logical vector (provided by the ggraph object g)
+	# where leaf is TRUE, those are the terminal nodes (the ones we want to draw)
+	# .data$group is the character vector from vertices
+	ggraph::geom_node_point(
+		ggplot2::aes(
+			filter = .data$leaf,
+			colour = .data$group,
+			fill = .data$group
+		),
+		size = 3
+	) +
+	# labels nodes and groups
 	ggraph::geom_node_text(
 		ggplot2::aes(label = .data$name),
+		# prevents text labels from overlapping
 		repel = TRUE,
 		size = 3
 	) +