Merge pull request #453 from Blosc/expand_dims

Extend functionality of expand_dims

Author: FrancescAlted

CMakeLists.txt

@@ -50,7 +50,7 @@ else()
     include(FetchContent)
     FetchContent_Declare(blosc2
         GIT_REPOSITORY https://github.com/Blosc/c-blosc2
-        GIT_TAG 110493ad34c3f9cc8bafd922a85333c0a46afd95  # v2.19.1
+        GIT_TAG 2d34ddb0b4832b58d68eae042da00c05a76b72fb  # v2.19.2dev0
     )
     FetchContent_MakeAvailable(blosc2)
     include_directories("${blosc2_SOURCE_DIR}/include")

examples/ndarray/animated_plot.py

@@ -33,9 +33,9 @@
 x = blosc2.linspace(0, n_frames, n_frames, dtype=dtype)
 y = blosc2.linspace(-4 * np.pi, 4 * np.pi, width, dtype=dtype)
 z = blosc2.linspace(-4 * np.pi, 4 * np.pi, height, dtype=dtype)
-X = blosc2.expand_dims(blosc2.expand_dims(x, 1), 2)  # Shape: (N, 1, 1)
-Y = blosc2.expand_dims(blosc2.expand_dims(y, 0), 2)  # Shape: (1, N, 1)
-Z = blosc2.expand_dims(blosc2.expand_dims(z, 0), 0)  # Shape: (1, 1, N)
+X = blosc2.expand_dims(x, (1, 2))  # Shape: (N, 1, 1)
+Y = blosc2.expand_dims(y, (0, 2))  # Shape: (1, N, 1)
+Z = blosc2.expand_dims(z, (0, 1))  # Shape: (1, 1, N)
 if not use_blosc2:
     # If not using Blosc2, convert to NumPy arrays
     # X, Y, Z = np.meshgrid(x, y, z)

src/blosc2/blosc2_ext.pyx

@@ -509,7 +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_expand_dims(const b2nd_array_t *array, b2nd_array_t ** view, const c_bool *axis, const uint8_t final_dims)
     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)
@@ -2971,13 +2971,17 @@ def concat(arr1: NDArray, arr2: NDArray, axis: int, **kwargs):
         # Return the first array, which now contains the concatenated data
         return arr1
 
-def expand_dims(arr1: NDArray, axis: int):
+def expand_dims(arr1: NDArray, axis_mask: list[bool], final_dims: int) -> blosc2.NDArray:
     """
     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")
+    cdef c_bool mask_[B2ND_MAX_DIM]
+    if final_dims > B2ND_MAX_DIM:
+        raise ValueError(f"Cannot expand dimensions beyond {B2ND_MAX_DIM} dimensions")
+    for i in range(final_dims):
+        mask_[i] = axis_mask[i]
+    _check_rc(b2nd_expand_dims(arr1.array, &view, mask_, final_dims),"Error while expanding 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

@@ -3730,13 +3730,33 @@ def concatenate(arrays: list[NDArray], /, axis=0, **kwargs: Any) -> NDArray:
 
 
 def expand_dims(array: NDArray, axis=0) -> NDArray:
+    """
+    Expand the shape of an array by adding new axes at the specified positions.
+
+    Parameters
+    ----------
+    array: :ref:`NDArray`
+        The array to be expanded.
+    axis: int or list of int, optional
+        Position in the expanded axes where the new axis (or axes) is placed. Default is 0.
+
+    Returns
+    -------
+    out: :ref:`NDArray`
+        A new NDArray with the expanded shape.
+    """
     if not isinstance(array, blosc2.NDArray):
         raise TypeError("Argument array must be instance of blosc2.NDArray")
-    if axis < 0:
-        axis += array.ndim + 1  # Adjust axis to be within the new stacked array's dimensions
-    if axis > array.ndim:
-        raise ValueError(f"Axis {axis} is out of bounds for expanded array of dimension {array.ndim + 1}.")
-    return blosc2_ext.expand_dims(array, axis=axis)
+    axis = [axis] if isinstance(axis, int) else axis
+    final_dims = array.ndim + len(axis)
+    mask = [False for i in range(final_dims)]
+    for a in axis:
+        if a < 0:
+            a += final_dims  # Adjust axis to be within the new stacked array's dimensions
+        if mask[a]:
+            raise ValueError("Axis values must be unique.")
+        mask[a] = True
+    return blosc2_ext.expand_dims(array, axis_mask=mask, final_dims=final_dims)
 
 
 def stack(arrays: list[NDArray], axis=0, **kwargs: Any) -> NDArray:

tests/ndarray/test_resize.py

@@ -37,7 +37,7 @@ def test_resize(shape, new_shape, chunks, blocks, fill_value):
         ((0,), 1, (0,), (0,), 1),
         ((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),
+        ((80, 51, 60), (-1, -2, 1), (20, 10, 33), (6, 6, 26), 3.333),
     ],
 )
 def test_expand_dims(shape, axis, chunks, blocks, fill_value):
@@ -49,6 +49,8 @@ def test_expand_dims(shape, axis, chunks, blocks, fill_value):
     np.testing.assert_array_equal(npa, b[:])
 
     # Repeated expansion
+    axis = (axis,) if isinstance(axis, int) else axis
+    axis = axis[0] if (len(axis) + b.ndim) > blosc2.MAX_DIM else axis
     b = blosc2.expand_dims(b, axis=axis)
     npa = np.expand_dims(npa, axis)
     assert npa.shape == b.shape