R/graph_visualization.R

In MetaNeighbor: Single cell replicability analysis

# Functions related to graph visualization of AUROCs.

#' Convert AUROC matrix into a graph.
#'
#' This representation is a useful alternative for heatmaps for large datasets
#' and sparse AUROC matrices (MetaNeighborUS with one_vs_best = TRUE)
#'
#' @param best_hits Matrix of AUROCs produced by MetaNeighborUS.
#' @param low_threshold AUROC threshold value. An edge is drawn between two
#' clusters only if their similarity exceeds low_threshold.
#' @param high_threshold AUROC threshold value. An edge is drawn between two
#' clusters only if their similarity is lower than high_threshold
#' (enables focusing on close calls).
#'
#' @return A graph in igraph format, where nodes are clusters and edges are
#' AUROC similarities.
#'
#' @export
makeClusterGraph <- function(best_hits, low_threshold = 0, high_threshold = 1) {
    filtered_hits <- best_hits
    filtered_hits[is.na(filtered_hits)] <- 0
    filtered_hits[best_hits>high_threshold | best_hits < low_threshold] <- 0
    result <- igraph::graph_from_adjacency_matrix(t(filtered_hits), weighted = TRUE)
    result <- igraph::simplify(result, remove.loops = TRUE)
    return(result)
}

#' Plot cluster graph generated with makeClusterGraph.
#'
#' In this visualization, edges are colored in black when AUROC > 0.5 and
#' orange when AUROC < 0.5, edge width scales linearly with AUROC.
#' Edges are oriented from training cluster towards
#' test cluster. A black bidirectional edge indicates that two clusters are
#' reciprocal top matches.
#' Node radius reflects cluster size (small: up to 10 cells,
#' medium: up to 100 cells, large: all other clusters).
#'
#' @param graph Graph in igraph format generated by makeClusterGraph.
#' @param study_id Vector with study IDs provided to MetaNeighborUS to compute
#' AUROCs stored in graph (used to compute cluster size). If NULL, all nodes
#' have medium size.
#' @param cell_type Vector with cell type labels provided to MetaNeighborUS
#' to compute AUROCs stored in graph (used to compute cluster size). If NULL,
#' all nodes have medium size.
#' @param size_factor Numeric value controling the size of nodes and edges.
#' @param label_cex Numeric value controling the size of cell type labels.
#' @param legend_cex Numeric value controling the size of the legend.
#' @param study_cols Named vector where values are RGB colors and names are
#' unique study identifiers. If NULL, a default color palette is used.
#'
#' @export
plotClusterGraph <- function(graph, study_id=NULL, cell_type=NULL,
                             size_factor=1, label_cex = 0.2*size_factor,
                             legend_cex = 2, study_cols=NULL) {
    vertex_colors <- study_cols
    if (is.null(vertex_colors)) {
        vertex_colors <- make_vertex_colors(graph)
    }
    graph <- color_graph(graph, vertex_colors, study_id, cell_type)
    plot(graph,
         vertex.size = igraph::V(graph)$size * size_factor,
         vertex.label.cex=label_cex,
         vertex.label.font=2,
         vertex.frame.color = NA,
         edge.width = igraph::E(graph)$width * size_factor,
         edge.arrow.size=.1*size_factor,
         edge.arrow.width=0.5*size_factor)
    graphics::legend(
        "topright", legend = names(vertex_colors), pt.bg = vertex_colors,
        pt.cex = legend_cex, cex = 0.5*legend_cex, bty="n", pch=21
    )
}

# Set node color (dataset) and size (# cells), edge color and size (AUROC)
color_graph <- function(graph, vertex_colors, study_id=NULL, cell_type=NULL) {
    vertex_study <- factor(getStudyId(igraph::V(graph)$name))
    levels(vertex_study) <- vertex_colors[levels(vertex_study)]
    
    full_label_list <- makeClusterName(study_id, cell_type)
    cluster_size <- table(full_label_list)
    cluster_size_class <- as.numeric(
        cut(cluster_size, breaks = c(0, 10, 100, Inf))
    )
    names(cluster_size_class) <- names(cluster_size)

    igraph::V(graph)$color <- as.character(vertex_study)
    igraph::V(graph)$label.color <- "black"
    igraph::V(graph)$label <- getCellType(igraph::V(graph)$name)
    node_size <- cluster_size_class[igraph::V(graph)$name]
    if (any(is.na(node_size))) {
        warning("Missing cluster size information for ", sum(is.na(node_size)),
                " cell types, setting node size to defaut (medium).")
    }
    node_size[is.na(node_size)] <- 2
    igraph::V(graph)$size <- node_size

    igraph::E(graph)$width <- igraph::E(graph)$weight
    igraph::E(graph)$color <- c("orange","darkgray")[as.numeric(igraph::E(graph)$weight >= 0.5) + 1]

    return(graph)
}

# Color palette for nodes
make_vertex_colors <- function(graph) {
    study_ids <- unique(getStudyId(igraph::V(graph)$name))
    result <- gg_color(length(study_ids))
    names(result) <- study_ids
    return(result)
}

#' Extend cluster set to nearest neighbors on cluster graph.
#'
#' Note that the graph is directed, i.e. neighbors are retrieved
#' by following arrows that start from the initial clusters.
#'
#' @param graph Graph in igraph format generated by makeClusterGraph.
#' @param initial_set Vector of cluster labels 
#' @param max_neighbor_distance Include more distantly related nodes
#' by performing neigbor extension max_neighbor_distance rounds.
#'
#' @return Character vector including initial cluster set and all
#' neighboring clusters (if any).
#'
#' @export
extendClusterSet <- function(graph, initial_set, max_neighbor_distance=2) {
    A <- as.matrix(igraph::as_adj(graph))
    A[A>0] <- 1
    diag(A) <- 1
    V <- as.numeric(igraph::V(graph)$name %in% initial_set)
    result <- V
    for (i in seq_len(max_neighbor_distance)) {
        result <- crossprod(A, result)
    }
    return(rownames(result)[as.logical(result)])
}

#' Subset cluster graph to clusters of interest.
#'
#' @param graph Graph in igraph format generated by makeClusterGraph.
#' @param vertices Vector of cluster labels 
#'
#' @return Graph in igraph format, where nodes have been restricted
#' to clusters of interests.
#'
#' @seealso \code{\link{extendClusterSet}}
#'
#' @export
subsetClusterGraph <- function(graph, vertices) {
    return(igraph::induced_subgraph(graph, vertices))
}