perf: Faster single bracket querying of a graph

R/indexing.R

@@ -53,6 +53,66 @@
 # - G[1:3,2,eid=TRUE]
 #               create an edge sequence
 
+get_adjacency_submatrix <- function(x, i = NULL, j = NULL, attr = NULL, sparse = TRUE) {
+  # If i or j is NULL, assume all nodes
+  # if not NULL make sure to handle duplicates correctly
+  if (is.null(i)) {
+    i <- i_unique <- i_map <- seq_len(vcount(x))
+  } else {
+    i_unique <- unique(i)
+    i_map <- match(i, i_unique)
+  }
+
+  if (is.null(j)) {
+    j <- j_unique <- j_map <- seq_len(vcount(x))
+  } else {
+    j_unique <- unique(j)
+    j_map <- match(j, j_unique)
+  }
+
+  adj <- adjacent_vertices(x, i_unique, mode = "out")
+  i_degree <- map_int(adj, length)
+
+  from_id <- rep(i_unique, i_degree)
+  to_id <- unlist(adj)
+
+  edge_list <- cbind(from_id, to_id)
+  edge_list <- edge_list[edge_list[, 2] %in% j_unique, , drop = FALSE]
+
+  row_indices <- edge_list[, 1]
+  col_indices <- edge_list[, 2]
+
+  values <- if (is.null(attr)) {
+    1
+  } else {
+    # get edge ids to locate edge attribute values
+    valid_edges <- get_edge_ids(x, c(t(edge_list)))
+    edge_attr(x, attr, valid_edges)
+  }
+
+
+  unique_res <- Matrix::sparseMatrix(
+    i = match(row_indices, i_unique),
+    j = match(col_indices, j_unique),
+    x = values,
+    dims = c(length(i_unique), length(j_unique))
+  )
+
+  # Expand to handle duplicated entries in i and j
+  res <- unique_res[i_map, j_map, drop = TRUE]
+
+  if (!sparse) {
+    res <- as.matrix(res)
+  }
+
+  if ("name" %in% vertex_attr_names(x) && !is.null(dim(res))) {
+    rownames(res) <- vertex_attr(x, "name", i)
+    colnames(res) <- vertex_attr(x, "name", j)
+  }
+
+  return(res)
+}
+
 
 #' Query and manipulate a graph as it were an adjacency matrix
 #'
@@ -152,11 +212,9 @@
 #' @method [ igraph
 #' @export
 `[.igraph` <- function(x, i, j, ..., from, to,
-                       sparse = igraph_opt("sparsematrices"),
-                       edges = FALSE, drop = TRUE,
-                       attr = if (is_weighted(x)) "weight" else NULL) {
-  ## TODO: make it faster, don't need the whole matrix usually
-
+    sparse = igraph_opt("sparsematrices"),
+    edges = FALSE, drop = TRUE,
+    attr = if (is_weighted(x)) "weight" else NULL) {
   ################################################################
   ## Argument checks
   if ((!missing(from) || !missing(to)) &&
@@ -192,31 +250,26 @@
     } else {
       res <- as.logical(res) + 0
     }
-    res
-  } else if (missing(i) && missing(j)) {
-    if (missing(edges)) {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr)
-    } else {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr, edges = edges)
-    }
-  } else if (missing(j)) {
-    if (missing(edges)) {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr)[i, , drop = drop]
-    } else {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr, edges = edges)[i, , drop = drop]
-    }
+    return(res)
+  }
+
+  # convert logical, character or negative i/j to proper vertex ids
+  if (!missing(i)) {
+    i <- as_igraph_vs(x, i)
+  }
+  if (!missing(j)) {
+    j <- as_igraph_vs(x, j)
+  }
+
+  if (missing(i) && missing(j)) {
+    return(as_adjacency_matrix(x, sparse = sparse, attr = attr))
+  }
+  if (missing(j)) {
+    get_adjacency_submatrix(x, i = i, j = NULL, attr = attr, sparse = sparse)
   } else if (missing(i)) {
-    if (missing(edges)) {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr)[, j, drop = drop]
-    } else {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr, edges = edges)[, j, drop = drop]
-    }
+    get_adjacency_submatrix(x, i = NULL, j = j, attr = attr, sparse = sparse)
   } else {
-    if (missing(edges)) {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr)[i, j, drop = drop]
-    } else {
-      as_adjacency_matrix(x, sparse = sparse, attr = attr, edges = edges)[i, j, drop = drop]
-    }
+    get_adjacency_submatrix(x, i = i, j = j, attr = attr, sparse = sparse)
   }
 }
 

tests/testthat/test-indexing.R

@@ -294,3 +294,49 @@ test_that("[[ handles from and to properly even if the graph has conflicting ver
   expect_true(is_igraph_vs(g[[1:3, 2:6]][[1]]))
   expect_true(is_igraph_vs(g[[from = 1:3, to = 2:6]][[1]]))
 })
+
+test_that("[ handles errors in input parameters well", {
+  g <- make_full_graph(10)
+  expect_error(g[from = 1, to = 1, i = 1, j = 1])
+  expect_error(g[from = 1])
+  expect_error(g[to = 1])
+  expect_error(g[from = NA, to = 2])
+  expect_error(g[from = 1, to = NA])
+  expect_error(g[from = 1, to = c(1, 2)])
+})
+
+test_that("[ handles all combinations of i and/or j", {
+  A <- matrix(
+    rep(
+      c(0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0),
+      c(
+        4L, 1L, 4L, 1L, 2L, 1L, 5L, 2L, 3L, 1L, 10L, 3L, 9L, 1L, 1L, 1L, 3L, 1L, 1L,
+        1L, 1L, 1L, 10L, 1L, 1L, 1L, 1L, 5L, 11L, 1L, 2L, 1L, 5L, 1L, 3L
+      )
+    ),
+    nrow = 10L,
+    ncol = 10L
+  )
+  g <- graph_from_adjacency_matrix(A, "directed")
+  expect_equal(canonicalize_matrix(g[1:3, ]), A[1:3, ])
+  expect_equal(canonicalize_matrix(g[, 4:7]), A[, 4:7])
+  expect_equal(canonicalize_matrix(g[1:3, 4:7]), A[1:3, 4:7])
+})
+
+test_that("[ handles duplicated i/j well", {
+  A <- matrix(
+    rep(
+      c(0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0),
+      c(
+        4L, 1L, 4L, 1L, 2L, 1L, 5L, 2L, 3L, 1L, 10L, 3L, 9L, 1L, 1L, 1L, 3L, 1L, 1L,
+        1L, 1L, 1L, 10L, 1L, 1L, 1L, 1L, 5L, 11L, 1L, 2L, 1L, 5L, 1L, 3L
+      )
+    ),
+    nrow = 10L,
+    ncol = 10L
+  )
+  g <- graph_from_adjacency_matrix(A, "directed")
+  expect_equal(canonicalize_matrix(g[c(1, 2, 2), ]), A[c(1, 2, 2), ])
+  expect_equal(canonicalize_matrix(g[, c(3, 3, 4, 4)]), A[, c(3, 3, 4, 4)])
+  expect_equal(canonicalize_matrix(g[c(1, 2, 2), c(3, 3, 4, 4)]), A[c(1, 2, 2), c(3, 3, 4, 4)])
+})