Add elemwise examples (#697)

Author: mtsokol

examples/__init__.py

@@ -0,0 +1,72 @@
+import importlib
+import operator
+import os
+
+import sparse
+
+from utils import benchmark
+
+import numpy as np
+import scipy.sparse as sps
+
+LEN = 10000
+DENSITY = 0.001
+ITERS = 3
+rng = np.random.default_rng(0)
+
+
+if __name__ == "__main__":
+    print("Elementwise Example:\n")
+
+    for func_name in ["multiply", "add", "greater_equal"]:
+        print(f"{func_name} benchmark:\n")
+
+        s1_sps = sps.random(LEN, LEN, format="csr", density=DENSITY, random_state=rng) * 10
+        s1_sps.sum_duplicates()
+        s2_sps = sps.random(LEN, LEN, format="csr", density=DENSITY, random_state=rng) * 10
+        s2_sps.sum_duplicates()
+
+        # ======= Finch =======
+        os.environ[sparse._ENV_VAR_NAME] = "Finch"
+        importlib.reload(sparse)
+
+        s1 = sparse.asarray(s1_sps.asformat("csc"), format="csc")
+        s2 = sparse.asarray(s2_sps.asformat("csc"), format="csc")
+
+        func = getattr(sparse, func_name)
+        # Compile
+        result_finch = func(s1, s2)
+        # Benchmark
+        benchmark(func, args=[s1, s2], info="Finch", iters=ITERS)
+
+        # ======= Numba =======
+        os.environ[sparse._ENV_VAR_NAME] = "Numba"
+        importlib.reload(sparse)
+
+        s1 = sparse.asarray(s1_sps)
+        s2 = sparse.asarray(s2_sps)
+
+        func = getattr(sparse, func_name)
+        # Compile
+        result_numba = func(s1, s2)
+        # Benchmark
+        benchmark(func, args=[s1, s2], info="Numba", iters=ITERS)
+
+        # ======= SciPy =======
+        s1 = s1_sps
+        s2 = s2_sps
+
+        if func_name == "multiply":
+            func, args = s1.multiply, [s2]
+        elif func_name == "add":
+            func, args = operator.add, [s1, s2]
+        elif func_name == "greater_equal":
+            func, args = operator.ge, [s1, s2]
+
+        result_scipy = func(*args)
+        # Benchmark
+        benchmark(func, args=args, info="SciPy", iters=ITERS)
+
+        np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
+        np.testing.assert_allclose(result_finch.todense(), result_numba.todense())
+        np.testing.assert_allclose(result_finch.todense(), result_scipy.toarray())

examples/elemwise_example.py

@@ -1,9 +1,10 @@
 import importlib
 import os
-import time
 
 import sparse
 
+from utils import benchmark
+
 import numpy as np
 
 I_ = 1000
@@ -15,15 +16,6 @@
 rng = np.random.default_rng(0)
 
 
-def benchmark(func, info, args):
-    print(info)
-    start = time.time()
-    for _ in range(ITERS):
-        func(*args)
-    elapsed = time.time() - start
-    print(f"Took {elapsed / ITERS} s.\n")
-
-
 if __name__ == "__main__":
     print("MTTKRP Example:\n")
 
@@ -45,9 +37,8 @@ def mttkrp_finch(B, D, C):
 
     # Compile
     result_finch = mttkrp_finch(B, D, C)
-    assert sparse.nonzero(result_finch)[0].size > 5
     # Benchmark
-    benchmark(mttkrp_finch, info="Finch", args=[B, D, C])
+    benchmark(mttkrp_finch, args=[B, D, C], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -63,6 +54,6 @@ def mttkrp_numba(B, D, C):
     # Compile
     result_numba = mttkrp_numba(B, D, C)
     # Benchmark
-    benchmark(mttkrp_numba, info="Numba", args=[B, D, C])
+    benchmark(mttkrp_numba, args=[B, D, C], info="Numba", iters=ITERS)
 
     np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/mttkrp_example.py

@@ -1,9 +1,10 @@
 import importlib
 import os
-import time
 
 import sparse
 
+from utils import benchmark
+
 import numpy as np
 import scipy.sparse as sps
 
@@ -13,15 +14,6 @@
 rng = np.random.default_rng(0)
 
 
-def benchmark(func, info, args):
-    print(info)
-    start = time.time()
-    for _ in range(ITERS):
-        func(*args)
-    elapsed = time.time() - start
-    print(f"Took {elapsed / ITERS} s.\n")
-
-
 if __name__ == "__main__":
     print("SDDMM Example:\n")
 
@@ -47,9 +39,8 @@ def sddmm_finch(s, a, b):
 
     # Compile
     result_finch = sddmm_finch(s, a, b)
-    assert sparse.nonzero(result_finch)[0].size > 5
     # Benchmark
-    benchmark(sddmm_finch, info="Finch", args=[s, a, b])
+    benchmark(sddmm_finch, args=[s, a, b], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -64,9 +55,8 @@ def sddmm_numba(s, a, b):
 
     # Compile
     result_numba = sddmm_numba(s, a, b)
-    assert sparse.nonzero(result_numba)[0].size > 5
     # Benchmark
-    benchmark(sddmm_numba, info="Numba", args=[s, a, b])
+    benchmark(sddmm_numba, args=[s, a, b], info="Numba", iters=ITERS)
 
     # ======= SciPy =======
     def sddmm_scipy(s, a, b):
@@ -78,7 +68,7 @@ def sddmm_scipy(s, a, b):
 
     result_scipy = sddmm_scipy(s, a, b)
     # Benchmark
-    benchmark(sddmm_scipy, info="SciPy", args=[s, a, b])
+    benchmark(sddmm_scipy, args=[s, a, b], info="SciPy", iters=ITERS)
 
     np.testing.assert_allclose(result_numba.todense(), result_scipy.toarray())
     np.testing.assert_allclose(result_finch.todense(), result_numba.todense())

examples/sddmm_example.py

@@ -1,9 +1,10 @@
 import importlib
 import os
-import time
 
 import sparse
 
+from utils import benchmark
+
 import numpy as np
 import scipy.sparse as sps
 
@@ -13,15 +14,6 @@
 rng = np.random.default_rng(0)
 
 
-def benchmark(func, info, args):
-    print(info)
-    start = time.time()
-    for _ in range(ITERS):
-        func(*args)
-    elapsed = time.time() - start
-    print(f"Took {elapsed / ITERS} s.\n")
-
-
 if __name__ == "__main__":
     print("SpMv_add Example:\n")
 
@@ -43,9 +35,8 @@ def spmv_finch(A, x, y):
 
     # Compile
     result_finch = spmv_finch(A, x, y)
-    assert sparse.nonzero(result_finch)[0].size > 5
     # Benchmark
-    benchmark(spmv_finch, info="Finch", args=[A, x, y])
+    benchmark(spmv_finch, args=[A, x, y], info="Finch", iters=ITERS)
 
     # ======= Numba =======
     os.environ[sparse._ENV_VAR_NAME] = "Numba"
@@ -62,7 +53,7 @@ def spmv_numba(A, x, y):
     result_numba = spmv_numba(A, x, y)
     assert sparse.nonzero(result_numba)[0].size > 5
     # Benchmark
-    benchmark(spmv_numba, info="Numba", args=[A, x, y])
+    benchmark(spmv_numba, args=[A, x, y], info="Numba", iters=ITERS)
 
     # ======= SciPy =======
     def spmv_scipy(A, x, y):
@@ -74,7 +65,7 @@ def spmv_scipy(A, x, y):
 
     result_scipy = spmv_scipy(A, x, y)
     # Benchmark
-    benchmark(spmv_scipy, info="SciPy", args=[A, x, y])
+    benchmark(spmv_scipy, args=[A, x, y], info="SciPy", iters=ITERS)
 
     np.testing.assert_allclose(result_numba, result_scipy)
     np.testing.assert_allclose(result_finch.todense(), result_numba)

examples/spmv_add_example.py

@@ -0,0 +1,12 @@
+import time
+from collections.abc import Callable, Iterable
+from typing import Any
+
+
+def benchmark(func: Callable, args: Iterable[Any], info: str, iters: int):
+    print(info)
+    start = time.time()
+    for _ in range(iters):
+        func(*args)
+    elapsed = time.time() - start
+    print(f"Took {elapsed / iters} s.\n")