Update benchmark and add oindex setitem

Author: Luke Shaw

bench/ndarray/fancy_index.py

@@ -14,123 +14,128 @@
 import time
 from memory_profiler import memory_usage, profile
 import matplotlib.pyplot as plt
+import zarr
+import h5py
+plt.rcParams.update({'text.usetex':True,'font.serif': ['cm'],'font.size':16})
+plt.rcParams['figure.dpi'] = 1000
+plt.rcParams['savefig.dpi'] = 1000
+plt.rc('text', usetex=True)
+plt.rc('font',**{'serif':['cm']})
+plt.style.use('seaborn-v0_8-paper')
 
-MEM_PROFILE = False
+NUMPY_BLOSC = False
 
-# @profile
-def fancyBlosc2(arr, fancyIdx):
+def genarray(r, ndims=1, verbose=True):
+    d = int((r*2**30/8)**(1/ndims))
+    shape = (d,) * ndims
+    chunks = (d // 4,) * ndims
+    blocks = (max(d // 10, 1),) * ndims
     t = time.time()
-    res = arr[fancyIdx]
-    dt = time.time() - t
-    return res, dt
-
-# @profile
-def fancyNumpy(arr, fancyIdx):
-    t = time.time()
-    res = arr[:]
-    res = res[fancyIdx]
-    dt = time.time() - t
-    return res, dt
-
-def genarray(d, verbose=False):
-    shape = (d,) * 4
-    chunks = (d // 4,) * 4
-    blocks = (max(d // 10, 1),) * 4
-    t = time.time()
-    arr = blosc2.ones(shape=shape, chunks=chunks, blocks=blocks)  # , urlpath=file, mode="w")
+    arr = blosc2.ones(shape=shape, chunks=chunks, blocks=blocks, dtype=np.int64)  # , urlpath=file, mode="w")
     t = time.time() - t
     if verbose:
         print(f"Array shape: {arr.shape}")
         print(f"Array size: {np.prod(arr.shape) * arr.dtype.itemsize / 2 ** 30:.6f} GB")
         print(f"Time to create array: {t:.6f} seconds")
     return arr
 
-if __name__ == '__main__':
-    blosc_times = []
-    np_times = []
-    sizes = []
-    # dims = np.int64(np.array([2**6.76, 2**6.8, 2**6.85, 2**6.9, 2**6.95]))
-    dims = np.int64(np.array([2**4.3, 2**5.25, 2**6.25, 2**6.75, 2**7, 2**7.25, 2**7.4, 2**7.5]))
-    rng = np.random.default_rng()
+blosc_times = []
+np_times = []
+zarr_times = []
+sizes = []
+dims = np.int64(np.array([1, 2, 4, 6, 8]))
+rng = np.random.default_rng()
+blosctimes = []
+nptimes = []
+zarrtimes = []
+h5pytimes = []
 
-    for d in dims:
-        arr = genarray(d)
-        arr_size = np.prod(arr.shape) * arr.dtype.itemsize / 2 ** 30
-        print(arr_size)
-        sizes.append(arr_size)
-        idx = rng.integers(low=0, high=d, size=(d,))
-        if MEM_PROFILE:
-            fancyIdx = np.s_[idx, :, :d // 2, d // 2:]
+for d in dims:
+    arr = genarray(d, ndims=2)
+    sizes.append(d)
+    idx = rng.integers(low=0, high=arr.shape[0], size=(arr.shape[0],))
+    row = np.arange(arr.shape[0])
+    col = row
+    mask = rng.integers(low=0, high=2, size=(d,)) == 1
 
-            interval = 0.001
-            offset = 0
-            for f in [fancyBlosc2, fancyNumpy]:
-                mem = memory_usage((f, (arr, fancyIdx)), interval=interval)
-                times = offset + interval * np.arange(len(mem))
-                offset = times[-1]
-                plt.plot(times, mem)
 
-            plt.xlabel('Time (s)')
-            plt.ylabel('Memory usage (MiB)')
-            plt.title('Memory usage fancy indexing')
-            plt.legend(['blosc2', 'numpy'])
-            plt.savefig(f'plots/MemoryUsagefancyIdx_d{d}.png', format="png")
+    ## Test fancy indexing for different use cases
+    m, M = np.min(idx), np.max(idx)
+    def timer(arr, skip_flag=True, row=row, col=col):
+        time_list = []
+        if not skip_flag:
+            t = time.time()
+            b = arr[row, col]
+            time_list += [time.time() - t]
+            t = time.time()
+            b = arr[slice(1, M // 2, 5), col]
+            time_list += [time.time() - t]
+            t = time.time()
+            b = arr[[[m // 2, M // 2], [m // 4, M // 4]]]
+            time_list += [time.time() - t]
+        t = time.time()
+        b = arr[[m, M//2, M]]
+        time_list += [time.time() - t]
+        t = time.time()
+        b = arr[m, col]
+        time_list += [time.time() - t]
+        return np.array(time_list)
 
-        row = idx
-        col = rng.permutation(idx)
-        mask = rng.integers(low=0, high=2, size=(d,)) == 1
+    if NUMPY_BLOSC:
+        blosctimes += [timer(arr, skip_flag=False, row=idx, col=idx)]
+        arr=arr[:]
+        nptimes += [timer(arr, skip_flag=False, row=idx, col=idx)]
+    else:
+        blosctimes += [timer(arr)]
+        arr = arr[:]
+        nptimes += [timer(arr)]
+        z_test = zarr.zeros(shape=arr.shape, dtype=arr.dtype)
+        z_test[:] = arr
+        # zarr is more limited, as must provide same number of coord arrays as dims of array
+        # also cannot mix with slices
+        zarrtimes += [timer(z_test)]
+        with h5py.File('my_hdf5_file.h5', 'w') as f:
+            dset = f.create_dataset("init", data=arr)
+            h5pytimes += [timer(dset)]
 
-        ## Test fancy indexing for different use cases
-        loc_blosc_times = []
-        loc_np_times = []
-        m, M = np.min(idx), np.max(idx)
-        for i, fancyIdx in enumerate([
-            [m, M//2, M],  # i)
-            [[[m//2, M//2], [m//4, M//4]]],  # ii)
-            [row, col],  # iii)
-            # [row[:, None], col],  # iv)
-            [m, col],  # v)
-            [slice(1, M//2, 5), col],  # vi)
-            # [row[:, None], mask],  # vii)
-        ]):
-            # print(f'\n(case {i + 1})')
-            try:
-                r, c = fancyIdx
-                idx = (r, c)
-            except:
-                r, c = fancyIdx, None
-                idx = r
-            b, blosctime = fancyBlosc2(arr,idx)
-            n, nptime = fancyNumpy(arr,idx)
-            slice_size = np.prod(b.shape) * b.dtype.itemsize / 2 ** 30
-            np.testing.assert_allclose(b, n)
-            loc_blosc_times.append(blosctime)
-            loc_np_times.append(nptime)
-        blosc_times.append(loc_blosc_times)
-        np_times.append(loc_np_times)
+x = np.arange(len(sizes))
+width = 0.2
+blosctimes = np.array(blosctimes)
+nptimes = np.array(nptimes)
+if NUMPY_BLOSC:
+    # Create bars for axis 0 plot
+    for i, r in enumerate((["Numpy", nptimes, -width], ["Blosc2", blosctimes, 0])):
+        label, times, w = r
+        c = ['b', 'r'][i]
+        plt.bar(x + w, times.mean(axis=1), width, color=c, alpha=0.5)
+        plt.bar(x + w, times.max(axis=1), width, color=c, alpha=0.25)
+        plt.bar(x + w, times.min(axis=1), width, label=label, color=c, alpha=1)
 
-    blosc_times = np.array(blosc_times)
-    np_times = np.array(np_times)
-    x = np.arange(len(sizes))
-    width = 0.25
+    plt.xlabel('Array size (GB)')
+    plt.legend()
+    plt.xticks(x, np.round(sizes, 2))
+    plt.ylabel("Time (s)")
+    plt.title('Fancy indexing NumPy vs Blosc2')
+    plt.savefig('plots/fancyIdxNumpyVsBlosc.png', format="png")
+    plt.show()
+else:
+    zarrtimes = np.array(zarrtimes)
+    h5pytimes = np.array(h5pytimes)
 
     # Create bars for axis 0 plot
-    plt.bar(x - width, np_times.mean(axis=1), width, color='b', alpha=0.5)
-    plt.bar(x - width, np_times.max(axis=1), width, color='b', alpha=0.25)
-    plt.bar(x - width, np_times.min(axis=1), width, label='NumPy', color='b', alpha=1)
-    plt.bar(x, blosc_times.mean(axis=1), width, color='r', alpha=0.5)
-    plt.bar(x, blosc_times.max(axis=1), width, color='r', alpha=0.25)
-    plt.bar(x, blosc_times.min(axis=1), width, label='Blosc2', color='r', alpha=1)
-    plt.bar([],[], label='max', color='k', alpha=0.25)
-    plt.bar([],[], label='min', color='k', alpha=1)
-    plt.bar([],[], label='mean', color='k', alpha=0.5)
+    for i, r in enumerate((["Numpy",nptimes,-2*width],["Blosc2",blosctimes, -width],["Zarr",zarrtimes, 0],["HDF5",h5pytimes, width])):
+        label,times,w = r
+        c = ['b', 'r', 'g', 'm'][i]
+        plt.bar(x + w, times.mean(axis=1), width, color=c, alpha=0.5)
+        plt.bar(x + w, times.max(axis=1), width, color=c, alpha=0.25)
+        plt.bar(x + w, times.min(axis=1), width, label=label, color=c, alpha=1)
 
     plt.xlabel('Array size (GB)')
     plt.legend()
     plt.xticks(x, np.round(sizes, 2))
     plt.ylabel("Time (s)")
-    plt.title('Fancy indexing, Blosc2 vs NumPy')
+    plt.title('Fancy indexing performance comparison')
     plt.savefig('plots/fancyIdx.png', format="png")
     plt.show()
-    ## slowest cases are the ones with broadcasting
-    ## in fact it's a memory problem for blosc2 seemingly
+
+print("Finished everything!")

src/blosc2/blosc2_ext.pyx

@@ -513,6 +513,9 @@ cdef extern from "b2nd.h":
     int b2nd_get_orthogonal_selection(const b2nd_array_t *array, int64_t ** selection,
                                       int64_t *selection_size, void *buffer,
                                       int64_t *buffershape, int64_t buffersize)
+    int b2nd_set_orthogonal_selection(const b2nd_array_t *array, int64_t ** selection,
+                                      int64_t *selection_size, void *buffer,
+                                      int64_t *buffershape, int64_t buffersize)
     int b2nd_from_schunk(blosc2_schunk *schunk, b2nd_array_t **array)
 
     void blosc2_unidim_to_multidim(uint8_t ndim, int64_t *shape, int64_t i, int64_t *index)
@@ -2426,7 +2429,6 @@ cdef class NDArray:
         """
         Orthogonal indexing. Key is a tuple of lists of integer indices.
         """
-        ## have to maybe edit key to be C-compatible
         if len(key) != self.array.ndim:
             raise ValueError(f"Key must have {self.array.ndim} dimensions, got {len(key)}.")
         cdef int64_t[B2ND_MAX_DIM] buffershape_
@@ -2455,6 +2457,38 @@ cdef class NDArray:
         free(key_)
         return arr
 
+    def set_oindex_numpy(self, key, arr):
+        """
+        Orthogonal indexing. Set elements of self with arr using key.
+        """
+        if len(key) != self.array.ndim:
+            raise ValueError(f"Key must have {self.array.ndim} dimensions, got {len(key)}.")
+        cdef int64_t[B2ND_MAX_DIM] buffershape_
+        cdef int64_t** key_
+        cdef int64_t buffersize_ = self.array.sc.typesize
+        cdef int64_t[B2ND_MAX_DIM] sel_size
+
+        key_ = <int64_t**> malloc(len(key) * sizeof(int64_t *))
+
+        for i in range(self.array.ndim):
+            buffershape_[i] = len(key[i])
+            buffersize_ *= buffershape_[i]
+            sel_size[i] = len(key[i])
+            key_[i] = <int64_t *> malloc(sel_size[i] * sizeof(int64_t))
+            for j in range(len(key[i])):
+                key_[i][j] = key[i][j]
+
+        cdef Py_buffer buf
+        PyObject_GetBuffer(arr, &buf, PyBUF_SIMPLE)
+
+        _check_rc(b2nd_set_orthogonal_selection(self.array, key_, sel_size, buf.buf,
+                                                buffershape_, buffersize_), "Error while getting orthogonal selection")
+        PyBuffer_Release(&buf)
+        for i in range(len(key)):
+            free(key_[i])  # Free the allocated memory for each key
+        free(key_)
+        return arr
+
 
     def get_slice(self, key, mask, **kwargs):
         start, stop = key

src/blosc2/ndarray.py

@@ -1469,9 +1469,11 @@ def get_oselection_numpy(self, key):
         shape = tuple(len(k) for k in key) + self.shape[len(key) :]
         # Create the array to store the result
         arr = np.empty(shape, dtype=self.dtype)
-
         return super().get_oindex_numpy(arr, key)
 
+    def set_oselection_numpy(self, key, arr):
+        return super().set_oindex_numpy(key, arr)
+
     def __getitem__(  # noqa: C901
         self,
         key: int | slice | Sequence[slice | int] | np.ndarray[np.bool_] | NDArray | blosc2.LazyExpr | str,
@@ -4386,10 +4388,12 @@ class OIndex:
     def __init__(self, array: NDArray):
         self.array = array
 
-    # TODO: add __setitem__
     def __getitem__(self, selection) -> np.ndarray:
         return self.array.get_oselection_numpy(selection)
 
+    def __setitem__(self, selection, input) -> np.ndarray:
+        return self.array.set_oselection_numpy(selection, input)
+
 
 class VIndex:
     def __init__(self, array: NDArray):

tests/ndarray/test_ndarray.py

@@ -269,6 +269,7 @@ def test_save():
 
 
 def test_oindex():
+    # Test Get
     ndim = 3
     shape = (10,) * ndim
     arr = blosc2.linspace(0, 100, num=np.prod(shape), shape=shape, dtype="i4")