dask-array 0.1.0__py3-none-any.whl

dask_array/_test_utils.py

@@ -0,0 +1,257 @@
+"""Test utilities for dask_array.
+
+These functions are used in tests to compare dask arrays with numpy arrays.
+"""
+
+from __future__ import annotations
+
+import itertools
+import math
+
+import numpy as np
+
+from dask.base import is_dask_collection
+
+from dask_array._utils import is_arraylike, is_cupy_type
+
+
+def normalize_to_array(x):
+    """Convert CuPy arrays to numpy arrays."""
+    if is_cupy_type(x):
+        return x.get()
+    else:
+        return x
+
+
+def allclose(a, b, equal_nan=False, **kwargs):
+    """Check if two arrays are element-wise equal within a tolerance."""
+    a = normalize_to_array(a)
+    b = normalize_to_array(b)
+    if getattr(a, "dtype", None) != "O":
+        if hasattr(a, "mask") or hasattr(b, "mask"):
+            return np.ma.allclose(a, b, masked_equal=True, **kwargs)
+        else:
+            return np.allclose(a, b, equal_nan=equal_nan, **kwargs)
+    if equal_nan:
+        return a.shape == b.shape and all(np.isnan(b) if np.isnan(a) else a == b for (a, b) in zip(a.flat, b.flat))
+    return (a == b).all()
+
+
+def same_keys(a, b):
+    """Check if two dask collections have the same keys in their graphs."""
+
+    def key(k):
+        if isinstance(k, str):
+            return (k, -1, -1, -1)
+        else:
+            return k
+
+    return sorted(a.dask, key=key) == sorted(b.dask, key=key)
+
+
+def _not_empty(x):
+    return x.shape and 0 not in x.shape
+
+
+def assert_eq_shape(a, b, check_ndim=True, check_nan=True):
+    """Assert that two shapes are equal, handling NaN values."""
+    if check_ndim:
+        assert len(a) == len(b)
+
+    for aa, bb in zip(a, b):
+        if math.isnan(aa) or math.isnan(bb):
+            if check_nan:
+                assert math.isnan(aa) == math.isnan(bb)
+        else:
+            assert aa == bb
+
+
+def _check_chunks(x, check_ndim=True, scheduler=None):
+    """Check that chunk shapes match expected shapes."""
+    x = x.persist(scheduler=scheduler)
+    dsk = x.dask  # Cache to avoid repeated graph materialization
+    for idx in itertools.product(*(range(len(c)) for c in x.chunks)):
+        chunk = dsk[(x.name,) + idx]
+        if hasattr(chunk, "result"):  # it's a future
+            chunk = chunk.result()
+        if not hasattr(chunk, "dtype"):
+            chunk = np.array(chunk, dtype="O")
+        expected_shape = tuple(c[i] for c, i in zip(x.chunks, idx))
+        assert_eq_shape(expected_shape, chunk.shape, check_ndim=check_ndim, check_nan=False)
+        assert chunk.dtype == x.dtype, "maybe you forgot to pass the scheduler to `assert_eq`?"
+    return x
+
+
+def _get_dt_meta_computed(
+    x,
+    check_shape=True,
+    check_graph=True,
+    check_chunks=True,
+    check_ndim=True,
+    scheduler=None,
+):
+    """Get dtype, meta, and computed value from an array-like."""
+    x_original = x
+    x_meta = None
+    x_computed = None
+
+    if is_dask_collection(x) and is_arraylike(x):
+        assert x.dtype is not None
+        adt = x.dtype
+        # Note: check_graph is ignored in array-expr mode as it triggers
+        # expensive graph regeneration. The HLG validation can be enabled
+        # when needed for debugging.
+        x_meta = getattr(x, "_meta", None)
+        if check_chunks:
+            # Replace x with persisted version to avoid computing it twice.
+            x = _check_chunks(x, check_ndim=check_ndim, scheduler=scheduler)
+        x = x.compute(scheduler=scheduler)
+        x_computed = x
+        if hasattr(x, "todense"):
+            x = x.todense()
+        if not hasattr(x, "dtype"):
+            x = np.array(x, dtype="O")
+        if _not_empty(x):
+            assert x.dtype == x_original.dtype
+        if check_shape:
+            assert_eq_shape(x_original.shape, x.shape, check_nan=False)
+    else:
+        if not hasattr(x, "dtype"):
+            x = np.array(x, dtype="O")
+        adt = getattr(x, "dtype", None)
+
+    return x, adt, x_meta, x_computed
+
+
+def assert_eq(
+    a,
+    b,
+    check_shape=True,
+    check_graph=True,
+    check_meta=True,
+    check_chunks=True,
+    check_ndim=True,
+    check_type=True,
+    check_dtype=True,
+    equal_nan=True,
+    scheduler="sync",
+    **kwargs,
+):
+    """Assert that two arrays are equal.
+
+    This function handles dask arrays, numpy arrays, and other array-likes.
+    It computes dask arrays before comparison and performs various checks.
+
+    Parameters
+    ----------
+    a, b : array-like
+        Arrays to compare
+    check_shape : bool
+        Whether to check that shapes match
+    check_graph : bool
+        Whether to validate the dask graph (currently not implemented locally)
+    check_meta : bool
+        Whether to check metadata consistency
+    check_chunks : bool
+        Whether to check chunk shapes
+    check_ndim : bool
+        Whether to check that ndims match
+    check_type : bool
+        Whether to check that types match
+    check_dtype : bool
+        Whether to check that dtypes match
+    equal_nan : bool
+        Whether to treat NaN values as equal
+    scheduler : str
+        Scheduler to use for computing dask arrays
+
+    Returns
+    -------
+    bool
+        True if arrays are equal
+    """
+    a_original = a
+    b_original = b
+
+    if isinstance(a, (list, int, float)):
+        a = np.array(a)
+    if isinstance(b, (list, int, float)):
+        b = np.array(b)
+
+    a, adt, a_meta, a_computed = _get_dt_meta_computed(
+        a,
+        check_shape=check_shape,
+        check_graph=check_graph,
+        check_chunks=check_chunks,
+        check_ndim=check_ndim,
+        scheduler=scheduler,
+    )
+    b, bdt, b_meta, b_computed = _get_dt_meta_computed(
+        b,
+        check_shape=check_shape,
+        check_graph=check_graph,
+        check_chunks=check_chunks,
+        check_ndim=check_ndim,
+        scheduler=scheduler,
+    )
+
+    if check_dtype and str(adt) != str(bdt):
+        raise AssertionError(f"a and b have different dtypes: (a: {adt}, b: {bdt})")
+
+    try:
+        assert a.shape == b.shape, f"a and b have different shapes (a: {a.shape}, b: {b.shape})"
+        if check_type:
+            _a = a if a.shape else a.item()
+            _b = b if b.shape else b.item()
+            assert type(_a) == type(_b), f"a and b have different types (a: {type(_a)}, b: {type(_b)})"
+        if check_meta:
+            if hasattr(a, "_meta") and hasattr(b, "_meta"):
+                assert_eq(a._meta, b._meta)
+            if hasattr(a_original, "_meta"):
+                msg = (
+                    f"compute()-ing 'a' changes its number of dimensions "
+                    f"(before: {a_original._meta.ndim}, after: {a.ndim})"
+                )
+                assert a_original._meta.ndim == a.ndim, msg
+                if a_meta is not None:
+                    msg = (
+                        f"compute()-ing 'a' changes its type (before: {type(a_original._meta)}, after: {type(a_meta)})"
+                    )
+                    assert type(a_original._meta) == type(a_meta), msg
+                    if not (np.isscalar(a_meta) or np.isscalar(a_computed)):
+                        msg = (
+                            f"compute()-ing 'a' results in a different type than implied by its metadata "
+                            f"(meta: {type(a_meta)}, computed: {type(a_computed)})"
+                        )
+                        assert type(a_meta) == type(a_computed), msg
+            if hasattr(b_original, "_meta"):
+                msg = (
+                    f"compute()-ing 'b' changes its number of dimensions "
+                    f"(before: {b_original._meta.ndim}, after: {b.ndim})"
+                )
+                assert b_original._meta.ndim == b.ndim, msg
+                if b_meta is not None:
+                    msg = (
+                        f"compute()-ing 'b' changes its type (before: {type(b_original._meta)}, after: {type(b_meta)})"
+                    )
+                    assert type(b_original._meta) == type(b_meta), msg
+                    if not (np.isscalar(b_meta) or np.isscalar(b_computed)):
+                        msg = (
+                            f"compute()-ing 'b' results in a different type than implied by its metadata "
+                            f"(meta: {type(b_meta)}, computed: {type(b_computed)})"
+                        )
+                        assert type(b_meta) == type(b_computed), msg
+        msg = "found values in 'a' and 'b' which differ by more than the allowed amount"
+        assert allclose(a, b, equal_nan=equal_nan, **kwargs), msg
+        return True
+    except TypeError:
+        pass
+
+    c = a == b
+
+    if isinstance(c, np.ndarray):
+        assert c.all()
+    else:
+        assert c
+
+    return True

dask_array/_ufunc.py

@@ -0,0 +1,385 @@
+from __future__ import annotations
+
+import functools
+from functools import partial
+from operator import getitem
+
+import numpy as np
+
+from dask import core as dask_core
+from dask_array._new_collection import new_collection
+from dask_array._collection import Array, asarray, blockwise, elemwise
+from dask_array._expr import ArrayExpr
+from dask_array._core_utils import apply_infer_dtype
+from dask.base import is_dask_collection
+from dask.tokenize import _tokenize_deterministic, normalize_token
+from dask.utils import derived_from, funcname
+
+
+def wrap_elemwise(numpy_ufunc, source=np):
+    """Wrap up numpy function into dask.array"""
+
+    def wrapped(*args, **kwargs):
+        dsk = [arg for arg in args if hasattr(arg, "_elemwise")]
+        if len(dsk) > 0:
+            return dsk[0]._elemwise(numpy_ufunc, *args, **kwargs)
+        else:
+            return numpy_ufunc(*args, **kwargs)
+
+    # functools.wraps cannot wrap ufunc in Python 2.x
+    wrapped.__name__ = numpy_ufunc.__name__
+    return derived_from(source)(wrapped)
+
+
+class da_frompyfunc:
+    """A serializable `frompyfunc` object"""
+
+    def __init__(self, func, nin, nout):
+        self._ufunc = np.frompyfunc(func, nin, nout)
+        self._func = func
+        self.nin = nin
+        self.nout = nout
+        self._name = funcname(func)
+        self.__name__ = f"frompyfunc-{self._name}"
+
+    def outer(self, *args, **kwargs):
+        """Outer product - tokenizable because da_frompyfunc has __dask_tokenize__."""
+        return self._ufunc.outer(*args, **kwargs)
+
+    def __repr__(self):
+        return f"da.frompyfunc<{self._name}, {self.nin}, {self.nout}>"
+
+    def __dask_tokenize__(self):
+        return (normalize_token(self._func), self.nin, self.nout)
+
+    def __reduce__(self):
+        return (da_frompyfunc, (self._func, self.nin, self.nout))
+
+    def __call__(self, *args, **kwargs):
+        return self._ufunc(*args, **kwargs)
+
+    def __getattr__(self, a):
+        if not a.startswith("_"):
+            return getattr(self._ufunc, a)
+        raise AttributeError(f"{type(self).__name__!r} object has no attribute {a!r}")
+
+    def __dir__(self):
+        o = set(dir(type(self)))
+        o.update(self.__dict__)
+        o.update(dir(self._ufunc))
+        return list(o)
+
+
+@derived_from(np)
+def frompyfunc(func, nin, nout):
+    if nout > 1:
+        raise NotImplementedError("frompyfunc with more than one output")
+    return ufunc(da_frompyfunc(func, nin, nout))
+
+
+class ufunc:
+    _forward_attrs = {
+        "nin",
+        "nargs",
+        "nout",
+        "ntypes",
+        "identity",
+        "signature",
+        "types",
+    }
+
+    def __init__(self, ufunc):
+        if not isinstance(ufunc, (np.ufunc, da_frompyfunc)):
+            raise TypeError(f"must be an instance of `ufunc` or `da_frompyfunc`, got `{type(ufunc).__name__}")
+        self._ufunc = ufunc
+        self.__name__ = ufunc.__name__
+        if isinstance(ufunc, np.ufunc):
+            derived_from(np)(self)
+
+    def __dask_tokenize__(self):
+        return self.__name__, normalize_token(self._ufunc)
+
+    def __getattr__(self, key):
+        if key in self._forward_attrs:
+            return getattr(self._ufunc, key)
+        raise AttributeError(f"{type(self).__name__!r} object has no attribute {key!r}")
+
+    def __dir__(self):
+        return list(self._forward_attrs.union(dir(type(self)), self.__dict__))
+
+    def __repr__(self):
+        return repr(self._ufunc)
+
+    def __call__(self, *args, **kwargs):
+        # Validate kwargs - only allow known ufunc kwargs
+        valid_kwargs = {"out", "where", "dtype"}
+        extra_kwargs = set(kwargs) - valid_kwargs
+        if extra_kwargs:
+            raise TypeError(f"{self.__name__} does not take the following keyword arguments {sorted(extra_kwargs)}")
+
+        dsks = [arg for arg in args if hasattr(arg, "_elemwise")]
+        if len(dsks) > 0:
+            for dsk in dsks:
+                result = dsk._elemwise(self._ufunc, *args, **kwargs)
+                if type(result) != type(NotImplemented):
+                    return result
+            raise TypeError(f"Parameters of such types are not supported by {self.__name__}")
+        else:
+            return self._ufunc(*args, **kwargs)
+
+    @derived_from(np.ufunc)
+    def outer(self, A, B, **kwargs):
+        if self.nin != 2:
+            raise ValueError("outer product only supported for binary functions")
+        if "out" in kwargs:
+            raise ValueError("`out` kwarg not supported")
+
+        A_is_dask = is_dask_collection(A)
+        B_is_dask = is_dask_collection(B)
+        if not A_is_dask and not B_is_dask:
+            return self._ufunc.outer(A, B, **kwargs)
+        elif A_is_dask and not isinstance(A, Array) or B_is_dask and not isinstance(B, Array):
+            raise NotImplementedError("Dask objects besides `dask.array.Array` are not supported at this time.")
+
+        A = asarray(A)
+        B = asarray(B)
+        ndim = A.ndim + B.ndim
+        out_inds = tuple(range(ndim))
+        A_inds = out_inds[: A.ndim]
+        B_inds = out_inds[A.ndim :]
+
+        dtype = apply_infer_dtype(self._ufunc.outer, [A, B], kwargs, "ufunc.outer", suggest_dtype=False)
+
+        if "dtype" in kwargs:
+            func = partial(self._ufunc.outer, dtype=kwargs.pop("dtype"))
+        else:
+            func = self._ufunc.outer
+
+        return blockwise(
+            func,
+            out_inds,
+            A,
+            A_inds,
+            B,
+            B_inds,
+            dtype=dtype,
+            name=self.__name__ + ".outer-" + _tokenize_deterministic(self._ufunc),
+            **kwargs,
+        )
+
+
+# ufuncs, copied from this page:
+# https://docs.scipy.org/doc/numpy/reference/ufuncs.html
+
+# math operations
+add = ufunc(np.add)
+subtract = ufunc(np.subtract)
+multiply = ufunc(np.multiply)
+divide = ufunc(np.divide)
+logaddexp = ufunc(np.logaddexp)
+logaddexp2 = ufunc(np.logaddexp2)
+true_divide = ufunc(np.true_divide)
+floor_divide = ufunc(np.floor_divide)
+negative = ufunc(np.negative)
+positive = ufunc(np.positive)
+power = ufunc(np.power)
+float_power = ufunc(np.float_power)
+remainder = ufunc(np.remainder)
+mod = ufunc(np.mod)
+# fmod: see below
+conj = conjugate = ufunc(np.conjugate)
+exp = ufunc(np.exp)
+exp2 = ufunc(np.exp2)
+log = ufunc(np.log)
+log2 = ufunc(np.log2)
+log10 = ufunc(np.log10)
+log1p = ufunc(np.log1p)
+expm1 = ufunc(np.expm1)
+sqrt = ufunc(np.sqrt)
+square = ufunc(np.square)
+cbrt = ufunc(np.cbrt)
+reciprocal = ufunc(np.reciprocal)
+
+# trigonometric functions
+sin = ufunc(np.sin)
+cos = ufunc(np.cos)
+tan = ufunc(np.tan)
+arcsin = ufunc(np.arcsin)
+arccos = ufunc(np.arccos)
+arctan = ufunc(np.arctan)
+arctan2 = ufunc(np.arctan2)
+hypot = ufunc(np.hypot)
+sinh = ufunc(np.sinh)
+cosh = ufunc(np.cosh)
+tanh = ufunc(np.tanh)
+arcsinh = ufunc(np.arcsinh)
+arccosh = ufunc(np.arccosh)
+arctanh = ufunc(np.arctanh)
+deg2rad = ufunc(np.deg2rad)
+rad2deg = ufunc(np.rad2deg)
+
+# comparison functions
+greater = ufunc(np.greater)
+greater_equal = ufunc(np.greater_equal)
+less = ufunc(np.less)
+less_equal = ufunc(np.less_equal)
+not_equal = ufunc(np.not_equal)
+equal = ufunc(np.equal)
+isneginf = partial(equal, -np.inf)
+isposinf = partial(equal, np.inf)
+logical_and = ufunc(np.logical_and)
+logical_or = ufunc(np.logical_or)
+logical_xor = ufunc(np.logical_xor)
+logical_not = ufunc(np.logical_not)
+maximum = ufunc(np.maximum)
+minimum = ufunc(np.minimum)
+fmax = ufunc(np.fmax)
+fmin = ufunc(np.fmin)
+
+# bitwise functions
+bitwise_and = ufunc(np.bitwise_and)
+bitwise_or = ufunc(np.bitwise_or)
+bitwise_xor = ufunc(np.bitwise_xor)
+bitwise_not = ufunc(np.bitwise_not)
+invert = bitwise_not
+left_shift = ufunc(np.left_shift)
+right_shift = ufunc(np.right_shift)
+
+# floating functions
+isfinite = ufunc(np.isfinite)
+isinf = ufunc(np.isinf)
+isnan = ufunc(np.isnan)
+signbit = ufunc(np.signbit)
+copysign = ufunc(np.copysign)
+nextafter = ufunc(np.nextafter)
+spacing = ufunc(np.spacing)
+# modf: see below
+ldexp = ufunc(np.ldexp)
+# frexp: see below
+fmod = ufunc(np.fmod)
+floor = ufunc(np.floor)
+ceil = ufunc(np.ceil)
+trunc = ufunc(np.trunc)
+
+# more math routines, from this page:
+# https://docs.scipy.org/doc/numpy/reference/routines.math.html
+degrees = ufunc(np.degrees)
+radians = ufunc(np.radians)
+rint = ufunc(np.rint)
+fabs = ufunc(np.fabs)
+sign = ufunc(np.sign)
+absolute = ufunc(np.absolute)
+abs = absolute
+
+# non-ufunc elementwise functions
+clip = wrap_elemwise(np.clip)
+isreal = wrap_elemwise(np.isreal)
+iscomplex = wrap_elemwise(np.iscomplex)
+real = wrap_elemwise(np.real)
+imag = wrap_elemwise(np.imag)
+fix = wrap_elemwise(np.fix)
+i0 = wrap_elemwise(np.i0)
+sinc = wrap_elemwise(np.sinc)
+nan_to_num = wrap_elemwise(np.nan_to_num)
+
+
+@derived_from(np)
+def angle(x, deg=0):
+    deg = bool(deg)
+    if hasattr(x, "_elemwise"):
+        return x._elemwise(np.angle, x, deg)
+    return np.angle(x, deg=deg)
+
+
+class DoubleOutputs(ArrayExpr):
+    _parameters = ["array", "index", "meta", "name", "func"]
+
+    @functools.cached_property
+    def _meta(self):
+        meta = self.operand("meta")
+        a = np.empty_like(meta, shape=(1,) * meta.ndim, dtype=meta.dtype)
+        result = self.operand("func")(a)
+        return result[self.operand("index")]
+
+    @functools.cached_property
+    def chunks(self):
+        return self.array.chunks
+
+    @functools.cached_property
+    def _name(self):
+        return self.operand("name") + _tokenize_deterministic(*self.operands)
+
+    def _layer(self) -> dict:
+        return {
+            (self._name,) + key[1:]: (getitem, key, self.operand("index"))
+            for key in dask_core.flatten(self.array.__dask_keys__())
+        }
+
+
+@derived_from(np)
+def frexp(x):
+    # Not actually object dtype, just need to specify something
+    tmp = elemwise(np.frexp, x, dtype=object)
+    left = DoubleOutputs(tmp, 0, getattr(x, "_meta", x), "mantissa-", np.frexp)
+    right = DoubleOutputs(tmp, 1, getattr(x, "_meta", x), "exponent-", np.frexp)
+    return new_collection(left), new_collection(right)
+
+
+@derived_from(np)
+def modf(x):
+    # Not actually object dtype, just need to specify something
+    tmp = elemwise(np.modf, x, dtype=object)
+    left = DoubleOutputs(tmp, 0, getattr(x, "_meta", x), "modf1-", np.modf)
+    right = DoubleOutputs(tmp, 1, getattr(x, "_meta", x), "modf2-", np.modf)
+    return new_collection(left), new_collection(right)
+
+
+@derived_from(np)
+def divmod(x, y):
+    res1 = x // y
+    res2 = x % y
+    return res1, res2
+
+
+def round(a, decimals=0):
+    """Round an array to the given number of decimals."""
+    a = asarray(a)
+    return elemwise(np.round, a, dtype=a.dtype, decimals=decimals)
+
+
+@derived_from(np)
+def around(x, decimals=0):
+    """Evenly round to the given number of decimals."""
+    return round(x, decimals=decimals)
+
+
+def _asarray_isnull(values):
+    import pandas as pd
+
+    return np.asarray(pd.isnull(values))
+
+
+def isnull(values):
+    """pandas.isnull for dask arrays"""
+    # eagerly raise ImportError, if pandas isn't available
+    import pandas as pd  # noqa: F401
+
+    return elemwise(_asarray_isnull, values, dtype="bool")
+
+
+def notnull(values):
+    """pandas.notnull for dask arrays"""
+    return ~isnull(values)
+
+
+@derived_from(np)
+def isclose(arr1, arr2, rtol=1e-5, atol=1e-8, equal_nan=False):
+    """Returns a boolean array where two arrays are element-wise equal within a tolerance."""
+    func = partial(np.isclose, rtol=rtol, atol=atol, equal_nan=equal_nan)
+    return elemwise(func, arr1, arr2, dtype="bool")
+
+
+@derived_from(np)
+def allclose(arr1, arr2, rtol=1e-5, atol=1e-8, equal_nan=False):
+    """Returns True if two arrays are element-wise equal within a tolerance."""
+    return isclose(arr1, arr2, rtol=rtol, atol=atol, equal_nan=equal_nan).all()