dask-array 0.1.0__py3-none-any.whl

dask_array/slicing/_blocks.py

@@ -0,0 +1,138 @@
+"""Block indexing expression for x.blocks[...] access."""
+
+from __future__ import annotations
+
+import functools
+import math
+from itertools import product
+from numbers import Number
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask._task_spec import Alias
+from dask_array._expr import ArrayExpr
+from dask_array.slicing._utils import normalize_index
+
+
+class BlockView:
+    """An array-like interface to the blocks of an array.
+
+    BlockView provides an array-like interface to the blocks of a dask array.
+    Numpy-style indexing of a BlockView returns a selection of blocks as a
+    new dask array.
+
+    You can index BlockView like a numpy array of shape equal to the number
+    of blocks in each dimension (available as array.blocks.size). The
+    dimensionality of the output array matches the dimension of this array,
+    even if integer indices are passed. Slicing with np.newaxis or multiple
+    lists is not supported.
+    """
+
+    __slots__ = ("_array",)
+
+    def __init__(self, array):
+        self._array = array
+
+    def __getitem__(self, index):
+        return new_collection(blocks_getitem(self._array.expr, index))
+
+    def __eq__(self, other):
+        # Check if other is any BlockView type (including legacy)
+        if hasattr(other, "_array") and type(other).__name__ == "BlockView":
+            return self._array is other._array
+        return NotImplemented
+
+    @property
+    def size(self):
+        """The total number of blocks in the array."""
+        return math.prod(self.shape)
+
+    @property
+    def shape(self):
+        """The number of blocks per axis. Alias of dask.array.numblocks."""
+        return self._array.numblocks
+
+    def ravel(self):
+        """Return a flattened list of all the blocks in the array in C order."""
+        return [self[idx] for idx in np.ndindex(self.shape)]
+
+
+class Blocks(ArrayExpr):
+    """Expression for block-based indexing (x.blocks[...]).
+
+    This expression allows accessing array blocks by block index rather than
+    element index. The index is normalized to always use slices (never integers)
+    to preserve dimensionality.
+
+    Parameters
+    ----------
+    array : ArrayExpr
+        The source array expression
+    index : tuple
+        Normalized block indices (after converting integers to length-1 slices)
+    """
+
+    _parameters = ["array", "index"]
+
+    @functools.cached_property
+    def _name(self):
+        return f"blocks-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _meta(self):
+        return self.array._meta
+
+    @functools.cached_property
+    def chunks(self):
+        """Compute chunks by selecting from the source array's chunks."""
+        return tuple(tuple(np.array(c)[idx].tolist()) for c, idx in zip(self.array.chunks, self.index))
+
+    def _layer(self) -> dict:
+        """Generate the task graph layer.
+
+        Each output block is an alias to the corresponding input block.
+        """
+        # Pre-compute index mappings for each dimension
+        index_maps = [np.arange(n)[idx] for n, idx in zip(self.array.numblocks, self.index)]
+
+        dsk = {}
+        for out_key in product(*(range(len(c)) for c in self.chunks)):
+            in_key = tuple(int(m[i]) for m, i in zip(index_maps, out_key))
+            out_name = (self._name,) + out_key
+            in_name = (self.array._name,) + in_key
+            dsk[out_name] = Alias(out_name, in_name)
+
+        return dsk
+
+
+def blocks_getitem(array, index):
+    """Create a Blocks expression for block indexing.
+
+    Parameters
+    ----------
+    array : ArrayExpr
+        The source array expression
+    index : tuple
+        The block index (may contain integers, slices, or lists)
+
+    Returns
+    -------
+    Blocks
+        The blocks expression
+    """
+    if not isinstance(index, tuple):
+        index = (index,)
+
+    if sum(isinstance(ind, (np.ndarray, list)) for ind in index) > 1:
+        raise ValueError("Can only slice with a single list")
+    if any(ind is None for ind in index):
+        raise ValueError("Slicing with np.newaxis or None is not supported")
+
+    # Normalize index to array's numblocks
+    index = normalize_index(index, array.numblocks)
+
+    # Convert integers to length-1 slices to preserve dimensionality
+    index = tuple(slice(k, k + 1) if isinstance(k, Number) else k for k in index)
+
+    return Blocks(array, index)

dask_array/slicing/_bool_index.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import functools
+import warnings
+from operator import getitem
+
+import numpy as np
+
+from dask._task_spec import Alias
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+
+
+def getitem_variadic(x, *index):
+    """Helper function for boolean indexing."""
+    return x[index]
+
+
+def slice_with_bool_dask_array(x, index):
+    """Slice x with one or more dask arrays of bools.
+
+    This is a helper function of :meth:`Array.__getitem__`.
+
+    Parameters
+    ----------
+    x: Array
+    index: tuple with as many elements as x.ndim, among which there are
+           one or more Array's with dtype=bool
+
+    Returns
+    -------
+    tuple of (sliced x, new index)
+
+    where the new index is the same as the input, but with slice(None)
+    replaced to the original slicer when a filter has been applied.
+
+    Note: The sliced x will have nan chunks on the sliced axes.
+    """
+    from dask_array._collection import (
+        Array,
+        blockwise,
+        elemwise,
+        new_collection,
+    )
+    from dask_array._expr import ChunksOverride
+
+    out_index = [slice(None) if isinstance(ind, Array) and ind.dtype == bool else ind for ind in index]
+
+    # Case 1: Full-dimensional boolean mask
+    if len(index) == 1 and index[0].ndim == x.ndim:
+        if not np.isnan(x.shape).any() and not np.isnan(index[0].shape).any():
+            x = x.ravel()
+            index = tuple(i.ravel() for i in index)
+        elif x.ndim > 1:
+            warnings.warn(
+                "When slicing a Dask array of unknown chunks with a boolean mask "
+                "Dask array, the output array may have a different ordering "
+                "compared to the equivalent NumPy operation. This will raise an "
+                "error in a future release of Dask.",
+                stacklevel=3,
+            )
+        # Use elemwise to apply getitem across blocks
+        y = elemwise(getitem, x, index[0], dtype=x.dtype)
+        # Trigger eager chunk validation to match legacy behavior
+        # This will raise if x and index have incompatible chunks
+        _ = y.chunks
+        result = BooleanIndexFlattened(y.expr)
+        return new_collection(result), out_index
+
+    # Case 2: 1D boolean arrays on specific dimensions
+    if any(isinstance(ind, Array) and ind.dtype == bool and ind.ndim != 1 for ind in index):
+        raise NotImplementedError(
+            "Slicing with dask.array of bools only permitted when "
+            "the indexer has only one dimension or when "
+            "it has the same dimension as the sliced "
+            "array"
+        )
+
+    indexes = [ind if isinstance(ind, Array) and ind.dtype == bool else slice(None) for ind in index]
+
+    # Track which dimension indices have boolean arrays
+    dsk_ind = []
+
+    from toolz import concat
+
+    arginds = []
+    i = 0
+    for dim, ind in enumerate(indexes):
+        if isinstance(ind, Array) and ind.dtype == bool:
+            dsk_ind.append(dim)
+            new = (ind, tuple(range(i, i + ind.ndim)))
+            i += x.ndim
+        else:
+            new = (slice(None), None)
+            i += 1
+        arginds.append(new)
+
+    arginds = list(concat(arginds))
+
+    out = blockwise(
+        getitem_variadic,
+        tuple(range(x.ndim)),
+        x,
+        tuple(range(x.ndim)),
+        *arginds,
+        dtype=x.dtype,
+    )
+
+    # For boolean indexing, override chunks on boolean-indexed dimensions
+    # with nan values since the output size is unknown
+    new_chunks = tuple(
+        tuple(np.nan for _ in range(len(c))) if dim in dsk_ind else c for dim, c in enumerate(out.chunks)
+    )
+    result = ChunksOverride(out.expr, new_chunks)
+    return new_collection(result), tuple(out_index)
+
+
+class BooleanIndexFlattened(ArrayExpr):
+    """Flattens the output of a full-dimensional boolean index operation."""
+
+    _parameters = ["array"]
+
+    @functools.cached_property
+    def _name(self):
+        return f"getitem-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.array._meta, ndim=1)
+
+    @functools.cached_property
+    def chunks(self):
+        # Total number of blocks = product of numblocks
+        from functools import reduce
+        from operator import mul
+
+        nblocks = reduce(mul, self.array.numblocks, 1)
+        return ((np.nan,) * nblocks,)
+
+    def _layer(self) -> dict:
+        from dask.base import flatten
+
+        # Flatten the keys from the elemwise result
+        keys = list(flatten(self.array.__dask_keys__()))
+        return {(self._name, i): Alias((self._name, i), k) for i, k in enumerate(keys)}

dask_array/slicing/_setitem.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import functools
+import math
+from itertools import product
+
+import numpy as np
+
+from dask._task_spec import Alias, List, Task, TaskRef
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+from dask_array._core_utils import concatenate3 as concatenate_shaped
+from dask_array.slicing._utils import parse_assignment_indices, setitem
+from dask.base import is_dask_collection
+from dask.core import flatten
+from dask.utils import cached_cumsum
+
+
+def setitem_array_expr(out_name, array, indices, value):
+    """Array-expr version of setitem_array that generates Task objects directly.
+
+    This function creates a new dask graph that assigns values to each block
+    that is touched by the indices, leaving other blocks unchanged.
+    """
+    array_shape = array.shape
+    value_shape = value.shape
+    value_ndim = len(value_shape)
+
+    # Reformat input indices
+    indices, implied_shape, reverse, implied_shape_positions = parse_assignment_indices(indices, array_shape)
+
+    # Empty slices can only be assigned size 1 values
+    if 0 in implied_shape and value_shape and max(value_shape) > 1:
+        raise ValueError(
+            f"shape mismatch: value array of shape {value_shape} "
+            "could not be broadcast to indexing result "
+            f"of shape {tuple(implied_shape)}"
+        )
+
+    # Set variables needed when creating the part of the assignment value
+    offset = len(implied_shape) - value_ndim
+    if offset >= 0:
+        array_common_shape = implied_shape[offset:]
+        value_common_shape = value_shape
+        value_offset = 0
+        reverse = [i - offset for i in reverse if i >= offset]
+    else:
+        value_offset = -offset
+        array_common_shape = implied_shape
+        value_common_shape = value_shape[value_offset:]
+        offset = 0
+        if value_shape[:value_offset] != (1,) * value_offset:
+            raise ValueError(
+                f"could not broadcast input array from shape{value_shape} into shape {tuple(implied_shape)}"
+            )
+
+    base_value_indices = []
+    non_broadcast_dimensions = []
+
+    for i, (a, b, j) in enumerate(zip(array_common_shape, value_common_shape, implied_shape_positions)):
+        index = indices[j]
+        if is_dask_collection(index) and index.dtype == bool:
+            if math.isnan(b) or b <= index.size:
+                base_value_indices.append(None)
+                non_broadcast_dimensions.append(i)
+            else:
+                raise ValueError(
+                    f"shape mismatch: value array dimension size of {b} is "
+                    "greater then corresponding boolean index size of "
+                    f"{index.size}"
+                )
+            continue
+
+        if b == 1:
+            base_value_indices.append(slice(None))
+        elif a == b:
+            base_value_indices.append(None)
+            non_broadcast_dimensions.append(i)
+        elif math.isnan(a):
+            base_value_indices.append(None)
+            non_broadcast_dimensions.append(i)
+        else:
+            raise ValueError(
+                f"shape mismatch: value array of shape {value_shape} "
+                "could not be broadcast to indexing result of shape "
+                f"{tuple(implied_shape)}"
+            )
+
+    # Translate chunks tuple to array locations
+    chunks = array.chunks
+    cumdims = [cached_cumsum(bds, initial_zero=True) for bds in chunks]
+    array_locations = [[(s, s + dim) for s, dim in zip(starts, shapes)] for starts, shapes in zip(cumdims, chunks)]
+    array_locations = product(*array_locations)
+
+    in_keys = list(flatten(array.__dask_keys__()))
+
+    # Build graph with Task objects
+    dsk = {}
+    out_name_tuple = (out_name,)
+
+    # Helper closures for index handling (simplified from legacy)
+    def block_index_from_1d_index(index, loc0, loc1, is_bool):
+        if is_bool:
+            return index[loc0:loc1]
+        elif is_dask_collection(index):
+            i = np.where((loc0 <= index) & (index < loc1), index, loc1)
+            return i - loc0
+        else:
+            i = np.where((loc0 <= index) & (index < loc1))[0]
+            return index[i] - loc0
+
+    def block_index_shape_from_1d_bool_index(index, loc0, loc1):
+        return np.sum(index[loc0:loc1])
+
+    def n_preceding_from_1d_bool_index(index, loc0):
+        return np.sum(index[:loc0])
+
+    def value_indices_from_1d_int_index(index, vsize, loc0, loc1):
+        if is_dask_collection(index):
+            if np.isnan(index.size):
+                i = np.where((loc0 <= index) & (index < loc1), True, False)
+                i = concatenate_array_chunks_expr(i)
+                i._chunks = ((vsize,),)
+            else:
+                i = np.where((loc0 <= index) & (index < loc1))[0]
+                i = concatenate_array_chunks_expr(i)
+        else:
+            i = np.where((loc0 <= index) & (index < loc1))[0]
+        return i
+
+    for in_key, locations in zip(in_keys, array_locations):
+        block_indices = []
+        block_indices_shape = []
+        block_preceding_sizes = []
+        overlaps = True
+        dim_1d_int_index = None
+
+        for dim, (index, (loc0, loc1)) in enumerate(zip(indices, locations)):
+            integer_index = isinstance(index, int)
+            if isinstance(index, slice):
+                stop = loc1 - loc0
+                if index.stop < loc1:
+                    stop -= loc1 - index.stop
+                start = index.start - loc0
+                if start < 0:
+                    start %= index.step
+                if start >= stop:
+                    overlaps = False
+                    break
+                step = index.step
+                block_index = slice(start, stop, step)
+                block_index_size, rem = divmod(stop - start, step)
+                if rem:
+                    block_index_size += 1
+                pre = index.indices(loc0)
+                n_preceding, rem = divmod(pre[1] - pre[0], step)
+                if rem:
+                    n_preceding += 1
+            elif integer_index:
+                if not loc0 <= index < loc1:
+                    overlaps = False
+                    break
+                block_index = index - loc0
+            else:
+                is_bool = index.dtype == bool
+                block_index = block_index_from_1d_index(index, loc0, loc1, is_bool)
+                if is_bool:
+                    block_index_size = block_index_shape_from_1d_bool_index(index, loc0, loc1)
+                    n_preceding = n_preceding_from_1d_bool_index(index, loc0)
+                else:
+                    block_index_size = None
+                    n_preceding = None
+                    dim_1d_int_index = dim
+                    loc0_loc1 = loc0, loc1
+
+                if not is_dask_collection(index) and not block_index.size:
+                    overlaps = False
+                    break
+
+            block_indices.append(block_index)
+            if not integer_index:
+                block_indices_shape.append(block_index_size)
+                block_preceding_sizes.append(n_preceding)
+
+        out_key = out_name_tuple + in_key[1:]
+
+        if not overlaps:
+            dsk[out_key] = Alias(out_key, in_key)
+            continue
+
+        # Build value indices for this block
+        value_indices = base_value_indices[:]
+        for i in non_broadcast_dimensions:
+            j = i + offset
+            if j == dim_1d_int_index:
+                value_indices[i] = value_indices_from_1d_int_index(
+                    indices[j], value_shape[i + value_offset], *loc0_loc1
+                )
+            else:
+                start = block_preceding_sizes[j]
+                value_indices[i] = slice(start, start + block_indices_shape[j])
+
+        for i in reverse:
+            size = value_common_shape[i]
+            start, stop, step = value_indices[i].indices(size)
+            size -= 1
+            start = size - start
+            stop = size - stop
+            if stop < 0:
+                stop = None
+            value_indices[i] = slice(start, stop, -1)
+
+        if value_ndim > len(indices):
+            value_indices.insert(0, Ellipsis)
+
+        # Get the value slice and concatenate to single chunk
+        v = value[tuple(value_indices)]
+        v = concatenate_array_chunks_expr(v)
+        v_key = next(flatten(v.__dask_keys__()))
+
+        # Merge value's graph into dsk
+        dsk.update(dict(v.__dask_graph__()))
+
+        # Convert block_indices to use TaskRef for any dask keys
+        task_block_indices = []
+        for idx in block_indices:
+            if is_dask_collection(idx):
+                idx = concatenate_array_chunks_expr(idx)
+                idx_key = next(flatten(idx.__dask_keys__()))
+                dsk.update(dict(idx.__dask_graph__()))
+                task_block_indices.append(TaskRef(idx_key))
+            else:
+                task_block_indices.append(idx)
+
+        # Create Task with proper TaskRef wrappers
+        dsk[out_key] = Task(
+            out_key,
+            setitem,
+            TaskRef(in_key),
+            TaskRef(v_key),
+            List(*task_block_indices),
+        )
+
+    return dsk
+
+
+class SetItem(ArrayExpr):
+    """Expression for array assignment (setitem)."""
+
+    _parameters = ["array", "index", "value"]
+
+    @functools.cached_property
+    def _name(self):
+        return f"setitem-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _meta(self):
+        meta = meta_from_array(self.array._meta, ndim=self.array.ndim)
+        if np.isscalar(meta):
+            meta = np.array(meta)
+        return meta
+
+    @property
+    def chunks(self):
+        return self.array.chunks
+
+    def _layer(self) -> dict:
+        from dask_array._collection import Array
+
+        # Wrap expressions as Array for setitem_array_expr
+        array = Array(self.array)
+        value = Array(self.value) if hasattr(self.value, "_meta") else self.value
+
+        return setitem_array_expr(self._name, array, self.index, value)
+
+
+class ConcatenateArrayChunks(ArrayExpr):
+    """Concatenate all chunks of an array into a single chunk.
+
+    This is an array-expr version of dask.array.slicing.concatenate_array_chunks.
+    """
+
+    _parameters = ["array"]
+
+    @functools.cached_property
+    def _name(self):
+        return f"concatenate-shaped-{self.deterministic_token}"
+
+    @functools.cached_property
+    def _meta(self):
+        return meta_from_array(self.array._meta, ndim=self.array.ndim)
+
+    @functools.cached_property
+    def chunks(self):
+        # Single chunk containing all the data
+        shape = self.array.shape
+        if not shape:
+            return ((1,),)
+        return tuple((s,) for s in shape)
+
+    def _layer(self) -> dict:
+        from dask.base import flatten
+
+        # Get all keys from the input array as TaskRefs
+        keys = [TaskRef(k) for k in flatten(self.array.__dask_keys__())]
+        # Output key has ndim indices, all 0 since we have a single chunk
+        out_key = (self._name,) + (0,) * self.array.ndim
+
+        return {
+            out_key: Task(
+                out_key,
+                concatenate_shaped,
+                List(*keys),
+                self.array.numblocks,
+            )
+        }
+
+
+def concatenate_array_chunks_expr(x):
+    """Concatenate all chunks of an array into a single chunk.
+
+    Array-expr version of dask.array.slicing.concatenate_array_chunks.
+    """
+    from dask_array._new_collection import new_collection
+
+    if x.npartitions == 1:
+        return x
+
+    return new_collection(ConcatenateArrayChunks(x.expr))