Remove refs to vindex and finalise bench and tests. Add oindex setitem.

Author: lshaw8317

bench/ndarray/fancy_index.py

@@ -12,130 +12,143 @@
 import ndindex
 import blosc2
 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})
+import pickle
+plt.rcParams.update({'text.usetex':False,'font.serif': ['cm'],'font.size':16})
 plt.rcParams['figure.dpi'] = 1000
 plt.rcParams['savefig.dpi'] = 1000
-plt.rc('text', usetex=True)
+plt.rc('text', usetex=False)
 plt.rc('font',**{'serif':['cm']})
 plt.style.use('seaborn-v0_8-paper')
 
-NUMPY_BLOSC = False
+NUMPY_BLOSC = True
+NUMPY_BLOSC_ZARR = False
 
 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()
-    arr = blosc2.ones(shape=shape, chunks=chunks, blocks=blocks, dtype=np.int64)  # , urlpath=file, mode="w")
+    arr = blosc2.ones(shape=shape, dtype=np.int64)
     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
 
-blosc_times = []
-np_times = []
-zarr_times = []
-sizes = []
-dims = np.int64(np.array([1, 2, 4, 6, 8]))
+
+sizes = np.int64(np.array([1, 2, 4, 8, 16]))
 rng = np.random.default_rng()
 blosctimes = []
 nptimes = []
 zarrtimes = []
 h5pytimes = []
-
-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
+x = np.arange(len(sizes))
+width = 0.2
+labs = 'NumpyBlosc2' if NUMPY_BLOSC else 'NumpyBlosc2ZarrHDF5' 
+labs = 'NumpyBlosc2Zarr' if NUMPY_BLOSC_ZARR else labs
+try:
+    with open(f"results{labs}.pkl", 'rb') as f:
+        result_tuple = pickle.load(f)
+    labs = ''
+    for i, r in enumerate(result_tuple):
+        if r[1].shape != (0,):
+            label,times,w = r
+            c = ['b', 'r', 'g', 'm'][i]
+            mean = times.mean(axis=1)
+            err = [mean - times.min(axis=1), times.max(axis=1)-mean]
+            plt.bar(x + w, mean , width, color=c, label=label, yerr=err, capsize=5, ecolor='k',
+            error_kw=dict(lw=2, capthick=2, ecolor='k'))
+            labs+=label
+except:
+    for d in sizes:
+        arr = genarray(d, ndims=2)
+        idx = rng.integers(low=0, high=arr.shape[0], size=(arr.shape[0],))
+        row = np.sort(np.unique(idx))
+        col = np.sort(np.unique(rng.integers(low=0, high=arr.shape[0], size=(arr.shape[0],))))
+        mask = rng.integers(low=0, high=2, size=(arr.shape[0],)) == 1
 
 
-    ## 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]
+        ## Test fancy indexing for different use cases
+        m, M = np.min(idx), np.max(idx)
+        def timer(arr, row=row, col=col):
+            time_list = []
+            if NUMPY_BLOSC or NUMPY_BLOSC_ZARR:
+                t = time.time()
+                b = arr[row, col]
+                time_list += [time.time() - t]
+            if NUMPY_BLOSC:
+                t = time.time()
+                b = arr[slice(1, M // 2, 5), col]
+                time_list += [time.time() - t]
+                t = time.time()
+                b = arr[[[row], [col]]]
+                time_list += [time.time() - t]
+                t = time.time()
+                b = arr[row[:10, None], col[:10]]
+                time_list += [time.time() - t]
+                t = time.time()
+                b = arr[row[:10, None], mask]
+                time_list += [time.time() - t]
             t = time.time()
-            b = arr[slice(1, M // 2, 5), col]
+            b = arr[row]
             time_list += [time.time() - t]
             t = time.time()
-            b = arr[[[m // 2, M // 2], [m // 4, M // 4]]]
+            b = arr[m, col]
             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)
+            return np.array(time_list)
 
-    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)]
+        if NUMPY_BLOSC or NUMPY_BLOSC_ZARR:
+            blosctimes += [timer(arr, row=idx, col=idx)]
+            arr=arr[:]
+            nptimes += [timer(arr, row=idx, col=idx)]
+            if NUMPY_BLOSC_ZARR:
+                z_test = zarr.zeros(shape=arr.shape, dtype=arr.dtype)
+                z_test[:] = arr
+                zarrtimes += [timer(z_test, 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
+            zarrtimes += [timer(z_test)]
+            with h5py.File('my_hdf5_file.h5', 'w') as f:
+                dset = f.create_dataset("init", data=arr)
+                h5pytimes += [timer(dset)]
 
-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)
-
-    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:
+    blosctimes = np.array(blosctimes)
+    nptimes = np.array(nptimes)
     zarrtimes = np.array(zarrtimes)
     h5pytimes = np.array(h5pytimes)
+    labs=''
+    result_tuple = (["Numpy",nptimes,-2*width],["Blosc2",blosctimes, -width],["Zarr",zarrtimes, 0],["HDF5",h5pytimes, width])
+
+    # Create barplot for Numpy vs Blosc vs Zarr vs H5py
+    for i, r in enumerate(result_tuple):
+        if r[1].shape != (0,):
+            label, times, w = r
+            c = ['b', 'r', 'g', 'm'][i]
+            mean = times.mean(axis=1)
+            err = (mean - times.min(axis=1), times.max(axis=1)-mean)
+            plt.bar(x + w, mean , width, color=c, label=label, yerr=err, capsize=5, ecolor='k', 
+            error_kw=dict(lw=2, capthick=2, ecolor='k'))            
+            labs+=label
 
-    # Create bars for axis 0 plot
-    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)
+    with open(f"results{labs}.pkl", 'wb') as f:
+        pickle.dump(result_tuple, f)
 
-    plt.xlabel('Array size (GB)')
-    plt.legend()
-    plt.xticks(x, np.round(sizes, 2))
-    plt.ylabel("Time (s)")
-    plt.title('Fancy indexing performance comparison')
-    plt.savefig('plots/fancyIdx.png', format="png")
-    plt.show()
+plt.xlabel('Array size (GB)')
+plt.legend()
+plt.xticks(x-width, np.round(sizes, 2))
+plt.ylabel("Time (s)")
+plt.title('Fancy indexing performance comparison')
+plt.ylim([0,10])
+plt.savefig(f'plots/fancyIdx{labs}.png', format="png")
+plt.show()
 
 print("Finished everything!")

src/blosc2/ndarray.py

@@ -1429,7 +1429,7 @@ def oindex(self) -> OIndex:
 
     @property
     def vindex(self) -> VIndex:
-        """Shortcut for fancy indexing, see :func:`get_vselection_numpy`"""
+        """Shortcut for vectorised indexing. Not yet supported."""
         return VIndex(self)
 
     @property
@@ -1439,7 +1439,7 @@ def T(self):
             raise ValueError("This property only works for 2-dimensional arrays.")
         return permute_dims(self)
 
-    def get_vselection_numpy(self, key):
+    def get_fselection_numpy(self, key):
         # TODO: Make this faster for broadcasted keys
         shape = self.shape
         if math.prod(shape) * self.dtype.itemsize < blosc2.MAX_FAST_PATH_SIZE:
@@ -1466,19 +1466,27 @@ def get_vselection_numpy(self, key):
         return out
 
     def get_oselection_numpy(self, key):
+        '''
+        Select independently from self along axes specified in key. Key must be same length as self shape.
+        See Zarr https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#orthogonal-indexing.
+        '''
         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):
+    def set_oselection_numpy(self, key, arr: np.ndarray):
+        '''
+        Select independently from self along axes specified in key and set to entries in arr. Key must be same length as self shape.
+        See Zarr https://zarr.readthedocs.io/en/stable/user-guide/arrays.html#orthogonal-indexing.
+        '''
         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,
     ) -> np.ndarray | blosc2.LazyExpr:
-        """Retrieve a (multidimensional) slice as specified by the key.
+        """Retrieve a (multidimensional) slice as specified by the key. Matches NumPy fancy indexing behaviour.
 
         Parameters
         ----------
@@ -1521,7 +1529,7 @@ def __getitem__(  # noqa: C901
             key_ = (slice(key, key + 1), *(slice(None),) * (self.ndim - 1))
         elif isinstance(key, tuple):
             if builtins.any(isinstance(k, (list, np.ndarray)) for k in key):
-                return self.get_vselection_numpy(key)  # fancy index
+                return self.get_fselection_numpy(key)  # fancy index
             if builtins.sum(isinstance(k, (list, builtins.slice)) for k in key) != self.ndim:
                 key_, mask = process_key(key, self.shape)
         elif isinstance(key, (list, np.ndarray, NDArray)):
@@ -1545,7 +1553,7 @@ def __getitem__(  # noqa: C901
             if key.dtype == np.int64 and key.ndim == 1 and self.ndim == 1:
                 return extract_values(self, key)
             else:
-                return self.get_vselection_numpy(key)  # fancy index
+                return self.get_fselection_numpy(key)  # fancy index
         else:
             # The more general case (this is quite slow)
             # If the key is a LazyExpr, decorate with ``where`` and return it
@@ -4399,6 +4407,9 @@ class VIndex:
     def __init__(self, array: NDArray):
         self.array = array
 
-    # TODO: add __setitem__
+    # TODO: all this
     def __getitem__(self, selection) -> np.ndarray:
-        return self.array.get_vselection_numpy(selection)
+        return NotImplementedError 
+
+    def __setitem__(self, selection, input) -> np.ndarray:
+        return NotImplementedError