dask-array 0.1.0__py3-none-any.whl

dask_array/_xarray.py

@@ -0,0 +1,337 @@
+"""
+xarray ChunkManager integration for dask-array expressions.
+
+This module registers a ChunkManagerEntrypoint under the entry point name
+"dask" — the same name used by xarray's built-in DaskManager.  We *must*
+replace the built-in rather than coexist alongside it because:
+
+1. ``dask_array.Array`` is a dask collection (implements ``__dask_graph__``)
+   and a duck array, so xarray's built-in ``DaskManager.is_chunked_array``
+   recognises it via ``is_duck_dask_array``.
+2. If two managers both claim the same array type, xarray's
+   ``get_chunked_array_type`` raises
+   ``"Multiple ChunkManagers recognise type ..."``.
+3. Therefore only one "dask"-flavoured manager can be active at a time.
+
+Because both xarray and dask-array register an entry point named "dask",
+the winner of ``importlib.metadata.entry_points()`` iteration is
+non-deterministic (it depends on filesystem enumeration order).  To make
+the result reproducible, ``_ensure_registered`` mutates the cached dict
+returned by ``list_chunkmanagers()`` at import time so that our manager
+is always the one stored under the "dask" key.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint
+
+if TYPE_CHECKING:
+    from dask_array._collection import Array
+
+
+class DaskArrayExprManager(ChunkManagerEntrypoint["Array"]):
+    """
+    ChunkManager for dask-array expressions.
+
+    This integrates dask_array.Array with xarray's chunked array interface,
+    enabling expression-based optimizations for xarray operations.
+    """
+
+    array_cls: type[Array]
+    available: bool = True
+
+    def __init__(self) -> None:
+        from dask_array._collection import Array
+
+        self.array_cls = Array
+
+    def is_chunked_array(self, data: Any) -> bool:
+        return isinstance(data, self.array_cls)
+
+    def chunks(self, data: Array) -> tuple[tuple[int, ...], ...]:
+        return data.chunks
+
+    def normalize_chunks(
+        self,
+        chunks: Any,
+        shape: tuple[int, ...] | None = None,
+        limit: int | None = None,
+        dtype: np.dtype[Any] | None = None,
+        previous_chunks: tuple[tuple[int, ...], ...] | None = None,
+    ) -> tuple[tuple[int, ...], ...]:
+        from dask_array._core_utils import normalize_chunks
+
+        return normalize_chunks(
+            chunks,
+            shape=shape,
+            limit=limit,
+            dtype=dtype,
+            previous_chunks=previous_chunks,
+        )
+
+    def from_array(
+        self,
+        data: Any,
+        chunks: Any,
+        **kwargs: Any,
+    ) -> Array:
+        import dask_array as da
+        from xarray.core.indexing import ImplicitToExplicitIndexingAdapter
+
+        if isinstance(data, ImplicitToExplicitIndexingAdapter):
+            # Lazily loaded backend arrays should use NumPy arrays for meta.
+            kwargs["meta"] = np.ndarray
+
+        return da.from_array(data, chunks, **kwargs)
+
+    def rechunk(
+        self,
+        data: Array,
+        chunks: Any,
+        **kwargs: Any,
+    ) -> Array:
+        return data.rechunk(chunks, **kwargs)
+
+    def compute(
+        self,
+        *data: Array | Any,
+        **kwargs: Any,
+    ) -> tuple[np.ndarray[Any, Any], ...]:
+        from dask import compute
+
+        return compute(*data, **kwargs)
+
+    def persist(
+        self,
+        *data: Array | Any,
+        **kwargs: Any,
+    ) -> tuple[Array | Any, ...]:
+        from dask import persist
+
+        return persist(*data, **kwargs)
+
+    @property
+    def array_api(self) -> Any:
+        import dask_array as da
+
+        return da
+
+    def reduction(
+        self,
+        arr: Array,
+        func: Callable[..., Any],
+        combine_func: Callable[..., Any] | None = None,
+        aggregate_func: Callable[..., Any] | None = None,
+        axis: int | Sequence[int] | None = None,
+        dtype: np.dtype[Any] | None = None,
+        keepdims: bool = False,
+    ) -> Array:
+        from dask_array import reduction
+
+        return reduction(
+            arr,
+            chunk=func,
+            combine=combine_func,
+            aggregate=aggregate_func,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+        )
+
+    def scan(
+        self,
+        func: Callable[..., Any],
+        binop: Callable[..., Any],
+        ident: float,
+        arr: Array,
+        axis: int | None = None,
+        dtype: np.dtype[Any] | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        from dask_array import cumreduction
+
+        return cumreduction(
+            func,
+            binop,
+            ident,
+            arr,
+            axis=axis,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    def apply_gufunc(
+        self,
+        func: Callable[..., Any],
+        signature: str,
+        *args: Any,
+        axes: Sequence[tuple[int, ...]] | None = None,
+        axis: int | None = None,
+        keepdims: bool = False,
+        output_dtypes: Sequence[np.dtype[Any]] | None = None,
+        output_sizes: dict[str, int] | None = None,
+        vectorize: bool | None = None,
+        allow_rechunk: bool = False,
+        meta: tuple[np.ndarray[Any, Any], ...] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        from dask_array import apply_gufunc
+
+        return apply_gufunc(
+            func,
+            signature,
+            *args,
+            axes=axes,
+            axis=axis,
+            keepdims=keepdims,
+            output_dtypes=output_dtypes,
+            output_sizes=output_sizes,
+            vectorize=vectorize,
+            allow_rechunk=allow_rechunk,
+            meta=meta,
+            **kwargs,
+        )
+
+    def map_blocks(
+        self,
+        func: Callable[..., Any],
+        *args: Any,
+        dtype: np.dtype[Any] | None = None,
+        chunks: tuple[int, ...] | None = None,
+        drop_axis: int | Sequence[int] | None = None,
+        new_axis: int | Sequence[int] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        from dask_array import map_blocks
+
+        return map_blocks(
+            func,
+            *args,
+            dtype=dtype,
+            chunks=chunks,
+            drop_axis=drop_axis,
+            new_axis=new_axis,
+            **kwargs,
+        )
+
+    def map_blocks_multi_output(
+        self,
+        func: Callable[..., Any],
+        input_exprs: Sequence[Any],
+        input_indices: Sequence[Iterable[Any]],
+        shared_indices: Iterable[Any],
+        block_specs: Mapping[tuple[int, ...], Any],
+        outputs: Sequence[Mapping[str, Any]],
+        *,
+        token: str,
+    ) -> list[Array]:
+        from dask_array._map_blocks import map_blocks_multi_output
+
+        return map_blocks_multi_output(
+            func,
+            input_exprs,
+            input_indices,
+            shared_indices,
+            block_specs,
+            outputs,
+            token=token,
+        )
+
+    def blockwise(
+        self,
+        func: Callable[..., Any],
+        out_ind: Iterable[Any],
+        *args: Any,
+        name: str | None = None,
+        token: Any | None = None,
+        dtype: np.dtype[Any] | None = None,
+        adjust_chunks: dict[Any, Callable[..., Any]] | None = None,
+        new_axes: dict[Any, int] | None = None,
+        align_arrays: bool = True,
+        concatenate: bool | None = None,
+        meta: tuple[np.ndarray[Any, Any], ...] | None = None,
+        **kwargs: Any,
+    ) -> Array:
+        from dask_array import blockwise
+
+        return blockwise(
+            func,
+            out_ind,
+            *args,
+            name=name,
+            token=token,
+            dtype=dtype,
+            adjust_chunks=adjust_chunks,
+            new_axes=new_axes,
+            align_arrays=align_arrays,
+            concatenate=concatenate,
+            meta=meta,
+            **kwargs,
+        )
+
+    def unify_chunks(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> tuple[dict[str, tuple[tuple[int, ...], ...]], list[Array]]:
+        from dask_array import unify_chunks
+
+        return unify_chunks(*args, **kwargs)
+
+    def store(
+        self,
+        sources: Array | Sequence[Array],
+        targets: Any,
+        **kwargs: Any,
+    ) -> Any:
+        from dask_array import store
+
+        return store(
+            sources=sources,
+            targets=targets,
+            **kwargs,
+        )
+
+    def shuffle(
+        self,
+        x: Array,
+        indexer: list[list[int]],
+        axis: int,
+        chunks: Any,
+    ) -> Array:
+        from dask_array import shuffle
+
+        if chunks is None:
+            chunks = "auto"
+        if chunks != "auto":
+            raise NotImplementedError("Only chunks='auto' is supported at present.")
+        return shuffle(x, indexer, axis, chunks="auto")
+
+    def get_auto_chunk_size(self) -> int:
+        from dask import config as dask_config
+        from dask.utils import parse_bytes
+
+        return parse_bytes(dask_config.get("array.chunk-size"))
+
+
+def _ensure_registered() -> None:
+    """Ensure DaskArrayExprManager is the "dask" chunk manager in xarray.
+
+    Both xarray and this package register an entry point named "dask" under
+    the ``xarray.chunkmanagers`` group.  ``list_chunkmanagers`` builds a dict
+    from those entry points, so the *last* one enumerated wins.  Because
+    ``importlib.metadata.entry_points`` iteration order is non-deterministic,
+    we fix the race here by replacing the cached value after the fact.
+    """
+    try:
+        from xarray.namedarray.parallelcompat import list_chunkmanagers
+    except ImportError:
+        return
+
+    managers = list_chunkmanagers()
+    if not isinstance(managers.get("dask"), DaskArrayExprManager):
+        managers["dask"] = DaskArrayExprManager()

dask_array/core/__init__.py

@@ -0,0 +1,34 @@
+"""Core array types and wrapping functions.
+
+This module re-exports the core Array class and conversion functions.
+"""
+
+from dask_array.core._blockwise_funcs import blockwise, elemwise
+from dask_array.core._conversion import (
+    array,
+    asanyarray,
+    asarray,
+    from_array,
+)
+from dask_array.core._from_graph import from_graph
+
+
+def __getattr__(name):
+    """Lazy import of Array to avoid circular imports."""
+    if name == "Array":
+        from dask_array._collection import Array
+
+        return Array
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "Array",
+    "from_array",
+    "from_graph",
+    "asarray",
+    "asanyarray",
+    "array",
+    "blockwise",
+    "elemwise",
+]

dask_array/core/_blockwise_funcs.py

@@ -0,0 +1,312 @@
+"""Blockwise and elemwise function wrappers."""
+
+from __future__ import annotations
+
+import numpy as np
+import toolz
+
+from dask_array._new_collection import new_collection
+from dask_array._blockwise import Blockwise, Elemwise
+from dask_array._core_utils import is_scalar_for_elemwise
+
+
+def blockwise(
+    func,
+    out_ind,
+    *args,
+    name=None,
+    token=None,
+    dtype=None,
+    adjust_chunks=None,
+    new_axes=None,
+    align_arrays=True,
+    concatenate=None,
+    meta=None,
+    **kwargs,
+):
+    """Tensor operation: Generalized inner and outer products
+
+    A broad class of blocked algorithms and patterns can be specified with a
+    concise multi-index notation.  The ``blockwise`` function applies an in-memory
+    function across multiple blocks of multiple inputs in a variety of ways.
+    Many dask.array operations are special cases of blockwise including
+    elementwise, broadcasting, reductions, tensordot, and transpose.
+
+    Parameters
+    ----------
+    func : callable
+        Function to apply to individual tuples of blocks
+    out_ind : iterable
+        Block pattern of the output, something like 'ijk' or (1, 2, 3)
+    *args : sequence of Array, index pairs
+        You may also pass literal arguments, accompanied by None index
+        e.g. (x, 'ij', y, 'jk', z, 'i', some_literal, None)
+    **kwargs : dict
+        Extra keyword arguments to pass to function
+    dtype : np.dtype
+        Datatype of resulting array.
+    concatenate : bool, keyword only
+        If true concatenate arrays along dummy indices, else provide lists
+    adjust_chunks : dict
+        Dictionary mapping index to function to be applied to chunk sizes
+    new_axes : dict, keyword only
+        New indexes and their dimension lengths
+    align_arrays: bool
+        Whether or not to align chunks along equally sized dimensions when
+        multiple arrays are provided.  This allows for larger chunks in some
+        arrays to be broken into smaller ones that match chunk sizes in other
+        arrays such that they are compatible for block function mapping. If
+        this is false, then an error will be thrown if arrays do not already
+        have the same number of blocks in each dimension.
+
+    Examples
+    --------
+    2D embarrassingly parallel operation from two arrays, x, and y.
+
+    >>> import operator, numpy as np, dask.array as da
+    >>> x = da.from_array([[1, 2],
+    ...                    [3, 4]], chunks=(1, 2))
+    >>> y = da.from_array([[10, 20],
+    ...                    [0, 0]])
+    >>> z = blockwise(operator.add, 'ij', x, 'ij', y, 'ij', dtype='f8')
+    >>> z.compute()
+    array([[11, 22],
+           [ 3,  4]])
+
+    Outer product multiplying a by b, two 1-d vectors
+
+    >>> a = da.from_array([0, 1, 2], chunks=1)
+    >>> b = da.from_array([10, 50, 100], chunks=1)
+    >>> z = blockwise(np.outer, 'ij', a, 'i', b, 'j', dtype='f8')
+    >>> z.compute()
+    array([[  0,   0,   0],
+           [ 10,  50, 100],
+           [ 20, 100, 200]])
+
+    z = x.T
+
+    >>> z = blockwise(np.transpose, 'ji', x, 'ij', dtype=x.dtype)
+    >>> z.compute()
+    array([[1, 3],
+           [2, 4]])
+
+    The transpose case above is illustrative because it does transposition
+    both on each in-memory block by calling ``np.transpose`` and on the order
+    of the blocks themselves, by switching the order of the index ``ij -> ji``.
+
+    We can compose these same patterns with more variables and more complex
+    in-memory functions
+
+    z = X + Y.T
+
+    >>> z = blockwise(lambda x, y: x + y.T, 'ij', x, 'ij', y, 'ji', dtype='f8')
+    >>> z.compute()
+    array([[11,  2],
+           [23,  4]])
+
+    Any index, like ``i`` missing from the output index is interpreted as a
+    contraction (note that this differs from Einstein convention; repeated
+    indices do not imply contraction.)  In the case of a contraction the passed
+    function should expect an iterable of blocks on any array that holds that
+    index.  To receive arrays concatenated along contracted dimensions instead
+    pass ``concatenate=True``.
+
+    Inner product multiplying a by b, two 1-d vectors
+
+    >>> def sequence_dot(a_blocks, b_blocks):
+    ...     result = 0
+    ...     for a, b in zip(a_blocks, b_blocks):
+    ...         result += a.dot(b)
+    ...     return result
+
+    >>> z = blockwise(sequence_dot, '', a, 'i', b, 'i', dtype='f8')
+    >>> z.compute()
+    250
+
+    Add new single-chunk dimensions with the ``new_axes=`` keyword, including
+    the length of the new dimension.  New dimensions will always be in a single
+    chunk.
+
+    >>> def f(a):
+    ...     return a[:, None] * np.ones((1, 5))
+
+    >>> z = blockwise(f, 'az', a, 'a', new_axes={'z': 5}, dtype=a.dtype)
+
+    New dimensions can also be multi-chunk by specifying a tuple of chunk
+    sizes.  This has limited utility as is (because the chunks are all the
+    same), but the resulting graph can be modified to achieve more useful
+    results (see ``da.map_blocks``).
+
+    >>> z = blockwise(f, 'az', a, 'a', new_axes={'z': (5, 5)}, dtype=x.dtype)
+    >>> z.chunks
+    ((1, 1, 1), (5, 5))
+
+    If the applied function changes the size of each chunk you can specify this
+    with a ``adjust_chunks={...}`` dictionary holding a function for each index
+    that modifies the dimension size in that index.
+
+    >>> def double(x):
+    ...     return np.concatenate([x, x])
+
+    >>> y = blockwise(double, 'ij', x, 'ij',
+    ...               adjust_chunks={'i': lambda n: 2 * n}, dtype=x.dtype)
+    >>> y.chunks
+    ((2, 2), (2,))
+
+    Include literals by indexing with None
+
+    >>> z = blockwise(operator.add, 'ij', x, 'ij', 1234, None, dtype=x.dtype)
+    >>> z.compute()
+    array([[1235, 1236],
+           [1237, 1238]])
+    """
+    from dask_array.core import asanyarray
+
+    new_axes = new_axes or {}
+
+    # Normalize dtype to numpy dtype (handles cases like dtype=float)
+    if dtype is not None:
+        dtype = np.dtype(dtype)
+
+    # Input Validation
+    if len(set(out_ind)) != len(out_ind):
+        raise ValueError(
+            "Repeated elements not allowed in output index",
+            [k for k, v in toolz.frequencies(out_ind).items() if v > 1],
+        )
+    new = set(out_ind) - {a for arg in args[1::2] if arg is not None for a in arg} - set(new_axes or ())
+    if new:
+        raise ValueError("Unknown dimension", new)
+
+    # Convert scalars with empty tuple index to 0-d dask arrays
+    # This mirrors what traditional unify_chunks does
+    normalized_args = []
+    for arg, ind in toolz.partition(2, args):
+        if ind == () and not hasattr(arg, "chunks"):
+            arg = asanyarray(arg)
+        normalized_args.extend([arg, ind])
+
+    return new_collection(
+        Blockwise(
+            func,
+            out_ind,
+            name,
+            token,
+            dtype,
+            adjust_chunks,
+            new_axes,
+            align_arrays,
+            concatenate,
+            meta,
+            kwargs,
+            *normalized_args,
+        )
+    )
+
+
+def elemwise(op, *args, out=None, where=True, dtype=None, name=None, **kwargs):
+    """Apply an elementwise ufunc-like function blockwise across arguments.
+
+    Like numpy ufuncs, broadcasting rules are respected.
+
+    Parameters
+    ----------
+    op : callable
+        The function to apply. Should be numpy ufunc-like in the parameters
+        that it accepts.
+    *args : Any
+        Arguments to pass to `op`. Non-dask array-like objects are first
+        converted to dask arrays, then all arrays are broadcast together before
+        applying the function blockwise across all arguments. Any scalar
+        arguments are passed as-is following normal numpy ufunc behavior.
+    out : dask array, optional
+        If out is a dask.array then this overwrites the contents of that array
+        with the result.
+    where : array_like, optional
+        An optional boolean mask marking locations where the ufunc should be
+        applied. Can be a scalar, dask array, or any other array-like object.
+        Mirrors the ``where`` argument to numpy ufuncs, see e.g. ``numpy.add``
+        for more information.
+    dtype : dtype, optional
+        If provided, overrides the output array dtype.
+    name : str, optional
+        A unique key name to use when building the backing dask graph. If not
+        provided, one will be automatically generated based on the input
+        arguments.
+    **kwargs : dict
+        Additional keyword arguments to pass to `op`.
+
+    Examples
+    --------
+    >>> elemwise(add, x, y)  # doctest: +SKIP
+    >>> elemwise(sin, x)  # doctest: +SKIP
+    >>> elemwise(sin, x, out=dask_array)  # doctest: +SKIP
+
+    See Also
+    --------
+    blockwise
+    """
+    # Lazy import to avoid circular dependency
+    from dask_array.core import asanyarray
+
+    # Normalize where parameter
+    if where is True:
+        pass  # keep as True
+    elif where is False or where is None:
+        where = False
+    else:
+        # Convert to dask array
+        where = asanyarray(where)
+
+    # Normalize out parameter
+    out = _normalize_out(out)
+
+    args = [np.asarray(a) if isinstance(a, (list, tuple)) else a for a in args]
+
+    # Only convert non-scalar arguments to dask arrays
+    # Scalars are kept as-is to preserve proper dtype behavior (e.g., 2.0 * float32_array = float32)
+    args = [asanyarray(a) if not is_scalar_for_elemwise(a) else a for a in args]
+
+    user_kwargs = dict(kwargs) if kwargs else None
+
+    result = new_collection(Elemwise(op, dtype, name, where, out, user_kwargs, *args))
+
+    return _handle_out(out, result)
+
+
+def _normalize_out(out):
+    """Normalize out parameter for elemwise operations."""
+    from dask_array._collection import Array
+
+    if isinstance(out, tuple):
+        if len(out) == 1:
+            out = out[0]
+        elif len(out) > 1:
+            raise NotImplementedError("The out parameter is not fully supported")
+        else:
+            out = None
+    if not (out is None or isinstance(out, Array)):
+        raise NotImplementedError(
+            f"The out parameter is not fully supported. Received type {type(out).__name__}, expected Dask Array"
+        )
+    return out
+
+
+def _handle_out(out, result):
+    """Handle out parameters for array-expr.
+
+    If out is a dask Array then this overwrites the contents of that array with
+    the result by replacing its internal expression.
+    """
+    from dask_array._collection import Array
+
+    if isinstance(out, Array):
+        if out.shape != result.shape:
+            raise ValueError(
+                f"Mismatched shapes between result and out parameter. out={out.shape}, result={result.shape}"
+            )
+        # Modify the out array in-place by replacing its expression
+        out._expr = result._expr
+        return out
+    else:
+        return result