dask-array 0.1.0__py3-none-any.whl

dask_array/random/_random_state.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import contextlib
+import importlib
+import numbers
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask_array.creation import arange
+from dask_array._backends_array import array_creation_dispatch
+from dask.utils import derived_from, typename
+
+from ._utils import _wrap_func
+
+
+class RandomState:
+    """
+    Mersenne Twister pseudo-random number generator
+
+    This object contains state to deterministically generate pseudo-random
+    numbers from a variety of probability distributions.  It is identical to
+    ``np.random.RandomState`` except that all functions also take a ``chunks=``
+    keyword argument.
+
+    Parameters
+    ----------
+    seed: Number
+        Object to pass to RandomState to serve as deterministic seed
+    RandomState: Callable[seed] -> RandomState
+        A callable that, when provided with a ``seed`` keyword provides an
+        object that operates identically to ``np.random.RandomState`` (the
+        default).  This might also be a function that returns a
+        ``mkl_random``, or ``cupy.random.RandomState`` object.
+
+    Examples
+    --------
+    >>> import dask_array as da
+    >>> state = da.random.RandomState(1234)  # a seed
+    >>> x = state.normal(10, 0.1, size=3, chunks=(2,))
+    >>> x.compute()
+    array([10.01867852, 10.04812289,  9.89649746])
+
+    See Also
+    --------
+    np.random.RandomState
+    """
+
+    def __init__(self, seed=None, RandomState=None):
+        self._numpy_state = np.random.RandomState(seed)
+        self._RandomState = array_creation_dispatch.RandomState if RandomState is None else RandomState
+
+    @property
+    def _backend(self):
+        # Assumes typename(self._RandomState) starts with
+        # an importable array-library name (e.g. "numpy" or "cupy")
+        _backend_name = typename(self._RandomState).split(".")[0]
+        return importlib.import_module(_backend_name)
+
+    def seed(self, seed=None):
+        self._numpy_state.seed(seed)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def beta(self, a, b, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "beta", a, b, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def binomial(self, n, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "binomial", n, p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def chisquare(self, df, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "chisquare", df, size=size, chunks=chunks, **kwargs)
+
+    with contextlib.suppress(AttributeError):
+
+        @derived_from(np.random.RandomState, skipblocks=1)
+        def choice(self, a, size=None, replace=True, p=None, chunks="auto"):
+            from ._choice import RandomChoice, _choice_validate_params
+
+            (
+                a_val,
+                a_expr,
+                size,
+                replace,
+                p_expr,
+                axis,  # np.random.RandomState.choice does not use axis
+                chunks,
+                meta,
+            ) = _choice_validate_params(self, a, size, replace, p, 0, chunks)
+
+            return new_collection(RandomChoice(a_val, a_expr, chunks, meta, self._numpy_state, replace, p_expr))
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def exponential(self, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "exponential", scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def f(self, dfnum, dfden, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "f", dfnum, dfden, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def gamma(self, shape, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "gamma", shape, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def geometric(self, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "geometric", p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def gumbel(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "gumbel", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def hypergeometric(self, ngood, nbad, nsample, size=None, chunks="auto", **kwargs):
+        return _wrap_func(
+            self,
+            "hypergeometric",
+            ngood,
+            nbad,
+            nsample,
+            size=size,
+            chunks=chunks,
+            **kwargs,
+        )
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def laplace(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "laplace", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def logistic(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "logistic", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def lognormal(self, mean=0.0, sigma=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "lognormal", mean, sigma, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def logseries(self, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "logseries", p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def multinomial(self, n, pvals, size=None, chunks="auto", **kwargs):
+        return _wrap_func(
+            self,
+            "multinomial",
+            n,
+            pvals,
+            size=size,
+            chunks=chunks,
+            extra_chunks=((len(pvals),),),
+            **kwargs,
+        )
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def negative_binomial(self, n, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "negative_binomial", n, p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def noncentral_chisquare(self, df, nonc, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "noncentral_chisquare", df, nonc, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def noncentral_f(self, dfnum, dfden, nonc, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "noncentral_f", dfnum, dfden, nonc, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def normal(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "normal", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def pareto(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "pareto", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def permutation(self, x):
+        from dask_array.slicing._utils import shuffle_slice
+
+        if isinstance(x, numbers.Number):
+            x = arange(x, chunks="auto")
+
+        index = np.arange(len(x))
+        self._numpy_state.shuffle(index)
+        return shuffle_slice(x, index)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def poisson(self, lam=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "poisson", lam, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def power(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "power", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def randint(self, low, high=None, size=None, chunks="auto", dtype="l", **kwargs):
+        return _wrap_func(self, "randint", low, high, size=size, chunks=chunks, dtype=dtype, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def random_integers(self, low, high=None, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "random_integers", low, high, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def random_sample(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "random_sample", size=size, chunks=chunks, **kwargs)
+
+    random = random_sample
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def rayleigh(self, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "rayleigh", scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def standard_cauchy(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_cauchy", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def standard_exponential(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_exponential", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def standard_gamma(self, shape, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_gamma", shape, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def standard_normal(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_normal", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def standard_t(self, df, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_t", df, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def tomaxint(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "tomaxint", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def triangular(self, left, mode, right, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "triangular", left, mode, right, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def uniform(self, low=0.0, high=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "uniform", low, high, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def vonmises(self, mu, kappa, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "vonmises", mu, kappa, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def wald(self, mean, scale, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "wald", mean, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def weibull(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "weibull", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.RandomState, skipblocks=1)
+    def zipf(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "zipf", a, size=size, chunks=chunks, **kwargs)

dask_array/random/_utils.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import importlib
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask_array._collection import Array
+from dask_array._core_utils import normalize_chunks
+from dask.utils import typename
+
+
+def _rng_from_bitgen(bitgen):
+    # Assumes typename(bitgen) starts with importable
+    # library name (e.g. "numpy" or "cupy")
+    backend_name = typename(bitgen).split(".")[0]
+    backend_lib = importlib.import_module(backend_name)
+    return backend_lib.random.default_rng(bitgen)
+
+
+def _shuffle(bit_generator, x, axis=0):
+    """Shuffle array in place and advance bit generator state."""
+    state_data = bit_generator.state
+    new_bitgen = type(bit_generator)()
+    new_bitgen.state = state_data
+    state = _rng_from_bitgen(new_bitgen)
+    state.shuffle(x, axis=axis)
+    # Copy advanced state back to original so subsequent calls get different results
+    bit_generator.state = new_bitgen.state
+
+
+def _broadcast_array_arg(arg, size, target_chunks):
+    """Broadcast and rechunk an array argument to match output shape."""
+    from dask_array._broadcast import broadcast_to
+    from dask_array.core._conversion import from_array
+
+    if isinstance(arg, np.ndarray) and arg.shape:
+        arg = from_array(arg, chunks=arg.shape)
+        arg = broadcast_to(arg, size).rechunk(target_chunks)
+    elif isinstance(arg, Array):
+        arg = broadcast_to(arg, size).rechunk(target_chunks)
+    return arg
+
+
+def _wrap_func(rng, funcname, *args, size=None, chunks="auto", extra_chunks=(), **kwargs):
+    from ._expr import RandomNormal, RandomPoisson
+
+    if size is not None and not isinstance(size, (tuple, list)):
+        size = (size,)
+
+    # Collect shapes from array arguments for broadcasting
+    shapes = []
+    for arg in args:
+        if isinstance(arg, (np.ndarray, Array)) and arg.shape:
+            shapes.append(arg.shape)
+    for v in kwargs.values():
+        if isinstance(v, (np.ndarray, Array)) and v.shape:
+            shapes.append(v.shape)
+
+    # Validate that all shapes can be broadcast together with size
+    if size is not None and shapes:
+        np.broadcast_shapes(*shapes, size)  # Raises ValueError if incompatible
+    elif size is None and shapes:
+        size = np.broadcast_shapes(*shapes)
+
+    # Broadcast and rechunk array arguments to match output shape/chunks
+    if size is not None and shapes:
+        target_chunks = normalize_chunks(chunks, size, dtype=kwargs.get("dtype", np.float64))
+        args = tuple(_broadcast_array_arg(arg, size, target_chunks) for arg in args)
+        kwargs = {k: _broadcast_array_arg(v, size, target_chunks) for k, v in kwargs.items()}
+
+    # Dispatch to specific subclass if available
+    if funcname == "normal":
+        loc = kwargs.pop("loc", args[0] if len(args) > 0 else 0.0)
+        scale = kwargs.pop("scale", args[1] if len(args) > 1 else 1.0)
+        return new_collection(RandomNormal(rng, size, chunks, extra_chunks, loc, scale))
+    elif funcname == "poisson":
+        lam = args[0] if len(args) > 0 else kwargs.pop("lam", 1.0)
+        return new_collection(RandomPoisson(rng, size, chunks, extra_chunks, lam))
+
+    # Fallback: use generic Random with args/kwargs tuples
+    from ._expr import Random
+
+    return new_collection(Random(rng, funcname, size, chunks, extra_chunks, args, kwargs))

dask_array/reductions/__init__.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dask_array.reductions._arg_reduction import arg_reduction
+from dask_array.reductions._common import (
+    all,
+    any,
+    argmax,
+    argmin,
+    nannumel,
+    nanargmax,
+    nanargmin,
+    max,
+    mean,
+    median,
+    min,
+    moment,
+    nanmax,
+    nanmean,
+    nanmedian,
+    nanmin,
+    nanprod,
+    nanquantile,
+    nanstd,
+    nansum,
+    nanvar,
+    numel,
+    prod,
+    quantile,
+    std,
+    sum,
+    var,
+)
+from dask_array.reductions._cumulative import (
+    cumprod,
+    cumreduction,
+    cumsum,
+    nancumprod,
+    nancumsum,
+)
+from dask_array.reductions._reduction import (
+    _tree_reduce,
+    reduction,
+)
+from dask_array.reductions._trace import trace
+from dask_array.reductions._percentile import nanpercentile, percentile
+
+__all__ = [
+    "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",
+    "_tree_reduce",
+]

dask_array/reductions/_arg_reduction.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import operator
+from itertools import product, repeat
+from numbers import Integral
+
+import numpy as np
+from tlz import accumulate, pluck
+
+from dask_array._expr import ArrayExpr
+from dask_array._utils import is_arraylike, validate_axis
+from dask.tokenize import _tokenize_deterministic
+from dask.utils import cached_property
+
+
+class ArgChunk(ArrayExpr):
+    """Expression for the initial chunk step of arg reductions (argmin/argmax).
+
+    Maps the chunk function across all blocks, tracking offsets to compute
+    global indices.
+    """
+
+    _parameters = ["array", "chunk_func", "axis", "ravel"]
+
+    @cached_property
+    def _name(self):
+        return "arg-chunk-" + _tokenize_deterministic(self.array, self.chunk_func, self.axis, self.ravel)
+
+    @cached_property
+    def _meta(self):
+        # The chunk function returns a structured array or dict with 'vals' and 'arg'
+        # fields. The dtype comes from argmin on the meta.
+        from dask_array._utils import asarray_safe, meta_from_array
+
+        dtype = np.argmin(asarray_safe([1], like=meta_from_array(self.array)))
+        if is_arraylike(dtype):
+            return dtype
+        # Return a small array with the correct dtype
+        return np.array([], dtype=np.intp)
+
+    @cached_property
+    def chunks(self):
+        # After the chunk step, each block is reduced to size 1 along the axis
+        return tuple((1,) * len(c) if i in self.axis else c for (i, c) in enumerate(self.array.chunks))
+
+    def _layer(self):
+        x = self.array
+        axis = self.axis
+        ravel = self.ravel
+
+        keys = list(product(*map(range, x.numblocks)))
+        offsets = list(product(*(accumulate(operator.add, bd[:-1], 0) for bd in x.chunks)))
+        if ravel:
+            offset_info = list(zip(offsets, repeat(x.shape)))
+        else:
+            offset_info = list(pluck(axis[0], offsets))
+
+        dsk = {}
+        for k, off in zip(keys, offset_info):
+            dsk[(self._name,) + tuple(k)] = (
+                self.chunk_func,
+                (x.name,) + tuple(k),
+                axis,
+                off,
+            )
+        return dsk
+
+
+def arg_reduction(x, chunk, combine, agg, axis=None, keepdims=False, split_every=None, out=None):
+    """Generic function for arg reductions in array-expr.
+
+    Parameters
+    ----------
+    x : Array
+    chunk : callable
+        Partialed ``arg_chunk``.
+    combine : callable
+        Partialed ``arg_combine``.
+    agg : callable
+        Partialed ``arg_agg``.
+    axis : int, optional
+    split_every : int or dict, optional
+    """
+    from dask_array.core._blockwise_funcs import _handle_out
+    from dask_array._utils import asarray_safe, meta_from_array
+
+    if axis is None:
+        axis = tuple(range(x.ndim))
+        ravel = True
+    elif isinstance(axis, Integral):
+        axis = validate_axis(axis, x.ndim)
+        axis = (axis,)
+        ravel = x.ndim == 1
+    else:
+        raise TypeError(f"axis must be either `None` or int, got '{axis}'")
+
+    for ax in axis:
+        chunks = x.chunks[ax]
+        if len(chunks) > 1 and np.isnan(chunks).any():
+            raise ValueError(
+                "Arg-reductions do not work with arrays that have "
+                "unknown chunksizes. At some point in your computation "
+                "this array lost chunking information.\n\n"
+                "A possible solution is with \n"
+                "  x.compute_chunk_sizes()"
+            )
+
+    # Create the ArgChunk expression for the initial chunk step
+    tmp = ArgChunk(x.expr, chunk, axis, ravel)
+
+    # Determine dtype
+    dtype = np.argmin(asarray_safe([1], like=meta_from_array(x)))
+    if hasattr(dtype, "dtype"):
+        dtype = dtype.dtype
+    else:
+        dtype = np.dtype(type(dtype))
+
+    # Import _tree_reduce from the same package
+    from dask_array.reductions._reduction import _tree_reduce
+
+    result = _tree_reduce(
+        tmp,
+        agg,
+        axis,
+        keepdims=keepdims,
+        dtype=dtype,
+        split_every=split_every,
+        combine=combine,
+    )
+    return _handle_out(out, result)