Add plots to benchmark and test for broadcasting

Author: Luke Shaw

bench/ndarray/slice-expr.py

@@ -2,7 +2,8 @@
 import numpy as np
 import blosc2
 import time
-
+from memory_profiler import memory_usage
+import matplotlib.pyplot as plt
 
 file = "dset-ones.b2nd"
 # a = blosc2.open(file)
@@ -25,27 +26,46 @@
 print(f"Time to create expression: {time.time() - t:.6f} seconds")
 
 # dim0
-t = time.time()
-res = expr[1]
-t0 = time.time() - t
-print(f"Time to access dim0: {t0:.6f} seconds")
+def slice_dim0():
+    t = time.time()
+    res = expr[1]
+    t0 = time.time() - t
+    print(f"Time to access dim0: {t0:.6f} seconds")
 
 # dim1
-t = time.time()
-res = expr[:,1]
-t1 = time.time() - t
-print(f"Time to access dim1: {t1:.6f} seconds")
+def slice_dim1():
+    t = time.time()
+    res = expr[:,1]
+    t1 = time.time() - t
+    print(f"Time to access dim1: {t1:.6f} seconds")
 
 # dim2
-t = time.time()
-res = expr[:,:,1]
-t2 = time.time() - t
-print(f"Time to access dim2: {t2:.6f} seconds")
+def slice_dim2():
+    t = time.time()
+    res = expr[:,:,1]
+    t2 = time.time() - t
+    print(f"Time to access dim2: {t2:.6f} seconds")
 
 # dim3
-t = time.time()
-res = expr[:,:,:,1]
-#res = expr[1]
-t3 = time.time() - t
+def slice_dim3():
+    t = time.time()
+    res = expr[:,:,:,1]
+    #res = expr[1]
+    t3 = time.time() - t
+
+    print(f"Time to access dim3: {t3:.6f} seconds")
+
+fig = plt.figure()
+interval = 0.001
+offset = 0
+for f in [slice_dim0, slice_dim1, slice_dim2, slice_dim3]:
+    mem = memory_usage((f,), interval=interval)
+    times = offset + interval * np.arange(len(mem))
+    offset = times[-1]
+    plt.plot(times, mem)
 
-print(f"Time to access dim3: {t3:.6f} seconds")
+plt.xlabel('Time (s)')
+plt.ylabel('Memory usage (MiB)')
+plt.title('Memory usage over time for slicing operations, slice-expr.py')
+plt.legend(['dim0', 'dim1', 'dim2', 'dim3'])
+plt.savefig('plots/slice-expr.png', format="png")

tests/ndarray/test_lazyexpr.py

@@ -1076,7 +1076,7 @@ def test_eval_getitem(array_fixture):
     np.testing.assert_allclose(expr[:10], nres[:10])
     np.testing.assert_allclose(expr[0:10:2], nres[0:10:2])
 
-    # Small test
+    # Small test for non-isomorphic shape
     shape = (2, 10, 5)
     test_arr = blosc2.linspace(0, 10, np.prod(shape), shape=shape)
     expr = test_arr * 30
@@ -1085,6 +1085,15 @@ def test_eval_getitem(array_fixture):
     np.testing.assert_allclose(expr[:10], nres[:10])
     np.testing.assert_allclose(expr[0:10:2], nres[0:10:2])
 
+    # Small test for broadcasting
+    shape = (2, 10, 5)
+    test_arr = blosc2.linspace(0, 10, np.prod(shape), shape=shape)
+    expr = test_arr + test_arr.slice(slice(1, 2))
+    nres = test_arr[:] + test_arr[1]
+    np.testing.assert_allclose(expr[0], nres[0])
+    np.testing.assert_allclose(expr[:10], nres[:10])
+    np.testing.assert_allclose(expr[0:10:2], nres[0:10:2])
+
 
 # Test lazyexpr's slice method
 def test_eval_slice(array_fixture):