From 22e317cfadf17b078daedc3c1eac714681aa7c5d Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 30 Nov 2023 16:11:43 -0600 Subject: [PATCH] nx-cugraph: add `ancestors` and `descendants` --- python/nx-cugraph/_nx_cugraph/__init__.py | 2 + .../nx_cugraph/algorithms/__init__.py | 5 +- .../nx-cugraph/nx_cugraph/algorithms/core.py | 6 +- .../nx-cugraph/nx_cugraph/algorithms/dag.py | 61 +++++++++++++++++++ python/nx-cugraph/nx_cugraph/classes/graph.py | 16 +++++ 5 files changed, 87 insertions(+), 3 deletions(-) create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/dag.py diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py index 1fd436bb845..b0ab9881660 100644 --- a/python/nx-cugraph/_nx_cugraph/__init__.py +++ b/python/nx-cugraph/_nx_cugraph/__init__.py @@ -29,6 +29,7 @@ # "description": "TODO", "functions": { # BEGIN: functions + "ancestors", "barbell_graph", "betweenness_centrality", "bull_graph", @@ -44,6 +45,7 @@ "davis_southern_women_graph", "degree_centrality", "desargues_graph", + "descendants", "diamond_graph", "dodecahedral_graph", "edge_betweenness_centrality", diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py index 63841b15bd5..9f95354793a 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py @@ -15,13 +15,14 @@ centrality, community, components, - shortest_paths, link_analysis, + shortest_paths, ) from .bipartite import complete_bipartite_graph from .centrality import * from .components import * from .core import * +from .dag import * from .isolate import * -from .shortest_paths import * from .link_analysis import * +from .shortest_paths import * diff --git a/python/nx-cugraph/nx_cugraph/algorithms/core.py b/python/nx-cugraph/nx_cugraph/algorithms/core.py index 2219388bc58..390598d070e 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/core.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/core.py @@ -31,7 +31,11 @@ def k_truss(G, k): if is_nx := isinstance(G, nx.Graph): G = nxcg.from_networkx(G, preserve_all_attrs=True) if nxcg.number_of_selfloops(G) > 0: - raise nx.NetworkXError( + if nx.__version__[:3] <= "3.2": + exc_class = nx.NetworkXError + else: + exc_class = nx.NetworkXNotImplemented + raise exc_class( "Input graph has self loops which is not permitted; " "Consider using G.remove_edges_from(nx.selfloop_edges(G))." ) diff --git a/python/nx-cugraph/nx_cugraph/algorithms/dag.py b/python/nx-cugraph/nx_cugraph/algorithms/dag.py new file mode 100644 index 00000000000..2014d4c3d27 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/algorithms/dag.py @@ -0,0 +1,61 @@ +# 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 graph.") + src_index = source if G.key_to_id is None else G.key_to_id[source] + distances, predecessors, node_ids = plc.bfs( + # XXX: why can't I pass arguments as keywords?! + plc.ResourceHandle(), + G._get_plc_graph(switch_indices=is_ancestors), + cp.array([src_index], dtype=index_dtype), + False, + -1, + False, + False, + # resource_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) diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py index e32f93d8bfe..8af7b585190 100644 --- a/python/nx-cugraph/nx_cugraph/classes/graph.py +++ b/python/nx-cugraph/nx_cugraph/classes/graph.py @@ -458,6 +458,22 @@ def has_edge(self, u: NodeKey, v: NodeKey) -> bool: return False return bool(((self.src_indices == u) & (self.dst_indices == v)).any()) + def _neighbors(self, n: NodeKey) -> cp.ndarray[NodeValue]: + if n not in self: + hash(n) # To raise TypeError if appropriate + raise nx.NetworkXError(f"The node {n} is not in the graph.") + if self.key_to_id is not None: + n = self.key_to_id[n] + nbrs = self.dst_indices[self.src_indices == n] + if self.is_multigraph(): + nbrs = cp.unique(nbrs) + return nbrs + + @networkx_api + def neighbors(self, n: NodeKey) -> Iterator[NodeKey]: + nbrs = self._neighbors(n) + return iter(self._nodeiter_to_iter(nbrs.tolist())) + @networkx_api def has_node(self, n: NodeKey) -> bool: return n in self