Adds benchmark for louvain_communities using medium size graphs.

rapidsai · Sep 27, 2023 · 15b28ee · 15b28ee
1 parent 9248af0
commit 15b28ee
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/benchmarks/nx-cugraph/pytest-based/bench_algos.py b/benchmarks/nx-cugraph/pytest-based/bench_algos.py
@@ -221,3 +221,15 @@ def bench_louvain_communities_small(benchmark,
         G = small_graph_obj
     result = benchmark(backend_selector(nx.community.louvain_communities), G)
     assert type(result) is list
+
+
+def bench_louvain_communities_medium(benchmark,
+                                     medium_graph_obj,
+                                     backend_selector):
+    # The cugraph backend for louvain_communities only supports undirected graphs
+    if isinstance(medium_graph_obj, nx.DiGraph):
+        G = medium_graph_obj.to_undirected()
+    else:
+        G = medium_graph_obj
+    result = benchmark(backend_selector(nx.community.louvain_communities), G)
+    assert type(result) is list