Add an eye constructor

Author: FrancescAlted

doc/reference/ndarray.rst

@@ -82,4 +82,5 @@ Constructors
     full
     arange
     linspace
+    eye
     reshape

examples/ndarray/eye-constructor.py

@@ -0,0 +1,45 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+# This example shows how to use the `eye()` constructor to create a blosc2 array.
+
+import math
+from time import time
+
+import numpy as np
+
+import blosc2
+
+N = 20_000
+
+shape = (N, N)
+print(f"*** Creating a blosc2 eye array with shape: {shape} ***")
+t0 = time()
+a = blosc2.eye(*shape, dtype=np.int8)
+cratio = a.schunk.nbytes / a.schunk.cbytes
+print(
+    f"Time: {time() - t0:.3f} s ({math.prod(shape) / (time() - t0) / 1e6:.2f} M/s)"
+    f"\tStorage required: {a.schunk.cbytes / 1e6:.2f} MB (cratio: {cratio:.2f}x)"
+)
+print(f"Last 3 elements:\n{a[-3:]}")
+
+# You can create rectangular arrays too
+shape = (N, N * 5)
+print(f"*** Creating a blosc2 eye array with shape: {shape} ***")
+t0 = time()
+a = blosc2.eye(*shape, dtype=np.int8)
+cratio = a.schunk.nbytes / a.schunk.cbytes
+print(
+    f"Time: {time() - t0:.3f} s ({math.prod(shape) / (time() - t0) / 1e6:.2f} M/s)"
+    f"\tStorage required: {a.schunk.cbytes / 1e6:.2f} MB (cratio: {cratio:.2f}x)"
+)
+print(f"First 3 elements:\n{a[:3]}")
+
+
+# In conclusion, you can use blosc2 eye() to create blosc2 arrays requiring much less storage
+# than numpy arrays.

src/blosc2/__init__.py

@@ -215,6 +215,7 @@ class Tuner(Enum):
     are_partitions_behaved,
     arange,
     linspace,
+    eye,
     asarray,
     indices,
     sort,

src/blosc2/lazyexpr.py

@@ -23,7 +23,7 @@
 from enum import Enum
 from pathlib import Path
 from queue import Empty, Queue
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from numpy.exceptions import ComplexWarning
 
@@ -197,7 +197,7 @@ def sort(self, order: str | list[str] | None = None) -> blosc2.LazyArray:
         pass
 
     @abstractmethod
-    def compute(self, item: slice | list[slice] | None = None, **kwargs: dict) -> blosc2.NDArray:
+    def compute(self, item: slice | list[slice] | None = None, **kwargs: Any) -> blosc2.NDArray:
         """
         Return an :ref:`NDArray` containing the evaluation of the :ref:`LazyArray`.
 
@@ -207,7 +207,7 @@ def compute(self, item: slice | list[slice] | None = None, **kwargs: dict) -> bl
             If not None, only the chunks that intersect with the slices
             in items will be evaluated.
 
-        kwargs: dict, optional
+        kwargs: Any, optional
             Keyword arguments that are supported by the :func:`empty` constructor.
             These arguments will be set in the resulting :ref:`NDArray`.
 
@@ -284,13 +284,13 @@ def __getitem__(self, item: int | slice | Sequence[slice]) -> blosc2.NDArray:
         pass
 
     @abstractmethod
-    def save(self, **kwargs: dict) -> None:
+    def save(self, **kwargs: Any) -> None:
         """
         Save the :ref:`LazyArray` on disk.
 
         Parameters
         ----------
-        kwargs: dict, optional
+        kwargs: Any, optional
             Keyword arguments that are supported by the :func:`empty` constructor.
             The `urlpath` must always be provided.
 
@@ -384,7 +384,7 @@ def info(self) -> InfoReporter:
 
 
 def convert_inputs(inputs):
-    if len(inputs) == 0:
+    if not inputs or len(inputs) == 0:
         return []
     inputs_ = []
     for obj in inputs:
@@ -595,7 +595,10 @@ def validate_inputs(inputs: dict, out=None) -> tuple:  # noqa: C901
                 "You really want to pass at least one input or one output for building a LazyArray."
                 "  Maybe you want blosc2.empty() instead?"
             )
-        return out.shape, out.chunks, out.blocks, True
+        if isinstance(out, blosc2.NDArray):
+            return out.shape, out.chunks, out.blocks, True
+        else:
+            return out.shape, None, None, True
 
     inputs = [input for input in inputs.values() if hasattr(input, "shape") and input is not np]
     shape = compute_broadcast_shape(inputs)
@@ -878,7 +881,7 @@ def fast_eval(  # noqa: C901
     getitem: bool, optional
         Indicates whether the expression is being evaluated for a getitem operation or eval().
         Default is False.
-    kwargs: dict, optional
+    kwargs: Any, optional
         Additional keyword arguments supported by the :func:`empty` constructor.
 
     Returns
@@ -892,6 +895,9 @@ def fast_eval(  # noqa: C901
     if isinstance(out, blosc2.NDArray):
         # If 'out' has been passed, and is a NDArray, use it as the base array
         basearr = out
+    elif isinstance(out, np.ndarray):
+        # If 'out' is a NumPy array, create a NDArray with the same shape and dtype
+        basearr = blosc2.empty(out.shape, dtype=out.dtype, **kwargs)
     else:
         # Otherwise, find the operand with the 'chunks' attribute and the longest shape
         operands_with_chunks = [o for o in operands.values() if hasattr(o, "chunks")]
@@ -1044,7 +1050,7 @@ def slices_eval(  # noqa: C901
     _slice: slice, list of slices, optional
         If provided, only the chunks that intersect with this slice
         will be evaluated.
-    kwargs: dict, optional
+    kwargs: Any, optional
         Additional keyword arguments that are supported by the :func:`empty` constructor.
 
     Returns
@@ -1293,7 +1299,7 @@ def reduce_slices(  # noqa: C901
     _slice: slice, list of slices, optional
         If provided, only the chunks that intersect with this slice
         will be evaluated.
-    kwargs: dict, optional
+    kwargs: Any, optional
         Additional keyword arguments supported by the :func:`empty` constructor.
 
     Returns
@@ -1530,7 +1536,7 @@ def chunked_eval(  # noqa: C901
         A dictionary containing the operands for the expression.
     item: int, slice or sequence of slices, optional
         The slice(s) to be retrieved. Note that step parameter is not honored yet.
-    kwargs: dict, optional
+    kwargs: Any, optional
         Additional keyword arguments supported by the :func:`empty` constructor.  In addition,
         the following keyword arguments are supported:
         _getitem: bool, optional
@@ -2680,11 +2686,11 @@ def save(self, **kwargs):
 
 def lazyudf(
     func: Callable[[tuple, np.ndarray, tuple[int]], None],
-    inputs: tuple | list,
+    inputs: tuple | list | None,
     dtype: np.dtype,
-    shape: tuple[int] | None = None,
+    shape: tuple | list | None = None,
     chunked_eval: bool = True,
-    **kwargs: dict,
+    **kwargs: Any,
 ) -> LazyUDF:
     """
     Get a LazyUDF from a python user-defined function.
@@ -2698,7 +2704,7 @@ def lazyudf(
         in :paramref:`inputs`.
         - `output`: The buffer to be filled as a multidimensional numpy.ndarray.
         - `offset`: The multidimensional offset corresponding to the start of the block being computed.
-    inputs: tuple or list
+    inputs: tuple or list or None
         The sequence of inputs. Supported inputs are:
         NumPy.ndarray, :ref:`NDArray`, :ref:`NDField`, :ref:`C2Array`.
         Any other object is supported too, and will be passed as is to the user-defined function.
@@ -2709,7 +2715,7 @@ def lazyudf(
         The shape of the resulting array. If None, the shape will be guessed from inputs.
     chunked_eval: bool, optional
         Whether to evaluate the function in chunks or not (blocks).
-    kwargs: dict, optional
+    kwargs: Any, optional
         Keyword arguments that are supported by the :func:`empty` constructor.
         These arguments will be used by the :meth:`LazyArray.__getitem__` and
         :meth:`LazyArray.eval` methods. The

src/blosc2/ndarray.py

@@ -3200,6 +3200,55 @@ def linspace_fill(inputs, output, offset):
     return reshape(larr, shape, c_order=c_order, **kwargs)
 
 
+def eye(N, M=None, k=0, dtype=np.float64, **kwargs: Any):
+    """Return a 2-D array with ones on the diagonal and zeros elsewhere.
+
+    Parameters
+    ----------
+    N: int
+        Number of rows in the output.
+    M: int, optional
+        Number of columns in the output. If None, defaults to `N`.
+    k: int, optional
+        Index of the diagonal: 0 (the default) refers to the main diagonal,
+        a positive value refers to an upper diagonal, and a negative value
+        to a lower diagonal.
+    dtype: np.dtype or str
+        The data type of the array elements in NumPy format. Default is `np.float64`.
+
+    Returns
+    -------
+    out: :ref:`NDArray`
+        A :ref:`NDArray` is returned.
+
+    Examples
+    --------
+    >>> import blosc2
+    >>> import numpy as np
+    >>> array = blosc2.eye(2, 3, dtype=np.int32)
+    >>> print(array[:])
+    [[1 0 0]
+     [0 1 0]]
+    """
+
+    def fill_eye(inputs, output: np.array, offset: tuple):
+        out_k = offset[0] - offset[1] + inputs[0]
+        output[:] = np.eye(*output.shape, out_k, dtype=output.dtype)
+
+    if M is None:
+        M = N
+    shape = (N, M)
+    # Check dtype
+    dtype = np.dtype(dtype)
+
+    if is_inside_new_expr():
+        # We already have the dtype and shape, so return immediately
+        return blosc2.zeros(shape, dtype=dtype)
+
+    lazyarr = blosc2.lazyudf(fill_eye, (k,), dtype=dtype, shape=shape)
+    return lazyarr.compute(**kwargs)
+
+
 def fromiter(iterable, shape, dtype, c_order=True, **kwargs) -> NDArray:
     """Create a new array from an iterable object.
 

tests/ndarray/test_ndarray.py

@@ -189,6 +189,18 @@ def test_linspace(ss, shape, dtype, chunks, blocks, endpoint, c_order):
         pass
 
 
+@pytest.mark.parametrize(("N", "M"), [(10, None), (10, 20), (20, 10)])
+@pytest.mark.parametrize("k", [-1, 0, 1, 2, 3])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int32])
+@pytest.mark.parametrize("chunks", [(5, 6), (10, 9)])
+def test_eye(k, N, M, dtype, chunks):
+    a = np.eye(N, M, k, dtype=dtype)
+    b = blosc2.eye(N, M, k, dtype=dtype, chunks=chunks)
+    assert a.shape == b.shape
+    assert a.dtype == b.dtype
+    np.testing.assert_allclose(a, b[:])
+
+
 @pytest.mark.parametrize(
     ("it", "shape", "dtype", "chunks", "blocks"),
     [