nx-cugraph: add ancestors and descendants (#4029)

Authors:
  - Erik Welch (https://github.com/eriknw)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Rick Ratzel (https://github.com/rlratzel)

URL: #4029

ci/test_python.sh

@@ -111,6 +111,11 @@ popd
 rapids-logger "pytest networkx using nx-cugraph backend"
 pushd python/nx-cugraph
 ./run_nx_tests.sh
+# Individually run tests that are skipped above b/c they may run out of memory
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestDAG and test_antichains"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestMultiDiGraph_DAGLCA and test_all_pairs_lca_pairs_without_lca"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestDAGLCA and test_all_pairs_lca_pairs_without_lca"
+PYTEST_NO_SKIP=True ./run_nx_tests.sh --cov-append -k "TestEfficiency and test_using_ego_graph"
 # run_nx_tests.sh outputs coverage data, so check that total coverage is >0.0%
 # in case nx-cugraph failed to load but fallback mode allowed the run to pass.
 _coverage=$(coverage report|grep "^TOTAL")

python/nx-cugraph/_nx_cugraph/__init__.py

@@ -29,8 +29,14 @@
     # "description": "TODO",
     "functions": {
         # BEGIN: functions
+        "ancestors",
         "barbell_graph",
         "betweenness_centrality",
+        "bfs_edges",
+        "bfs_layers",
+        "bfs_predecessors",
+        "bfs_successors",
+        "bfs_tree",
         "bull_graph",
         "caveman_graph",
         "chvatal_graph",
@@ -44,6 +50,8 @@
         "davis_southern_women_graph",
         "degree_centrality",
         "desargues_graph",
+        "descendants",
+        "descendants_at_distance",
         "diamond_graph",
         "dodecahedral_graph",
         "edge_betweenness_centrality",
@@ -53,6 +61,7 @@
         "from_pandas_edgelist",
         "from_scipy_sparse_array",
         "frucht_graph",
+        "generic_bfs_edges",
         "heawood_graph",
         "hits",
         "house_graph",
@@ -99,9 +108,14 @@
     "extra_docstrings": {
         # BEGIN: extra_docstrings
         "betweenness_centrality": "`weight` parameter is not yet supported.",
+        "bfs_edges": "`sort_neighbors` parameter is not yet supported.",
+        "bfs_predecessors": "`sort_neighbors` parameter is not yet supported.",
+        "bfs_successors": "`sort_neighbors` parameter is not yet supported.",
+        "bfs_tree": "`sort_neighbors` parameter is not yet supported.",
         "edge_betweenness_centrality": "`weight` parameter is not yet supported.",
         "eigenvector_centrality": "`nstart` parameter is not used, but it is checked for validity.",
         "from_pandas_edgelist": "cudf.DataFrame inputs also supported.",
+        "generic_bfs_edges": "`neighbors` and `sort_neighbors` parameters are not yet supported.",
         "k_truss": (
             "Currently raises `NotImplementedError` for graphs with more than one connected\n"
             "component when k >= 3. We expect to fix this soon."

python/nx-cugraph/nx_cugraph/algorithms/__init__.py

@@ -17,11 +17,14 @@
     components,
     link_analysis,
     shortest_paths,
+    traversal,
 )
 from .bipartite import complete_bipartite_graph
 from .centrality import *
 from .components import *
 from .core import *
+from .dag import *
 from .isolate import *
 from .link_analysis import *
 from .shortest_paths import *
+from .traversal import *

python/nx-cugraph/nx_cugraph/algorithms/core.py

@@ -15,7 +15,12 @@
 import pylibcugraph as plc
 
 import nx_cugraph as nxcg
-from nx_cugraph.utils import _get_int_dtype, networkx_algorithm, not_implemented_for
+from nx_cugraph.utils import (
+    _get_int_dtype,
+    index_dtype,
+    networkx_algorithm,
+    not_implemented_for,
+)
 
 __all__ = ["k_truss"]
 
@@ -81,10 +86,8 @@ def k_truss(G, k):
         edge_values = {key: val[edge_indices] for key, val in G.edge_values.items()}
         edge_masks = {key: val[edge_indices] for key, val in G.edge_masks.items()}
     # Renumber step 2: edge indices
-    mapper = cp.zeros(len(G), src_indices.dtype)
-    mapper[node_indices] = cp.arange(node_indices.size, dtype=mapper.dtype)
-    src_indices = mapper[src_indices]
-    dst_indices = mapper[dst_indices]
+    src_indices = cp.searchsorted(node_indices, src_indices).astype(index_dtype)
+    dst_indices = cp.searchsorted(node_indices, dst_indices).astype(index_dtype)
     # Renumber step 3: node values
     node_values = {key: val[node_indices] for key, val in G.node_values.items()}
     node_masks = {key: val[node_indices] for key, val in G.node_masks.items()}

python/nx-cugraph/nx_cugraph/algorithms/dag.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import cupy as cp
+import networkx as nx
+import numpy as np
+import pylibcugraph as plc
+
+from nx_cugraph.convert import _to_graph
+from nx_cugraph.utils import index_dtype, networkx_algorithm
+
+__all__ = [
+    "descendants",
+    "ancestors",
+]
+
+
+def _ancestors_and_descendants(G, source, *, is_ancestors):
+    G = _to_graph(G)
+    if source not in G:
+        hash(source)  # To raise TypeError if appropriate
+        raise nx.NetworkXError(
+            f"The node {source} is not in the {G.__class__.__name__.lower()}."
+        )
+    src_index = source if G.key_to_id is None else G.key_to_id[source]
+    distances, predecessors, node_ids = plc.bfs(
+        handle=plc.ResourceHandle(),
+        graph=G._get_plc_graph(switch_indices=is_ancestors),
+        sources=cp.array([src_index], dtype=index_dtype),
+        direction_optimizing=False,
+        depth_limit=-1,
+        compute_predecessors=False,
+        do_expensive_check=False,
+    )
+    mask = (distances != np.iinfo(distances.dtype).max) & (distances != 0)
+    return G._nodearray_to_set(node_ids[mask])
+
+
+@networkx_algorithm
+def descendants(G, source):
+    return _ancestors_and_descendants(G, source, is_ancestors=False)
+
+
+@networkx_algorithm
+def ancestors(G, source):
+    return _ancestors_and_descendants(G, source, is_ancestors=True)

python/nx-cugraph/nx_cugraph/algorithms/traversal/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .breadth_first_search import *