From 3a3aa897f5a06926fe45624da6ab98dabacb6bae Mon Sep 17 00:00:00 2001 From: Ishika Roy Date: Wed, 3 Feb 2021 15:15:09 -0600 Subject: [PATCH] add batched ego test --- python/cugraph/__init__.py | 1 + python/cugraph/community/__init__.py | 1 + python/cugraph/community/egonet.py | 4 ++-- python/cugraph/structure/number_map.py | 2 +- python/cugraph/tests/test_egonet.py | 23 ++++++++++++++++++++++- 5 files changed, 27 insertions(+), 4 deletions(-) diff --git a/python/cugraph/__init__.py b/python/cugraph/__init__.py index d0629933544..46013903a38 100644 --- a/python/cugraph/__init__.py +++ b/python/cugraph/__init__.py @@ -25,6 +25,7 @@ subgraph, triangles, ego_graph, + batched_ego_graphs, ) from cugraph.structure import ( diff --git a/python/cugraph/community/__init__.py b/python/cugraph/community/__init__.py index 60d1a14ad07..9cc92637e20 100644 --- a/python/cugraph/community/__init__.py +++ b/python/cugraph/community/__init__.py @@ -26,3 +26,4 @@ from cugraph.community.ktruss_subgraph import ktruss_subgraph from cugraph.community.ktruss_subgraph import k_truss from cugraph.community.egonet import ego_graph +from cugraph.community.egonet import batched_ego_graphs diff --git a/python/cugraph/community/egonet.py b/python/cugraph/community/egonet.py index b057a0fa50e..728adc49008 100644 --- a/python/cugraph/community/egonet.py +++ b/python/cugraph/community/egonet.py @@ -149,7 +149,7 @@ def batched_ego_graphs( df, offsets = egonet_wrapper.egonet(G, seeds, radius) if G.renumbered: - df = G.unrenumber(df, "src") - df = G.unrenumber(df, "dst") + df = G.unrenumber(df, "src", preserve_order=True) + df = G.unrenumber(df, "dst", preserve_order=True) return _convert_df_series_to_output_type(df, offsets, input_type) diff --git a/python/cugraph/structure/number_map.py b/python/cugraph/structure/number_map.py index f1b8949eb5d..4b714f7a84b 100644 --- a/python/cugraph/structure/number_map.py +++ b/python/cugraph/structure/number_map.py @@ -895,7 +895,7 @@ def unrenumber(self, df, column_name, preserve_order=False): if preserve_order: df = df.sort_values( index_name - ).drop(index_name).reset_index(drop=True) + ).drop(columns=index_name).reset_index(drop=True) if type(df) is dask_cudf.DataFrame: return df.map_partitions( diff --git a/python/cugraph/tests/test_egonet.py b/python/cugraph/tests/test_egonet.py index aa875561d66..009fd1252f1 100644 --- a/python/cugraph/tests/test_egonet.py +++ b/python/cugraph/tests/test_egonet.py @@ -54,7 +54,10 @@ def test_ego_graph_nx(graph_file, seed, radius): assert nx.is_isomorphic(ego_nx, ego_cugraph) -def batched_ego_graphs(G, seeds, radius=1): +@pytest.mark.parametrize("graph_file", utils.DATASETS) +@pytest.mark.parametrize("seeds", [[0, 5, 13]]) +@pytest.mark.parametrize("radius", [1, 2, 3]) +def test_batched_ego_graphs(graph_file, seeds, radius): """ Compute the induced subgraph of neighbors for each node in seeds within a given radius. @@ -78,3 +81,21 @@ def batched_ego_graphs(G, seeds, radius=1): Series containing the starting offset in the returned edge list for each seed. """ + gc.collect() + + # Nx + df = utils.read_csv_for_nx(graph_file, read_weights_in_sp=True) + Gnx = nx.from_pandas_edgelist( + df, create_using=nx.Graph(), source="0", target="1", edge_attr="weight" + ) + + # cugraph + df, offsets = cugraph.batched_ego_graphs(Gnx, seeds, radius=radius) + for i in range(len(seeds)): + ego_nx = nx.ego_graph(Gnx, seeds[i], radius=radius) + ego_df = df[offsets[i]:offsets[i+1]] + ego_cugraph = nx.from_pandas_edgelist(ego_df, + source="src", + target="dst", + edge_attr="weight") + assert nx.is_isomorphic(ego_nx, ego_cugraph)