dask-array 0.1.0__py3-none-any.whl

dask_array/creation/_pad.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray, concatenate
+from dask_array._utils import meta_from_array
+from dask.utils import derived_from
+
+from ._utils import (
+    expand_pad_value,
+    get_pad_shapes_chunks,
+    linear_ramp_chunk,
+    wrapped_pad_func,
+)
+
+
+def _pad_reuse_expr(array, pad_width, mode, **kwargs):
+    """
+    Helper function for padding boundaries with values in the array.
+
+    Handles the cases where the padding is constructed from values in
+    the array. Namely by reflecting them or tiling them to create periodic
+    boundary constraints.
+    """
+    from dask_array._collection import block
+    from dask_array.manipulation._flip import flip
+
+    if mode in {"reflect", "symmetric"}:
+        reflect_type = kwargs.get("reflect", "even")
+        if reflect_type == "odd":
+            raise NotImplementedError("`pad` does not support `reflect_type` of `odd`.")
+        if reflect_type != "even":
+            raise ValueError("unsupported value for reflect_type, must be one of (`even`, `odd`)")
+
+    result = np.empty(array.ndim * (3,), dtype=object)
+    for idx in np.ndindex(result.shape):
+        select = []
+        flip_axes = []
+        for axis, (i, s, pw) in enumerate(zip(idx, array.shape, pad_width)):
+            if mode == "wrap":
+                pw = pw[::-1]
+
+            if i < 1:
+                if mode == "reflect":
+                    select.append(slice(1, pw[0] + 1, None))
+                else:
+                    select.append(slice(None, pw[0], None))
+            elif i > 1:
+                if mode == "reflect":
+                    select.append(slice(s - pw[1] - 1, s - 1, None))
+                else:
+                    select.append(slice(s - pw[1], None, None))
+            else:
+                select.append(slice(None))
+
+            if i != 1 and mode in ["reflect", "symmetric"]:
+                flip_axes.append(axis)
+
+        select = tuple(select)
+
+        if mode == "wrap":
+            idx = tuple(2 - i for i in idx)
+
+        chunk = array[select]
+        # Apply flips for each axis that needs reversal
+        for axis in flip_axes:
+            chunk = flip(chunk, axis)
+        result[idx] = chunk
+
+    result = block(result.tolist())
+
+    return result
+
+
+def _pad_stats_expr(array, pad_width, mode, stat_length):
+    """
+    Helper function for padding boundaries with statistics from the array.
+
+    In cases where the padding requires computations of statistics from part
+    or all of the array, this function helps compute those statistics as
+    requested and then adds those statistics onto the boundaries of the array.
+    """
+    from dask_array._collection import block, broadcast_to
+    from dask_array._ufunc import rint
+
+    if mode == "median":
+        raise NotImplementedError("`pad` does not support `mode` of `median`.")
+
+    stat_length = expand_pad_value(array, stat_length)
+
+    result = np.empty(array.ndim * (3,), dtype=object)
+    for idx in np.ndindex(result.shape):
+        axes = []
+        select = []
+        pad_shape = []
+        pad_chunks = []
+        for d, (i, s, c, w, l) in enumerate(zip(idx, array.shape, array.chunks, pad_width, stat_length)):
+            if i < 1:
+                axes.append(d)
+                select.append(slice(None, l[0], None))
+                pad_shape.append(w[0])
+                pad_chunks.append(w[0])
+            elif i > 1:
+                axes.append(d)
+                select.append(slice(s - l[1], None, None))
+                pad_shape.append(w[1])
+                pad_chunks.append(w[1])
+            else:
+                select.append(slice(None))
+                pad_shape.append(s)
+                pad_chunks.append(c)
+
+        axes = tuple(axes)
+        select = tuple(select)
+        pad_shape = tuple(pad_shape)
+        pad_chunks = tuple(pad_chunks)
+
+        result_idx = array[select]
+        if axes:
+            stat_funcs = {"maximum": "max", "mean": "mean", "minimum": "min"}
+            result_idx = getattr(result_idx, stat_funcs[mode])(axis=axes, keepdims=True)
+            result_idx = broadcast_to(result_idx, pad_shape, chunks=pad_chunks)
+
+            if mode == "mean":
+                if np.issubdtype(array.dtype, np.integer):
+                    result_idx = rint(result_idx)
+                result_idx = result_idx.astype(array.dtype)
+
+        result[idx] = result_idx
+
+    result = block(result.tolist())
+
+    return result
+
+
+def _pad_udf_expr(array, pad_width, mode, **kwargs):
+    """
+    Helper function for padding boundaries with a user defined function.
+
+    In cases where the padding requires a custom user defined function be
+    applied to the array, this function assists in the prepping and
+    application of this function to the Dask Array to construct the desired
+    boundaries.
+    """
+    result = _pad_edge_expr(array, pad_width, "constant", constant_values=0)
+
+    chunks = result.chunks
+    for d in range(result.ndim):
+        result = result.rechunk(chunks[:d] + (result.shape[d : d + 1],) + chunks[d + 1 :])
+
+        result = result.map_blocks(
+            wrapped_pad_func,
+            name="pad",
+            dtype=result.dtype,
+            pad_func=mode,
+            iaxis_pad_width=pad_width[d],
+            iaxis=d,
+            pad_func_kwargs=kwargs,
+        )
+
+        result = result.rechunk(chunks)
+
+    return result
+
+
+def _pad_edge_expr(array, pad_width, mode, **kwargs):
+    """
+    Helper function for padding edges - array-expr version.
+
+    Handles the cases where the only the values on the edge are needed.
+    """
+    from dask_array._collection import broadcast_to
+    from dask_array._utils import asarray_safe
+
+    from ._ones_zeros import empty_like
+
+    kwargs = {k: expand_pad_value(array, v) for k, v in kwargs.items()}
+
+    result = array
+    for d in range(array.ndim):
+        pad_shapes, pad_chunks = get_pad_shapes_chunks(result, pad_width, (d,), mode=mode)
+        pad_arrays = [result, result]
+
+        if mode == "constant":
+            constant_values = kwargs["constant_values"][d]
+            constant_values = [
+                asarray_safe(c, like=meta_from_array(array), dtype=result.dtype) for c in constant_values
+            ]
+
+            pad_arrays = [broadcast_to(v, s, c) for v, s, c in zip(constant_values, pad_shapes, pad_chunks)]
+        elif mode in ["edge", "linear_ramp"]:
+            pad_slices = [result.ndim * [slice(None)], result.ndim * [slice(None)]]
+            pad_slices[0][d] = slice(None, 1, None)
+            pad_slices[1][d] = slice(-1, None, None)
+            pad_slices = [tuple(sl) for sl in pad_slices]
+
+            pad_arrays = [result[sl] for sl in pad_slices]
+
+            if mode == "edge":
+                pad_arrays = [broadcast_to(a, s, c) for a, s, c in zip(pad_arrays, pad_shapes, pad_chunks)]
+            elif mode == "linear_ramp":
+                end_values = kwargs["end_values"][d]
+
+                pad_arrays = [
+                    a.map_blocks(
+                        linear_ramp_chunk,
+                        ev,
+                        pw,
+                        chunks=c,
+                        dtype=result.dtype,
+                        dim=d,
+                        step=(2 * i - 1),
+                    )
+                    for i, (a, ev, pw, c) in enumerate(zip(pad_arrays, end_values, pad_width[d], pad_chunks))
+                ]
+        elif mode == "empty":
+            pad_arrays = [
+                empty_like(array, shape=s, dtype=array.dtype, chunks=c) for s, c in zip(pad_shapes, pad_chunks)
+            ]
+
+        result = concatenate([pad_arrays[0], result, pad_arrays[1]], axis=d)
+
+    return result
+
+
+@derived_from(np)
+def pad(array, pad_width, mode="constant", **kwargs):
+    array = asarray(array)
+
+    pad_width = expand_pad_value(array, pad_width)
+
+    if callable(mode):
+        return _pad_udf_expr(array, pad_width, mode, **kwargs)
+
+    # Make sure that no unsupported keywords were passed for the current mode
+    allowed_kwargs = {
+        "empty": [],
+        "edge": [],
+        "wrap": [],
+        "constant": ["constant_values"],
+        "linear_ramp": ["end_values"],
+        "maximum": ["stat_length"],
+        "mean": ["stat_length"],
+        "median": ["stat_length"],
+        "minimum": ["stat_length"],
+        "reflect": ["reflect_type"],
+        "symmetric": ["reflect_type"],
+    }
+    try:
+        unsupported_kwargs = set(kwargs) - set(allowed_kwargs[mode])
+    except KeyError as e:
+        raise ValueError(f"mode '{mode}' is not supported") from e
+    if unsupported_kwargs:
+        raise ValueError(f"unsupported keyword arguments for mode '{mode}': {unsupported_kwargs}")
+
+    if mode in {"maximum", "mean", "median", "minimum"}:
+        stat_length = kwargs.get("stat_length", tuple((n, n) for n in array.shape))
+        return _pad_stats_expr(array, pad_width, mode, stat_length)
+    elif mode == "constant":
+        kwargs.setdefault("constant_values", 0)
+        return _pad_edge_expr(array, pad_width, mode, **kwargs)
+    elif mode == "linear_ramp":
+        kwargs.setdefault("end_values", 0)
+        return _pad_edge_expr(array, pad_width, mode, **kwargs)
+    elif mode in {"edge", "empty"}:
+        return _pad_edge_expr(array, pad_width, mode)
+    elif mode in ["reflect", "symmetric", "wrap"]:
+        return _pad_reuse_expr(array, pad_width, mode, **kwargs)
+
+    raise RuntimeError("unreachable")

dask_array/creation/_repeat.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from numbers import Integral
+
+import numpy as np
+from tlz import sliding_window
+
+from dask_array._collection import concatenate
+from dask.utils import cached_cumsum, derived_from
+
+
+@derived_from(np)
+def repeat(a, repeats, axis=None):
+    if axis is None:
+        if a.ndim == 1:
+            axis = 0
+        else:
+            raise NotImplementedError("Must supply an integer axis value")
+
+    if not isinstance(repeats, Integral):
+        raise NotImplementedError("Only integer valued repeats supported")
+
+    if -a.ndim <= axis < 0:
+        axis += a.ndim
+    elif not 0 <= axis <= a.ndim - 1:
+        raise ValueError(f"axis(={axis}) out of bounds")
+
+    if repeats == 0:
+        return a[tuple(slice(None) if d != axis else slice(0) for d in range(a.ndim))]
+    elif repeats == 1:
+        return a
+
+    cchunks = cached_cumsum(a.chunks[axis], initial_zero=True)
+    slices = []
+    for c_start, c_stop in sliding_window(2, cchunks):
+        ls = np.linspace(c_start, c_stop, repeats).round(0)
+        for ls_start, ls_stop in sliding_window(2, ls):
+            if ls_start != ls_stop:
+                slices.append(slice(ls_start, ls_stop))
+
+    all_slice = slice(None, None, None)
+    slices = [(all_slice,) * axis + (s,) + (all_slice,) * (a.ndim - axis - 1) for s in slices]
+
+    slabs = [a[slc] for slc in slices]
+
+    out = []
+    for slab in slabs:
+        chunks = list(slab.chunks)
+        assert len(chunks[axis]) == 1
+        chunks[axis] = (chunks[axis][0] * repeats,)
+        chunks = tuple(chunks)
+        result = slab.map_blocks(np.repeat, repeats, axis=axis, chunks=chunks, dtype=slab.dtype)
+        out.append(result)
+
+    return concatenate(out, axis=axis)

dask_array/creation/_tile.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def tile(A, reps):
+    from dask_array._collection import block
+
+    from ._ones_zeros import empty
+
+    try:
+        tup = tuple(reps)
+    except TypeError:
+        tup = (reps,)
+    if any(i < 0 for i in tup):
+        raise ValueError("Negative `reps` are not allowed.")
+    c = asarray(A)
+
+    if all(tup):
+        for nrep in tup[::-1]:
+            c = nrep * [c]
+        return block(c)
+
+    d = len(tup)
+    if d < c.ndim:
+        tup = (1,) * (c.ndim - d) + tup
+    if c.ndim < d:
+        shape = (1,) * (d - c.ndim) + c.shape
+    else:
+        shape = c.shape
+    shape_out = tuple(s * t for s, t in zip(shape, tup))
+    return empty(shape=shape_out, dtype=c.dtype)

dask_array/creation/_tri.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import numpy as np
+
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def tri(N, M=None, k=0, dtype=float, chunks="auto", *, like=None):
+    from dask_array._ufunc import greater_equal
+    from dask_array._core_utils import normalize_chunks
+
+    from ._arange import arange
+
+    if M is None:
+        M = N
+
+    chunks = normalize_chunks(chunks, shape=(N, M), dtype=dtype)
+
+    m = greater_equal(
+        arange(N, chunks=chunks[0][0], like=like).reshape(1, N).T,
+        arange(-k, M - k, chunks=chunks[1][0], like=like),
+    )
+
+    # Avoid making a copy if the requested type is already bool
+    m = m.astype(dtype, copy=False)
+
+    return m

dask_array/creation/_utils.py

@@ -0,0 +1,296 @@
+"""Helper functions for array creation operations.
+
+Copied/adapted from dask.array.creation and dask.array.wrap to reduce
+imports from dask.array.* modules.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from numbers import Number
+
+import numpy as np
+from tlz import curry
+
+from dask_array._core_utils import normalize_chunks
+from dask.base import tokenize
+from dask.utils import funcname
+
+
+def _parse_wrap_args(func, args, kwargs, shape):
+    """Parse arguments for wrap functions (ones, zeros, full, empty).
+
+    Parameters
+    ----------
+    func : callable
+        The numpy function (e.g., np.ones_like)
+    args : tuple
+        Positional arguments after shape
+    kwargs : dict
+        Keyword arguments (may include name, chunks, dtype)
+    shape : tuple or int
+        The desired shape
+
+    Returns
+    -------
+    dict with keys: shape, dtype, kwargs, chunks, name
+    """
+    if isinstance(shape, np.ndarray):
+        shape = shape.tolist()
+
+    if not isinstance(shape, (tuple, list)):
+        shape = (shape,)
+
+    name = kwargs.pop("name", None)
+    chunks = kwargs.pop("chunks", "auto")
+
+    dtype = kwargs.pop("dtype", None)
+    if dtype is None:
+        dtype = func(shape, *args, **kwargs).dtype
+    dtype = np.dtype(dtype)
+
+    chunks = normalize_chunks(chunks, shape, dtype=dtype)
+
+    name = name or funcname(func) + "-" + tokenize(func, shape, chunks, dtype, args, kwargs)
+
+    return {
+        "shape": shape,
+        "dtype": dtype,
+        "kwargs": kwargs,
+        "chunks": chunks,
+        "name": name,
+    }
+
+
+@curry
+def _broadcast_trick_inner(func, shape, meta=(), *args, **kwargs):
+    """Inner function for broadcast_trick.
+
+    cupy-specific hack. numpy is happy with hardcoded shape=().
+    """
+    null_shape = () if shape == () else 1
+    return np.broadcast_to(func(meta, *args, shape=null_shape, **kwargs), shape)
+
+
+def broadcast_trick(func):
+    """Provide a decorator to wrap common numpy function with a broadcast trick.
+
+    Dask arrays are currently immutable; thus when we know an array is uniform,
+    we can replace the actual data by a single value and have all elements point
+    to it, thus reducing the size.
+
+    >>> x = np.broadcast_to(1, (100,100,100))
+    >>> x.base.nbytes
+    8
+
+    Those array are not only more efficient locally, but dask serialisation is
+    aware of the _real_ size of those array and thus can send them around
+    efficiently and schedule accordingly.
+
+    Note that those array are read-only and numpy will refuse to assign to them,
+    so should be safe.
+    """
+    inner = _broadcast_trick_inner(func)
+    inner.__doc__ = func.__doc__
+    inner.__name__ = func.__name__
+    return inner
+
+
+def _get_like_function_shapes_chunks(a, chunks, shape):
+    """Helper function for finding shapes and chunks for *_like() array creation functions.
+
+    Parameters
+    ----------
+    a : dask array
+        The input array to get shape/chunks from
+    chunks : tuple or None
+        Desired chunks (None means use a's chunks)
+    shape : tuple or None
+        Desired shape (None means use a's shape)
+
+    Returns
+    -------
+    shape, chunks : tuple, tuple
+    """
+    if shape is None:
+        shape = a.shape
+        if chunks is None:
+            chunks = a.chunks
+    elif chunks is None:
+        chunks = "auto"
+    return shape, chunks
+
+
+def expand_pad_value(array, pad_value):
+    """Expand pad_value to a per-dimension format.
+
+    Parameters
+    ----------
+    array : dask array
+        The array to be padded (used to get ndim)
+    pad_value : various
+        The pad value in various formats
+
+    Returns
+    -------
+    tuple of tuples
+        Normalized pad_value as ((before, after), ...) for each dimension
+    """
+    if isinstance(pad_value, Number) or getattr(pad_value, "ndim", None) == 0:
+        pad_value = array.ndim * ((pad_value, pad_value),)
+    elif isinstance(pad_value, Sequence) and all(isinstance(pw, Number) for pw in pad_value) and len(pad_value) == 1:
+        pad_value = array.ndim * ((pad_value[0], pad_value[0]),)
+    elif isinstance(pad_value, Sequence) and len(pad_value) == 2 and all(isinstance(pw, Number) for pw in pad_value):
+        pad_value = array.ndim * (tuple(pad_value),)
+    elif (
+        isinstance(pad_value, Sequence)
+        and len(pad_value) == array.ndim
+        and all(isinstance(pw, Sequence) for pw in pad_value)
+        and all((len(pw) == 2) for pw in pad_value)
+        and all(all(isinstance(w, Number) for w in pw) for pw in pad_value)
+    ):
+        pad_value = tuple(tuple(pw) for pw in pad_value)
+    elif (
+        isinstance(pad_value, Sequence)
+        and len(pad_value) == 1
+        and isinstance(pad_value[0], Sequence)
+        and len(pad_value[0]) == 2
+        and all(isinstance(pw, Number) for pw in pad_value[0])
+    ):
+        pad_value = array.ndim * (tuple(pad_value[0]),)
+    else:
+        raise TypeError("`pad_value` must be composed of integral typed values.")
+
+    return pad_value
+
+
+def get_pad_shapes_chunks(array, pad_width, axes, mode):
+    """Helper function for finding shapes and chunks of end pads.
+
+    Parameters
+    ----------
+    array : dask array
+        The array to be padded
+    pad_width : tuple of tuples
+        The pad widths as ((before, after), ...) for each dimension
+    axes : tuple of ints
+        Which axes to compute pad info for
+    mode : str
+        The padding mode
+
+    Returns
+    -------
+    pad_shapes : list of tuples
+        Shape for [before, after] pads
+    pad_chunks : list of tuples
+        Chunks for [before, after] pads
+    """
+    pad_shapes = [list(array.shape), list(array.shape)]
+    pad_chunks = [list(array.chunks), list(array.chunks)]
+
+    for d in axes:
+        for i in range(2):
+            pad_shapes[i][d] = pad_width[d][i]
+            if mode != "constant" or pad_width[d][i] == 0:
+                pad_chunks[i][d] = (pad_width[d][i],)
+            else:
+                pad_chunks[i][d] = normalize_chunks((max(pad_chunks[i][d]),), (pad_width[d][i],))[0]
+
+    pad_shapes = [tuple(s) for s in pad_shapes]
+    pad_chunks = [tuple(c) for c in pad_chunks]
+
+    return pad_shapes, pad_chunks
+
+
+def linear_ramp_chunk(start, stop, num, dim, step):
+    """Helper function to find the linear ramp for a chunk.
+
+    Parameters
+    ----------
+    start : array
+        Starting values (shape has size 1 in dim)
+    stop : scalar
+        End value for ramp
+    num : int
+        Number of points in ramp
+    dim : int
+        Dimension along which to ramp
+    step : int
+        Direction (1 or -1)
+
+    Returns
+    -------
+    array with linear ramp values
+    """
+    num1 = num + 1
+
+    shape = list(start.shape)
+    shape[dim] = num
+    shape = tuple(shape)
+
+    dtype = np.dtype(start.dtype)
+
+    result = np.empty_like(start, shape=shape, dtype=dtype)
+    for i in np.ndindex(start.shape):
+        j = list(i)
+        j[dim] = slice(None)
+        j = tuple(j)
+
+        result[j] = np.linspace(start[i], stop, num1, dtype=dtype)[1:][::step]
+
+    return result
+
+
+def wrapped_pad_func(array, pad_func, iaxis_pad_width, iaxis, pad_func_kwargs):
+    """Wrapper to apply a user-defined pad function along an axis.
+
+    Parameters
+    ----------
+    array : ndarray
+        The input array chunk
+    pad_func : callable
+        User-defined padding function
+    iaxis_pad_width : tuple
+        (before, after) pad widths for this axis
+    iaxis : int
+        The axis index
+    pad_func_kwargs : dict
+        Keyword arguments to pass to pad_func
+
+    Returns
+    -------
+    array with padding applied
+    """
+    result = np.empty_like(array)
+    for i in np.ndindex(array.shape[:iaxis] + array.shape[iaxis + 1 :]):
+        i = i[:iaxis] + (slice(None),) + i[iaxis:]
+        result[i] = pad_func(array[i], iaxis_pad_width, iaxis, pad_func_kwargs)
+
+    return result
+
+
+def to_backend(x, backend: str | None = None, **kwargs):
+    """Move an Array collection to a new backend.
+
+    Parameters
+    ----------
+    x : Array
+        The input Array collection.
+    backend : str, Optional
+        The name of the new backend to move to. The default
+        is the current "array.backend" configuration.
+
+    Returns
+    -------
+    Array
+        A new Array collection with the backend specified
+        by ``backend``.
+    """
+    from dask_array._backends_array import array_creation_dispatch
+
+    # Get desired backend
+    backend = backend or array_creation_dispatch.backend
+    # Check that "backend" has a registered entrypoint
+    backend_entrypoint = array_creation_dispatch.dispatch(backend)
+    # Call `ArrayBackendEntrypoint.to_backend`
+    return backend_entrypoint.to_backend(x, **kwargs)