From 48dcf73a245946d7f1d9f827c0684c69ffb8d8cc Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Thu, 12 Oct 2023 22:01:06 -0500 Subject: [PATCH 1/3] Add multigraph support to nx-cugraph --- python/nx-cugraph/lint.yaml | 14 +- .../nx-cugraph/nx_cugraph/classes/__init__.py | 2 + python/nx-cugraph/nx_cugraph/classes/graph.py | 22 +- .../nx_cugraph/classes/multidigraph.py | 30 +++ .../nx_cugraph/classes/multigraph.py | 251 ++++++++++++++++++ python/nx-cugraph/nx_cugraph/convert.py | 167 ++++++++---- python/nx-cugraph/nx_cugraph/interface.py | 12 +- .../nx_cugraph/tests/bench_convert.py | 16 +- .../nx_cugraph/tests/test_convert.py | 39 ++- 9 files changed, 468 insertions(+), 85 deletions(-) create mode 100644 python/nx-cugraph/nx_cugraph/classes/multidigraph.py create mode 100644 python/nx-cugraph/nx_cugraph/classes/multigraph.py diff --git a/python/nx-cugraph/lint.yaml b/python/nx-cugraph/lint.yaml index 6a462a6af79..d6ed0c501d3 100644 --- a/python/nx-cugraph/lint.yaml +++ b/python/nx-cugraph/lint.yaml @@ -11,7 +11,7 @@ default_language_version: python: python3 repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 + rev: v4.5.0 hooks: - id: check-added-large-files - id: check-case-conflict @@ -26,7 +26,7 @@ repos: - id: mixed-line-ending - id: trailing-whitespace - repo: https://github.com/abravalheri/validate-pyproject - rev: v0.14 + rev: v0.15 hooks: - id: validate-pyproject name: Validate pyproject.toml @@ -40,7 +40,7 @@ repos: hooks: - id: isort - repo: https://github.com/asottile/pyupgrade - rev: v3.13.0 + rev: v3.15.0 hooks: - id: pyupgrade args: [--py39-plus] @@ -50,7 +50,7 @@ repos: - id: black # - id: black-jupyter - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.291 + rev: v0.0.292 hooks: - id: ruff args: [--fix-only, --show-fixes] @@ -70,18 +70,18 @@ repos: - id: yesqa additional_dependencies: *flake8_dependencies - repo: https://github.com/codespell-project/codespell - rev: v2.2.5 + rev: v2.2.6 hooks: - id: codespell types_or: [python, rst, markdown] additional_dependencies: [tomli] files: ^(nx_cugraph|docs)/ - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.0.291 + rev: v0.0.292 hooks: - id: ruff - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 + rev: v4.5.0 hooks: - id: no-commit-to-branch args: [-p, "^branch-2....$"] diff --git a/python/nx-cugraph/nx_cugraph/classes/__init__.py b/python/nx-cugraph/nx_cugraph/classes/__init__.py index e47641ae812..9916bcbe241 100644 --- a/python/nx-cugraph/nx_cugraph/classes/__init__.py +++ b/python/nx-cugraph/nx_cugraph/classes/__init__.py @@ -11,5 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from .graph import Graph +from .multigraph import MultiGraph from .digraph import DiGraph # isort:skip +from .multidigraph import MultiDiGraph # isort:skip diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py index 1432f68c752..40581fd1906 100644 --- a/python/nx-cugraph/nx_cugraph/classes/graph.py +++ b/python/nx-cugraph/nx_cugraph/classes/graph.py @@ -43,7 +43,8 @@ class Graph: # Tell networkx to dispatch calls with this object to nx-cugraph - __networkx_plugin__: ClassVar[str] = "cugraph" + __networkx_backend__: ClassVar[str] = "cugraph" # nx >=3.2 + __networkx_plugin__: ClassVar[str] = "cugraph" # nx <3.2 # networkx properties graph: dict @@ -58,7 +59,7 @@ class Graph: node_values: dict[AttrKey, cp.ndarray[NodeValue]] node_masks: dict[AttrKey, cp.ndarray[bool]] key_to_id: dict[NodeKey, IndexValue] | None - _id_to_key: dict[IndexValue, NodeKey] | None + _id_to_key: list[NodeKey] | None _N: int #################### @@ -77,7 +78,7 @@ def from_coo( node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, *, key_to_id: dict[NodeKey, IndexValue] | None = None, - id_to_key: dict[IndexValue, NodeKey] | None = None, + id_to_key: list[NodeKey] | None = None, **attr, ) -> Graph: new_graph = object.__new__(cls) @@ -88,7 +89,7 @@ def from_coo( new_graph.node_values = {} if node_values is None else dict(node_values) new_graph.node_masks = {} if node_masks is None else dict(node_masks) new_graph.key_to_id = None if key_to_id is None else dict(key_to_id) - new_graph._id_to_key = None if id_to_key is None else dict(id_to_key) + new_graph._id_to_key = None if id_to_key is None else list(id_to_key) new_graph._N = op.index(N) # Ensure N is integral new_graph.graph = new_graph.graph_attr_dict_factory() new_graph.graph.update(attr) @@ -123,7 +124,7 @@ def from_csr( node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, *, key_to_id: dict[NodeKey, IndexValue] | None = None, - id_to_key: dict[IndexValue, NodeKey] | None = None, + id_to_key: list[NodeKey] | None = None, **attr, ) -> Graph: N = indptr.size - 1 @@ -155,7 +156,7 @@ def from_csc( node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, *, key_to_id: dict[NodeKey, IndexValue] | None = None, - id_to_key: dict[IndexValue, NodeKey] | None = None, + id_to_key: list[NodeKey] | None = None, **attr, ) -> Graph: N = indptr.size - 1 @@ -189,7 +190,7 @@ def from_dcsr( node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, *, key_to_id: dict[NodeKey, IndexValue] | None = None, - id_to_key: dict[IndexValue, NodeKey] | None = None, + id_to_key: list[NodeKey] | None = None, **attr, ) -> Graph: row_indices = cp.array( @@ -222,7 +223,7 @@ def from_dcsc( node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, *, key_to_id: dict[NodeKey, IndexValue] | None = None, - id_to_key: dict[IndexValue, NodeKey] | None = None, + id_to_key: list[NodeKey] | None = None, **attr, ) -> Graph: col_indices = cp.array( @@ -295,11 +296,11 @@ def node_dtypes(self) -> dict[AttrKey, Dtype]: return {key: val.dtype for key, val in self.node_values.items()} @property - def id_to_key(self) -> dict[IndexValue, NodeKey] | None: + def id_to_key(self) -> [NodeKey] | None: if self.key_to_id is None: return None if self._id_to_key is None: - self._id_to_key = {val: key for key, val in self.key_to_id.items()} + self._id_to_key = sorted(self.key_to_id, key=self.key_to_id.__getitem__) return self._id_to_key name = nx.Graph.name @@ -447,6 +448,7 @@ def to_undirected(self, as_view: bool = False) -> Graph: ################### def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False): + # DRY warning: see also MultiGraph._copy indptr = self.indptr row_indices = self.row_indices col_indices = self.col_indices diff --git a/python/nx-cugraph/nx_cugraph/classes/multidigraph.py b/python/nx-cugraph/nx_cugraph/classes/multidigraph.py new file mode 100644 index 00000000000..5629e2c9c06 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/classes/multidigraph.py @@ -0,0 +1,30 @@ +# 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 __future__ import annotations + +import networkx as nx + +import nx_cugraph as nxcg + +from .digraph import DiGraph +from .multigraph import MultiGraph + +__all__ = ["MultiDiGraph"] + +networkx_api = nxcg.utils.decorators.networkx_class(nx.MultiDiGraph) + + +class MultiDiGraph(MultiGraph, DiGraph): + @classmethod + def to_networkx_class(cls) -> type[nx.MultiDiGraph]: + return nx.MultiDiGraph diff --git a/python/nx-cugraph/nx_cugraph/classes/multigraph.py b/python/nx-cugraph/nx_cugraph/classes/multigraph.py new file mode 100644 index 00000000000..aad1c653757 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/classes/multigraph.py @@ -0,0 +1,251 @@ +# 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 __future__ import annotations + +from copy import deepcopy +from typing import TYPE_CHECKING, ClassVar + +import cupy as cp +import networkx as nx +import numpy as np + +import nx_cugraph as nxcg + +from .graph import Graph + +if TYPE_CHECKING: + from nx_cugraph.typing import ( + AttrKey, + EdgeKey, + EdgeValue, + IndexValue, + NodeKey, + NodeValue, + ) +__all__ = ["MultiGraph"] + +networkx_api = nxcg.utils.decorators.networkx_class(nx.MultiGraph) + + +class MultiGraph(Graph): + # networkx properties + edge_key_dict_factory: ClassVar[type] = dict + + # Not networkx properties + + # In a MultiGraph, each edge has a unique `(row, col, key)` key. + # By default, `key` is 0 if possible, else 1, else 2, etc. + # This key can be any hashable Python object in NetworkX. + # We don't use a dict for our data structure here, because + # that would require a `(row, col, key)` key. + # Instead, we keep `edge_keys` and/or `edge_indices`. + # `edge_keys` is the list of Python objects for each edge. + # `edge_indices` is for the common case of default multiedge keys, + # in which case we can store it as a cupy array. + # `edge_indices` is generally preferred. It is possible to provide + # both where edge_indices is the default and edge_keys is anything. + # It is also possible for them both to be None, which means the + # default edge indices has not yet been calculated. + edge_indices: cp.ndarray[IndexValue] | None + edge_keys: list[EdgeKey] | None + + #################### + # Creation methods # + #################### + + @classmethod + def from_coo( + cls, + N: int, + row_indices: cp.ndarray[IndexValue], + col_indices: cp.ndarray[IndexValue], + edge_indices: cp.ndarray[IndexValue] | None = None, + edge_values: dict[AttrKey, cp.ndarray[EdgeValue]] | None = None, + edge_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + node_values: dict[AttrKey, cp.ndarray[NodeValue]] | None = None, + node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + *, + key_to_id: dict[NodeKey, IndexValue] | None = None, + id_to_key: list[NodeKey] | None = None, + edge_keys: list[EdgeKey] | None = None, + **attr, + ) -> MultiGraph: + new_graph = super().from_coo( + N, + row_indices, + col_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + **attr, + ) + new_graph.edge_indices = edge_indices + new_graph.edge_keys = edge_keys + # Easy and fast sanity checks + if ( + new_graph.edge_keys is not None + and len(new_graph.edge_keys) != row_indices.size + ): + raise ValueError + return new_graph + + # TODO: + # from_csr + # from_csc + # from_dcsr + # from_dcsc + + def __new__(cls, incoming_graph_data=None, multigraph_input=None, **attr) -> Graph: + # TODO: handle multigraph_input + if incoming_graph_data is None: + new_graph = cls.from_coo(0, cp.empty(0, np.int32), cp.empty(0, np.int32)) + elif incoming_graph_data.__class__ is new_graph.__class__: + new_graph = incoming_graph_data.copy() + elif incoming_graph_data.__class__ is new_graph.to_networkx_class(): + new_graph = nxcg.from_networkx(incoming_graph_data, preserve_all_attrs=True) + else: + raise NotImplementedError + new_graph.graph.update(attr) + return new_graph + + ################# + # Class methods # + ################# + + @classmethod + @networkx_api + def is_directed(cls) -> bool: + return False + + @classmethod + @networkx_api + def is_multigraph(cls) -> bool: + return True + + @classmethod + @networkx_api + def to_directed_class(cls) -> type[nxcg.MultiDiGraph]: + return nxcg.MultiDiGraph + + @classmethod + def to_networkx_class(cls) -> type[nx.MultiGraph]: + return nx.MultiGraph + + @classmethod + @networkx_api + def to_undirected_class(cls) -> type[MultiGraph]: + return MultiGraph + + ########################## + # NetworkX graph methods # + ########################## + + @networkx_api + def clear(self) -> None: + super().clear() + self.edge_indices = None + self.edge_keys = None + + @networkx_api + def clear_edges(self) -> None: + super().clear_edges() + self.edge_indices = None + self.edge_keys = None + + # TODO: + # copy + # get_edge_data + + @networkx_api + def has_edge(self, u: NodeKey, v: NodeKey, key: EdgeKey | None = None) -> bool: + if self.key_to_id is not None: + try: + u = self.key_to_id[u] + v = self.key_to_id[v] + except KeyError: + return False + mask = (self.row_indices == u) & (self.col_indices == v) + if key is None or (self.edge_indices is None and self.edge_keys is None): + return bool(mask.any()) + if self.edge_keys is None: + return bool((mask & (self.edge_indices == key)).any()) + indices = cp.nonzero(mask)[0] + if indices.size == 0: + return False + edge_keys = self.edge_keys + return any(edge_keys[i] == key for i in indices.tolist()) + + @networkx_api + def to_directed(self, as_view: bool = False) -> nxcg.MultiDiGraph: + return self._copy(as_view, self.to_directed_class()) + + @networkx_api + def to_undirected(self, as_view: bool = False) -> MultiGraph: + # Does deep copy in networkx + return self.copy(as_view) + + ################### + # Private methods # + ################### + + def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False): + # DRY warning: see also Graph._copy + indptr = self.indptr + row_indices = self.row_indices + col_indices = self.col_indices + edge_indices = self.edge_indices + edge_values = self.edge_values + edge_masks = self.edge_masks + node_values = self.node_values + node_masks = self.node_masks + key_to_id = self.key_to_id + id_to_key = None if key_to_id is None else self._id_to_key + edge_keys = self.edge_keys + if not as_view: + indptr = indptr.copy() + row_indices = row_indices.copy() + col_indices = col_indices.copy() + edge_indices = edge_indices.copy() + edge_values = {key: val.copy() for key, val in edge_values.items()} + edge_masks = {key: val.copy() for key, val in edge_masks.items()} + node_values = {key: val.copy() for key, val in node_values.items()} + node_masks = {key: val.copy() for key, val in node_masks.items()} + if key_to_id is not None: + key_to_id = key_to_id.copy() + if id_to_key is not None: + id_to_key = id_to_key.copy() + if edge_keys is not None: + edge_keys = edge_keys.copy() + if reverse: + row_indices, col_indices = col_indices, row_indices + rv = cls.from_coo( + indptr, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + edge_keys=edge_keys, + ) + if as_view: + rv.graph = self.graph + else: + rv.graph.update(deepcopy(self.graph)) + return rv diff --git a/python/nx-cugraph/nx_cugraph/convert.py b/python/nx-cugraph/nx_cugraph/convert.py index 9be8cac7877..4f822564f46 100644 --- a/python/nx-cugraph/nx_cugraph/convert.py +++ b/python/nx-cugraph/nx_cugraph/convert.py @@ -122,8 +122,6 @@ def from_networkx( graph = G else: raise TypeError(f"Expected networkx.Graph; got {type(graph)}") - elif graph.is_multigraph(): - raise NotImplementedError("MultiGraph support is not yet implemented") if preserve_all_attrs: preserve_edge_attrs = True @@ -144,7 +142,7 @@ def from_networkx( else: node_attrs = {node_attrs: None} - if graph.__class__ in {nx.Graph, nx.DiGraph}: + if graph.__class__ in {nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph}: # This is a NetworkX private attribute, but is much faster to use adj = graph._adj else: @@ -163,7 +161,11 @@ def from_networkx( edge_attrs = None elif preserve_edge_attrs: # Using comprehensions should be just as fast starting in Python 3.11 - attr_sets = set(map(frozenset, concat(map(dict.values, adj.values())))) + it = concat(map(dict.values, adj.values())) + if graph.is_multigraph(): + it = concat(map(dict.values, it)) + # PERF: should we add `filter(None, ...)` to remove empty data dicts? + attr_sets = set(map(frozenset, it)) attrs = frozenset.union(*attr_sets) edge_attrs = dict.fromkeys(attrs, REQUIRED) if len(attr_sets) > 1: @@ -180,11 +182,19 @@ def from_networkx( if len(required) == 1: # Fast path for the common case of a single attribute with no default [attr] = required - it = ( - attr in edgedata - for rowdata in adj.values() - for edgedata in rowdata.values() - ) + if graph.is_multigraph(): + it = ( + attr in edgedata + for rowdata in adj.values() + for multiedges in rowdata.values() + for edgedata in multiedges.values() + ) + else: + it = ( + attr in edgedata + for rowdata in adj.values() + for edgedata in rowdata.values() + ) if next(it): if all(it): # All edges have data @@ -195,9 +205,10 @@ def from_networkx( del edge_attrs[attr] # Else some edges have attribute (default already None) else: - attr_sets = set( - map(required.intersection, concat(map(dict.values, adj.values()))) - ) + it = concat(map(dict.values, adj.values())) + if graph.is_multigraph(): + it = concat(map(dict.values, it)) + attr_sets = set(map(required.intersection, it)) for attr in required - frozenset.union(*attr_sets): # No edges have these attributes del edge_attrs[attr] @@ -254,7 +265,23 @@ def from_networkx( key_to_id = None else: col_iter = map(key_to_id.__getitem__, col_iter) - col_indices = cp.fromiter(col_iter, np.int32) + if graph.is_multigraph(): + col_indices = np.fromiter(col_iter, np.int32) + num_multiedges = np.fromiter( + map(len, concat(map(dict.values, adj.values()))), np.int32 + ) + # cp.repeat is slow to use here, so use numpy instead + col_indices = cp.array(np.repeat(col_indices, num_multiedges)) + # Determine edge keys and edge ids for multigraphs + edge_keys = list(concat(concat(map(dict.values, adj.values())))) + edge_indices = np.fromiter( + concat(map(range, map(len, concat(map(dict.values, adj.values()))))), + np.int32, + ) + if edge_keys == edge_indices.tolist(): + edge_keys = None # Prefer edge_indices + else: + col_indices = cp.fromiter(col_iter, np.int32) edge_values = {} edge_masks = {} @@ -268,16 +295,29 @@ def from_networkx( if edge_default is None: vals = [] append = vals.append - iter_mask = ( - append( - edgedata[edge_attr] - if (present := edge_attr in edgedata) - else False + if graph.is_multigraph(): + iter_mask = ( + append( + edgedata[edge_attr] + if (present := edge_attr in edgedata) + else False + ) + or present + for rowdata in adj.values() + for multiedges in rowdata.values() + for edgedata in multiedges.values() + ) + else: + iter_mask = ( + append( + edgedata[edge_attr] + if (present := edge_attr in edgedata) + else False + ) + or present + for rowdata in adj.values() + for edgedata in rowdata.values() ) - or present - for rowdata in adj.values() - for edgedata in rowdata.values() - ) edge_masks[edge_attr] = cp.fromiter(iter_mask, bool) edge_values[edge_attr] = cp.array(vals, dtype) # if vals.ndim > 1: ... @@ -289,8 +329,16 @@ def from_networkx( # for rowdata in adj.values() # for edgedata in rowdata.values() # ) - iter_values = map( - op.itemgetter(edge_attr), concat(map(dict.values, adj.values())) + it = concat(map(dict.values, adj.values())) + if graph.is_multigraph(): + it = concat(map(dict.values, it)) + iter_values = map(op.itemgetter(edge_attr), it) + elif graph.is_multigraph(): + iter_values = ( + edgedata.get(edge_attr, edge_default) + for rowdata in adj.values() + for multiedges in rowdata.values() + for edgedata in multiedges.values() ) else: iter_values = ( @@ -304,13 +352,13 @@ def from_networkx( edge_values[edge_attr] = cp.fromiter(iter_values, dtype) # if vals.ndim > 1: ... - row_indices = cp.array( - # cp.repeat is slow to use here, so use numpy instead - np.repeat( - np.arange(N, dtype=np.int32), - np.fromiter(map(len, adj.values()), np.int32), - ) + # cp.repeat is slow to use here, so use numpy instead + row_indices = np.repeat( + np.arange(N, dtype=np.int32), np.fromiter(map(len, adj.values()), np.int32) ) + if graph.is_multigraph(): + row_indices = np.repeat(row_indices, num_multiedges) + row_indices = cp.array(row_indices) node_values = {} node_masks = {} @@ -350,21 +398,38 @@ def from_networkx( else: node_values[node_attr] = cp.fromiter(iter_values, dtype) # if vals.ndim > 1: ... - - if graph.is_directed() or as_directed: - klass = nxcg.DiGraph + if graph.is_multigraph(): + if graph.is_directed() or as_directed: + klass = nxcg.MultiDiGraph + else: + klass = nxcg.MultiGraph + rv = klass.from_coo( + N, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + edge_keys=edge_keys, + ) else: - klass = nxcg.Graph - rv = klass.from_coo( - N, - row_indices, - col_indices, - edge_values, - edge_masks, - node_values, - node_masks, - key_to_id=key_to_id, - ) + if graph.is_directed() or as_directed: + klass = nxcg.DiGraph + else: + klass = nxcg.Graph + rv = klass.from_coo( + N, + row_indices, + col_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + ) if preserve_graph_attrs: rv.graph.update(graph.graph) # deepcopy? return rv @@ -427,7 +492,7 @@ def to_networkx(G: nxcg.Graph) -> nx.Graph: full_node_dicts = _iter_attr_dicts(node_values, node_masks) rv.add_nodes_from(zip(node_iter, full_node_dicts)) elif id_to_key is not None: - rv.add_nodes_from(id_to_key.values()) + rv.add_nodes_from(id_to_key) else: rv.add_nodes_from(range(len(G))) @@ -448,7 +513,17 @@ def to_networkx(G: nxcg.Graph) -> nx.Graph: if id_to_key is not None: row_iter = map(id_to_key.__getitem__, row_indices) col_iter = map(id_to_key.__getitem__, col_indices) - if edge_values: + if G.is_multigraph() and (G.edge_keys is not None or G.edge_indices is not None): + if G.edge_keys is not None: + edge_keys = G.edge_keys + else: + edge_keys = G.edge_indices.tolist() + if edge_values: + full_edge_dicts = _iter_attr_dicts(edge_values, edge_masks) + rv.add_edges_from(zip(row_iter, col_iter, edge_keys, full_edge_dicts)) + else: + rv.add_edges_from(zip(row_iter, col_iter, edge_keys)) + elif edge_values: full_edge_dicts = _iter_attr_dicts(edge_values, edge_masks) rv.add_edges_from(zip(row_iter, col_iter, full_edge_dicts)) else: diff --git a/python/nx-cugraph/nx_cugraph/interface.py b/python/nx-cugraph/nx_cugraph/interface.py index 2ad23acd940..124b86a4af2 100644 --- a/python/nx-cugraph/nx_cugraph/interface.py +++ b/python/nx-cugraph/nx_cugraph/interface.py @@ -174,6 +174,10 @@ def key(testpath): ): louvain_different, key("test_louvain.py:test_none_weight_param"): louvain_different, key("test_louvain.py:test_multigraph"): louvain_different, + # See networkx#6630 + key( + "test_louvain.py:test_undirected_selfloops" + ): "self-loops not handled in Louvain", } ) @@ -189,10 +193,4 @@ def can_run(cls, name, args, kwargs): This is a proposed API to add to networkx dispatching machinery and may change. """ - return ( - hasattr(cls, name) - and getattr(cls, name).can_run(*args, **kwargs) - # We don't support MultiGraphs yet - and not any(isinstance(x, nx.MultiGraph) for x in args) - and not any(isinstance(x, nx.MultiGraph) for x in kwargs.values()) - ) + return hasattr(cls, name) and getattr(cls, name).can_run(*args, **kwargs) diff --git a/python/nx-cugraph/nx_cugraph/tests/bench_convert.py b/python/nx-cugraph/nx_cugraph/tests/bench_convert.py index 7e6278661c2..2eb432230eb 100644 --- a/python/nx-cugraph/nx_cugraph/tests/bench_convert.py +++ b/python/nx-cugraph/nx_cugraph/tests/bench_convert.py @@ -39,6 +39,8 @@ gpubenchmark = pytest_benchmark.plugin.benchmark +CREATE_USING = [nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph] + def _bench_helper(gpubenchmark, N, attr_kind, create_using, method): G = method(N, create_using=create_using) @@ -103,7 +105,7 @@ def _bench_helper_scipy(gpubenchmark, N, attr_kind, create_using, method, fmt): None, ], ) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) def bench_cycle_graph(gpubenchmark, N, attr_kind, create_using): _bench_helper(gpubenchmark, N, attr_kind, create_using, nx.cycle_graph) @@ -111,7 +113,7 @@ def bench_cycle_graph(gpubenchmark, N, attr_kind, create_using): @pytest.mark.skipif("not cugraph") @pytest.mark.parametrize("N", [1, 10**6]) @pytest.mark.parametrize("attr_kind", ["full", None]) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) @pytest.mark.parametrize("do_renumber", [True, False]) def bench_cycle_graph_cugraph(gpubenchmark, N, attr_kind, create_using, do_renumber): if N == 1 and not do_renumber: @@ -124,7 +126,7 @@ def bench_cycle_graph_cugraph(gpubenchmark, N, attr_kind, create_using, do_renum @pytest.mark.skipif("not scipy") @pytest.mark.parametrize("N", [1, 10**6]) @pytest.mark.parametrize("attr_kind", ["full", None]) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) @pytest.mark.parametrize("fmt", ["coo", "csr"]) def bench_cycle_graph_scipy(gpubenchmark, N, attr_kind, create_using, fmt): _bench_helper_scipy(gpubenchmark, N, attr_kind, create_using, nx.cycle_graph, fmt) @@ -143,14 +145,14 @@ def bench_cycle_graph_scipy(gpubenchmark, N, attr_kind, create_using, fmt): None, ], ) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) def bench_complete_graph_edgedata(gpubenchmark, N, attr_kind, create_using): _bench_helper(gpubenchmark, N, attr_kind, create_using, nx.complete_graph) @pytest.mark.parametrize("N", [3000]) @pytest.mark.parametrize("attr_kind", [None]) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) def bench_complete_graph_noedgedata(gpubenchmark, N, attr_kind, create_using): _bench_helper(gpubenchmark, N, attr_kind, create_using, nx.complete_graph) @@ -158,7 +160,7 @@ def bench_complete_graph_noedgedata(gpubenchmark, N, attr_kind, create_using): @pytest.mark.skipif("not cugraph") @pytest.mark.parametrize("N", [1, 1500]) @pytest.mark.parametrize("attr_kind", ["full", None]) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) @pytest.mark.parametrize("do_renumber", [True, False]) def bench_complete_graph_cugraph(gpubenchmark, N, attr_kind, create_using, do_renumber): if N == 1 and not do_renumber: @@ -171,7 +173,7 @@ def bench_complete_graph_cugraph(gpubenchmark, N, attr_kind, create_using, do_re @pytest.mark.skipif("not scipy") @pytest.mark.parametrize("N", [1, 1500]) @pytest.mark.parametrize("attr_kind", ["full", None]) -@pytest.mark.parametrize("create_using", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize("create_using", CREATE_USING) @pytest.mark.parametrize("fmt", ["coo", "csr"]) def bench_complete_graph_scipy(gpubenchmark, N, attr_kind, create_using, fmt): _bench_helper_scipy( diff --git a/python/nx-cugraph/nx_cugraph/tests/test_convert.py b/python/nx-cugraph/nx_cugraph/tests/test_convert.py index ba3cd7aaee1..83820f7621f 100644 --- a/python/nx-cugraph/nx_cugraph/tests/test_convert.py +++ b/python/nx-cugraph/nx_cugraph/tests/test_convert.py @@ -18,7 +18,9 @@ from nx_cugraph import interface -@pytest.mark.parametrize("graph_class", [nx.Graph, nx.DiGraph]) +@pytest.mark.parametrize( + "graph_class", [nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph] +) @pytest.mark.parametrize( "kwargs", [ @@ -48,9 +50,10 @@ def test_convert_empty(graph_class, kwargs): assert G.graph == Gcg.graph == H.graph == {} -def test_convert(): +@pytest.mark.parametrize("graph_class", [nx.Graph, nx.MultiGraph]) +def test_convert(graph_class): # FIXME: can we break this into smaller tests? - G = nx.Graph() + G = graph_class() G.add_edge(0, 1, x=2) G.add_node(0, foo=10) G.add_node(1, foo=20, bar=100) @@ -88,7 +91,10 @@ def test_convert(): cp.testing.assert_array_equal(Gcg.node_values["foo"], [10, 20]) assert Gcg.edge_values == Gcg.edge_masks == {} H = nxcg.to_networkx(Gcg) - assert set(G.edges) == set(H.edges) == {(0, 1)} + if G.is_multigraph(): + assert set(G.edges) == set(H.edges) == {(0, 1, 0)} + else: + assert set(G.edges) == set(H.edges) == {(0, 1)} assert G.nodes == H.nodes # Fill completely missing attribute with default value @@ -100,6 +106,8 @@ def test_convert(): assert Gcg.edge_masks == Gcg.node_values == Gcg.node_masks == {} H = nxcg.to_networkx(Gcg) assert list(H.edges(data=True)) == [(0, 1, {"y": 0})] + if Gcg.is_multigraph(): + assert set(H.edges) == {(0, 1, 0)} # If attribute is completely missing (and no default), then just ignore it Gcg = nxcg.from_networkx(G, edge_attrs={"y": None}) @@ -108,6 +116,8 @@ def test_convert(): assert sorted(Gcg.edge_values) == sorted(Gcg.edge_masks) == [] H = nxcg.to_networkx(Gcg) assert list(H.edges(data=True)) == [(0, 1, {})] + if Gcg.is_multigraph(): + assert set(H.edges) == {(0, 1, 0)} G.add_edge(0, 2) # Some edges are missing 'x' attribute; need to use a mask @@ -120,6 +130,8 @@ def test_convert(): cp.testing.assert_array_equal(Gcg.edge_values["x"][Gcg.edge_masks["x"]], [2, 2]) H = nxcg.to_networkx(Gcg) assert list(H.edges(data=True)) == [(0, 1, {"x": 2}), (0, 2, {})] + if Gcg.is_multigraph(): + assert set(H.edges) == {(0, 1, 0), (0, 2, 0)} with pytest.raises(KeyError, match="x"): nxcg.from_networkx(G, edge_attrs={"x": nxcg.convert.REQUIRED}) @@ -131,7 +143,7 @@ def test_convert(): nxcg.from_networkx(G, node_attrs={"bar": ...}) # Now for something more complicated... - G = nx.Graph() + G = graph_class() G.add_edge(10, 20, x=1) G.add_edge(10, 30, x=2, y=1.5) G.add_node(10, foo=100) @@ -147,7 +159,7 @@ def test_convert(): {"edge_attrs": {"x": 0, "y": None}, "edge_dtypes": {"x": int, "y": float}}, ]: Gcg = nxcg.from_networkx(G, **kwargs) - assert Gcg.id_to_key == {0: 10, 1: 20, 2: 30} # Remap node IDs to 0, 1, ... + assert Gcg.id_to_key == [10, 20, 30] # Remap node IDs to 0, 1, ... cp.testing.assert_array_equal(Gcg.row_indices, [0, 0, 1, 2]) cp.testing.assert_array_equal(Gcg.col_indices, [1, 2, 0, 0]) cp.testing.assert_array_equal(Gcg.edge_values["x"], [1, 2, 1, 2]) @@ -168,7 +180,7 @@ def test_convert(): {"node_attrs": {"foo": 0, "bar": None, "missing": None}}, ]: Gcg = nxcg.from_networkx(G, **kwargs) - assert Gcg.id_to_key == {0: 10, 1: 20, 2: 30} # Remap node IDs to 0, 1, ... + assert Gcg.id_to_key == [10, 20, 30] # Remap node IDs to 0, 1, ... cp.testing.assert_array_equal(Gcg.row_indices, [0, 0, 1, 2]) cp.testing.assert_array_equal(Gcg.col_indices, [1, 2, 0, 0]) cp.testing.assert_array_equal(Gcg.node_values["foo"], [100, 200, 300]) @@ -189,7 +201,7 @@ def test_convert(): {"node_attrs": {"bar": 0, "foo": None}, "node_dtypes": int}, ]: Gcg = nxcg.from_networkx(G, **kwargs) - assert Gcg.id_to_key == {0: 10, 1: 20, 2: 30} # Remap node IDs to 0, 1, ... + assert Gcg.id_to_key == [10, 20, 30] # Remap node IDs to 0, 1, ... cp.testing.assert_array_equal(Gcg.row_indices, [0, 0, 1, 2]) cp.testing.assert_array_equal(Gcg.col_indices, [1, 2, 0, 0]) cp.testing.assert_array_equal(Gcg.node_values["bar"], [0, 1000, 0]) @@ -201,3 +213,14 @@ def test_convert(): interface.BackendInterface.convert_from_nx(G, edge_attrs={"x": 1}, weight="x") with pytest.raises(TypeError, match="Expected networkx.Graph"): nxcg.from_networkx({}) + + +@pytest.mark.parametrize("graph_class", [nx.MultiGraph, nx.MultiDiGraph]) +def test_multigraph(graph_class): + G = graph_class() + G.add_edge(0, 1, "key1", x=10) + G.add_edge(0, 1, "key2", y=20) + Gcg = nxcg.from_networkx(G, preserve_edge_attrs=True) + H = nxcg.to_networkx(Gcg) + assert type(G) is type(H) + assert nx.utils.graphs_equal(G, H) From 71f6fb576193df3e93776846a5ac95d2417e7343 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Fri, 13 Oct 2023 20:31:29 -0500 Subject: [PATCH 2/3] More methods and tests --- python/nx-cugraph/nx_cugraph/classes/graph.py | 10 +- .../nx_cugraph/classes/multigraph.py | 240 ++++++++++++++++-- python/nx-cugraph/nx_cugraph/convert.py | 2 +- .../nx_cugraph/tests/test_multigraph.py | 72 ++++++ python/nx-cugraph/pyproject.toml | 2 +- 5 files changed, 303 insertions(+), 23 deletions(-) create mode 100644 python/nx-cugraph/nx_cugraph/tests/test_multigraph.py diff --git a/python/nx-cugraph/nx_cugraph/classes/graph.py b/python/nx-cugraph/nx_cugraph/classes/graph.py index 40581fd1906..4c1e6eeee92 100644 --- a/python/nx-cugraph/nx_cugraph/classes/graph.py +++ b/python/nx-cugraph/nx_cugraph/classes/graph.py @@ -246,9 +246,9 @@ def from_dcsc( def __new__(cls, incoming_graph_data=None, **attr) -> Graph: if incoming_graph_data is None: new_graph = cls.from_coo(0, cp.empty(0, np.int32), cp.empty(0, np.int32)) - elif incoming_graph_data.__class__ is new_graph.__class__: + elif incoming_graph_data.__class__ is cls: new_graph = incoming_graph_data.copy() - elif incoming_graph_data.__class__ is new_graph.to_networkx_class(): + elif incoming_graph_data.__class__ is cls.to_networkx_class(): new_graph = nxcg.from_networkx(incoming_graph_data, preserve_all_attrs=True) else: raise NotImplementedError @@ -371,6 +371,12 @@ def get_edge_data( v = self.key_to_id[v] except KeyError: return default + else: + try: + if u < 0 or v < 0 or u >= self._N or v >= self._N: + return default + except TypeError: + return default index = cp.nonzero((self.row_indices == u) & (self.col_indices == v))[0] if index.size == 0: return default diff --git a/python/nx-cugraph/nx_cugraph/classes/multigraph.py b/python/nx-cugraph/nx_cugraph/classes/multigraph.py index aad1c653757..499ca7ce212 100644 --- a/python/nx-cugraph/nx_cugraph/classes/multigraph.py +++ b/python/nx-cugraph/nx_cugraph/classes/multigraph.py @@ -102,22 +102,162 @@ def from_coo( raise ValueError return new_graph - # TODO: - # from_csr - # from_csc - # from_dcsr - # from_dcsc - - def __new__(cls, incoming_graph_data=None, multigraph_input=None, **attr) -> Graph: - # TODO: handle multigraph_input - if incoming_graph_data is None: - new_graph = cls.from_coo(0, cp.empty(0, np.int32), cp.empty(0, np.int32)) - elif incoming_graph_data.__class__ is new_graph.__class__: - new_graph = incoming_graph_data.copy() - elif incoming_graph_data.__class__ is new_graph.to_networkx_class(): - new_graph = nxcg.from_networkx(incoming_graph_data, preserve_all_attrs=True) + @classmethod + def from_csr( + cls, + indptr: cp.ndarray[IndexValue], + col_indices: cp.ndarray[IndexValue], + edge_indices: cp.ndarray[IndexValue] | None = None, + edge_values: dict[AttrKey, cp.ndarray[EdgeValue]] | None = None, + edge_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + node_values: dict[AttrKey, cp.ndarray[NodeValue]] | None = None, + node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + *, + key_to_id: dict[NodeKey, IndexValue] | None = None, + id_to_key: list[NodeKey] | None = None, + edge_keys: list[EdgeKey] | None = None, + **attr, + ) -> MultiGraph: + N = indptr.size - 1 + row_indices = cp.array( + # cp.repeat is slow to use here, so use numpy instead + np.repeat(np.arange(N, dtype=np.int32), cp.diff(indptr).get()) + ) + return cls.from_coo( + N, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + edge_keys=edge_keys, + **attr, + ) + + @classmethod + def from_csc( + cls, + indptr: cp.ndarray[IndexValue], + row_indices: cp.ndarray[IndexValue], + edge_indices: cp.ndarray[IndexValue] | None = None, + edge_values: dict[AttrKey, cp.ndarray[EdgeValue]] | None = None, + edge_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + node_values: dict[AttrKey, cp.ndarray[NodeValue]] | None = None, + node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + *, + key_to_id: dict[NodeKey, IndexValue] | None = None, + id_to_key: list[NodeKey] | None = None, + edge_keys: list[EdgeKey] | None = None, + **attr, + ) -> MultiGraph: + N = indptr.size - 1 + col_indices = cp.array( + # cp.repeat is slow to use here, so use numpy instead + np.repeat(np.arange(N, dtype=np.int32), cp.diff(indptr).get()) + ) + return cls.from_coo( + N, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + edge_keys=edge_keys, + **attr, + ) + + @classmethod + def from_dcsr( + cls, + N: int, + compressed_rows: cp.ndarray[IndexValue], + indptr: cp.ndarray[IndexValue], + col_indices: cp.ndarray[IndexValue], + edge_indices: cp.ndarray[IndexValue] | None = None, + edge_values: dict[AttrKey, cp.ndarray[EdgeValue]] | None = None, + edge_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + node_values: dict[AttrKey, cp.ndarray[NodeValue]] | None = None, + node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + *, + key_to_id: dict[NodeKey, IndexValue] | None = None, + id_to_key: list[NodeKey] | None = None, + edge_keys: list[EdgeKey] | None = None, + **attr, + ) -> MultiGraph: + row_indices = cp.array( + # cp.repeat is slow to use here, so use numpy instead + np.repeat(compressed_rows.get(), cp.diff(indptr).get()) + ) + return cls.from_coo( + N, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + edge_keys=edge_keys, + **attr, + ) + + @classmethod + def from_dcsc( + cls, + N: int, + compressed_cols: cp.ndarray[IndexValue], + indptr: cp.ndarray[IndexValue], + row_indices: cp.ndarray[IndexValue], + edge_indices: cp.ndarray[IndexValue] | None = None, + edge_values: dict[AttrKey, cp.ndarray[EdgeValue]] | None = None, + edge_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + node_values: dict[AttrKey, cp.ndarray[NodeValue]] | None = None, + node_masks: dict[AttrKey, cp.ndarray[bool]] | None = None, + *, + key_to_id: dict[NodeKey, IndexValue] | None = None, + id_to_key: list[NodeKey] | None = None, + edge_keys: list[EdgeKey] | None = None, + **attr, + ) -> Graph: + col_indices = cp.array( + # cp.repeat is slow to use here, so use numpy instead + np.repeat(compressed_cols.get(), cp.diff(indptr).get()) + ) + return cls.from_coo( + N, + row_indices, + col_indices, + edge_indices, + edge_values, + edge_masks, + node_values, + node_masks, + key_to_id=key_to_id, + id_to_key=id_to_key, + edge_keys=edge_keys, + **attr, + ) + + def __new__( + cls, incoming_graph_data=None, multigraph_input=None, **attr + ) -> MultiGraph: + if isinstance(incoming_graph_data, dict) and multigraph_input is not False: + new_graph = nxcg.from_networkx( + nx.MultiGraph(incoming_graph_data, multigraph_input=multigraph_input), + preserve_all_attrs=True, + ) else: - raise NotImplementedError + new_graph = super().__new__(cls, incoming_graph_data) new_graph.graph.update(attr) return new_graph @@ -165,9 +305,68 @@ def clear_edges(self) -> None: self.edge_indices = None self.edge_keys = None - # TODO: - # copy - # get_edge_data + @networkx_api + def copy(self, as_view: bool = False) -> MultiGraph: + # Does shallow copy in networkx + return self._copy(as_view, self.__class__) + + @networkx_api + def get_edge_data( + self, + u: NodeKey, + v: NodeKey, + key: EdgeKey | None = None, + default: EdgeValue | None = None, + ): + if self.key_to_id is not None: + try: + u = self.key_to_id[u] + v = self.key_to_id[v] + except KeyError: + return default + else: + try: + if u < 0 or v < 0 or u >= self._N or v >= self._N: + return default + except TypeError: + return default + mask = (self.row_indices == u) & (self.col_indices == v) + if not mask.any(): + return default + if self.edge_keys is None: + if self.edge_indices is None: + self._calculate_edge_indices() + if key is not None: + try: + mask = mask & (self.edge_indices == key) + except TypeError: + return default + indices = cp.nonzero(mask)[0] + if indices.size == 0: + return default + edge_keys = self.edge_keys + if key is not None and edge_keys is not None: + mask[[i for i in indices.tolist() if edge_keys[i] != key]] = False + indices = cp.nonzero(mask)[0] + if indices.size == 0: + return default + if key is not None: + [index] = indices.tolist() + return { + k: v[index].tolist() + for k, v in self.edge_values.items() + if k not in self.edge_masks or self.edge_masks[k][index] + } + return { + edge_keys[index] + if edge_keys is not None + else index: { + k: v[index].tolist() + for k, v in self.edge_values.items() + if k not in self.edge_masks or self.edge_masks[k][index] + } + for index in indices.tolist() + } @networkx_api def has_edge(self, u: NodeKey, v: NodeKey, key: EdgeKey | None = None) -> bool: @@ -181,7 +380,10 @@ def has_edge(self, u: NodeKey, v: NodeKey, key: EdgeKey | None = None) -> bool: if key is None or (self.edge_indices is None and self.edge_keys is None): return bool(mask.any()) if self.edge_keys is None: - return bool((mask & (self.edge_indices == key)).any()) + try: + return bool((mask & (self.edge_indices == key)).any()) + except TypeError: + return False indices = cp.nonzero(mask)[0] if indices.size == 0: return False diff --git a/python/nx-cugraph/nx_cugraph/convert.py b/python/nx-cugraph/nx_cugraph/convert.py index 4f822564f46..83423c4ade8 100644 --- a/python/nx-cugraph/nx_cugraph/convert.py +++ b/python/nx-cugraph/nx_cugraph/convert.py @@ -274,7 +274,7 @@ def from_networkx( col_indices = cp.array(np.repeat(col_indices, num_multiedges)) # Determine edge keys and edge ids for multigraphs edge_keys = list(concat(concat(map(dict.values, adj.values())))) - edge_indices = np.fromiter( + edge_indices = cp.fromiter( concat(map(range, map(len, concat(map(dict.values, adj.values()))))), np.int32, ) diff --git a/python/nx-cugraph/nx_cugraph/tests/test_multigraph.py b/python/nx-cugraph/nx_cugraph/tests/test_multigraph.py new file mode 100644 index 00000000000..a8f189a4745 --- /dev/null +++ b/python/nx-cugraph/nx_cugraph/tests/test_multigraph.py @@ -0,0 +1,72 @@ +# 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 networkx as nx +import pytest + +import nx_cugraph as nxcg + + +@pytest.mark.parametrize("test_nxcugraph", [False, True]) +def test_get_edge_data(test_nxcugraph): + G = nx.MultiGraph() + G.add_edge(0, 1, 0, x=10) + G.add_edge(0, 1, 1, y=20) + G.add_edge(0, 2, "a", x=100) + G.add_edge(0, 2, "b", y=200) + G.add_edge(0, 3) + G.add_edge(0, 3) + if test_nxcugraph: + G = nxcg.MultiGraph(G) + default = object() + assert G.get_edge_data(0, 0, default=default) is default + assert G.get_edge_data("a", "b", default=default) is default + assert G.get_edge_data(0, 1, 2, default=default) is default + assert G.get_edge_data(-1, 1, default=default) is default + assert G.get_edge_data(0, 1, 0, default=default) == {"x": 10} + assert G.get_edge_data(0, 1, 1, default=default) == {"y": 20} + assert G.get_edge_data(0, 1, default=default) == {0: {"x": 10}, 1: {"y": 20}} + assert G.get_edge_data(0, 2, "a", default=default) == {"x": 100} + assert G.get_edge_data(0, 2, "b", default=default) == {"y": 200} + assert G.get_edge_data(0, 2, default=default) == {"a": {"x": 100}, "b": {"y": 200}} + assert G.get_edge_data(0, 3, 0, default=default) == {} + assert G.get_edge_data(0, 3, 1, default=default) == {} + assert G.get_edge_data(0, 3, 2, default=default) is default + assert G.get_edge_data(0, 3, default=default) == {0: {}, 1: {}} + assert G.has_edge(0, 1) + assert G.has_edge(0, 1, 0) + assert G.has_edge(0, 1, 1) + assert not G.has_edge(0, 1, 2) + assert not G.has_edge(0, 1, "a") + assert not G.has_edge(0, -1) + assert G.has_edge(0, 2) + assert G.has_edge(0, 2, "a") + assert G.has_edge(0, 2, "b") + assert not G.has_edge(0, 2, "c") + assert not G.has_edge(0, 2, 0) + assert G.has_edge(0, 3) + assert not G.has_edge(0, 0) + assert not G.has_edge(0, 0, 0) + + G = nx.MultiGraph() + G.add_edge(0, 1) + if test_nxcugraph: + G = nxcg.MultiGraph(G) + assert G.get_edge_data(0, 1, default=default) == {0: {}} + assert G.get_edge_data(0, 1, 0, default=default) == {} + assert G.get_edge_data(0, 1, 1, default=default) is default + assert G.get_edge_data(0, 1, "b", default=default) is default + assert G.get_edge_data(-1, 2, default=default) is default + assert G.has_edge(0, 1) + assert G.has_edge(0, 1, 0) + assert not G.has_edge(0, 1, 1) + assert not G.has_edge(0, 1, "a") diff --git a/python/nx-cugraph/pyproject.toml b/python/nx-cugraph/pyproject.toml index 98adbb439d8..7bc0650ae02 100644 --- a/python/nx-cugraph/pyproject.toml +++ b/python/nx-cugraph/pyproject.toml @@ -109,7 +109,7 @@ addopts = [ # "-ra", # Print summary of all fails/errors "--benchmark-warmup=off", "--benchmark-max-time=0", - "--benchmark-min-rounds=3", + "--benchmark-min-rounds=1", "--benchmark-columns=min,median,max", ] From a03964276c6fa41af14ab6adac20f2e8ed642a90 Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Tue, 17 Oct 2023 20:54:14 -0500 Subject: [PATCH 3/3] add numpy as a dependency to nx-cugraph --- dependencies.yaml | 3 ++- python/nx-cugraph/pyproject.toml | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/dependencies.yaml b/dependencies.yaml index 70fb54b4721..736aa2e019f 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -401,12 +401,13 @@ dependencies: - output_types: [conda, pyproject] packages: - networkx>=3.0 + - &numpy numpy>=1.21 python_run_cugraph_dgl: common: - output_types: [conda, pyproject] packages: - *numba - - &numpy numpy>=1.21 + - *numpy - output_types: [pyproject] packages: - &cugraph cugraph==23.12.* diff --git a/python/nx-cugraph/pyproject.toml b/python/nx-cugraph/pyproject.toml index 7bc0650ae02..2478a02df9b 100644 --- a/python/nx-cugraph/pyproject.toml +++ b/python/nx-cugraph/pyproject.toml @@ -32,6 +32,7 @@ classifiers = [ dependencies = [ "cupy-cuda11x>=12.0.0", "networkx>=3.0", + "numpy>=1.21", "pylibcugraph==23.12.*", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.