From 733f44ba0b95cc0c822e806b8a9f5646be037651 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= <mat646@gmail.com>
Date: Fri, 27 Sep 2024 12:31:32 +0000
Subject: [PATCH] ENH: Added counting triangles example

---
 examples/triangles_example.py | 53 +++++++++++++++++++++++++++++++++++
 pyproject.toml                |  1 +
 2 files changed, 54 insertions(+)
 create mode 100644 examples/triangles_example.py

diff --git a/examples/triangles_example.py b/examples/triangles_example.py
new file mode 100644
index 00000000..97ee3f24
--- /dev/null
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 662fb895..77fab122 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,6 +41,7 @@ extras = [
     "sparse[finch]",
     "scipy",
     "scikit-learn",
+    "networkx",
 ]
 tests = [
     "sparse[extras]",