From 6e765bb36e7a8d2d7120a3afd034bb005b48eda9 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Mon, 20 Nov 2023 17:36:50 -0600 Subject: [PATCH] nx-cugraph: add SSSP (unweighted) (#3976) There are many more traversal algorithms to implement, but these get us started! Authors: - Erik Welch (https://github.com/eriknw) - Brad Rees (https://github.com/BradReesWork) - Rick Ratzel (https://github.com/rlratzel) Approvers: - Brad Rees (https://github.com/BradReesWork) - Rick Ratzel (https://github.com/rlratzel) URL: https://github.com/rapidsai/cugraph/pull/3976 --- python/nx-cugraph/_nx_cugraph/__init__.py | 2 + .../nx_cugraph/algorithms/__init__.py | 3 +- .../algorithms/shortest_paths/__init__.py | 13 +++++ .../algorithms/shortest_paths/unweighted.py | 53 +++++++++++++++++++ python/nx-cugraph/nx_cugraph/classes/graph.py | 9 +++- python/nx-cugraph/nx_cugraph/interface.py | 2 + 6 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/__init__.py create mode 100644 python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py index 26638d1e735..910db1bc379 100644 --- a/python/nx-cugraph/_nx_cugraph/__init__.py +++ b/python/nx-cugraph/_nx_cugraph/__init__.py @@ -79,6 +79,8 @@ "path_graph", "petersen_graph", "sedgewick_maze_graph", + "single_source_shortest_path_length", + "single_target_shortest_path_length", "star_graph", "tadpole_graph", "tetrahedral_graph", diff --git a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py index 87b1967fa93..32cd6f31a47 100644 --- a/python/nx-cugraph/nx_cugraph/algorithms/__init__.py +++ b/python/nx-cugraph/nx_cugraph/algorithms/__init__.py @@ -10,9 +10,10 @@ # 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 . import bipartite, centrality, community, components +from . import bipartite, centrality, community, components, shortest_paths from .bipartite import complete_bipartite_graph from .centrality import * from .components import * from .core import * from .isolate import * +from .shortest_paths import * diff --git a/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/__init__.py b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/__init__.py new file mode 100644 index 00000000000..b7d6b742176 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/__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 .unweighted import * diff --git a/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py new file mode 100644 index 00000000000..3413a637b32 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/algorithms/shortest_paths/unweighted.py @@ -0,0 +1,53 @@ +# 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__ = ["single_source_shortest_path_length", "single_target_shortest_path_length"] + + +@networkx_algorithm +def single_source_shortest_path_length(G, source, cutoff=None): + return _single_shortest_path_length(G, source, cutoff, "Source") + + +@networkx_algorithm +def single_target_shortest_path_length(G, target, cutoff=None): + return _single_shortest_path_length(G, target, cutoff, "Target") + + +def _single_shortest_path_length(G, source, cutoff, kind): + G = _to_graph(G) + if source not in G: + raise nx.NodeNotFound(f"{kind} {source} is not in G") + if G.src_indices.size == 0: + return {source: 0} + if cutoff is None: + cutoff = -1 + 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=kind == "Target"), + sources=cp.array([src_index], index_dtype), + direction_optimizing=False, # True for undirected only; what's recommended? + depth_limit=cutoff, + compute_predecessors=False, + do_expensive_check=False, + ) + mask = distances != np.iinfo(distances.dtype).max + return G._nodearrays_to_dict(node_ids[mask], distances[mask]) diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py index 23004651fc5..fea318e036e 100644 --- a/python/nx-cugraph/nx_cugraph/classes/graph.py +++ b/python/nx-cugraph/nx_cugraph/classes/graph.py @@ -559,6 +559,7 @@ def _get_plc_graph( edge_dtype: Dtype | None = None, *, store_transposed: bool = False, + switch_indices: bool = False, edge_array: cp.ndarray[EdgeValue] | None = None, ): if edge_array is not None: @@ -613,14 +614,18 @@ def _get_plc_graph( elif edge_array.dtype not in self._plc_allowed_edge_types: raise TypeError(edge_array.dtype) # Should we cache PLC graph? + src_indices = self.src_indices + dst_indices = self.dst_indices + if switch_indices: + src_indices, dst_indices = dst_indices, src_indices return plc.SGGraph( resource_handle=plc.ResourceHandle(), graph_properties=plc.GraphProperties( is_multigraph=self.is_multigraph(), is_symmetric=not self.is_directed(), ), - src_or_offset_array=self.src_indices, - dst_or_index_array=self.dst_indices, + src_or_offset_array=src_indices, + dst_or_index_array=dst_indices, weight_array=edge_array, store_transposed=store_transposed, renumber=False, diff --git a/python/nx-cugraph/nx_cugraph/interface.py b/python/nx-cugraph/nx_cugraph/interface.py index 875f8621021..8903fdc541e 100644 --- a/python/nx-cugraph/nx_cugraph/interface.py +++ b/python/nx-cugraph/nx_cugraph/interface.py @@ -224,9 +224,11 @@ def key(testpath): ) too_slow = "Too slow to run" + maybe_oom = "out of memory in CI" skip = { key("test_tree_isomorphism.py:test_positive"): too_slow, key("test_tree_isomorphism.py:test_negative"): too_slow, + key("test_efficiency.py:TestEfficiency.test_using_ego_graph"): maybe_oom, } for item in items: