dask-array 0.1.0__py3-none-any.whl

dask_array/reductions/_trace.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+
+def trace(a, offset=0, axis1=0, axis2=1, dtype=None):
+    """
+    Return the sum along diagonals of the array.
+
+    This docstring was copied from numpy.trace.
+
+    Some inconsistencies with the Dask version may exist.
+
+    If `a` is 2-D, the sum along its diagonal with the given offset
+    is returned, i.e., the sum of elements ``a[i,i+offset]`` for all i.
+
+    If `a` has more than two dimensions, then the axes specified by axis1 and
+    axis2 are used to determine the 2-D sub-arrays whose traces are returned.
+    The shape of the resulting array is the same as that of `a` with `axis1`
+    and `axis2` removed.
+
+    Parameters
+    ----------
+    a : array_like
+        Input array, from which the diagonals are taken.
+    offset : int, optional
+        Offset of the diagonal from the main diagonal. Can be both positive
+        and negative. Defaults to 0.
+    axis1, axis2 : int, optional
+        Axes to be used as the first and second axis of the 2-D sub-arrays
+        from which the diagonals should be taken. Defaults are the first two
+        axes of `a`.
+    dtype : dtype, optional
+        Determines the data-type of the returned array and of the accumulator
+        where the elements are summed. If dtype has the value None and `a` is
+        of integer type of precision less than the default integer precision,
+        then the default integer precision is used. Otherwise, the precision
+        is the same as that of `a`.
+
+    Returns
+    -------
+    sum_along_diagonals : ndarray
+        If `a` is 2-D, the sum along the diagonal is returned.  If `a` has
+        larger dimensions, then an array of sums along diagonals is returned.
+
+    See Also
+    --------
+    diag, diagonal, diagflat
+
+    Examples
+    --------
+    >>> import dask_array as da
+    >>> da.trace(da.eye(3)).compute()  # doctest: +SKIP
+    3.0
+    """
+    from dask_array.creation import diagonal
+
+    return diagonal(a, offset=offset, axis1=axis1, axis2=axis2).sum(-1, dtype=dtype)

dask_array/routines/__init__.py

@@ -0,0 +1,133 @@
+"""Array routines for array-expr."""
+
+import numpy as np
+
+from dask.utils import derived_from
+
+# Direct imports from submodules
+# Re-exports from other modules
+from dask_array._blockwise import outer  # noqa: F401
+from dask_array._collection import asanyarray, asarray
+from dask_array._ufunc import (  # noqa: F401
+    allclose,
+    around,
+    isclose,
+    isnull,
+    notnull,
+    round,
+)
+from dask_array.routines._apply import apply_along_axis, apply_over_axes
+from dask_array.routines._bincount import bincount
+from dask_array.routines._broadcast import broadcast_arrays, unify_chunks
+from dask_array.routines._coarsen import aligned_coarsen_chunks, coarsen
+from dask_array.routines._diff import diff
+from dask_array.routines._gradient import gradient
+from dask_array.routines._indexing import ravel_multi_index, unravel_index
+from dask_array.routines._insert_delete import (
+    append,
+    delete,
+    ediff1d,
+    insert,
+)
+from dask_array.routines._misc import (
+    compress,
+    ndim,
+    result_type,
+    shape,
+    take,
+)
+from dask_array.routines._nonzero import (
+    argwhere,
+    count_nonzero,
+    flatnonzero,
+    isnonzero,
+    nonzero,
+)
+from dask_array.routines._search import isin, searchsorted
+from dask_array.routines._select import (
+    choose,
+    digitize,
+    extract,
+    piecewise,
+    select,
+)
+from dask_array.routines._statistics import average, corrcoef, cov
+from dask_array.routines._topk import argtopk, topk
+from dask_array.routines._triangular import (
+    tril,
+    tril_indices,
+    tril_indices_from,
+    triu,
+    triu_indices,
+    triu_indices_from,
+)
+from dask_array.routines._unique import union1d, unique
+from dask_array.routines._where import where
+
+
+@derived_from(np)
+def array(x, dtype=None, ndmin=None, *, like=None):
+    x = asarray(x, like=like)
+    while ndmin is not None and x.ndim < ndmin:
+        x = x[None, :]
+    if dtype is not None and x.dtype != dtype:
+        x = x.astype(dtype)
+    return x
+
+
+__all__ = [
+    "aligned_coarsen_chunks",
+    "allclose",
+    "append",
+    "apply_along_axis",
+    "array",
+    "apply_over_axes",
+    "argwhere",
+    "argtopk",
+    "around",
+    "average",
+    "bincount",
+    "broadcast_arrays",
+    "choose",
+    "coarsen",
+    "compress",
+    "corrcoef",
+    "count_nonzero",
+    "cov",
+    "delete",
+    "diff",
+    "digitize",
+    "ediff1d",
+    "extract",
+    "flatnonzero",
+    "gradient",
+    "insert",
+    "isclose",
+    "isin",
+    "isnonzero",
+    "isnull",
+    "ndim",
+    "nonzero",
+    "notnull",
+    "outer",
+    "piecewise",
+    "ravel_multi_index",
+    "result_type",
+    "round",
+    "searchsorted",
+    "select",
+    "shape",
+    "take",
+    "topk",
+    "tril",
+    "tril_indices",
+    "tril_indices_from",
+    "triu",
+    "triu_indices",
+    "triu_indices_from",
+    "unify_chunks",
+    "union1d",
+    "unique",
+    "unravel_index",
+    "where",
+]

dask_array/routines/_apply.py

@@ -0,0 +1,84 @@
+"""Apply functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask.utils import derived_from, funcname
+
+
+def _inner_apply_along_axis(arr, func1d, func1d_axis, func1d_args, func1d_kwargs):
+    return np.apply_along_axis(func1d, func1d_axis, arr, *func1d_args, **func1d_kwargs)
+
+
+@derived_from(np)
+def apply_along_axis(func1d, axis, arr, *args, dtype=None, shape=None, **kwargs):
+    """
+    This is a blocked variant of :func:`numpy.apply_along_axis` implemented via
+    :func:`dask.array.map_blocks`
+
+    Notes
+    -----
+    If either of `dtype` or `shape` are not provided, Dask attempts to
+    determine them by calling `func1d` on a dummy array. This may produce
+    incorrect values for `dtype` or `shape`, so we recommend providing them.
+    """
+    arr = asarray(arr)
+
+    # Verify that axis is valid and throw an error otherwise
+    axis = len(arr.shape[:axis])
+
+    # If necessary, infer dtype and shape of the output of func1d by calling it on test data.
+    if shape is None or dtype is None:
+        test_data = np.ones((1,), dtype=arr.dtype)
+        test_result = np.array(func1d(test_data, *args, **kwargs))
+        if shape is None:
+            shape = test_result.shape
+        if dtype is None:
+            dtype = test_result.dtype
+
+    # Rechunk so that func1d is applied over the full axis.
+    arr = arr.rechunk(arr.chunks[:axis] + (arr.shape[axis : axis + 1],) + arr.chunks[axis + 1 :])
+
+    # Map func1d over the data to get the result
+    # Adds other axes as needed.
+    result = arr.map_blocks(
+        _inner_apply_along_axis,
+        name=funcname(func1d) + "-along-axis",
+        dtype=dtype,
+        chunks=(arr.chunks[:axis] + shape + arr.chunks[axis + 1 :]),
+        drop_axis=axis,
+        new_axis=list(range(axis, axis + len(shape), 1)),
+        func1d=func1d,
+        func1d_axis=axis,
+        func1d_args=args,
+        func1d_kwargs=kwargs,
+    )
+
+    return result
+
+
+@derived_from(np)
+def apply_over_axes(func, a, axes):
+    """Apply a function repeatedly over multiple axes."""
+    a = asarray(a)
+    try:
+        axes = tuple(axes)
+    except TypeError:
+        axes = (axes,)
+
+    sl = a.ndim * (slice(None),)
+
+    # Compute using `apply_along_axis`.
+    result = a
+    for i in axes:
+        result = apply_along_axis(func, i, result, 0)
+
+        # Restore original dimensionality or error.
+        if result.ndim == (a.ndim - 1):
+            result = result[sl[:i] + (None,)]
+        elif result.ndim != a.ndim:
+            raise ValueError("func must either preserve dimensionality of the input or reduce it by one.")
+
+    return result

dask_array/routines/_bincount.py

@@ -0,0 +1,112 @@
+"""Bincount implementation for array-expr."""
+
+from __future__ import annotations
+
+from functools import cached_property, partial
+
+import numpy as np
+
+from dask._task_spec import Task, TaskRef
+from dask_array._collection import asarray, new_collection
+from dask_array._expr import ArrayExpr
+from dask.utils import derived_from
+
+
+def _bincount_chunk(x, weights, minlength):
+    """Apply bincount to a single chunk, wrapping result in extra dimension."""
+    if weights is not None:
+        result = np.bincount(x, weights=weights, minlength=minlength)
+    else:
+        result = np.bincount(x, minlength=minlength)
+    return result[np.newaxis]
+
+
+def _bincount_sum(bincounts, axis, keepdims, dtype=None):
+    """Sum bincount results, handling variable lengths when minlength=0."""
+    if not isinstance(bincounts, list):
+        bincounts = [bincounts]
+
+    n = max(b.shape[1] for b in bincounts)
+    out = np.zeros((1, n), dtype=bincounts[0].dtype)
+    for b in bincounts:
+        out[0, : b.shape[1]] += b[0]
+
+    if not keepdims:
+        return out[0]
+    return out
+
+
+class BincountChunked(ArrayExpr):
+    """Expression for per-chunk bincount computation."""
+
+    _parameters = ["x", "weights", "minlength", "output_size", "meta_provided"]
+    _defaults = {"weights": None}
+
+    @cached_property
+    def _meta(self):
+        return np.empty((0, 0), dtype=self.meta_provided.dtype)
+
+    @cached_property
+    def chunks(self):
+        nchunks = len(self.x.chunks[0])
+        return ((1,) * nchunks, (self.output_size,))
+
+    @cached_property
+    def _name(self):
+        return f"bincount-{self.deterministic_token}"
+
+    def _layer(self):
+        dsk = {}
+        minlen = self.minlength
+        for i in range(len(self.x.chunks[0])):
+            key = (self._name, i, 0)
+            x_ref = TaskRef((self.x._name, i))
+            w_ref = TaskRef((self.weights._name, i)) if self.weights is not None else None
+            dsk[key] = Task(key, _bincount_chunk, x_ref, w_ref, minlen)
+        return dsk
+
+    @property
+    def _dependencies(self):
+        deps = [self.x]
+        if self.weights is not None:
+            deps.append(self.weights)
+        return deps
+
+
+@derived_from(np)
+def bincount(x, weights=None, minlength=0, split_every=None):
+    """Count number of occurrences of each value in array of non-negative ints."""
+    from dask_array.reductions import _tree_reduce
+
+    x = asarray(x)
+    if x.ndim != 1:
+        raise ValueError("Input array must be one dimensional. Try using x.ravel()")
+    if weights is not None:
+        weights = asarray(weights)
+        if weights.chunks != x.chunks:
+            raise ValueError("Chunks of input array x and weights must match.")
+
+    if weights is not None:
+        meta = np.bincount([1], weights=np.array([1], dtype=weights.dtype))
+    else:
+        meta = np.bincount([])
+
+    if minlength == 0:
+        output_size = np.nan
+    else:
+        output_size = minlength
+
+    chunked_counts = new_collection(BincountChunked(x, weights, minlength, output_size, meta_provided=meta))
+
+    if minlength > 0:
+        return chunked_counts.sum(axis=0)
+    else:
+        return _tree_reduce(
+            chunked_counts,
+            aggregate=partial(_bincount_sum, dtype=meta.dtype),
+            axis=(0,),
+            keepdims=False,
+            dtype=meta.dtype,
+            split_every=split_every,
+            concatenate=False,
+        )

dask_array/routines/_broadcast.py

@@ -0,0 +1,111 @@
+"""Broadcasting utilities for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asanyarray, asarray
+from dask.utils import derived_from
+
+
+def unify_chunks(*args, **kwargs):
+    """
+    Unify chunks across a sequence of arrays
+
+    This utility function is used within other common operations like
+    :func:`dask.array.core.map_blocks` and :func:`dask.array.core.blockwise`.
+    It is not commonly used by end-users directly.
+
+    Parameters
+    ----------
+    *args: sequence of Array, index pairs
+        Sequence like (x, 'ij', y, 'jk', z, 'i')
+
+    Examples
+    --------
+    >>> import dask_array as da
+    >>> x = da.ones(10, chunks=((5, 2, 3),))
+    >>> y = da.ones(10, chunks=((2, 3, 5),))
+    >>> chunkss, arrays = unify_chunks(x, 'i', y, 'i')
+    >>> chunkss
+    {'i': (2, 3, 2, 3)}
+
+    Returns
+    -------
+    chunkss : dict
+        Map like {index: chunks}.
+    arrays : list
+        List of rechunked arrays.
+    """
+    from toolz import partition
+
+    from dask_array._new_collection import new_collection
+    from dask_array._expr import unify_chunks_expr
+
+    if not args:
+        return {}, []
+
+    arginds = [(asanyarray(a) if ind is not None else a, ind) for a, ind in partition(2, args)]
+
+    arrays, inds = zip(*arginds)
+    if all(ind is None for ind in inds):
+        return {}, list(arrays)
+
+    # Convert to expression-level args
+    expr_args = []
+    for a, ind in arginds:
+        if ind is not None:
+            expr_args.extend([a.expr, ind])
+        else:
+            expr_args.extend([a, ind])
+
+    warn = kwargs.pop("warn", True)
+    if kwargs:
+        raise TypeError(f"Unexpected keyword arguments: {kwargs}")
+    chunkss, expr_arrays, _ = unify_chunks_expr(*expr_args, warn=warn)
+
+    # Convert back to collections
+    result_arrays = []
+    for a, orig_a_ind in zip(expr_arrays, arginds):
+        orig_a, ind = orig_a_ind
+        if ind is None:
+            result_arrays.append(orig_a)
+        else:
+            result_arrays.append(new_collection(a))
+
+    return chunkss, result_arrays
+
+
+@derived_from(np)
+def broadcast_arrays(*args, subok=False):
+    """Broadcast any number of arrays against each other."""
+    from toolz import concat
+
+    from dask_array._collection import broadcast_to
+    from dask_array._core_utils import broadcast_chunks, broadcast_shapes
+    from dask_array._numpy_compat import NUMPY_GE_200
+
+    subok = bool(subok)
+
+    to_array = asanyarray if subok else asarray
+    args = tuple(to_array(e) for e in args)
+
+    if not args:
+        if NUMPY_GE_200:
+            return ()
+        return []
+
+    # Unify uneven chunking
+    inds = [list(reversed(range(x.ndim))) for x in args]
+    uc_args = list(concat(zip(args, inds)))
+    _, args = unify_chunks(*uc_args, warn=False)
+
+    shape = broadcast_shapes(*(e.shape for e in args))
+    chunks = broadcast_chunks(*(e.chunks for e in args))
+
+    if NUMPY_GE_200:
+        result = tuple(broadcast_to(e, shape=shape, chunks=chunks) for e in args)
+    else:
+        result = [broadcast_to(e, shape=shape, chunks=chunks) for e in args]
+
+    return result

dask_array/routines/_coarsen.py

@@ -0,0 +1,115 @@
+"""Coarsen implementation for array-expr."""
+
+from __future__ import annotations
+
+import itertools
+from functools import cached_property, partial
+
+import numpy as np
+
+from dask._task_spec import Task, TaskRef
+from dask_array._collection import asarray, new_collection
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+
+
+def _partition(n, size):
+    """Partition n into evenly distributed sizes."""
+    quotient, remainder = divmod(n, size)
+    return [size] * quotient, [remainder] if remainder else []
+
+
+def aligned_coarsen_chunks(chunks, multiple):
+    """Returns a new chunking aligned with the coarsening multiple."""
+    chunks = np.asarray(chunks)
+    overflow = chunks % multiple
+    excess = overflow.sum()
+    new_chunks = chunks - overflow
+    chunk_validity = new_chunks == chunks
+    valid_inds, invalid_inds = np.where(chunk_validity)[0], np.where(~chunk_validity)[0]
+    chunk_modification_order = [
+        *invalid_inds[np.argsort(new_chunks[invalid_inds])],
+        *valid_inds[np.argsort(new_chunks[valid_inds])],
+    ]
+    partitioned_excess, remainder = _partition(excess, multiple)
+    for idx, extra in enumerate(partitioned_excess):
+        new_chunks[chunk_modification_order[idx]] += extra
+    new_chunks = np.array([*new_chunks, *remainder])
+    new_chunks = new_chunks[new_chunks > 0]
+    return tuple(new_chunks)
+
+
+class Coarsen(ArrayExpr):
+    """Expression class for coarsen operation."""
+
+    _parameters = ["x", "reduction", "axes", "trim_excess", "kwargs"]
+    _defaults = {"trim_excess": False, "kwargs": None}
+
+    @cached_property
+    def _reduction(self):
+        reduction = self.reduction
+        # Handle dask.array or dask_array functions - use numpy equivalent
+        if reduction.__module__.startswith(("dask.", "dask_array")):
+            return getattr(np, reduction.__name__)
+        return reduction
+
+    @cached_property
+    def _kwargs(self):
+        return self.kwargs or {}
+
+    @cached_property
+    def _meta(self):
+        x = self.x
+        meta = self._reduction(np.empty((1,) * x.ndim, dtype=x.dtype), **self._kwargs)
+        return meta_from_array(meta, ndim=x.ndim)
+
+    @cached_property
+    def chunks(self):
+        x = self.x
+        axes = self.axes
+        coarsen_dim = lambda dim, ax: int(dim // axes.get(ax, 1))
+        return tuple(
+            tuple(coarsen_dim(bd, i) for bd in bds if coarsen_dim(bd, i) > 0) for i, bds in enumerate(x.chunks)
+        )
+
+    def _layer(self):
+        from dask_array import _chunk as chunk
+
+        x = self.x
+        axes = self.axes
+        name = self._name
+        dsk = {}
+
+        in_ranges = [range(len(c)) for c in x.chunks]
+        for in_idx in itertools.product(*in_ranges):
+            in_key = (x._name,) + in_idx
+            out_key = (name,) + in_idx
+            func = partial(
+                chunk.coarsen,
+                self._reduction,
+                axes=axes,
+                trim_excess=self.trim_excess,
+                **self._kwargs,
+            )
+            dsk[out_key] = Task(out_key, func, TaskRef(in_key))
+
+        return dsk
+
+
+def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
+    """Coarsen array by applying reduction to fixed size neighborhoods."""
+    x = asarray(x)
+
+    if not trim_excess and not all(x.shape[i] % div == 0 for i, div in axes.items()):
+        msg = f"Coarsening factors {axes} do not align with array shape {x.shape}."
+        raise ValueError(msg)
+
+    new_chunks = {}
+    for i, div in axes.items():
+        aligned = aligned_coarsen_chunks(x.chunks[i], div)
+        if aligned != x.chunks[i]:
+            new_chunks[i] = aligned
+    if new_chunks:
+        x = x.rechunk(new_chunks)
+
+    return new_collection(Coarsen(x, reduction, axes, trim_excess, kwargs or None))

dask_array/routines/_diff.py

@@ -0,0 +1,79 @@
+"""Difference operation."""
+
+from __future__ import annotations
+
+
+def diff(a, n=1, axis=-1, prepend=None, append=None):
+    """Calculate the n-th discrete difference along the given axis.
+
+    Parameters
+    ----------
+    a : array_like
+        Input array.
+    n : int, optional
+        The number of times values are differenced. Default is 1.
+    axis : int, optional
+        The axis along which the difference is taken. Default is -1.
+    prepend, append : array_like, optional
+        Values to prepend or append to a along axis prior to
+        performing the difference.
+
+    Returns
+    -------
+    diff : Array
+        The n-th differences.
+
+    See Also
+    --------
+    numpy.diff
+    """
+    # Lazy imports to avoid circular dependencies
+    from dask_array._broadcast import broadcast_to
+    from dask_array.core import asarray
+    from dask_array.stacking import concatenate
+
+    a = asarray(a)
+    n = int(n)
+    axis = int(axis)
+
+    if n == 0:
+        return a
+    if n < 0:
+        raise ValueError(f"order must be non-negative but got {n}")
+
+    combined = []
+    if prepend is not None:
+        prepend = asarray(prepend)
+        if prepend.ndim == 0:
+            shape = list(a.shape)
+            shape[axis] = 1
+            prepend = broadcast_to(prepend, tuple(shape))
+        combined.append(prepend)
+
+    combined.append(a)
+
+    if append is not None:
+        append = asarray(append)
+        if append.ndim == 0:
+            shape = list(a.shape)
+            shape[axis] = 1
+            append = broadcast_to(append, tuple(shape))
+        combined.append(append)
+
+    if len(combined) > 1:
+        a = concatenate(combined, axis)
+
+    sl_1 = a.ndim * [slice(None)]
+    sl_2 = a.ndim * [slice(None)]
+
+    sl_1[axis] = slice(1, None)
+    sl_2[axis] = slice(None, -1)
+
+    sl_1 = tuple(sl_1)
+    sl_2 = tuple(sl_2)
+
+    r = a
+    for _ in range(n):
+        r = r[sl_1] - r[sl_2]
+
+    return r