R/subset_plot_df.R

In baruuum/igplot: Fast(er) plotting for igraph objects

#' Extract subset from plot-ready list
#' 
#' @param x list of data.tables to subset (generated by igraph_plot_prep)
#' @param sub_var variable to subset on; either a logical vector or a character string
#' @param sub_val value to subset on
#' @param keep_edges what edges to keep
#' @details if \code{sub_var} is a logical vector, it has to be of the same length as the number of rows in \code{x$vertices}. If \code{sub_var} is a string, it has to be equal a column in \code{x$vertices}.
#' @export
subset_plot_df = function(
    x, 
    sub_var, 
    sub_val, 
    keep_edges = c("all", "within", "none")
) {
    
    if (class(x) != "igplotdat")
        stop("`x` has to be a `igplotdat` object")
    
    # set NULL values to avoid R CMD check notes
    name1 = name2 = NULL
    
    keep = match.arg(keep_edges, several.ok = FALSE)
    
    stopifnot(
        is.logical(sub_var) || is.character(sub_var),
        "name" %in% names(x$vertices),
        all(c("name1", "name2") %in% names(x$edges))
    )
    
    if (is.logical(sub_var))
        stopifnot(
            length(sub_var) == NROW(x$vertices)
        )
    
    if (!is.logical(sub_var))
        stopifnot(
            length(sub_var) == 1L,
            isTRUE(sub_var %in% names(x$vertices)),
            all(sub_val %in% unique(x$vertices[[sub_var]]))
        )
    
    # subset vertices
    if (is.logical(sub_var)) {
        
        x$vertices = x$vertices[sub_var, ]
        
    } else {
        
        x$vertices = x$vertices[get(sub_var) %in% sub_val]
        
    }
    
    # update edges
    if (keep == "none") {
        
        x$edges = x$edges[0L]
        
    } else if (keep == "all") {
        
        # get names
        v_names = x$vertices$name
        el = x$edges[
            (name1 %in% v_names) | (name2 %in% v_names)
        ]
        
        x$edges = el
        
    } else {
        
        # get names
        v_names = x$vertices$name
        el = x$edges[
            (name1 %in% v_names) & (name2 %in% v_names)
        ]
        x$edges = el
        
    }
    
    return(x)
    
}