dask-array 0.1.0__py3-none-any.whl

dask_array/__init__.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import numpy as np
+
+import dask_array._backends
+from dask_array import _chunk as chunk
+from dask_array._core_utils import PerformanceWarning
+from dask.base import compute
+from dask_array._chunk_types import register_chunk_type
+
+from dask_array import fft, random
+from dask_array._collection import (
+    Array,
+    array,
+    asanyarray,
+    asarray,
+    atleast_1d,
+    atleast_2d,
+    atleast_3d,
+    block,
+    blockwise,
+    broadcast_to,
+    concatenate,
+    dstack,
+    elemwise,
+    expand_dims,
+    flip,
+    fliplr,
+    flipud,
+    from_array,
+    hstack,
+    moveaxis,
+    ravel,
+    rechunk,
+    reshape,
+    reshape_blockwise,
+    roll,
+    rollaxis,
+    rot90,
+    squeeze,
+    stack,
+    swapaxes,
+    transpose,
+    vstack,
+)
+from dask_array._einsum import einsum
+from dask_array._gufunc import *
+from dask_array._histogram import histogram, histogram2d, histogramdd
+from dask_array._map_blocks import map_blocks
+from dask_array._overlap import map_overlap, overlap, trim_overlap
+from dask_array._routines import (
+    aligned_coarsen_chunks,
+    allclose,
+    append,
+    apply_along_axis,
+    apply_over_axes,
+    argtopk,
+    argwhere,
+    around,
+    average,
+    bincount,
+    broadcast_arrays,
+    choose,
+    coarsen,
+    compress,
+    corrcoef,
+    count_nonzero,
+    cov,
+    delete,
+    digitize,
+    ediff1d,
+    extract,
+    flatnonzero,
+    gradient,
+    insert,
+    isclose,
+    iscomplexobj,
+    isin,
+    isnull,
+    ndim,
+    nonzero,
+    notnull,
+    outer,
+    piecewise,
+    ptp,
+    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,
+)
+from dask_array._shuffle import shuffle
+from dask_array._ufunc import *
+from dask_array.creation import (
+    arange,
+    diag,
+    diagonal,
+    empty,
+    empty_like,
+    eye,
+    fromfunction,
+    full,
+    full_like,
+    indices,
+    linspace,
+    meshgrid,
+    ones,
+    ones_like,
+    pad,
+    repeat,
+    tile,
+    tri,
+    zeros,
+    zeros_like,
+)
+from dask_array.io import (
+    from_delayed,
+    from_npy_stack,
+    from_tiledb,
+    from_zarr,
+    store,
+    to_hdf5,
+    to_npy_stack,
+    to_tiledb,
+    to_zarr,
+)
+from dask_array.linalg import dot, matmul, tensordot, vdot
+from dask_array.reductions import (
+    _tree_reduce,
+    all,
+    any,
+    arg_reduction,
+    argmax,
+    argmin,
+    cumprod,
+    cumreduction,
+    cumsum,
+    max,
+    mean,
+    median,
+    min,
+    moment,
+    nanargmax,
+    nanargmin,
+    nancumprod,
+    nancumsum,
+    nanmax,
+    nanmean,
+    nanmedian,
+    nanmin,
+    nanpercentile,
+    nanprod,
+    nanquantile,
+    nanstd,
+    nansum,
+    nanvar,
+    percentile,
+    prod,
+    quantile,
+    reduction,
+    std,
+    sum,
+    trace,
+    var,
+)
+from dask_array.routines._diff import diff
+from dask_array.routines._where import where
+from dask_array._expr_flow import expr_flow
+from dask_array._visualize import expr_table
+
+
+def optimize(dsk, keys=None, **kwargs):
+    """Optimize a dask-array collection.
+
+    Low-level graphs are returned unchanged because this package optimizes
+    through Array expressions before graph materialization.
+    """
+    if isinstance(dsk, Array):
+        result = dsk.optimize()
+        if keys is not None:
+            return result.__dask_graph__()
+        return result
+    return dsk
+
+
+newaxis = None
+nan = np.nan
+inf = np.inf
+e = np.e
+pi = np.pi
+euler_gamma = np.euler_gamma
+
+bool = np.bool
+int8 = np.int8
+int16 = np.int16
+int32 = np.int32
+int64 = np.int64
+uint8 = np.uint8
+uint16 = np.uint16
+uint32 = np.uint32
+uint64 = np.uint64
+float32 = np.float32
+float64 = np.float64
+complex64 = np.complex64
+complex128 = np.complex128
+
+# Ensure our xarray ChunkManager replaces the built-in DaskManager
+# regardless of entry point enumeration order. See _xarray.py for details.
+try:
+    from dask_array._xarray import _ensure_registered
+
+    _ensure_registered()
+    del _ensure_registered
+except ImportError:
+    pass

dask_array/_backends.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import numpy as np
+
+from dask._dispatch import get_collection_type
+
+try:
+    import sparse
+
+    sparse_installed = True
+except ImportError:
+    sparse_installed = False
+
+
+try:
+    import scipy.sparse as sp
+
+    scipy_installed = True
+except ImportError:
+    scipy_installed = False
+
+
+def create_array_collection(expr):
+    """Create an Array collection from an expression."""
+    from dask_array._collection import Array
+    from dask_array._expr import ArrayExpr
+
+    if isinstance(expr, ArrayExpr):
+        return Array(expr)
+
+    # For non-ArrayExpr (e.g., from dask-dataframe), wrap in adapter
+    # This is a fallback - most cases should be ArrayExpr
+    raise TypeError(f"Expected ArrayExpr, got {type(expr)}")
+
+
+@get_collection_type.register(np.ndarray)
+def get_collection_type_array(_):
+    return create_array_collection
+
+
+if sparse_installed:
+
+    @get_collection_type.register(sparse.COO)
+    def get_collection_type_sparse(_):
+        return create_array_collection
+
+
+if scipy_installed:
+
+    @get_collection_type.register(sp.csr_matrix)
+    def get_collection_type_scipy(_):
+        return create_array_collection
+
+
+if scipy_installed and hasattr(sp, "sparray"):
+
+    @get_collection_type.register(sp.csr_array)
+    def get_collection_type_scipy_array(_):
+        return create_array_collection
+
+
+@get_collection_type.register(object)
+def get_collection_type_object(_):
+    return create_scalar_collection
+
+
+def create_scalar_collection(expr):
+    from dask_array._expr import ArrayExpr
+
+    if isinstance(expr, ArrayExpr):
+        from dask_array._collection import Array
+
+        return Array(expr)
+
+    # For other expressions, try to return something sensible
+    raise TypeError(f"Cannot create collection from {type(expr)}")

dask_array/_backends_array.py

@@ -0,0 +1,99 @@
+"""Array creation dispatch for backend-specific array creation.
+
+This module provides `array_creation_dispatch` which is used by random modules
+and creation functions to dispatch to backend-specific implementations.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask.backends import CreationDispatch, DaskBackendEntrypoint
+
+
+class ArrayBackendEntrypoint(DaskBackendEntrypoint):
+    """Dask-Array version of ``DaskBackendEntrypoint``
+
+    See Also
+    --------
+    NumpyBackendEntrypoint
+    """
+
+    @property
+    def RandomState(self):
+        """Return the backend-specific RandomState class
+
+        For example, the 'numpy' backend simply returns
+        ``numpy.random.RandomState``.
+        """
+        raise NotImplementedError
+
+    @property
+    def default_bit_generator(self):
+        """Return the default BitGenerator type"""
+        raise NotImplementedError
+
+    @staticmethod
+    def ones(shape, *, dtype=None, meta=None, **kwargs):
+        """Create an array of ones
+
+        Returns a new array having a specified shape and filled
+        with ones.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def zeros(shape, *, dtype=None, meta=None, **kwargs):
+        """Create an array of zeros
+
+        Returns a new array having a specified shape and filled
+        with zeros.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def empty(shape, *, dtype=None, meta=None, **kwargs):
+        """Create an empty array
+
+        Returns an uninitialized array having a specified shape.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def full(shape, fill_value, *, dtype=None, meta=None, **kwargs):
+        """Create a uniformly filled array
+
+        Returns a new array having a specified shape and filled
+        with fill_value.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def arange(start, /, stop=None, step=1, *, dtype=None, meta=None, **kwargs):
+        """Create an ascending or descending array
+
+        Returns evenly spaced values within the half-open interval
+        ``[start, stop)`` as a one-dimensional array.
+        """
+        raise NotImplementedError
+
+
+class NumpyBackendEntrypoint(ArrayBackendEntrypoint):
+    @property
+    def RandomState(self):
+        return np.random.RandomState
+
+    @property
+    def default_bit_generator(self):
+        return np.random.PCG64
+
+
+array_creation_dispatch = CreationDispatch(
+    module_name="array",
+    default="numpy",
+    entrypoint_class=ArrayBackendEntrypoint,
+    name="array_creation_dispatch",
+)
+
+
+array_creation_dispatch.register_backend("numpy", NumpyBackendEntrypoint())