dask-array 0.1.0__py3-none-any.whl

dask_array/routines/_select.py

@@ -0,0 +1,113 @@
+"""Selection functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray, elemwise
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def digitize(a, bins, right=False):
+    """Return the indices of the bins to which each value in input array belongs.
+
+    Parameters
+    ----------
+    a : dask array
+        Input array to be binned.
+    bins : array_like
+        Array of bins. Must be 1-dimensional and monotonic.
+    right : bool, optional
+        Indicating whether the intervals include the right or left bin edge.
+
+    Returns
+    -------
+    indices : dask array of ints
+        Output array of indices.
+    """
+    bins = np.asarray(bins)
+    if bins.ndim != 1:
+        raise ValueError("bins must be 1-dimensional")
+
+    dtype = np.digitize(np.asarray([0], like=bins), bins, right=right).dtype
+    return elemwise(np.digitize, a, dtype=dtype, bins=bins, right=right)
+
+
+def _variadic_choose(a, *choices):
+    return np.choose(a, choices)
+
+
+@derived_from(np)
+def choose(a, choices):
+    a = asarray(a)
+    choices = [asarray(c) for c in choices]
+    return elemwise(_variadic_choose, a, *choices)
+
+
+@derived_from(np)
+def extract(condition, arr):
+    from dask_array.routines._misc import compress
+
+    condition = asarray(condition).astype(bool)
+    arr = asarray(arr)
+    return compress(condition.ravel(), arr.ravel())
+
+
+def _int_piecewise(x, *condlist, funclist=None, func_args=(), func_kw=None):
+    return np.piecewise(x, list(condlist), funclist, *func_args, **(func_kw or {}))
+
+
+@derived_from(np)
+def piecewise(x, condlist, funclist, *args, **kw):
+    x = asarray(x)
+    return elemwise(
+        _int_piecewise,
+        x,
+        *condlist,
+        dtype=x.dtype,
+        name="piecewise",
+        funclist=funclist,
+        func_args=args,
+        func_kw=kw,
+    )
+
+
+def _select(*args, **kwargs):
+    split_at = len(args) // 2
+    condlist = args[:split_at]
+    choicelist = args[split_at:]
+    return np.select(condlist, choicelist, **kwargs)
+
+
+@derived_from(np)
+def select(condlist, choicelist, default=0):
+    from dask_array._collection import blockwise
+    from dask_array.routines._misc import result_type
+
+    if len(condlist) != len(choicelist):
+        raise ValueError("list of cases must be same length as list of conditions")
+
+    if len(condlist) == 0:
+        raise ValueError("select with an empty condition list is not possible")
+
+    choicelist = [asarray(choice) for choice in choicelist]
+
+    try:
+        intermediate_dtype = result_type(*choicelist)
+    except TypeError as e:
+        msg = "Choicelist elements do not have a common dtype."
+        raise TypeError(msg) from e
+
+    blockwise_shape = tuple(range(choicelist[0].ndim))
+    condargs = [arg for elem in condlist for arg in (elem, blockwise_shape)]
+    choiceargs = [arg for elem in choicelist for arg in (elem, blockwise_shape)]
+
+    return blockwise(
+        _select,
+        blockwise_shape,
+        *condargs,
+        *choiceargs,
+        dtype=intermediate_dtype,
+        default=default,
+    )

dask_array/routines/_statistics.py

@@ -0,0 +1,171 @@
+"""Statistical functions for array-expr."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+
+from dask_array._collection import (
+    array,
+    asanyarray,
+    asarray,
+    broadcast_to,
+    concatenate,
+)
+from dask.utils import derived_from
+
+
+def result_type(*arrays_and_dtypes):
+    """Returns the type from NumPy type promotion rules."""
+    from dask_array._collection import Array
+
+    args = [a.dtype if isinstance(a, Array) else a for a in arrays_and_dtypes]
+    return np.result_type(*args)
+
+
+def average(a, axis=None, weights=None, returned=False, keepdims=False):
+    """Compute the weighted average along the specified axis."""
+    a = asanyarray(a)
+
+    if weights is None:
+        avg = a.mean(axis, keepdims=keepdims)
+        scl = avg.dtype.type(a.size / avg.size)
+    else:
+        wgt = asanyarray(weights)
+
+        if issubclass(a.dtype.type, (np.integer, np.bool_)):
+            result_dtype = result_type(a.dtype, wgt.dtype, "f8")
+        else:
+            result_dtype = result_type(a.dtype, wgt.dtype)
+
+        if a.shape != wgt.shape:
+            if axis is None:
+                raise TypeError("Axis must be specified when shapes of a and weights differ.")
+            if wgt.ndim != 1:
+                raise TypeError("1D weights expected when shapes of a and weights differ.")
+            if wgt.shape[0] != a.shape[axis]:
+                raise ValueError("Length of weights not compatible with specified axis.")
+
+            wgt = broadcast_to(wgt, (a.ndim - 1) * (1,) + wgt.shape)
+            wgt = wgt.swapaxes(-1, axis)
+
+        scl = wgt.sum(axis=axis, dtype=result_dtype, keepdims=keepdims)
+        from dask_array._ufunc import multiply
+
+        avg = multiply(a, wgt, dtype=result_dtype).sum(axis, keepdims=keepdims) / scl
+
+    if returned:
+        if scl.shape != avg.shape:
+            scl = broadcast_to(scl, avg.shape)
+        return avg, scl
+    else:
+        return avg
+
+
+@derived_from(np)
+def cov(
+    m,
+    y=None,
+    rowvar=True,
+    bias=False,
+    ddof=None,
+    fweights=None,
+    aweights=None,
+    *,
+    dtype=None,
+):
+    """Estimate a covariance matrix."""
+    from dask_array._ufunc import true_divide
+    from dask_array.linalg import dot
+
+    if ddof is not None and ddof != int(ddof):
+        raise ValueError("ddof must be integer")
+
+    m = asarray(m)
+
+    if y is None:
+        dtype = result_type(m, np.float64)
+    else:
+        y = asarray(y)
+        dtype = result_type(m, y, np.float64)
+
+    if m.ndim > 2:
+        raise ValueError("m has more than 2 dimensions")
+    if y is not None and y.ndim > 2:
+        raise ValueError("y has more than 2 dimensions")
+
+    X = array(m, ndmin=2, dtype=dtype)
+
+    if ddof is None:
+        ddof = 1 if bias == 0 else 0
+
+    if not rowvar and m.ndim != 1:
+        X = X.T
+    if X.shape[0] == 0:
+        return array([]).reshape(0, 0)
+    if y is not None:
+        y = array(y, ndmin=2, dtype=dtype)
+        if not rowvar and y.shape[0] != 1:
+            y = y.T
+        X = concatenate((X, y), axis=0)
+
+    w = None
+    if fweights is not None:
+        fweights = asarray(fweights, dtype=float)
+        if fweights.ndim > 1:
+            raise RuntimeError("cannot handle multidimensional fweights")
+        if fweights.shape[0] != X.shape[1]:
+            raise RuntimeError("incompatible numbers of samples and fweights")
+        w = fweights
+    if aweights is not None:
+        aweights = asarray(aweights, dtype=float)
+        if aweights.ndim > 1:
+            raise RuntimeError("cannot handle multidimensional aweights")
+        if aweights.shape[0] != X.shape[1]:
+            raise RuntimeError("incompatible numbers of samples and aweights")
+        if w is None:
+            w = aweights
+        else:
+            w *= aweights
+
+    avg, w_sum = average(X, axis=1, weights=w, returned=True)
+    w_sum = w_sum[0]
+
+    if w is None:
+        fact = X.shape[1] - ddof
+    elif ddof == 0:
+        fact = w_sum
+    elif aweights is None:
+        fact = w_sum - ddof
+    else:
+        fact = w_sum - ddof * (w * aweights).sum() / w_sum
+
+    if fact <= 0:
+        warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning)
+        fact = 0.0
+
+    X -= avg[:, None]
+    if w is None:
+        X_T = X.T
+    else:
+        X_T = (X * w).T
+    c = dot(X, X_T.conj())
+    c *= true_divide(1, fact)
+    return c.squeeze()
+
+
+@derived_from(np)
+def corrcoef(x, y=None, rowvar=1):
+    """Return Pearson product-moment correlation coefficients."""
+    from dask_array._ufunc import sqrt
+    from dask_array.creation import diag
+
+    c = cov(x, y, rowvar)
+
+    if c.shape == ():
+        return c / c
+    d = diag(c)
+    d = d.reshape((d.shape[0], 1))
+    sqr_d = sqrt(d)
+    return (c / sqr_d) / sqr_d.T

dask_array/routines/_topk.py

@@ -0,0 +1,82 @@
+"""Top-k selection for array-expr."""
+
+from __future__ import annotations
+
+from functools import partial
+from numbers import Number
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask_array._utils import validate_axis
+
+
+def topk(a, k, axis=-1, split_every=None):
+    """Extract the k largest elements from a on the given axis.
+
+    Returns them sorted from largest to smallest. If k is negative,
+    extract the -k smallest elements instead, and return them sorted
+    from smallest to largest.
+    """
+    from dask_array import _chunk as chunk
+    from dask_array.reductions import reduction
+
+    a = asarray(a)
+    axis = validate_axis(axis, a.ndim)
+
+    chunk_combine = partial(chunk.topk, k=k)
+    aggregate = partial(chunk.topk_aggregate, k=k)
+
+    return reduction(
+        a,
+        chunk=chunk_combine,
+        combine=chunk_combine,
+        aggregate=aggregate,
+        axis=axis,
+        keepdims=True,
+        dtype=a.dtype,
+        split_every=split_every,
+        output_size=abs(k),
+    )
+
+
+def argtopk(a, k, axis=-1, split_every=None):
+    """Extract the indices of the k largest elements from a on the given axis.
+
+    Returns them sorted from largest to smallest. If k is negative,
+    extract the indices of the -k smallest elements instead.
+    """
+    from dask_array import _chunk as chunk
+    from dask_array.creation import arange
+    from dask_array.reductions import reduction
+
+    a = asarray(a)
+    axis = validate_axis(axis, a.ndim)
+
+    idx = arange(a.shape[axis], chunks=(a.chunks[axis],), dtype=np.intp)
+    idx = idx[tuple(slice(None) if i == axis else np.newaxis for i in range(a.ndim))]
+    a_plus_idx = a.map_blocks(chunk.argtopk_preprocess, idx, dtype=object)
+
+    chunk_combine = partial(chunk.argtopk, k=k)
+    aggregate = partial(chunk.argtopk_aggregate, k=k)
+
+    if isinstance(axis, Number):
+        naxis = 1
+    else:
+        naxis = len(axis)
+
+    meta = a._meta.astype(np.intp).reshape((0,) * (a.ndim - naxis + 1))
+
+    return reduction(
+        a_plus_idx,
+        chunk=chunk_combine,
+        combine=chunk_combine,
+        aggregate=aggregate,
+        axis=axis,
+        keepdims=True,
+        dtype=np.intp,
+        split_every=split_every,
+        concatenate=False,
+        output_size=abs(k),
+        meta=meta,
+    )

dask_array/routines/_triangular.py

@@ -0,0 +1,74 @@
+"""Triangular matrix functions for array-expr."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask_array._collection import asarray
+from dask.utils import derived_from
+
+
+@derived_from(np)
+def tril(m, k=0):
+    from dask_array.creation import tri
+    from dask_array.routines._where import where
+    from dask_array._utils import meta_from_array
+
+    m = asarray(m)
+    mask = tri(
+        *m.shape[-2:],
+        k=k,
+        dtype=bool,
+        chunks=m.chunks[-2:],
+        like=meta_from_array(m),
+    )
+
+    return where(mask, m, np.zeros_like(m._meta, shape=(1,)))
+
+
+@derived_from(np)
+def triu(m, k=0):
+    from dask_array.creation import tri
+    from dask_array.routines._where import where
+    from dask_array._utils import meta_from_array
+
+    m = asarray(m)
+    mask = tri(
+        *m.shape[-2:],
+        k=k - 1,
+        dtype=bool,
+        chunks=m.chunks[-2:],
+        like=meta_from_array(m),
+    )
+
+    return where(mask, np.zeros_like(m._meta, shape=(1,)), m)
+
+
+@derived_from(np)
+def tril_indices(n, k=0, m=None, chunks="auto"):
+    from dask_array.creation import tri
+    from dask_array.routines._nonzero import nonzero
+
+    return nonzero(tri(n, m, k=k, dtype=bool, chunks=chunks))
+
+
+@derived_from(np)
+def tril_indices_from(arr, k=0):
+    if arr.ndim != 2:
+        raise ValueError("input array must be 2-d")
+    return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1], chunks=arr.chunks)
+
+
+@derived_from(np)
+def triu_indices(n, k=0, m=None, chunks="auto"):
+    from dask_array.creation import tri
+    from dask_array.routines._nonzero import nonzero
+
+    return nonzero(~tri(n, m, k=k - 1, dtype=bool, chunks=chunks))
+
+
+@derived_from(np)
+def triu_indices_from(arr, k=0):
+    if arr.ndim != 2:
+        raise ValueError("input array must be 2-d")
+    return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1], chunks=arr.chunks)

dask_array/routines/_unique.py

@@ -0,0 +1,232 @@
+"""Unique implementation for array-expr."""
+
+from __future__ import annotations
+
+from functools import cached_property
+
+import numpy as np
+
+from dask._task_spec import List, Task, TaskRef
+from dask_array._collection import (
+    asarray,
+    concatenate,
+    new_collection,
+    ravel,
+)
+from dask_array._expr import ArrayExpr
+from dask_array._utils import meta_from_array
+from dask.utils import derived_from
+
+
+def _unique_internal(ar, indices, counts, return_inverse=False):
+    """Helper function for np.unique with structured array output."""
+    return_index = indices is not None
+    return_counts = counts is not None
+
+    u = np.unique(ar)
+
+    dt = [("values", u.dtype)]
+    if return_index:
+        dt.append(("indices", np.intp))
+    if return_inverse:
+        dt.append(("inverse", np.intp))
+    if return_counts:
+        dt.append(("counts", np.intp))
+
+    r = np.empty(u.shape, dtype=dt)
+    r["values"] = u
+    if return_inverse:
+        r["inverse"] = np.arange(len(r), dtype=np.intp)
+    if return_index or return_counts:
+        for i, v in enumerate(r["values"]):
+            m = ar == v
+            if return_index:
+                indices[m].min(keepdims=True, out=r["indices"][i : i + 1])
+            if return_counts:
+                counts[m].sum(keepdims=True, out=r["counts"][i : i + 1])
+
+    return r
+
+
+class UniqueChunked(ArrayExpr):
+    """Expression for per-chunk unique computation."""
+
+    _parameters = ["x", "indices", "counts", "out_dtype"]
+    _defaults = {"indices": None, "counts": None}
+
+    @cached_property
+    def _meta(self):
+        return np.empty((0,), dtype=self.out_dtype)
+
+    @cached_property
+    def chunks(self):
+        nchunks = len(self.x.chunks[0])
+        return ((np.nan,) * nchunks,)
+
+    @cached_property
+    def _name(self):
+        return f"unique-chunk-{self.deterministic_token}"
+
+    def _layer(self):
+        dsk = {}
+        for i in range(len(self.x.chunks[0])):
+            key = (self._name, i)
+            x_ref = TaskRef((self.x._name, i))
+            idx_ref = TaskRef((self.indices._name, i)) if self.indices is not None else None
+            cnt_ref = TaskRef((self.counts._name, i)) if self.counts is not None else None
+            dsk[key] = Task(key, _unique_internal, x_ref, idx_ref, cnt_ref, False)
+        return dsk
+
+    @property
+    def _dependencies(self):
+        deps = [self.x]
+        if self.indices is not None:
+            deps.append(self.indices)
+        if self.counts is not None:
+            deps.append(self.counts)
+        return deps
+
+
+def _unique_aggregate_func(chunks, return_inverse):
+    """Aggregate unique results from multiple chunks."""
+    combined = np.concatenate(chunks)
+    return_index = "indices" in combined.dtype.names
+    return_counts = "counts" in combined.dtype.names
+
+    return _unique_internal(
+        combined["values"],
+        combined["indices"] if return_index else None,
+        combined["counts"] if return_counts else None,
+        return_inverse=return_inverse,
+    )
+
+
+class UniqueAggregate(ArrayExpr):
+    """Expression for aggregating unique results from all chunks."""
+
+    _parameters = ["chunked", "return_inverse", "out_dtype"]
+    _defaults = {"return_inverse": False}
+
+    @cached_property
+    def _meta(self):
+        return np.empty((0,), dtype=self.out_dtype)
+
+    @cached_property
+    def chunks(self):
+        return ((np.nan,),)
+
+    @cached_property
+    def _name(self):
+        return f"unique-aggregate-{self.deterministic_token}"
+
+    def _layer(self):
+        chunk_keys = [(self.chunked._name, i) for i in range(len(self.chunked.chunks[0]))]
+        key = (self._name, 0)
+        chunks_list = List(*[TaskRef(k) for k in chunk_keys])
+        dsk = {key: Task(key, _unique_aggregate_func, chunks_list, self.return_inverse)}
+        return dsk
+
+    @property
+    def _dependencies(self):
+        return [self.chunked]
+
+
+def unique_no_structured_arr(ar, return_index=False, return_inverse=False, return_counts=False):
+    """Simplified version of unique for arrays that don't support structured arrays."""
+    from dask_array._blockwise import Blockwise
+    from dask_array._expr import ChunksOverride
+    from dask_array.reductions import _tree_reduce
+
+    if return_index or return_inverse or return_counts:
+        raise ValueError(
+            "dask.array.unique does not support `return_index`, `return_inverse` "
+            "or `return_counts` with array types that don't support structured arrays."
+        )
+
+    ar = ravel(ar)
+    out = Blockwise(np.unique, "i", ar, "i", dtype=ar.dtype)
+    chunked = new_collection(out)
+    chunked = new_collection(ChunksOverride(chunked.expr, ((np.nan,) * len(ar.chunks[0]),)))
+
+    def _unique_agg(arrays, axis, keepdims):
+        if not isinstance(arrays, list):
+            arrays = [arrays]
+        return np.unique(np.concatenate(arrays))
+
+    return _tree_reduce(
+        chunked.expr,
+        aggregate=_unique_agg,
+        axis=(0,),
+        keepdims=False,
+        dtype=ar.dtype,
+        concatenate=False,
+    )
+
+
+@derived_from(np)
+def unique(ar, return_index=False, return_inverse=False, return_counts=False):
+    """Find the unique elements of an array."""
+    from dask_array.creation import arange, ones
+    from dask_array._numpy_compat import NUMPY_GE_200
+
+    try:
+        meta = meta_from_array(ar)
+        np.empty_like(meta, dtype=[("a", int), ("b", float)])
+    except TypeError:
+        return unique_no_structured_arr(
+            ar,
+            return_index=return_index,
+            return_inverse=return_inverse,
+            return_counts=return_counts,
+        )
+
+    orig_shape = ar.shape
+    ar = ravel(ar)
+
+    out_dtype = [("values", ar.dtype)]
+    indices_arr = None
+    counts_arr = None
+
+    if return_index:
+        indices_arr = arange(ar.shape[0], dtype=np.intp, chunks=ar.chunks[0])
+        out_dtype.append(("indices", np.intp))
+    if return_counts:
+        counts_arr = ones((ar.shape[0],), dtype=np.intp, chunks=ar.chunks[0])
+        out_dtype.append(("counts", np.intp))
+
+    out_dtype = np.dtype(out_dtype)
+    chunked = UniqueChunked(
+        ar.expr,
+        indices_arr.expr if indices_arr is not None else None,
+        counts_arr.expr if counts_arr is not None else None,
+        out_dtype,
+    )
+
+    final_dtype = out_dtype if not return_inverse else np.dtype(list(out_dtype.descr) + [("inverse", np.intp)])
+    aggregated = new_collection(UniqueAggregate(chunked, return_inverse, final_dtype))
+
+    result = [aggregated["values"]]
+    if return_index:
+        result.append(aggregated["indices"])
+    if return_inverse:
+        matches = (ar[:, None] == aggregated["values"][None, :]).astype(np.intp)
+        inverse = (matches * aggregated["inverse"]).sum(axis=1)
+        if NUMPY_GE_200:
+            from dask_array._reshape import reshape
+
+            inverse = reshape(inverse, orig_shape)
+        result.append(inverse)
+    if return_counts:
+        result.append(aggregated["counts"])
+
+    if len(result) == 1:
+        return result[0]
+    return tuple(result)
+
+
+@derived_from(np)
+def union1d(ar1, ar2):
+    """Find the union of two arrays."""
+    ar1 = asarray(ar1)
+    ar2 = asarray(ar2)
+    return unique(concatenate((ravel(ar1), ravel(ar2))))