dask-array 0.1.0__py3-none-any.whl

dask_array/_broadcast.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+import functools
+from itertools import product
+
+import numpy as np
+
+from dask._task_spec import Task, TaskRef
+from dask_array._expr import ArrayExpr
+from dask_array._core_utils import normalize_chunks
+from dask_array._utils import meta_from_array
+
+
+class BroadcastTo(ArrayExpr):
+    """Broadcast an array to a new shape."""
+
+    _parameters = ["array", "_shape", "_chunks", "_meta_override"]
+    _defaults = {"_meta_override": None}
+
+    @functools.cached_property
+    def _name(self):
+        return f"broadcast_to-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _meta(self):
+        meta_override = self.operand("_meta_override")
+        # Only use meta_override if it has the correct ndim
+        if meta_override is not None and hasattr(meta_override, "ndim") and meta_override.ndim == len(self._shape):
+            return meta_override
+        return meta_from_array(self.array._meta, ndim=len(self._shape))
+
+    @functools.cached_property
+    def chunks(self):
+        return self._chunks
+
+    def _layer(self) -> dict:
+        x = self.array
+        shape = self._shape
+        chunks = self._chunks
+        ndim_new = len(shape) - x.ndim
+
+        dsk = {}
+        enumerated_chunks = product(*(enumerate(bds) for bds in chunks))
+        for ec in enumerated_chunks:
+            new_index, chunk_shape = zip(*ec)
+            old_index = tuple(0 if bd == (1,) else i for bd, i in zip(x.chunks, new_index[ndim_new:]))
+            old_key = (x._name,) + old_index
+            new_key = (self._name,) + new_index
+            dsk[new_key] = Task(new_key, np.broadcast_to, TaskRef(old_key), chunk_shape)
+
+        return dsk
+
+    def _simplify_up(self, parent, dependents):
+        """Allow slice and shuffle operations to push through BroadcastTo."""
+        from dask_array._shuffle import Shuffle
+        from dask_array.slicing import SliceSlicesIntegers
+
+        if isinstance(parent, SliceSlicesIntegers):
+            return self._accept_slice(parent)
+        if isinstance(parent, Shuffle):
+            return self._accept_shuffle(parent)
+        return None
+
+    def _accept_shuffle(self, shuffle_expr):
+        """Accept a shuffle being pushed through BroadcastTo.
+
+        - Shuffle on a new dimension (added by broadcast): can't push through
+        - Shuffle on dimension broadcast from size 1: no-op, return self
+        - Shuffle on dimension with real data: push through to input
+        """
+        from dask_array._shuffle import Shuffle
+
+        axis = shuffle_expr.axis
+        ndim_new = len(self._shape) - self.array.ndim
+
+        # Shuffle on a new dimension (added by broadcast) - can't push through
+        if axis < ndim_new:
+            return None
+
+        # Map to input axis
+        input_axis = axis - ndim_new
+        input_size = self.array.shape[input_axis]
+
+        # If input dimension is size 1 (broadcasted), shuffle is a no-op
+        if input_size == 1:
+            return self
+
+        # Push shuffle through to input
+        shuffled_input = Shuffle(
+            self.array,
+            shuffle_expr.indexer,
+            input_axis,
+            shuffle_expr.operand("name"),
+        )
+        return BroadcastTo(shuffled_input, self._shape, self._chunks, self._meta)
+
+    def _accept_slice(self, slice_expr):
+        """Accept a slice being pushed through BroadcastTo.
+
+        For broadcast_to(x, shape)[slices]:
+        - Dimensions added by broadcast (new dims): affect output shape only
+        - Dimensions from input with size > 1: push slice to input
+        - Dimensions from input with size == 1: affect output shape only
+        """
+        from numbers import Integral
+
+        from dask_array._new_collection import new_collection
+
+        input_arr = self.array
+        output_shape = self._shape
+        index = slice_expr.index
+
+        # Pad index to full length
+        full_index = index + (slice(None),) * (len(output_shape) - len(index))
+
+        # For now, only handle simple slices
+        if any(isinstance(idx, Integral) for idx in full_index):
+            return None
+        if any(idx is None for idx in full_index):
+            return None
+
+        ndim_new = len(output_shape) - input_arr.ndim
+
+        # Compute new output shape and input slices
+        new_output_shape = []
+        input_slices = []
+
+        for out_dim, idx in enumerate(full_index):
+            if not isinstance(idx, slice):
+                return None
+
+            out_size = output_shape[out_dim]
+            start, stop, step = idx.indices(out_size)
+            if step != 1:
+                return None
+            new_dim_size = max(0, stop - start)
+            new_output_shape.append(new_dim_size)
+
+            # Check if this dimension maps to input
+            if out_dim >= ndim_new:
+                in_dim = out_dim - ndim_new
+                in_size = input_arr.shape[in_dim]
+
+                if in_size == 1:
+                    # Broadcasted from size 1 - can't push, just take full slice
+                    input_slices.append(slice(None))
+                else:
+                    # Real dimension - push the slice
+                    input_slices.append(idx)
+
+        # Slice the input array
+        if input_slices:
+            sliced_input = new_collection(input_arr)[tuple(input_slices)]
+        else:
+            sliced_input = new_collection(input_arr)
+
+        # Compute new chunks for the output
+        # For dimensions from input: use input's (sliced) chunks
+        # For new dimensions: use the new output shape (single chunk)
+        old_chunks = self._chunks
+        new_chunks = []
+        for out_dim, old_chunk in enumerate(old_chunks):
+            if out_dim >= ndim_new:
+                in_dim = out_dim - ndim_new
+                in_size = input_arr.shape[in_dim]
+                if in_size == 1:
+                    # Broadcasted - compute new chunks from old
+                    idx = full_index[out_dim]
+                    start, stop, _ = idx.indices(output_shape[out_dim])
+                    new_chunks.append(self._slice_chunks(old_chunk, start, stop - start))
+                else:
+                    # Use sliced input's chunks
+                    new_chunks.append(sliced_input.expr.chunks[in_dim])
+            else:
+                # New dimension - compute from slice
+                idx = full_index[out_dim]
+                start, stop, _ = idx.indices(output_shape[out_dim])
+                new_chunks.append(self._slice_chunks(old_chunk, start, stop - start))
+
+        # Compute meta for the new broadcast
+        new_meta = meta_from_array(sliced_input.expr._meta)
+
+        # Create new BroadcastTo
+        return BroadcastTo(
+            sliced_input.expr,
+            tuple(new_output_shape),
+            tuple(new_chunks),
+            new_meta,
+        )
+
+    def _slice_chunks(self, chunks, start, length):
+        """Compute new chunks after slicing."""
+        if length == 0:
+            return (0,)
+
+        result = []
+        pos = 0
+        remaining = length
+        for chunk_size in chunks:
+            chunk_start = pos
+            chunk_end = pos + chunk_size
+            pos = chunk_end
+
+            if chunk_end <= start:
+                continue
+            if chunk_start >= start + length:
+                break
+
+            # Overlap with the slice
+            overlap_start = max(chunk_start, start)
+            overlap_end = min(chunk_end, start + length)
+            overlap_size = overlap_end - overlap_start
+
+            if overlap_size > 0:
+                result.append(overlap_size)
+                remaining -= overlap_size
+
+        return tuple(result)
+
+
+def broadcast_to(x, shape, chunks=None, meta=None):
+    """Broadcast an array to a new shape.
+
+    Parameters
+    ----------
+    x : array_like
+        The array to broadcast.
+    shape : tuple
+        The shape of the desired array.
+    chunks : tuple, optional
+        If provided, then the result will use these chunks instead of the same
+        chunks as the source array.
+    meta : empty ndarray, optional
+        empty ndarray created with same NumPy backend, ndim and dtype as the
+        Dask Array being created
+
+    Returns
+    -------
+    Array
+    """
+    from dask_array._new_collection import new_collection
+    from dask_array._collection import asarray
+
+    x = asarray(x)
+    shape = tuple(shape)
+
+    if meta is None:
+        meta = meta_from_array(x._meta)
+
+    # Identity case
+    if x.shape == shape and (chunks is None or chunks == x.chunks):
+        return x
+
+    ndim_new = len(shape) - x.ndim
+    if ndim_new < 0 or any(new != old for new, old in zip(shape[ndim_new:], x.shape) if old != 1):
+        raise ValueError(f"cannot broadcast shape {x.shape} to shape {shape}")
+
+    if chunks is None:
+        chunks = tuple((s,) for s in shape[:ndim_new]) + tuple(
+            bd if old > 1 else (new,) for bd, old, new in zip(x.chunks, x.shape, shape[ndim_new:])
+        )
+    else:
+        chunks = normalize_chunks(chunks, shape, dtype=x.dtype, previous_chunks=x.chunks)
+        for old_bd, new_bd in zip(x.chunks, chunks[ndim_new:]):
+            if old_bd != new_bd and old_bd != (1,):
+                raise ValueError(
+                    f"cannot broadcast chunks {x.chunks} to chunks {chunks}: "
+                    "new chunks must either be along a new "
+                    "dimension or a dimension of size 1"
+                )
+
+    return new_collection(BroadcastTo(x.expr, shape, chunks, meta))

dask_array/_chunk.py

@@ -0,0 +1,445 @@
+"""A set of NumPy functions to apply per chunk"""
+
+from __future__ import annotations
+
+import itertools
+from collections.abc import Container, Iterable, Sequence
+from functools import wraps
+from numbers import Integral
+
+import numpy as np
+
+
+def concat(seqs):
+    """Concatenate zero or more iterables, any of which may be infinite.
+
+    An infinite sequence will prevent the rest of the arguments from
+    being included.
+
+    We use chain.from_iterable rather than ``chain(*seqs)`` so that seqs
+    can be a generator.
+
+    >>> list(concat([[], [1], [2, 3]]))
+    [1, 2, 3]
+
+    See also:
+        itertools.chain.from_iterable  equivalent
+    """
+    return itertools.chain.from_iterable(seqs)
+
+
+def flatten(seq, container=list):
+    """Flatten nested sequences.
+
+    >>> list(flatten([1]))
+    [1]
+
+    >>> list(flatten([[1, 2], [1, 2]]))
+    [1, 2, 1, 2]
+
+    >>> list(flatten([[[1], [2]], [[1], [2]]]))
+    [1, 2, 1, 2]
+
+    >>> list(flatten(((1, 2), (1, 2)))) # Don't flatten tuples
+    [(1, 2), (1, 2)]
+
+    >>> list(flatten((1, 2, [3, 4]))) # support heterogeneous
+    [1, 2, 3, 4]
+    """
+    if isinstance(seq, str):
+        yield seq
+    else:
+        for item in seq:
+            if isinstance(item, container):
+                yield from flatten(item, container=container)
+            else:
+                yield item
+
+
+def astype(x, astype_dtype=None, **kwargs):
+    """Change array dtype."""
+    return x.astype(astype_dtype, **kwargs)
+
+
+def view(x, dtype, order="C"):
+    """View array as different dtype."""
+    if order == "C":
+        try:
+            x = np.ascontiguousarray(x, like=x)
+        except TypeError:
+            x = np.ascontiguousarray(x)
+        return x.view(dtype)
+    else:
+        try:
+            x = np.asfortranarray(x, like=x)
+        except TypeError:
+            x = np.asfortranarray(x)
+        return x.T.view(dtype).T
+
+
+def trim(x, axes=None):
+    """Trim boundaries off of array.
+
+    >>> x = np.arange(24).reshape((4, 6))
+    >>> trim(x, axes={0: 0, 1: 1})
+    array([[ 1,  2,  3,  4],
+           [ 7,  8,  9, 10],
+           [13, 14, 15, 16],
+           [19, 20, 21, 22]])
+
+    >>> trim(x, axes={0: 1, 1: 1})
+    array([[ 7,  8,  9, 10],
+           [13, 14, 15, 16]])
+    """
+    if isinstance(axes, Integral):
+        axes = [axes] * x.ndim
+    if isinstance(axes, dict):
+        axes = [axes.get(i, 0) for i in range(x.ndim)]
+
+    return x[tuple(slice(ax, -ax if ax else None) for ax in axes)]
+
+
+def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
+    """Coarsen array by applying reduction to fixed size neighborhoods.
+
+    Parameters
+    ----------
+    reduction: function
+        Function like np.sum, np.mean, etc...
+    x: np.ndarray
+        Array to be coarsened
+    axes: dict
+        Mapping of axis to coarsening factor
+
+    Examples
+    --------
+    >>> x = np.array([1, 2, 3, 4, 5, 6])
+    >>> coarsen(np.sum, x, {0: 2})
+    array([ 3,  7, 11])
+    >>> coarsen(np.max, x, {0: 3})
+    array([3, 6])
+    """
+    # Insert singleton dimensions if they don't exist already
+    for i in range(x.ndim):
+        if i not in axes:
+            axes[i] = 1
+
+    if trim_excess:
+        ind = tuple(slice(0, -(d % axes[i])) if d % axes[i] else slice(None, None) for i, d in enumerate(x.shape))
+        x = x[ind]
+
+    # (10, 10) -> (5, 2, 5, 2)
+    newshape = tuple(concat([(x.shape[i] // axes[i], axes[i]) for i in range(x.ndim)]))
+
+    return reduction(x.reshape(newshape), axis=tuple(range(1, x.ndim * 2, 2)), **kwargs)
+
+
+def keepdims_wrapper(a_callable):
+    """
+    A wrapper for functions that don't provide keepdims to ensure that they do.
+    """
+
+    @wraps(a_callable)
+    def keepdims_wrapped_callable(x, axis=None, keepdims=None, *args, **kwargs):
+        r = a_callable(x, *args, axis=axis, **kwargs)
+
+        if not keepdims:
+            return r
+
+        axes = axis
+
+        if axes is None:
+            axes = range(x.ndim)
+
+        if not isinstance(axes, (Container, Iterable, Sequence)):
+            axes = [axes]
+
+        r_slice = tuple()
+        for each_axis in range(x.ndim):
+            if each_axis in axes:
+                r_slice += (None,)
+            else:
+                r_slice += (slice(None),)
+
+        r = r[r_slice]
+
+        return r
+
+    return keepdims_wrapped_callable
+
+
+# Wrap NumPy functions to ensure they provide keepdims.
+sum = np.sum
+prod = np.prod
+min = np.min
+max = np.max
+argmin = keepdims_wrapper(np.argmin)
+nanargmin = keepdims_wrapper(np.nanargmin)
+argmax = keepdims_wrapper(np.argmax)
+nanargmax = keepdims_wrapper(np.nanargmax)
+any = np.any
+all = np.all
+nansum = np.nansum
+nanprod = np.nanprod
+
+nancumprod = np.nancumprod
+nancumsum = np.nancumsum
+
+nanmin = np.nanmin
+nanmax = np.nanmax
+mean = np.mean
+nanmean = np.nanmean
+
+var = np.var
+nanvar = np.nanvar
+
+std = np.std
+nanstd = np.nanstd
+
+
+def topk(a, k, axis, keepdims):
+    """Chunk and combine function of topk
+
+    Extract the k largest elements from a on the given axis.
+    If k is negative, extract the -k smallest elements instead.
+    Note that, unlike in the parent function, the returned elements
+    are not sorted internally.
+    """
+    assert keepdims is True
+    axis = axis[0]
+    if abs(k) >= a.shape[axis]:
+        return a
+
+    a = np.partition(a, -k, axis=axis)
+    k_slice = slice(-k, None) if k > 0 else slice(-k)
+    return a[tuple(k_slice if i == axis else slice(None) for i in range(a.ndim))]
+
+
+def topk_aggregate(a, k, axis, keepdims):
+    """Final aggregation function of topk
+
+    Invoke topk one final time and then sort the results internally.
+    """
+    assert keepdims is True
+    a = topk(a, k, axis, keepdims)
+    axis = axis[0]
+    a = np.sort(a, axis=axis)
+    if k < 0:
+        return a
+    return a[tuple(slice(None, None, -1) if i == axis else slice(None) for i in range(a.ndim))]
+
+
+def argtopk_preprocess(a, idx):
+    """Preparatory step for argtopk
+
+    Put data together with its original indices in a tuple.
+    """
+    return a, idx
+
+
+def argtopk(a_plus_idx, k, axis, keepdims):
+    """Chunk and combine function of argtopk
+
+    Extract the indices of the k largest elements from a on the given axis.
+    If k is negative, extract the indices of the -k smallest elements instead.
+    Note that, unlike in the parent function, the returned elements
+    are not sorted internally.
+    """
+    assert keepdims is True
+    axis = axis[0]
+
+    if isinstance(a_plus_idx, list):
+        a_plus_idx = list(flatten(a_plus_idx))
+        a = np.concatenate([ai for ai, _ in a_plus_idx], axis)
+        idx = np.concatenate([np.broadcast_to(idxi, ai.shape) for ai, idxi in a_plus_idx], axis)
+    else:
+        a, idx = a_plus_idx
+
+    if abs(k) >= a.shape[axis]:
+        return a_plus_idx
+
+    idx2 = np.argpartition(a, -k, axis=axis)
+    k_slice = slice(-k, None) if k > 0 else slice(-k)
+    idx2 = idx2[tuple(k_slice if i == axis else slice(None) for i in range(a.ndim))]
+
+    return np.take_along_axis(a, idx2, axis), np.take_along_axis(idx, idx2, axis)
+
+
+def argtopk_aggregate(a_plus_idx, k, axis, keepdims):
+    """Final aggregation function of argtopk
+
+    Invoke argtopk one final time, sort the results internally, and drop the data.
+    """
+    assert keepdims is True
+    a_plus_idx = a_plus_idx if len(a_plus_idx) > 1 else a_plus_idx[0]
+    a, idx = argtopk(a_plus_idx, k, axis, keepdims)
+    axis = axis[0]
+    idx2 = np.argsort(a, axis=axis)
+
+    idx = np.take_along_axis(idx, idx2, axis)
+    if k < 0:
+        return idx
+    return idx[tuple(slice(None, None, -1) if i == axis else slice(None) for i in range(idx.ndim))]
+
+
+def getitem(obj, index):
+    """Getitem function
+
+    This function creates a copy of the desired selection for array-like
+    inputs when the selection is smaller than half of the original array. This
+    avoids excess memory usage when extracting a small portion from a large array.
+    For more information, see
+    https://numpy.org/doc/stable/reference/arrays.indexing.html#basic-slicing-and-indexing.
+
+    Parameters
+    ----------
+    obj: ndarray, string, tuple, list
+        Object to get item from.
+    index: int, list[int], slice()
+        Desired selection to extract from obj.
+
+    Returns
+    -------
+    Selection obj[index]
+
+    """
+    try:
+        result = obj[index]
+    except IndexError as e:
+        raise ValueError("Array chunk size or shape is unknown. Possible solution with x.compute_chunk_sizes()") from e
+
+    try:
+        if not result.flags.owndata and obj.size >= 2 * result.size:
+            result = result.copy()
+    except AttributeError:
+        pass
+
+    return result
+
+
+def arange(start, stop, step, length, dtype, like=None):
+    """Arange wrapper for chunk generation.
+
+    Creates an arange and truncates if needed to match expected length.
+    """
+    if like is None:
+        res = np.arange(start, stop, step, dtype=dtype)
+    else:
+        try:
+            res = np.arange(start, stop, step, dtype=dtype, like=like)
+        except TypeError:
+            res = np.arange(start, stop, step, dtype=dtype)
+    return res[:-1] if len(res) > length else res
+
+
+def linspace(start, stop, num, endpoint=True, dtype=None):
+    """Linspace wrapper for chunk generation."""
+    return np.linspace(start, stop, num, endpoint=endpoint, dtype=dtype)
+
+
+def slice_with_int_dask_array(x, idx, offset, x_size, axis):
+    """Chunk function of `slice_with_int_dask_array_on_axis`.
+    Slice one chunk of x by one chunk of idx.
+
+    Parameters
+    ----------
+    x: ndarray, any dtype, any shape
+        i-th chunk of x
+    idx: ndarray, ndim=1, dtype=any integer
+        j-th chunk of idx (cartesian product with the chunks of x)
+    offset: ndarray, shape=(1, ), dtype=int64
+        Index of the first element along axis of the current chunk of x
+    x_size: int
+        Total size of the x da.Array along axis
+    axis: int
+        normalized axis to take elements from (0 <= axis < x.ndim)
+
+    Returns
+    -------
+    x sliced along axis, using only the elements of idx that fall inside the
+    current chunk.
+    """
+    from dask_array._utils import meta_from_array
+
+    # asarray_safe functionality - convert to appropriate array type
+    if hasattr(x, "__array_namespace__"):
+        try:
+            xp = x.__array_namespace__()
+            idx = xp.asarray(idx)
+        except (AttributeError, TypeError):
+            pass
+
+    idx = np.asarray(idx)
+    if not np.issubdtype(idx.dtype, np.integer):
+        idx = meta_from_array(x)
+
+    # Needed when idx is unsigned
+    idx = idx.astype(np.int64)
+
+    # Normalize negative indices
+    idx = np.where(idx < 0, idx + x_size, idx)
+
+    # A chunk of the offset dask Array is a numpy array with shape (1, ).
+    # It indicates the index of the first element along axis of the current
+    # chunk of x.
+    idx = idx - offset
+
+    # Drop elements of idx that do not fall inside the current chunk of x
+    idx_filter = (idx >= 0) & (idx < x.shape[axis])
+    idx = idx[idx_filter]
+
+    # np.take does not support slice indices
+    # return np.take(x, idx, axis)
+    return x[tuple(idx if i == axis else slice(None) for i in range(x.ndim))]
+
+
+def slice_with_int_dask_array_aggregate(idx, chunk_outputs, x_chunks, axis):
+    """Final aggregation function of `slice_with_int_dask_array_on_axis`.
+    Aggregate all chunks of x by one chunk of idx, reordering the output of
+    `slice_with_int_dask_array`.
+
+    Note that there is no combine function, as a recursive aggregation (e.g.
+    with split_every) would not give any benefit.
+
+    Parameters
+    ----------
+    idx: ndarray, ndim=1, dtype=any integer
+        j-th chunk of idx
+    chunk_outputs: ndarray
+        concatenation along axis of the outputs of `slice_with_int_dask_array`
+        for all chunks of x and the j-th chunk of idx
+    x_chunks: tuple
+        dask chunks of the x da.Array along axis, e.g. ``(3, 3, 2)``
+    axis: int
+        normalized axis to take elements from (0 <= axis < x.ndim)
+
+    Returns
+    -------
+    Selection from all chunks of x for the j-th chunk of idx, in the correct
+    order
+    """
+    # Needed when idx is unsigned
+    idx = idx.astype(np.int64)
+
+    # Normalize negative indices
+    idx = np.where(idx < 0, idx + sum(x_chunks), idx)
+
+    x_chunk_offset = 0
+    chunk_output_offset = 0
+
+    # Assemble the final index that picks from the output of the previous
+    # kernel by adding together one layer per chunk of x
+    # FIXME: this could probably be reimplemented with a faster search-based
+    # algorithm
+    idx_final = np.zeros_like(idx)
+    for x_chunk in x_chunks:
+        idx_filter = (idx >= x_chunk_offset) & (idx < x_chunk_offset + x_chunk)
+        idx_cum = np.cumsum(idx_filter)
+        idx_final += np.where(idx_filter, idx_cum - 1 + chunk_output_offset, 0)
+        x_chunk_offset += x_chunk
+        if idx_cum.size > 0:
+            chunk_output_offset += idx_cum[-1]
+
+    # np.take does not support slice indices
+    # return np.take(chunk_outputs, idx_final, axis)
+    return chunk_outputs[tuple(idx_final if i == axis else slice(None) for i in range(chunk_outputs.ndim))]