dask-array 0.1.0__py3-none-any.whl

dask_array/routines/_gradient.py

@@ -0,0 +1,158 @@
+"""Gradient implementation for array-expr."""
+
+from __future__ import annotations
+
+import math
+from numbers import Integral, Real
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask_array._overlap import map_overlap
+from dask_array._utils import validate_axis
+from dask.base import is_dask_collection
+
+
+def _gradient_kernel(x, block_id, coord, axis, array_locs, grad_kwargs):
+    """
+    x: nd-array
+        array of one block
+    coord: 1d-array or scalar
+        coordinate along which the gradient is computed.
+    axis: int
+        axis along which the gradient is computed
+    array_locs:
+        actual location along axis. None if coordinate is scalar
+    grad_kwargs:
+        keyword to be passed to np.gradient
+    """
+    block_loc = block_id[axis]
+    if array_locs is not None:
+        coord = coord[array_locs[0][block_loc] : array_locs[1][block_loc]]
+    grad = np.gradient(x, coord, axis=axis, **grad_kwargs)
+    return grad
+
+
+def gradient(f, *varargs, axis=None, **kwargs):
+    """
+    Return the gradient of an N-dimensional array.
+
+    This docstring was copied from numpy.gradient.
+
+    Some inconsistencies with the Dask version may exist.
+
+    The gradient is computed using second order accurate central differences
+    in the interior points and either first or second order accurate one-sides
+    (forward or backwards) differences at the boundaries.
+    The returned gradient hence has the same shape as the input array.
+
+    Parameters
+    ----------
+    f : array_like
+        An N-dimensional array containing samples of a scalar function.
+    varargs : list of scalar or array, optional
+        Spacing between f values. Default unitary spacing for all dimensions.
+        Spacing can be specified using:
+
+        1. single scalar to specify a sample distance for all dimensions.
+        2. N scalars to specify a constant sample distance for each dimension.
+           i.e. `dx`, `dy`, `dz`, ...
+        3. N arrays to specify the coordinates of the values along each
+           dimension of F. The length of the array must match the size of
+           the corresponding dimension
+        4. Any combination of N scalars/arrays with the meaning of 2. and 3.
+
+        If `axis` is given, the number of varargs must equal the number of axes.
+        Default: 1.
+    axis : None or int or tuple of ints, optional
+        Gradient is calculated only along the given axis or axes.
+        The default (axis = None) is to calculate the gradient for all the axes
+        of the input array. axis may be negative, in which case it counts from
+        the last to the first axis.
+
+    Returns
+    -------
+    gradient : ndarray or list of ndarray
+        A list of ndarrays (or a single ndarray if there is only one dimension)
+        corresponding to the derivatives of f with respect to each dimension.
+        Each derivative has the same shape as f.
+
+    Other Parameters
+    ----------------
+    edge_order : {1, 2}, optional
+        Gradient is calculated using N-th order accurate differences
+        at the boundaries. Default: 1.
+    """
+    f = asarray(f)
+
+    kwargs["edge_order"] = math.ceil(kwargs.get("edge_order", 1))
+    if kwargs["edge_order"] > 2:
+        raise ValueError("edge_order must be less than or equal to 2.")
+
+    drop_result_list = False
+    if axis is None:
+        axis = tuple(range(f.ndim))
+    elif isinstance(axis, Integral):
+        drop_result_list = True
+        axis = (axis,)
+
+    axis = validate_axis(axis, f.ndim)
+
+    if len(axis) != len(set(axis)):
+        raise ValueError("duplicate axes not allowed")
+
+    axis = tuple(ax % f.ndim for ax in axis)
+
+    if varargs == ():
+        varargs = (1,)
+    if len(varargs) == 1:
+        varargs = len(axis) * varargs
+    if len(varargs) != len(axis):
+        raise TypeError("Spacing must either be a single scalar, or a scalar / 1d-array per axis")
+
+    if issubclass(f.dtype.type, (np.bool_, Integral)):
+        f = f.astype(float)
+    elif issubclass(f.dtype.type, Real) and f.dtype.itemsize < 4:
+        f = f.astype(float)
+
+    results = []
+    for i, ax in enumerate(axis):
+        for c in f.chunks[ax]:
+            if np.min(c) < kwargs["edge_order"] + 1:
+                raise ValueError(
+                    "Chunk size must be larger than edge_order + 1. "
+                    f"Minimum chunk for axis {ax} is {np.min(c)}. Rechunk to "
+                    "proceed."
+                )
+
+        if np.isscalar(varargs[i]):
+            array_locs = None
+        else:
+            if is_dask_collection(varargs[i]):
+                raise NotImplementedError("dask array coordinated is not supported.")
+            # coordinate position for each block taking overlap into account
+            chunk = np.array(f.chunks[ax])
+            array_loc_stop = np.cumsum(chunk) + 1
+            array_loc_start = array_loc_stop - chunk - 2
+            array_loc_stop[-1] -= 1
+            array_loc_start[0] = 0
+            array_locs = (array_loc_start, array_loc_stop)
+
+        results.append(
+            map_overlap(
+                _gradient_kernel,
+                f,
+                dtype=f.dtype,
+                depth={j: 1 if j == ax else 0 for j in range(f.ndim)},
+                boundary="none",
+                coord=varargs[i],
+                axis=ax,
+                array_locs=array_locs,
+                grad_kwargs=kwargs,
+            )
+        )
+
+    if drop_result_list:
+        results = results[0]
+
+    return results

dask_array/routines/_indexing.py

@@ -0,0 +1,65 @@
+"""Index manipulation functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray, stack
+from dask.base import is_dask_collection
+from dask.utils import derived_from
+
+
+def _unravel_index_kernel(indices, func_kwargs):
+    return np.stack(np.unravel_index(indices, **func_kwargs))
+
+
+@derived_from(np)
+def unravel_index(indices, shape, order="C"):
+    from dask_array.creation import empty
+
+    indices = asarray(indices)
+    if shape and indices.size:
+        unraveled_indices = tuple(
+            indices.map_blocks(
+                _unravel_index_kernel,
+                dtype=np.intp,
+                chunks=(((len(shape),),) + indices.chunks),
+                new_axis=0,
+                func_kwargs={"shape": shape, "order": order},
+            )
+        )
+    else:
+        unraveled_indices = tuple(empty((0,), dtype=np.intp, chunks=1) for i in shape)
+
+    return unraveled_indices
+
+
+@derived_from(np)
+def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
+    from dask_array.routines._broadcast import broadcast_arrays
+
+    if np.isscalar(dims):
+        dims = (dims,)
+    if is_dask_collection(dims) or any(is_dask_collection(d) for d in dims):
+        raise NotImplementedError(f"Dask types are not supported in the `dims` argument: {dims!r}")
+
+    if hasattr(multi_index, "ndim") and multi_index.ndim > 0:
+        # It's an array-like
+        index_stack = asarray(multi_index)
+    else:
+        multi_index_arrs = broadcast_arrays(*multi_index)
+        index_stack = stack(multi_index_arrs)
+
+    if not np.isnan(index_stack.shape).any() and len(index_stack) != len(dims):
+        raise ValueError(f"parameter multi_index must be a sequence of length {len(dims)}")
+    if not np.issubdtype(index_stack.dtype, np.signedinteger):
+        raise TypeError("only int indices permitted")
+    return index_stack.map_blocks(
+        np.ravel_multi_index,
+        dtype=np.intp,
+        chunks=index_stack.chunks[1:],
+        drop_axis=0,
+        dims=dims,
+        mode=mode,
+        order=order,
+    )

dask_array/routines/_insert_delete.py

@@ -0,0 +1,132 @@
+"""Insert/delete operations for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import (
+    asanyarray,
+    asarray,
+    broadcast_to,
+    concatenate,
+    ravel,
+)
+from dask_array._utils import validate_axis
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def append(arr, values, axis=None):
+    """Append values to the end of an array."""
+    arr = asanyarray(arr)
+    if axis is None:
+        if arr.ndim != 1:
+            arr = ravel(arr)
+        values = ravel(asanyarray(values))
+        axis = arr.ndim - 1
+    return concatenate((arr, values), axis=axis)
+
+
+@derived_from(np)
+def ediff1d(ary, to_end=None, to_begin=None):
+    """Compute the differences between consecutive elements of an array."""
+    ary = asarray(ary)
+
+    aryf = ary.flatten()
+    r = aryf[1:] - aryf[:-1]
+
+    r = [r]
+    if to_begin is not None:
+        r = [asarray(to_begin).flatten()] + r
+    if to_end is not None:
+        r = r + [asarray(to_end).flatten()]
+    r = concatenate(r)
+
+    return r
+
+
+def _split_at_breaks(array, breaks, axis=0):
+    """Split an array into a list of arrays (using slices) at the given breaks
+
+    >>> _split_at_breaks(np.arange(6), [3, 5])
+    [array([0, 1, 2]), array([3, 4]), array([5])]
+    """
+    from tlz import concat, sliding_window
+
+    padded_breaks = list(concat([[None], breaks, [None]]))
+    slices = [slice(i, j) for i, j in sliding_window(2, padded_breaks)]
+    preslice = (slice(None),) * axis
+    split_array = [array[preslice + (s,)] for s in slices]
+    return split_array
+
+
+@derived_from(np)
+def insert(arr, obj, values, axis):
+    """Insert values along the given axis before the given indices."""
+    from tlz import interleave
+
+    # axis is a required argument here to avoid needing to deal with the numpy
+    # default case (which reshapes the array to make it flat)
+    arr = asarray(arr)
+    axis = validate_axis(axis, arr.ndim)
+
+    if isinstance(obj, slice):
+        obj = np.arange(*obj.indices(arr.shape[axis]))
+    obj = np.asarray(obj)
+    scalar_obj = obj.ndim == 0
+    if scalar_obj:
+        obj = np.atleast_1d(obj)
+
+    obj = np.where(obj < 0, obj + arr.shape[axis], obj)
+    if (np.diff(obj) < 0).any():
+        raise NotImplementedError("da.insert only implemented for monotonic ``obj`` argument")
+
+    split_arr = _split_at_breaks(arr, np.unique(obj), axis)
+
+    if getattr(values, "ndim", 0) == 0:
+        # we need to turn values into a dask array
+        values = asarray(values)
+
+        values_shape = tuple(len(obj) if axis == n else s for n, s in enumerate(arr.shape))
+        values = broadcast_to(values, values_shape)
+    elif scalar_obj:
+        values = values[(slice(None),) * axis + (None,)]
+
+    values = asarray(values)
+    values_chunks = tuple(
+        values_bd if axis == n else arr_bd for n, (arr_bd, values_bd) in enumerate(zip(arr.chunks, values.chunks))
+    )
+    values = values.rechunk(values_chunks)
+
+    counts = np.bincount(obj)[:-1]
+    values_breaks = np.cumsum(counts[counts > 0])
+    split_values = _split_at_breaks(values, values_breaks, axis)
+
+    interleaved = list(interleave([split_arr, split_values]))
+    interleaved = [i for i in interleaved if i.size]
+    return concatenate(interleaved, axis=axis)
+
+
+@derived_from(np)
+def delete(arr, obj, axis):
+    """Remove elements from an array along an axis."""
+    # axis is a required argument here to avoid needing to deal with the numpy
+    # default case (which reshapes the array to make it flat)
+    arr = asarray(arr)
+    axis = validate_axis(axis, arr.ndim)
+
+    if isinstance(obj, slice):
+        tmp = np.arange(*obj.indices(arr.shape[axis]))
+        obj = tmp[::-1] if obj.step and obj.step < 0 else tmp
+    else:
+        obj = np.asarray(obj)
+        obj = np.where(obj < 0, obj + arr.shape[axis], obj)
+        obj = np.unique(obj)
+
+    target_arr = _split_at_breaks(arr, obj, axis)
+
+    target_arr = [
+        (arr[tuple(slice(1, None) if axis == n else slice(None) for n in range(arr.ndim))] if i != 0 else arr)
+        for i, arr in enumerate(target_arr)
+    ]
+    return concatenate(target_arr, axis=axis)

dask_array/routines/_misc.py

@@ -0,0 +1,122 @@
+"""Miscellaneous utility functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import Array, asarray, elemwise
+from dask_array._utils import validate_axis
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def ndim(a):
+    """Return the number of dimensions of an array."""
+    a = asarray(a)
+    return a.ndim
+
+
+@derived_from(np)
+def shape(a):
+    """Return the shape of an array."""
+    a = asarray(a)
+    return a.shape
+
+
+def result_type(*arrays_and_dtypes):
+    """Returns the type from NumPy type promotion rules."""
+    args = [a.dtype if isinstance(a, Array) else a for a in arrays_and_dtypes]
+    return np.result_type(*args)
+
+
+def compress(condition, a, axis=None):
+    """
+    Return selected slices of an array along given axis.
+
+    This docstring was copied from numpy.compress.
+
+    Some inconsistencies with the Dask version may exist.
+
+    Parameters
+    ----------
+    condition : 1-D array of bools
+        Array that selects which entries to return. If len(condition)
+        is less than the size of a along the given axis, then output is
+        truncated to the length of the condition array.
+    a : array_like
+        Array from which to extract a part.
+    axis : int, optional
+        Axis along which to take slices. If None (default), work on the
+        flattened array.
+
+    Returns
+    -------
+    compressed_array : ndarray
+        A copy of a without the slices along axis for which condition
+        is false.
+    """
+    from dask_array._utils import is_arraylike
+
+    if not is_arraylike(condition):
+        condition = np.asarray(condition)
+    condition = condition.astype(bool)
+    a = asarray(a)
+
+    if condition.ndim != 1:
+        raise ValueError("Condition must be one dimensional")
+
+    if axis is None:
+        a = a.ravel()
+        axis = 0
+    axis = validate_axis(axis, a.ndim)
+
+    # Treat `condition` as filled with `False` (if it is too short)
+    a = a[tuple(slice(None, len(condition)) if i == axis else slice(None) for i in range(a.ndim))]
+
+    # Use `condition` to select along 1 dimension
+    a = a[tuple(condition if i == axis else slice(None) for i in range(a.ndim))]
+
+    return a
+
+
+def _take_constant(indices, a, axis):
+    """Take from a constant array using indices."""
+    return np.take(a, indices, axis)
+
+
+def _take_dask_array_from_numpy(a, indices, axis):
+    """Take from a numpy array using a dask array of indices."""
+    assert isinstance(a, np.ndarray)
+    assert isinstance(indices, Array)
+
+    return elemwise(_take_constant, indices, dtype=a.dtype, a=a, axis=axis)
+
+
+@derived_from(np)
+def take(a, indices, axis=0):
+    """
+    Take elements from an array along an axis.
+
+    This docstring was copied from numpy.take.
+
+    Parameters
+    ----------
+    a : dask array
+        The source array.
+    indices : array_like
+        The indices of the values to extract.
+    axis : int, optional
+        The axis over which to select values.
+
+    Returns
+    -------
+    out : dask array
+        The returned array has the same type as a.
+    """
+    a = asarray(a)
+    axis = validate_axis(axis, a.ndim)
+
+    if isinstance(a, np.ndarray) and isinstance(indices, Array):
+        return _take_dask_array_from_numpy(a, indices, axis)
+    else:
+        return a[(slice(None),) * axis + (indices,)]

dask_array/routines/_nonzero.py

@@ -0,0 +1,72 @@
+"""Nonzero-related functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray, elemwise, stack
+from dask.utils import derived_from
+
+
+def _isnonzero_vec(v):
+    return bool(np.count_nonzero(v))
+
+
+_isnonzero_vec = np.vectorize(_isnonzero_vec, otypes=[bool])
+
+
+def _isnonzero(a):
+    # Output of np.vectorize can't be pickled
+    return _isnonzero_vec(a)
+
+
+def isnonzero(a):
+    """Handle special cases where conversion to bool does not work correctly.
+    xref: https://github.com/numpy/numpy/issues/9479
+    """
+    try:
+        np.zeros([], dtype=a.dtype).astype(bool)
+    except ValueError:
+        return elemwise(_isnonzero, a, dtype=bool)
+    else:
+        return a.astype(bool)
+
+
+@derived_from(np)
+def count_nonzero(a, axis=None):
+    """Counts the number of non-zero values in the array."""
+    return isnonzero(asarray(a)).astype(np.intp).sum(axis=axis)
+
+
+@derived_from(np)
+def argwhere(a):
+    """Find the indices of array elements that are non-zero."""
+    from dask_array._routines import compress
+    from dask_array.creation import indices
+
+    a = asarray(a)
+
+    nz = isnonzero(a).flatten()
+
+    ind = indices(a.shape, dtype=np.intp, chunks=a.chunks)
+    if ind.ndim > 1:
+        ind = stack([ind[i].ravel() for i in range(len(ind))], axis=1)
+    ind = compress(nz, ind, axis=0)
+
+    return ind
+
+
+@derived_from(np)
+def flatnonzero(a):
+    """Return indices that are non-zero in the flattened array."""
+    return argwhere(asarray(a).ravel())[:, 0]
+
+
+@derived_from(np)
+def nonzero(a):
+    """Return the indices of the elements that are non-zero."""
+    ind = argwhere(a)
+    if ind.ndim > 1:
+        return tuple(ind[:, i] for i in range(ind.shape[1]))
+    else:
+        return (ind,)

dask_array/routines/_search.py

@@ -0,0 +1,123 @@
+"""Search functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask.utils import derived_from
+
+
+def _searchsorted_block(x, y, side):
+    res = np.searchsorted(x, y, side=side)
+    # 0 is only correct for the first block of a, but blockwise doesn't have a way
+    # of telling which block is being operated on (unlike map_blocks),
+    # so set all 0 values to a special value and set back at the end of searchsorted
+    res[res == 0] = -1
+    return res[np.newaxis, :]
+
+
+def searchsorted(a, v, side="left", sorter=None):
+    """
+    Find indices where elements should be inserted to maintain order.
+
+    This docstring was copied from numpy.searchsorted.
+
+    Some inconsistencies with the Dask version may exist.
+
+    Find the indices into a sorted array `a` such that, if the
+    corresponding elements in `v` were inserted before the indices, the
+    order of `a` would be preserved.
+
+    Parameters
+    ----------
+    a : 1-D array_like
+        Input array. If `sorter` is None, then it must be sorted in
+        ascending order, otherwise `sorter` must be an array of indices
+        that sort it.
+    v : array_like
+        Values to insert into `a`.
+    side : {'left', 'right'}, optional
+        If 'left', the index of the first suitable location found is given.
+        If 'right', return the last such index.  If there is no suitable
+        index, return either 0 or N (where N is the length of `a`).
+    sorter : 1-D array_like, optional
+        Optional array of integer indices that sort array a into ascending
+        order. They are typically the result of argsort.
+
+    Returns
+    -------
+    indices : int or array of ints
+        Array of insertion points with the same shape as `v`,
+        or an integer if `v` is a scalar.
+    """
+    from dask_array._collection import blockwise
+    from dask_array.routines._where import where
+    from dask_array._utils import array_safe, meta_from_array
+
+    if a.ndim != 1:
+        raise ValueError("Input array a must be one dimensional")
+
+    if sorter is not None:
+        raise NotImplementedError("da.searchsorted with a sorter argument is not supported")
+
+    # call np.searchsorted for each pair of blocks in a and v
+    meta = np.searchsorted(a._meta, v._meta)
+    out = blockwise(
+        _searchsorted_block,
+        list(range(v.ndim + 1)),
+        a,
+        [0],
+        v,
+        list(range(1, v.ndim + 1)),
+        side,
+        None,
+        meta=meta,
+        adjust_chunks={0: 1},  # one row for each block in a
+    )
+
+    # add offsets to take account of the position of each block within the array a
+    a_chunk_sizes = array_safe((0, *a.chunks[0]), like=meta_from_array(a))
+    a_chunk_offsets = np.cumsum(a_chunk_sizes)[:-1]
+    a_chunk_offsets = a_chunk_offsets[(Ellipsis,) + v.ndim * (np.newaxis,)]
+    a_offsets = asarray(a_chunk_offsets, chunks=1)
+    out = where(out < 0, out, out + a_offsets)
+
+    # combine the results from each block (of a)
+    out = out.max(axis=0)
+
+    # fix up any -1 values
+    out[out == -1] = 0
+
+    return out
+
+
+def _isin_kernel(element, test_elements, assume_unique=False):
+    values = np.isin(element.ravel(), test_elements, assume_unique=assume_unique)
+    return values.reshape(element.shape + (1,) * test_elements.ndim)
+
+
+@derived_from(np)
+def isin(element, test_elements, assume_unique=False, invert=False):
+    from dask_array._collection import blockwise
+
+    element = asarray(element)
+    test_elements = asarray(test_elements)
+    element_axes = tuple(range(element.ndim))
+    test_axes = tuple(i + element.ndim for i in range(test_elements.ndim))
+    mapped = blockwise(
+        _isin_kernel,
+        element_axes + test_axes,
+        element,
+        element_axes,
+        test_elements,
+        test_axes,
+        adjust_chunks={axis: lambda _: 1 for axis in test_axes},
+        dtype=bool,
+        assume_unique=assume_unique,
+    )
+
+    result = mapped.any(axis=test_axes)
+    if invert:
+        result = ~result
+    return result