dask-array 0.1.0__py3-none-any.whl

dask_array/creation/_arange.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import functools
+from functools import partial
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task
+from dask_array._expr import ArrayExpr
+from dask_array._chunk import arange as _arange
+from dask_array._core_utils import normalize_chunks
+from dask_array._utils import meta_from_array
+
+
+class Arange(ArrayExpr):
+    _parameters = ["start", "stop", "step", "chunks", "like", "dtype", "kwargs"]
+    _defaults = {"chunks": "auto", "like": None, "dtype": None, "kwargs": None}
+
+    @functools.cached_property
+    def num_rows(self):
+        return int(max(np.ceil((self.stop - self.start) / self.step), 0))
+
+    @functools.cached_property
+    def dtype(self):
+        # Use type(x)(0) to determine dtype without overflow issues
+        # when start/stop are very large integers
+        dt = self.operand("dtype")
+        if dt is not None:
+            return np.dtype(dt)
+        return np.arange(type(self.start)(0), type(self.stop)(0), self.step).dtype
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.like, ndim=1, dtype=self.dtype)
+
+    @functools.cached_property
+    def chunks(self):
+        return normalize_chunks(self.operand("chunks"), (self.num_rows,), dtype=self.dtype)
+
+    def _layer(self) -> dict:
+        dsk = {}
+        elem_count = 0
+        start, step = self.start, self.step
+        like = self.like
+        func = partial(_arange, like=like)
+
+        for i, bs in enumerate(self.chunks[0]):
+            blockstart = start + (elem_count * step)
+            blockstop = start + ((elem_count + bs) * step)
+            task = Task(
+                (self._name, i),
+                func,
+                blockstart,
+                blockstop,
+                step,
+                bs,
+                self.dtype,
+            )
+            dsk[(self._name, i)] = task
+            elem_count += bs
+        return dsk
+
+
+_arange_sentinel = object()
+
+
+def arange(start=_arange_sentinel, stop=None, step=1, *, chunks="auto", like=None, dtype=None):
+    """
+    Return evenly spaced values from `start` to `stop` with step size `step`.
+
+    The values are half-open [start, stop), so including start and excluding
+    stop. This is basically the same as python's range function but for dask
+    arrays.
+
+    When using a non-integer step, such as 0.1, the results will often not be
+    consistent. It is better to use linspace for these cases.
+
+    Parameters
+    ----------
+    start : int, optional
+        The starting value of the sequence. The default is 0.
+    stop : int
+        The end of the interval, this value is excluded from the interval.
+    step : int, optional
+        The spacing between the values. The default is 1 when not specified.
+    chunks :  int
+        The number of samples on each block. Note that the last block will have
+        fewer samples if ``len(array) % chunks != 0``.
+        Defaults to "auto" which will automatically determine chunk sizes.
+    dtype : numpy.dtype
+        Output dtype. Omit to infer it from start, stop, step
+        Defaults to ``None``.
+    like : array type or ``None``
+        Array to extract meta from. Defaults to ``None``.
+
+    Returns
+    -------
+    samples : dask array
+
+    See Also
+    --------
+    dask.array.linspace
+    """
+    if start is _arange_sentinel:
+        if stop is None:
+            raise TypeError("arange() requires stop to be specified.")
+        # Only stop was provided as a keyword argument
+        start = 0
+    elif stop is None:
+        # Only start was provided, treat it as stop
+        stop = start
+        start = 0
+
+    # Avoid loss of precision calculating blockstart and blockstop
+    # when start is a very large int (~2**63) and step is a small float
+    if start != 0 and not np.isclose(start + step - start, step, atol=0):
+        r = arange(0, stop - start, step, chunks=chunks, dtype=dtype, like=like)
+        return r + start
+
+    return new_collection(Arange(start, stop, step, chunks, like, dtype))

dask_array/creation/_diag.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import functools
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task, TaskRef
+from dask_array._collection import asarray
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+from dask.utils import derived_from
+
+
+class Diag1D(ArrayExpr):
+    """Create a diagonal matrix from a 1D array (k=0 case only)."""
+
+    _parameters = ["x"]
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.x, ndim=2)
+
+    @functools.cached_property
+    def dtype(self):
+        return self.x.dtype
+
+    @functools.cached_property
+    def chunks(self):
+        chunks_1d = self.x.chunks[0]
+        return (chunks_1d, chunks_1d)
+
+    def _layer(self) -> dict:
+        dsk = {}
+        x = self.x
+        chunks_1d = x.chunks[0]
+        meta = self._meta
+
+        for i, m in enumerate(chunks_1d):
+            for j, n in enumerate(chunks_1d):
+                key = (self._name, i, j)
+                if i == j:
+                    dsk[key] = Task(key, np.diag, TaskRef((x._name, i)))
+                else:
+                    dsk[key] = Task(key, np.zeros_like, meta, shape=(m, n))
+        return dsk
+
+
+class Diag2DSimple(ArrayExpr):
+    """Extract diagonal from a 2D array with square chunks (k=0 case only)."""
+
+    _parameters = ["x"]
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.x, ndim=1)
+
+    @functools.cached_property
+    def dtype(self):
+        return self.x.dtype
+
+    @functools.cached_property
+    def chunks(self):
+        return (self.x.chunks[0],)
+
+    def _layer(self) -> dict:
+        dsk = {}
+        x = self.x
+        x_keys = x.__dask_keys__()
+
+        for i, row in enumerate(x_keys):
+            key = (self._name, i)
+            dsk[key] = Task(key, np.diag, TaskRef(row[i]))
+        return dsk
+
+
+@derived_from(np)
+def diag(v, k=0):
+    from dask_array._collection import Array
+
+    from ._diagonal import diagonal
+    from ._pad import pad
+
+    if not isinstance(v, np.ndarray) and not isinstance(v, Array):
+        raise TypeError(f"v must be a dask array or numpy array, got {type(v)}")
+
+    # Handle numpy arrays - wrap and return
+    if isinstance(v, np.ndarray) or (hasattr(v, "__array_function__") and not isinstance(v, Array)):
+        if v.ndim == 1:
+            abs(k)
+            result = np.diag(v, k)
+            return asarray(result)
+        elif v.ndim == 2:
+            result = np.diag(v, k)
+            return asarray(result)
+        else:
+            raise ValueError("Array must be 1d or 2d only")
+
+    v = asarray(v)
+
+    if v.ndim != 1:
+        if v.ndim != 2:
+            raise ValueError("Array must be 1d or 2d only")
+        # 2D case: extract diagonal
+        if k == 0 and v.chunks[0] == v.chunks[1]:
+            return new_collection(Diag2DSimple(v.expr))
+        else:
+            return diagonal(v, k)
+
+    # 1D case: create diagonal matrix
+    if k == 0:
+        return new_collection(Diag1D(v.expr))
+    elif k > 0:
+        return pad(diag(v), [[0, k], [k, 0]], mode="constant")
+    else:  # k < 0
+        return pad(diag(v), [[-k, 0], [0, -k]], mode="constant")

dask_array/creation/_diagonal.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import functools
+from functools import partial
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task, TaskRef
+from dask_array._collection import asarray
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+from dask.utils import derived_from
+
+
+class Diagonal(ArrayExpr):
+    """Extract a diagonal from a multi-dimensional array."""
+
+    _parameters = ["x", "offset", "axis1", "axis2"]
+    _defaults = {"offset": 0, "axis1": 0, "axis2": 1}
+
+    @functools.cached_property
+    def _axis1_normalized(self):
+        axis = self.axis1
+        if axis < 0:
+            axis = self.x.ndim + axis
+        return axis
+
+    @functools.cached_property
+    def _axis2_normalized(self):
+        axis = self.axis2
+        if axis < 0:
+            axis = self.x.ndim + axis
+        return axis
+
+    @functools.cached_property
+    def _effective_axes(self):
+        """Return (axis1, axis2, k) with axis1 < axis2."""
+        axis1, axis2 = self._axis1_normalized, self._axis2_normalized
+        k = self.offset
+        if axis1 > axis2:
+            axis1, axis2 = axis2, axis1
+            k = -self.offset
+        return axis1, axis2, k
+
+    @functools.cached_property
+    def _diag_info(self):
+        """Compute diagonal metadata."""
+        from itertools import product
+
+        x = self.x
+        axis1, axis2, k = self._effective_axes
+
+        kdiag_row_start = max(0, -k)
+        kdiag_col_start = max(0, k)
+        kdiag_row_stop = min(x.shape[axis1], x.shape[axis2] - k)
+        len_kdiag = kdiag_row_stop - kdiag_row_start
+
+        free_axes = set(range(x.ndim)) - {axis1, axis2}
+        free_indices = list(product(*(range(x.numblocks[i]) for i in free_axes)))
+        ndims_free = len(free_axes)
+
+        return {
+            "axis1": axis1,
+            "axis2": axis2,
+            "k": k,
+            "len_kdiag": len_kdiag,
+            "kdiag_row_start": kdiag_row_start,
+            "kdiag_col_start": kdiag_col_start,
+            "kdiag_row_stop": kdiag_row_stop,
+            "free_axes": free_axes,
+            "free_indices": free_indices,
+            "ndims_free": ndims_free,
+        }
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.x, ndim=self._diag_info["ndims_free"] + 1)
+
+    @functools.cached_property
+    def dtype(self):
+        return self.x.dtype
+
+    @functools.cached_property
+    def chunks(self):
+        info = self._diag_info
+        x = self.x
+        axis1, axis2 = info["axis1"], info["axis2"]
+
+        def pop_axes(chunks, axis1, axis2):
+            chunks = list(chunks)
+            chunks.pop(axis2)
+            chunks.pop(axis1)
+            return tuple(chunks)
+
+        if info["len_kdiag"] <= 0:
+            return pop_axes(x.chunks, axis1, axis2) + ((0,),)
+
+        # Compute diagonal chunks by following the diagonal through blocks
+        info["k"]
+        kdiag_row_start = info["kdiag_row_start"]
+        kdiag_col_start = info["kdiag_col_start"]
+
+        row_stops_ = np.cumsum(x.chunks[axis1])
+        row_starts = np.roll(row_stops_, 1)
+        row_starts[0] = 0
+
+        col_stops_ = np.cumsum(x.chunks[axis2])
+        col_starts = np.roll(col_stops_, 1)
+        col_starts[0] = 0
+
+        row_blockid = np.arange(x.numblocks[axis1])
+        col_blockid = np.arange(x.numblocks[axis2])
+
+        row_filter = (row_starts <= kdiag_row_start) & (kdiag_row_start < row_stops_)
+        col_filter = (col_starts <= kdiag_col_start) & (kdiag_col_start < col_stops_)
+        (I,) = row_blockid[row_filter]
+        (J,) = col_blockid[col_filter]
+
+        kdiag_chunks = ()
+        kdiag_r_start = kdiag_row_start
+        kdiag_c_start = kdiag_col_start
+        curr_I, curr_J = I, J
+
+        while kdiag_r_start < x.shape[axis1] and kdiag_c_start < x.shape[axis2]:
+            nrows, ncols = x.chunks[axis1][curr_I], x.chunks[axis2][curr_J]
+            local_r_start = kdiag_r_start - row_starts[curr_I]
+            local_c_start = kdiag_c_start - col_starts[curr_J]
+            local_k = -local_r_start if local_r_start > 0 else local_c_start
+            kdiag_row_end = min(nrows, ncols - local_k)
+            kdiag_len = kdiag_row_end - local_r_start
+            kdiag_chunks += (kdiag_len,)
+
+            kdiag_r_start = kdiag_row_end + row_starts[curr_I]
+            kdiag_c_start = min(ncols, nrows + local_k) + col_starts[curr_J]
+            curr_I = curr_I + 1 if kdiag_r_start == row_stops_[curr_I] else curr_I
+            curr_J = curr_J + 1 if kdiag_c_start == col_stops_[curr_J] else curr_J
+
+        return pop_axes(x.chunks, axis1, axis2) + (kdiag_chunks,)
+
+    def _layer(self) -> dict:
+        from dask_array._utils import is_cupy_type
+
+        dsk = {}
+        info = self._diag_info
+        x = self.x
+        axis1, axis2, _k = info["axis1"], info["axis2"], info["k"]
+        free_indices = info["free_indices"]
+        ndims_free = info["ndims_free"]
+
+        if info["len_kdiag"] <= 0:
+            xp = np
+            if is_cupy_type(x._meta):
+                import cupy
+
+                xp = cupy
+
+            out_chunks = self.chunks
+            for free_idx in free_indices:
+                shape = tuple(out_chunks[axis][free_idx[axis]] for axis in range(ndims_free))
+                key = (self._name,) + free_idx + (0,)
+                dsk[key] = Task(key, partial(xp.empty, dtype=x.dtype), shape + (0,))  # type: ignore[misc]
+            return dsk
+
+        # Follow k-diagonal through chunks
+        kdiag_row_start = info["kdiag_row_start"]
+        kdiag_col_start = info["kdiag_col_start"]
+
+        row_stops_ = np.cumsum(x.chunks[axis1])
+        row_starts = np.roll(row_stops_, 1)
+        row_starts[0] = 0
+
+        col_stops_ = np.cumsum(x.chunks[axis2])
+        col_starts = np.roll(col_stops_, 1)
+        col_starts[0] = 0
+
+        row_blockid = np.arange(x.numblocks[axis1])
+        col_blockid = np.arange(x.numblocks[axis2])
+
+        row_filter = (row_starts <= kdiag_row_start) & (kdiag_row_start < row_stops_)
+        col_filter = (col_starts <= kdiag_col_start) & (kdiag_col_start < col_stops_)
+        (I,) = row_blockid[row_filter]
+        (J,) = col_blockid[col_filter]
+
+        i = 0
+        kdiag_r_start = kdiag_row_start
+        kdiag_c_start = kdiag_col_start
+
+        while kdiag_r_start < x.shape[axis1] and kdiag_c_start < x.shape[axis2]:
+            nrows, ncols = x.chunks[axis1][I], x.chunks[axis2][J]
+            local_r_start = kdiag_r_start - row_starts[I]
+            local_c_start = kdiag_c_start - col_starts[J]
+            local_k = -local_r_start if local_r_start > 0 else local_c_start
+            kdiag_row_end = min(nrows, ncols - local_k)
+
+            for free_idx in free_indices:
+                input_idx = free_idx[:axis1] + (I,) + free_idx[axis1 : axis2 - 1] + (J,) + free_idx[axis2 - 1 :]
+                output_idx = free_idx + (i,)
+                key = (self._name,) + output_idx
+                dsk[key] = Task(
+                    key,
+                    np.diagonal,
+                    TaskRef((x._name,) + input_idx),
+                    local_k,
+                    axis1,
+                    axis2,
+                )
+
+            i += 1
+            kdiag_r_start = kdiag_row_end + row_starts[I]
+            kdiag_c_start = min(ncols, nrows + local_k) + col_starts[J]
+            I = I + 1 if kdiag_r_start == row_stops_[I] else I
+            J = J + 1 if kdiag_c_start == col_stops_[J] else J
+
+        return dsk
+
+
+@derived_from(np)
+def diagonal(a, offset=0, axis1=0, axis2=1):
+    from dask_array._numpy_compat import AxisError
+
+    if a.ndim < 2:
+        raise ValueError("diag requires an array of at least two dimensions")
+
+    def _axis_fmt(axis, name, ndim):
+        if axis < 0:
+            t = ndim + axis
+            if t < 0:
+                msg = "{}: axis {} is out of bounds for array of dimension {}"
+                raise AxisError(msg.format(name, axis, ndim))
+            axis = t
+        return axis
+
+    axis1_norm = _axis_fmt(axis1, "axis1", a.ndim)
+    axis2_norm = _axis_fmt(axis2, "axis2", a.ndim)
+
+    if axis1_norm == axis2_norm:
+        raise ValueError("axis1 and axis2 cannot be the same")
+
+    a = asarray(a)
+    return new_collection(Diagonal(a.expr, offset, axis1, axis2))

dask_array/creation/_eye.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import functools
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task
+from dask_array._expr import ArrayExpr
+from dask_array._core_utils import normalize_chunks
+
+
+class Eye(ArrayExpr):
+    _parameters = ["N", "M", "k", "dtype", "chunks"]
+    _defaults = {"M": None, "k": 0, "dtype": float, "chunks": "auto"}
+
+    @functools.cached_property
+    def _M(self):
+        return self.M if self.M is not None else self.N
+
+    @functools.cached_property
+    def dtype(self):
+        return np.dtype(self.operand("dtype") or float)
+
+    @functools.cached_property
+    def _meta(self):
+        return np.empty((0, 0), dtype=self.dtype)
+
+    @functools.cached_property
+    def chunks(self):
+        vchunks, hchunks = normalize_chunks(self.operand("chunks"), shape=(self.N, self._M), dtype=self.dtype)
+        return (vchunks, hchunks)
+
+    @functools.cached_property
+    def _chunk_size(self):
+        # Use the first vertical chunk size for diagonal positioning logic
+        return self.chunks[0][0]
+
+    def _layer(self) -> dict:
+        dsk = {}
+        vchunks, hchunks = self.chunks
+        chunk_size = self._chunk_size
+        k = self.k
+        dtype = self.dtype
+
+        for i, vchunk in enumerate(vchunks):
+            for j, hchunk in enumerate(hchunks):
+                key = (self._name, i, j)
+                # Check if this block contains part of the k-diagonal
+                if (j - i - 1) * chunk_size <= k <= (j - i + 1) * chunk_size:
+                    local_k = k - (j - i) * chunk_size
+                    task = Task(
+                        key,
+                        np.eye,
+                        vchunk,
+                        hchunk,
+                        local_k,
+                        dtype,
+                    )
+                else:
+                    task = Task(key, np.zeros, (vchunk, hchunk), dtype)
+                dsk[key] = task
+        return dsk
+
+
+def eye(N, chunks="auto", M=None, k=0, dtype=float):
+    """
+    Return a 2-D Array with ones on the diagonal and zeros elsewhere.
+
+    Parameters
+    ----------
+    N : int
+      Number of rows in the output.
+    chunks : int, str
+        How to chunk the array. Must be one of the following forms:
+
+        -   A blocksize like 1000.
+        -   A size in bytes, like "100 MiB" which will choose a uniform
+            block-like shape
+        -   The word "auto" which acts like the above, but uses a configuration
+            value ``array.chunk-size`` for the chunk size
+    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 : data-type, optional
+      Data-type of the returned array.
+
+    Returns
+    -------
+    I : Array of shape (N,M)
+      An array where all elements are equal to zero, except for the `k`-th
+      diagonal, whose values are equal to one.
+    """
+    if dtype is None:
+        dtype = float
+
+    if not isinstance(chunks, (int, str)):
+        raise ValueError("chunks must be an int or string")
+
+    return new_collection(Eye(N, M, k, dtype, chunks))

dask_array/creation/_linspace.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import functools
+from functools import partial
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Task
+from dask_array._chunk import linspace as _linspace
+
+from ._arange import Arange
+
+
+class Linspace(Arange):
+    _parameters = ["start", "stop", "num", "endpoint", "chunks", "dtype"]
+    _defaults = {"num": 50, "endpoint": True, "chunks": "auto", "dtype": None}
+    like = None
+
+    @functools.cached_property
+    def num_rows(self):
+        return self.operand("num")
+
+    @functools.cached_property
+    def dtype(self):
+        dt = self.operand("dtype")
+        if dt is not None:
+            return np.dtype(dt)
+        return np.linspace(0, 1, 1).dtype
+
+    @functools.cached_property
+    def step(self):
+        range_ = self.stop - self.start
+
+        div = (self.num_rows - 1) if self.endpoint else self.num_rows
+        if div == 0:
+            div = 1
+
+        return float(range_) / div
+
+    def _layer(self) -> dict:
+        dsk = {}
+        blockstart = self.start
+        func = partial(_linspace, endpoint=self.endpoint, dtype=self.dtype)
+
+        for i, bs in enumerate(self.chunks[0]):
+            bs_space = bs - 1 if self.endpoint else bs
+            blockstop = blockstart + (bs_space * self.step)
+            task = Task(
+                (self._name, i),
+                func,
+                blockstart,
+                blockstop,
+                bs,
+            )
+            blockstart = blockstart + (self.step * bs)
+            dsk[task.key] = task
+        return dsk
+
+
+def linspace(start, stop, num=50, endpoint=True, retstep=False, chunks="auto", dtype=None):
+    """
+    Return `num` evenly spaced values over the closed interval [`start`,
+    `stop`].
+
+    Parameters
+    ----------
+    start : scalar
+        The starting value of the sequence.
+    stop : scalar
+        The last value of the sequence.
+    num : int, optional
+        Number of samples to include in the returned dask array, including the
+        endpoints. Default is 50.
+    endpoint : bool, optional
+        If True, ``stop`` is the last sample. Otherwise, it is not included.
+        Default is True.
+    retstep : bool, optional
+        If True, return (samples, step), where step is the spacing between
+        samples. Default is False.
+    chunks :  int
+        The number of samples on each block. Note that the last block will have
+        fewer samples if `num % blocksize != 0`
+    dtype : dtype, optional
+        The type of the output array.
+
+    Returns
+    -------
+    samples : dask array
+    step : float, optional
+        Only returned if ``retstep`` is True. Size of spacing between samples.
+
+    See Also
+    --------
+    dask.array.arange
+    """
+    num = int(num)
+    result = new_collection(Linspace(start, stop, num, endpoint, chunks, dtype))
+    if retstep:
+        return result, result.expr.step
+    else:
+        return result