Merge pull request #781 from pydata/triangles-example

ENH: Add counting triangles example

Author: mtsokol

examples/triangles_example.py

@@ -0,0 +1,53 @@
+import importlib
+import os
+
+import sparse
+
+import networkx as nx
+from utils import benchmark
+
+import numpy as np
+
+ITERS = 3
+
+
+if __name__ == "__main__":
+    print("Counting Triangles Example:\n")
+
+    G = nx.gnp_random_graph(n=200, p=0.2)
+
+    # ======= Finch =======
+    os.environ[sparse._ENV_VAR_NAME] = "Finch"
+    importlib.reload(sparse)
+
+    a_sps = nx.to_scipy_sparse_array(G)
+    a = sparse.asarray(a_sps)
+
+    # @sparse.compiled NOTE: blocked by https://github.com/willow-ahrens/Finch.jl/issues/615
+    def count_triangles_finch(a):
+        return sparse.sum(a @ a * a) / 6
+
+    # Compile & Benchmark
+    result_finch = benchmark(count_triangles_finch, args=[a], info="Finch", iters=ITERS)
+
+    # ======= SciPy =======
+    def count_triangles_scipy(a):
+        return (a @ a * a).sum() / 6
+
+    a = nx.to_scipy_sparse_array(G)
+
+    # Compile & Benchmark
+    result_scipy = benchmark(count_triangles_scipy, args=[a], info="SciPy", iters=ITERS)
+
+    # ======= NetworkX =======
+    def count_triangles_networkx(a):
+        return sum(nx.triangles(a).values()) / 3
+
+    a = G
+
+    # Compile & Benchmark
+    result_networkx = benchmark(count_triangles_networkx, args=[a], info="NetworkX", iters=ITERS)
+
+    np.testing.assert_equal(result_finch.todense(), result_scipy)
+    np.testing.assert_equal(result_finch.todense(), result_networkx)
+    assert result_networkx == result_scipy

pyproject.toml

@@ -41,6 +41,7 @@ extras = [
     "sparse[finch]",
     "scipy",
     "scikit-learn",
+    "networkx",
 ]
 tests = [
     "sparse[extras]",