Incorporate expand_dims and add blosc2.stack

Author: Luke Shaw

CMakeLists.txt

@@ -50,7 +50,7 @@ else()
     include(FetchContent)
     FetchContent_Declare(blosc2
         GIT_REPOSITORY https://github.com/Blosc/c-blosc2
-        GIT_TAG 1c2f8bb0c914c43e23b751fbcf6642cd7aec09db  # v2.18.0 (concatenate added)
+        GIT_TAG 22b828b412391a8cbeeb97415777c5beecef73ff  # v2.18.1 (expand_dims added)
     )
     FetchContent_MakeAvailable(blosc2)
     include_directories("${blosc2_SOURCE_DIR}/include")

doc/reference/ndarray.rst

@@ -74,6 +74,7 @@ Constructors
     copy
     concatenate
     empty
+    expand_dims
     frombuffer
     fromiter
     nans

src/blosc2/__init__.py

@@ -239,6 +239,7 @@ class Tuner(Enum):
     reshape,
     copy,
     concatenate,
+    expand_dims,
     empty,
     frombuffer,
     fromiter,

src/blosc2/blosc2_ext.pyx

@@ -509,6 +509,7 @@ cdef extern from "b2nd.h":
     int b2nd_copy(b2nd_context_t *ctx, b2nd_array_t *src, b2nd_array_t **array)
     int b2nd_concatenate(b2nd_context_t *ctx, b2nd_array_t *src1, b2nd_array_t *src2,
                          int8_t axis, c_bool copy, b2nd_array_t **array)
+    int b2nd_expand_dims(const b2nd_array_t *array, b2nd_array_t ** view, const int8_t axis)
     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)
@@ -2892,3 +2893,14 @@ def concatenate(arr1: NDArray, arr2: NDArray, axis: int, **kwargs):
     else:
         # Return the first array, which now contains the concatenated data
         return arr1
+
+def expand_dims(arr1: NDArray, axis: int):
+    """
+    Add new dummy axis to NDArray object at specified dimension.
+    """
+    cdef b2nd_array_t *view
+    _check_rc(b2nd_expand_dims(arr1.array, &view, axis),
+              "Error while concatenating the arrays")
+
+    return blosc2.NDArray(_schunk=PyCapsule_New(view.sc, <char *> "blosc2_schunk*", NULL),
+                              _array=PyCapsule_New(view, <char *> "b2nd_array_t*", NULL))

src/blosc2/ndarray.py

@@ -3588,8 +3588,12 @@ def concatenate(arrays: list[NDArray], /, axis=0, **kwargs: Any) -> NDArray:  #
     return arr1
 
 
+def expand_dims(array: NDArray, axis=0) -> NDArray:
+    return blosc2_ext.expand_dims(array, axis=axis)
+
+
 def stack(arrays: list[NDArray], axis=0, **kwargs: Any) -> NDArray:
-    """Stack two arrays, creating a new  axis.
+    """Stack multiple arrays, creating a new axis.
 
     Parameters
     ----------
@@ -3621,8 +3625,8 @@ def stack(arrays: list[NDArray], axis=0, **kwargs: Any) -> NDArray:
     if axis < 0:
         axis += arrays[0].ndim + 1  # Adjust axis to be within the new stacked array's dimensions
     newarrays = []
-    for i, arr in enumerate(arrays):
-        newarrays[i] = blosc2_ext.expand_dims(arr, axis=axis)
+    for arr in arrays:
+        newarrays += [blosc2.expand_dims(arr, axis=axis)]
     return blosc2.concatenate(newarrays, axis, **kwargs)
 
 

tests/ndarray/test_concatenate.py

@@ -64,6 +64,22 @@ def test_concat3(shape1, shape2, shape3, dtype, axis):
     np.testing.assert_almost_equal(result[:], nparray)
 
 
+@pytest.mark.parametrize(
+    ("shape", "dtype", "axis"),
+    [
+        ([521], "i2", 0),
+        ([521, 121], "u4", 0),
+        ([52, 21], "i8", 1),
+        ([521, 121, 10], "f4", 0),
+        ([121, 521, 10], "f8", 1),
+        ([121, 121, 101], "i4", 2),
+        # 4-dimensional arrays
+        ([21, 121, 101, 10], "f4", 0),
+        ([121, 21, 101, 10], "i8", 1),
+        ([121, 121, 10, 10], "i8", 2),
+        ([121, 121, 101, 2], "i8", -1),
+    ],
+)
 def test_stack(shape, dtype, axis):
     ndarr1 = blosc2.arange(0, int(np.prod(shape)), 1, dtype=dtype, shape=shape)
     ndarr2 = blosc2.arange(0, int(np.prod(shape)), 1, dtype=dtype, shape=shape)

tests/ndarray/test_resize.py

@@ -29,3 +29,20 @@ def test_resize(shape, new_shape, chunks, blocks, fill_value):
     slices = tuple(slice(s) for s in shape)
     for i in np.nditer(a[slices]):
         assert i == fill_value
+
+
+@pytest.mark.parametrize(
+    ("shape", "axis", "chunks", "blocks", "fill_value"),
+    [
+        ((100, 1230), 1, (200, 100), (55, 3), b"0123"),
+        ((23, 34), 0, (20, 20), (10, 10), 1234),
+        ((80, 51, 60), -1, (20, 10, 33), (6, 6, 26), 3.333),
+    ],
+)
+def test_expand_dims(shape, axis, chunks, blocks, fill_value):
+    a = blosc2.full(shape, fill_value=fill_value, chunks=chunks, blocks=blocks)
+
+    b = blosc2.expand_dims(a, axis=axis)
+    npa = np.expand_dims(a, axis)
+    assert npa.shape == b.shape
+    np.testing.assert_array_equal(npa, b[:])