Merge pull request #518 from Blosc/squeeze_view

Squeeze view

Author: lshaw8317

CMakeLists.txt

@@ -50,7 +50,7 @@ else()
     include(FetchContent)
     FetchContent_Declare(blosc2
         GIT_REPOSITORY https://github.com/Blosc/c-blosc2
-        GIT_TAG 0c853a639ba97997e33e29db9eed202459ebc6f0  # v2.21.3
+        GIT_TAG 77ab4b583eb7e03a6f2a411e1a728500e269540c  # v2.22.0
     )
     FetchContent_MakeAvailable(blosc2)
     include_directories("${blosc2_SOURCE_DIR}/include")

src/blosc2/blosc2_ext.pyx

@@ -503,8 +503,8 @@ cdef extern from "b2nd.h":
     int b2nd_to_cframe(const b2nd_array_t *array, uint8_t ** cframe, int64_t *cframe_len,
                        c_bool *needs_free);
 
-    int b2nd_squeeze(b2nd_array_t *array)
-    int b2nd_squeeze_index(b2nd_array_t *array, const c_bool *index)
+    int b2nd_squeeze(b2nd_array_t *array, b2nd_array_t **view)
+    int b2nd_squeeze_index(b2nd_array_t *array, b2nd_array_t **view, const c_bool *index)
     int b2nd_resize(b2nd_array_t *array, const int64_t *new_shape, const int64_t *start)
     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,
@@ -2530,7 +2530,7 @@ cdef class NDArray:
         cdef c_bool mask_[B2ND_MAX_DIM]
         for i in range(ndim):
             mask_[i] = mask[i]
-        _check_rc(b2nd_squeeze_index(array, mask_), "Error while squeezing sliced array")
+        _check_rc(b2nd_squeeze_index(array, &array, mask_), "Error while squeezing sliced array")
         ndarray = blosc2.NDArray(_schunk=PyCapsule_New(array.sc, <char *> "blosc2_schunk*", NULL),
                                  _array=PyCapsule_New(array, <char *> "b2nd_array_t*", NULL))
 
@@ -2604,19 +2604,6 @@ cdef class NDArray:
         _check_rc(b2nd_resize(self.array, new_shape_, NULL),
                   "Error while resizing the array")
 
-    def squeeze(self, mask=None):
-        cdef c_bool mask_[B2ND_MAX_DIM]
-        if mask is None:
-            _check_rc(b2nd_squeeze(self.array), "Error while performing the squeeze")
-        else:
-            for i in range(self.ndim):
-                mask_[i] = mask[i]
-            _check_rc(b2nd_squeeze_index(self.array, mask_), "Error while squeezing array")
-
-        #this squeezes even if not asked for by mask - may have to use in future though
-        #if self.array.shape[0] == 1 and self.ndim == 1:
-        #    self.array.ndim = 0
-
     def as_ffi_ptr(self):
         return PyCapsule_New(self.array, <char *> "b2nd_array_t*", NULL)
 
@@ -3001,3 +2988,22 @@ def expand_dims(arr1: NDArray, axis_mask: list[bool], final_dims: int) -> blosc2
     new_base = arr1 if arr1.base is None else arr1.base
     return blosc2.NDArray(_schunk=PyCapsule_New(view.sc, <char *> "blosc2_schunk*", NULL),
                           _array=PyCapsule_New(view, <char *> "b2nd_array_t*", NULL), _base=new_base)
+
+def squeeze(arr1: NDArray, axis_mask: list[bool]) -> blosc2.NDArray:
+    """
+    Remove axis from NDArray object at specified dimensions.
+    """
+    cdef b2nd_array_t *view
+    cdef c_bool mask_[B2ND_MAX_DIM]
+    for i in range(arr1.ndim):
+        mask_[i] = axis_mask[i]
+    _check_rc(b2nd_squeeze_index(arr1.array, &view, mask_), "Error while squeezing array")
+
+    # this squeezes even if not asked for by mask - may have to use in future though
+    # if arr1.array.shape[0] == 1 and arr1.ndim == 1:
+    #     arr1.array.ndim = 0
+
+    # create view with reference to self to hold onto
+    new_base = arr1 if arr1.base is None else arr1.base
+    return blosc2.NDArray(_schunk=PyCapsule_New(view.sc, <char *> "blosc2_schunk*", NULL),
+                        _array=PyCapsule_New(view, <char *> "b2nd_array_t*", NULL), _base=new_base)

src/blosc2/linalg.py

@@ -127,9 +127,9 @@ def matmul(x1: blosc2.Array, x2: blosc2.NDArray, **kwargs: Any) -> blosc2.NDArra
                         result[chunk + (slice(row, row_end), slice(col, col_end))] += np.matmul(bx1, bx2)
 
     if x1_is_vector:
-        result.squeeze(axis=-2)
+        result = result.squeeze(axis=-2)
     if x2_is_vector:
-        result.squeeze(axis=-1)
+        result = result.squeeze(axis=-1)
 
     return result
 

src/blosc2/ndarray.py

@@ -4681,8 +4681,7 @@ def slice(self, key: int | slice | Sequence[slice], **kwargs: Any) -> NDArray:
                 for order, nchunk in enumerate(aligned_chunks):
                     chunk = self.schunk.get_chunk(nchunk)
                     newarr.schunk.update_chunk(order, chunk)
-                newarr.squeeze(axis=np.where(mask)[0])  # remove any dummy dims introduced
-                return newarr
+                return newarr.squeeze(axis=np.where(mask)[0])  # remove any dummy dims introduced
 
         key = (start, stop)
         ndslice = super().get_slice(key, mask, **kwargs)
@@ -4724,16 +4723,7 @@ def squeeze(self, axis: int | Sequence[int]) -> NDArray:
         >>> a.shape
         (23, 11)
         """
-        axis = [axis] if isinstance(axis, int) else axis
-        mask = [False for i in range(self.ndim)]
-        for a in axis:
-            if a < 0:
-                a += self.ndim  # Adjust axis to be within the array's dimensions
-            if mask[a]:
-                raise ValueError("Axis values must be unique.")
-            mask[a] = True
-        super().squeeze(mask=mask)
-        return self
+        return blosc2.squeeze(self, axis=axis)
 
     def indices(self, order: str | list[str] | None = None, **kwargs: Any) -> NDArray:
         """
@@ -4797,8 +4787,15 @@ def squeeze(x: Array, axis: int | Sequence[int]) -> NDArray:
     >>> b.shape
     (23, 11)
     """
-    # TODO: implement squeeze as a view
-    return x.squeeze(axis)
+    axis = [axis] if isinstance(axis, int) else axis
+    mask = [False for i in range(x.ndim)]
+    for a in axis:
+        if a < 0:
+            a += x.ndim  # Adjust axis to be within the array's dimensions
+        if mask[a]:
+            raise ValueError("Axis values must be unique.")
+        mask[a] = True
+    return blosc2_ext.squeeze(x, axis_mask=mask)
 
 
 def array_from_ffi_ptr(array_ptr) -> NDArray:

tests/ndarray/test_lazyexpr.py

@@ -1714,9 +1714,6 @@ def test_lazylinalg():
     npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape
     np.testing.assert_array_almost_equal(out[()], npres)
-    # refresh D since squeeze is in-place
-    s = shapes["D"]
-    D = blosc2.linspace(0, np.prod(s), shape=s)
     out = blosc2.lazyexpr("D.squeeze(axis=-1)")
     npres = np.squeeze(npD, -1)
     assert out.shape == npres.shape

tests/ndarray/test_squeeze.py

@@ -25,8 +25,8 @@ def test_squeeze(shape, chunks, blocks, fill_value, axis):
     a = blosc2.full(shape, fill_value=fill_value, chunks=chunks, blocks=blocks)
 
     b = np.squeeze(a[...], axis)
-    a_ = a.squeeze(axis)
+    a_ = blosc2.squeeze(a, axis)
 
     assert a_.shape == b.shape
-    # TODO: this would work if squeeze returns a view
-    # assert a_.shape != a.shape
+    # Confirm squeeze returns a view (does not modify original array)
+    assert a_.shape != a.shape